Skip to main content

rig_core/operation/
completion.rs

1//! Generating an assistant turn. A completion decoder writes its reply one
2//! provider item at a time, keyed by the item's wire index: the item opens
3//! when its first event arrives and takes the next position in the reply,
4//! its fragments grow it, and closing it finalizes its block with the
5//! provider's item as the block's native. Nothing regroups: blocks keep the
6//! order their items opened in. The provider's end of the reply is a
7//! [`Finish`], which the fold needs to produce the response.
8//!
9//! ```
10//! use rig_core::operation::Finish;
11//! use rig_core::completion::FinishReason;
12//!
13//! let finish = Finish {
14//!     reason: Some(FinishReason::Stop),
15//!     ..Finish::default()
16//! };
17//! assert_eq!(finish.reason, Some(FinishReason::Stop));
18//! ```
19
20use std::collections::{BTreeMap, HashMap, HashSet};
21
22use crate::completion::{CompletionRequest, CompletionResponse, FinishReason, Usage};
23use crate::error::ProviderError;
24use crate::message::citation;
25use crate::message::{
26    Api, AssistantContent, CallId, Image, LocalCallId, Opaque, Origin, Reasoning, Text, ToolCall,
27    ToolFunction, ToolName,
28};
29use crate::streaming::{Item, Part, PartKind, StreamEvent};
30use crate::telemetry::{GenAiOperation, SpanBuilder, SpanCombinator};
31use crate::wire::{
32    Assembled, Call, Descriptor, Emit, Fold, Mode, Operation, Out, Reply, Shared, WireCitation,
33};
34
35/// Generating an assistant turn, unary or streamed.
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub struct Completion;
38
39impl Operation for Completion {
40    type Request = CompletionRequest;
41    type Event = StreamEvent;
42    type End = Finish;
43    type Response = CompletionResponse;
44    type Fold = ReplyFold;
45    type Emit = Assembled;
46
47    /// The call's span names the model the request overrides to, when it
48    /// names one: every wire honours the override on encode.
49    fn fold(request: &Self::Request, call: &mut Call<'_>) -> Self::Fold {
50        let telemetry = call
51            .wire
52            .telemetry
53            .map_or(GenAiOperation::Chat, |telemetry| telemetry(call.mode));
54        debug_assert!(telemetry.is_completion());
55        // A wire that picks its API per request records the one it used.
56        let replay = call
57            .wire
58            .replay
59            .map(|target| target.route(request).unwrap_or(target));
60        let model = request
61            .model
62            .as_deref()
63            .or(call.wire.model)
64            .or_else(|| replay.map(|target| target.model()))
65            .unwrap_or_default();
66        let span = SpanBuilder::new(call.wire.name, model, telemetry)
67            .streaming(matches!(call.mode, Mode::Streaming))
68            .system_instructions(
69                request.system_instructions(),
70                request.record_telemetry_content,
71            )
72            .build();
73        call.instrument(span.clone());
74        let api = replay.map_or_else(
75            || Api::from(call.wire.name.to_owned()),
76            |target| target.api(),
77        );
78        let mut origin = Origin::new(api, call.wire.name, model);
79        // Only a wire that binds items to the context needs it recorded.
80        if replay.is_some_and(|target| target.binds_context(model)) {
81            origin.context =
82                replay.map(|target| crate::completion::history::context_of(request, target, model));
83        }
84        ReplyFold {
85            span,
86            facts: call.wire.facts.cloned().unwrap_or_default(),
87            wire: replay.is_some_and(|target| target.states_finish_reason()),
88            call_id_slot: replay.and_then(|target| target.call_id_slot()),
89            catalog_priced: replay.map(|target| target.catalog_priced()),
90            blob_slots: replay.map_or(&[], |target| target.blob_slots()),
91            accept_unknown_finish: request.accept_unknown_finish_reasons,
92            ..ReplyFold::new(origin)
93        }
94    }
95
96    /// The reply's normalized usage.
97    fn observed_usage(finish: &Finish) -> Option<&crate::completion::Usage> {
98        Some(&finish.usage)
99    }
100
101    /// Resolve the model the request addresses once, as the request's
102    /// `model`: the one it names, else the wire's. The history is checked
103    /// with [`CompletionRequest::validate_message_content`], then
104    /// [`adapt`](crate::completion::adapt)ed for that model on the wire's
105    /// replay target and checked again. Encoders, the fold and replay all
106    /// read the one resolved model. Each option the routed target refuses
107    /// is reported first, through
108    /// [`options::check`](crate::completion::options::check).
109    fn prepare(
110        mut request: Self::Request,
111        wire: &Descriptor<'_>,
112    ) -> Result<Self::Request, ProviderError> {
113        let Some(target) = wire.replay else {
114            return Err(ProviderError::request(format!(
115                "completion wire `{}` names no replay target",
116                wire.name
117            )));
118        };
119        // The caller's history is checked as written, so a rejection names
120        // its own messages; adapting never empties a message it keeps.
121        request.validate_message_content()?;
122        // A wire that addresses no model (an interaction read back) leaves
123        // it to the reply, which names the model the turn is from.
124        request.model = request
125            .model
126            .take()
127            .filter(|model| !model.is_empty())
128            .or_else(|| Some(target.model().to_owned()).filter(|model| !model.is_empty()));
129        let target = target.route(&request).unwrap_or(target);
130        // Every option the routed wire cannot honour is reported here, before
131        // any wire encodes, whatever its encoder does.
132        crate::completion::options::check(target, &mut request)?;
133        // Message cache points name the history as written; they are moved
134        // with it once it is adapted.
135        let written = request.cache_points.iter().any(|point| {
136            matches!(
137                point,
138                crate::completion::CachePoint::Messages(_)
139                    | crate::completion::CachePoint::Within { .. }
140            )
141        });
142        let written = written.then(|| {
143            let documents = (!target.takes_documents())
144                .then(|| request.normalized_documents())
145                .flatten();
146            (request.chat_history.clone(), documents)
147        });
148        let stored = target.continues_stored(&request);
149        let model = request.model.as_deref().unwrap_or(target.model());
150        // Only a target that binds items to the context reads it.
151        let binds = target.binds_context(model) && !target.drops_unbound_items(&request);
152        let shape = crate::completion::history::Request {
153            model: request.model.as_deref(),
154            stored,
155            foreign_reasoning: request.foreign_reasoning,
156            // A conversation the provider stores holds its own tools, so a
157            // continuation that declares none still calls them.
158            tools: stored || target.declares_tools(&request),
159            context: binds.then(|| crate::completion::history::context_of(&request, target, model)),
160        };
161        // The history moves into the adapter. Documents join it first, so
162        // the adapter's rules apply to them, unless the encoder sends them
163        // itself.
164        let mut history = std::mem::take(&mut request.chat_history);
165        if !target.takes_documents() {
166            if let Some(documents) = request.normalized_documents() {
167                crate::completion::request::insert_after_leading_system(&mut history, documents);
168            }
169            request.documents.clear();
170        }
171        // What the model's last request left; another request waits for it.
172        let mut kept = wire.prepared.and_then(|prepared| prepared.0.lock().ok());
173        let mut fresh = Default::default();
174        let memo = kept.as_deref_mut().unwrap_or(&mut fresh);
175        let adapted = crate::completion::history::adapt_for(history, target, &shape, memo);
176        drop(kept);
177        if let Some((written, documents)) = written {
178            request.cache_points = crate::completion::history::cache_points(
179                &request.cache_points,
180                &written,
181                documents.as_ref(),
182                &adapted,
183                target,
184                &shape,
185            );
186        }
187        request.chat_history = adapted;
188        request.validate_message_content()?;
189        Ok(request)
190    }
191}
192
193impl crate::wire::reply::Closing<Completion> for Assembled {
194    fn close(shared: &mut Shared<Completion>) {
195        let Shared { fold, items, .. } = shared;
196        fold.close_open(items);
197    }
198}
199
200impl Emit<Completion> for Assembled {}
201
202/// What the provider sends when it ends a completion reply.
203#[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)]
204pub struct Finish {
205    /// Token usage the provider reported. A counter it did not report is
206    /// `None`. A cost the provider reported goes in `usage.cost` and is
207    /// kept; without one the fold prices the counters at the catalog's
208    /// [`Pricing`](crate::catalog::Pricing) for the request's model, and
209    /// leaves the cost `None` when the catalog has no price for it.
210    pub usage: Usage,
211    /// Why the model stopped, when the provider said.
212    #[serde(default, skip_serializing_if = "Option::is_none")]
213    pub reason: Option<FinishReason>,
214    /// The response id.
215    pub response_id: Option<String>,
216    /// The model the provider reports.
217    pub model: Option<String>,
218    /// The provider's report that the turn failed, such as a refusal's
219    /// explanation. The turn is then never replayed.
220    #[serde(default, skip_serializing_if = "Option::is_none")]
221    pub error: Option<String>,
222}
223
224/// What a provider item becomes when it opens.
225#[derive(Debug, Clone, PartialEq)]
226pub enum Block {
227    /// Answer text.
228    Text,
229    /// Reasoning, possibly withheld by the provider.
230    Reasoning {
231        /// Whether the provider withheld the text.
232        redacted: bool,
233    },
234    /// An image, its canonical fields stated at open; fragments append
235    /// base64 data.
236    Image(Image),
237    /// A tool call with its id and name.
238    Call {
239        /// The call's id.
240        id: CallId,
241        /// The tool's name.
242        name: ToolName,
243    },
244    /// An item with no canonical meaning.
245    Opaque {
246        /// Whether the item goes back to the model that produced it.
247        replay: bool,
248    },
249}
250
251/// The completion fold. The driver and the handler writer build it; no other
252/// code can feed one.
253///
254/// It is both sides of one reply: the writer state a decoder writes through
255/// (the open items by wire index) and the blocks the consumer has taken, in
256/// their position.
257pub struct ReplyFold {
258    span: tracing::Span,
259    /// Who the reply is from; the end adds the provider's model and id.
260    origin: Origin,
261    /// The facts of the wire that answers, which price the reply.
262    facts: crate::catalog::ModelFacts,
263    // The writer.
264    open: BTreeMap<usize, Draft>,
265    /// The position of the block each wire index last closed, whose item
266    /// [`Out::edit`] may still amend until the reply ends.
267    ended: HashMap<usize, usize>,
268    next_part: u32,
269    call_ids: HashSet<CallId>,
270    /// Whether the reply's end must name a finish reason: a wire's fold
271    /// whose target states one (a relayed or written reply states its own).
272    wire: bool,
273    /// Where the wire's call items hold their id ([`ReplayTarget::call_id_slot`]).
274    ///
275    /// [`ReplayTarget::call_id_slot`]: crate::completion::ReplayTarget::call_id_slot
276    call_id_slot: Option<&'static str>,
277    /// Whether the catalog prices the reply, as its wire's target answers
278    /// ([`ReplayTarget::catalog_priced`]). `None` for a relayed or written
279    /// reply, which the registered dialect of its origin's provider answers.
280    ///
281    /// [`ReplayTarget::catalog_priced`]: crate::completion::ReplayTarget::catalog_priced
282    catalog_priced: Option<bool>,
283    /// Where the wire's items keep a reasoning blob ([`ReplayTarget::blob_slots`]).
284    ///
285    /// [`ReplayTarget::blob_slots`]: crate::completion::ReplayTarget::blob_slots
286    blob_slots: &'static [&'static str],
287    /// Whether the request accepts an unknown finish reason as a normal stop
288    /// ([`CompletionRequest::accept_unknown_finish_reasons`]).
289    ///
290    /// [`CompletionRequest::accept_unknown_finish_reasons`]: field@CompletionRequest::accept_unknown_finish_reasons
291    accept_unknown_finish: bool,
292    /// The first position whose item closed without the provider stating it
293    /// complete. An item there may be the partner a later one needs, so
294    /// blocks from it on replay from their canonical fields.
295    first_incomplete: Option<usize>,
296    /// The calls still open when the provider ended the reply, in the
297    /// order they opened ([`CompletionResponse::unfinished_calls`]).
298    unfinished: Vec<CallId>,
299    /// The call the last index-less fragment went to.
300    last_call: Option<usize>,
301    /// The index of the block a boundary-less wire is streaming.
302    run: Option<usize>,
303    next_auto: usize,
304    /// The wire index of each position, for a wire whose choice follows
305    /// its indices ([`Out::order_by_index`]).
306    by_index: Option<BTreeMap<usize, usize>>,
307    /// The position of the block the choice holds first ([`Out::lead`]).
308    lead: Option<usize>,
309    // The fold.
310    choice: Vec<Option<AssistantContent>>,
311    /// The text the consumer took of text and reasoning parts still open,
312    /// by position.
313    open_text: BTreeMap<usize, AssistantContent>,
314}
315
316/// One open provider item.
317struct Draft {
318    part: Part,
319    /// Whether its start was emitted.
320    started: bool,
321    /// The provider's item as assembled so far; `Null` for none.
322    item: serde_json::Value,
323    body: Body,
324    /// The citations of a text item, resolved when it closes.
325    citations: Vec<WireCitation>,
326}
327
328/// Whether a closing item keeps its provider item as the block's native.
329#[derive(Clone, Copy, PartialEq, Eq)]
330enum Closing {
331    /// The provider stated the item complete: it becomes the native.
332    Complete,
333    /// The item never completed: the block has no native, and an opaque
334    /// item does not replay.
335    Incomplete,
336}
337
338enum Body {
339    Text(String),
340    Image(Image),
341    Reasoning {
342        text: String,
343        redacted: bool,
344    },
345    Call {
346        id: Option<CallId>,
347        name: String,
348        arguments: Arguments,
349    },
350    Opaque {
351        replay: bool,
352    },
353}
354
355impl Body {
356    fn kind(&self) -> PartKind {
357        match self {
358            Self::Text(_) => PartKind::Text,
359            Self::Reasoning { .. } => PartKind::Reasoning,
360            Self::Image(_) => PartKind::Image,
361            Self::Call { .. } => PartKind::ToolCall,
362            Self::Opaque { .. } => PartKind::Opaque,
363        }
364    }
365}
366
367/// A tool call's argument text as it arrives.
368#[derive(Default)]
369struct Arguments {
370    text: String,
371    /// How many bytes of `text` the stream has carried.
372    streamed: usize,
373    overflowed: bool,
374    /// Whether any fragment carried a non-blank byte.
375    substantive: bool,
376    /// Arguments the provider announced when the call opened, used only
377    /// when no fragment arrives.
378    announced: Option<serde_json::Value>,
379}
380
381/// The most argument bytes one call accumulates.
382const MAX_TOOL_INPUT_BYTES: usize = 32 * 1024 * 1024;
383
384impl Arguments {
385    fn push(&mut self, fragment: &str, name: &str) {
386        self.substantive |= !fragment.trim().is_empty();
387        // Some OpenAI-compatible gateways send a literal `null` before the
388        // real fragments; a non-blank fragment supersedes it.
389        if self.text.trim() == "null" && !fragment.trim().is_empty() {
390            self.text.clear();
391        }
392        if self.text.len().saturating_add(fragment.len()) > MAX_TOOL_INPUT_BYTES {
393            if !self.overflowed {
394                self.overflowed = true;
395                tracing::warn!(
396                    tool = name,
397                    "streamed tool-call input exceeded the accumulation bound; truncating"
398                );
399            }
400        } else {
401            self.text.push_str(fragment);
402        }
403    }
404
405    /// The call to `name` these arguments make: the announced ones when no
406    /// fragment arrived, else the text, read by [`ToolFunction::parse`].
407    fn function(&self, name: ToolName) -> ToolFunction {
408        if self.text.is_empty()
409            && let Some(announced) = &self.announced
410        {
411            return ToolFunction::new(name, announced.clone());
412        }
413        let mut function = ToolFunction::parse(name, &self.text);
414        if self.overflowed {
415            function.invalid_arguments = Some(self.text.clone());
416        }
417        function
418    }
419
420    /// The text the stream has not carried yet, now marked carried. A text
421    /// that may still be a gateway's placeholder `null` waits, since the
422    /// first real fragment replaces it.
423    fn unstreamed(&mut self) -> Option<String> {
424        if "null".starts_with(self.text.trim()) {
425            return None;
426        }
427        let fragment = self.text.get(self.streamed..)?;
428        if fragment.is_empty() {
429            return None;
430        }
431        let fragment = fragment.to_owned();
432        self.streamed = self.text.len();
433        Some(fragment)
434    }
435
436    /// Whether the text is a complete JSON object.
437    fn complete(&self) -> bool {
438        !self.overflowed && crate::json_utils::is_whole_object(&self.text)
439    }
440}
441
442/// Whether `id` names a call: some providers send `""` or `"null"` for a
443/// call they gave no id.
444fn stated(id: &str) -> bool {
445    !id.is_empty() && id != "null"
446}
447
448/// What one fragment of a buffered tool call carries.
449#[derive(Debug, Clone, Copy, Default)]
450pub struct CallFragment<'a> {
451    /// The provider's id for the call.
452    pub id: Option<&'a str>,
453    /// The tool's name.
454    pub name: Option<&'a str>,
455    /// A fragment of the argument JSON.
456    pub arguments: Option<&'a str>,
457}
458
459/// The first index the writer hands out itself ([`Out::fresh_index`]); a
460/// provider's wire indices stay below it.
461pub const AUTO_INDEX: usize = 1 << 30;
462
463fn not_open(index: usize) -> ProviderError {
464    ProviderError::Response(format!(
465        "the reply wrote to item {index}, which is not open"
466    ))
467}
468
469/// Carry what the open call `draft` has not streamed yet: its start once it
470/// names a tool, then the argument text since the last fragment. A call
471/// that never names a tool stays unseen, as its close drops it.
472fn stream_call(items: &mut Items, draft: &mut Draft) {
473    let Body::Call {
474        name, arguments, ..
475    } = &mut draft.body
476    else {
477        return;
478    };
479    if !draft.started {
480        let Ok(name) = ToolName::new(name.clone()) else {
481            return;
482        };
483        draft.started = true;
484        emit(
485            items,
486            StreamEvent::Start {
487                part: draft.part,
488                kind: PartKind::ToolCall,
489                name: Some(name),
490            },
491        );
492    }
493    if let Some(json) = arguments.unstreamed() {
494        emit(
495            items,
496            StreamEvent::Arguments {
497                part: draft.part,
498                json,
499            },
500        );
501    }
502}
503
504impl ReplyFold {
505    /// The writer of a reply from `origin`.
506    pub(crate) fn new(origin: Origin) -> Self {
507        Self {
508            span: tracing::Span::none(),
509            origin,
510            facts: crate::catalog::ModelFacts::default(),
511            open: BTreeMap::new(),
512            ended: HashMap::new(),
513            next_part: 0,
514            call_ids: HashSet::new(),
515            wire: false,
516            call_id_slot: None,
517            catalog_priced: None,
518            blob_slots: &[],
519            accept_unknown_finish: false,
520            first_incomplete: None,
521            unfinished: Vec::new(),
522            last_call: None,
523            run: None,
524            next_auto: AUTO_INDEX,
525            by_index: None,
526            lead: None,
527            choice: Vec::new(),
528            open_text: BTreeMap::new(),
529        }
530    }
531
532    /// The fold of a stream relayed from another fold, which built its
533    /// events: it only collects them, and takes its origin from the relay
534    /// ([`Self::set_origin`]). Until then it names the relay `label`.
535    pub(crate) fn relayed(label: impl Into<String>) -> Self {
536        let label = label.into();
537        Self::new(Origin::new(label.clone(), label.clone(), label))
538    }
539
540    /// Who the reply is from, as the relay or the writer's owner states it.
541    pub(crate) fn set_origin(&mut self, origin: Origin) {
542        self.origin = origin;
543    }
544
545    /// The pricing `facts` give `origin`'s model (the spec the wire was
546    /// connected to, else its catalog's entry), or `None` when they list
547    /// none or the catalog does not price the wire.
548    fn pricing(&self, origin: &Origin) -> Option<crate::catalog::Pricing> {
549        let priced = self.catalog_priced.unwrap_or_else(|| {
550            crate::providers::registry::ProviderId::catalog(&origin.provider)
551                .is_none_or(|id| id.catalog_priced())
552        });
553        if !priced {
554            return None;
555        }
556        self.facts
557            .for_model(&origin.provider, &origin.model)?
558            .pricing
559    }
560
561    /// The next position in the choice.
562    fn next(&mut self) -> Part {
563        let part = Part::new(self.next_part);
564        self.next_part += 1;
565        part
566    }
567
568    fn fresh_index(&mut self) -> usize {
569        let index = self.next_auto;
570        self.next_auto += 1;
571        index
572    }
573
574    fn draft(&mut self, index: usize) -> Result<&mut Draft, ProviderError> {
575        self.open.get_mut(&index).ok_or_else(|| not_open(index))
576    }
577
578    pub(crate) fn open_item(
579        &mut self,
580        items: &mut Items,
581        index: usize,
582        block: Block,
583        item: serde_json::Value,
584    ) -> Result<(), ProviderError> {
585        if self.open.contains_key(&index) {
586            return Err(ProviderError::Response(format!(
587                "the reply opened item {index} twice"
588            )));
589        }
590        let body = match block {
591            Block::Text => Body::Text(String::new()),
592            Block::Image(image) => Body::Image(image),
593            Block::Reasoning { redacted } => Body::Reasoning {
594                text: String::new(),
595                redacted,
596            },
597            Block::Call { id, name } => Body::Call {
598                id: Some(id),
599                name: name.into(),
600                arguments: Arguments::default(),
601            },
602            Block::Opaque { replay } => Body::Opaque { replay },
603        };
604        self.insert(index, body, item);
605        stream_call(items, self.draft(index)?);
606        Ok(())
607    }
608
609    fn insert(&mut self, index: usize, body: Body, item: serde_json::Value) {
610        self.ended.remove(&index);
611        let part = self.next();
612        if let Some(by_index) = &mut self.by_index {
613            by_index.insert(part.index(), index);
614        }
615        let started = false;
616        self.open.insert(
617            index,
618            Draft {
619                part,
620                started,
621                item,
622                body,
623                citations: Vec::new(),
624            },
625        );
626    }
627
628    pub(crate) fn push_item(
629        &mut self,
630        items: &mut Items,
631        index: usize,
632        fragment: &str,
633    ) -> Result<(), ProviderError> {
634        if fragment.is_empty() {
635            return Ok(());
636        }
637        let draft = self.draft(index)?;
638        let event = match &mut draft.body {
639            Body::Text(text) => {
640                text.push_str(fragment);
641                StreamEvent::Text {
642                    part: draft.part,
643                    text: fragment.to_owned(),
644                }
645            }
646            Body::Reasoning { text, .. } => {
647                text.push_str(fragment);
648                StreamEvent::Reasoning {
649                    part: draft.part,
650                    text: fragment.to_owned(),
651                }
652            }
653            Body::Call {
654                name, arguments, ..
655            } => {
656                arguments.push(fragment, name);
657                stream_call(items, draft);
658                return Ok(());
659            }
660            Body::Image(image) => {
661                if let crate::message::DocumentSourceKind::Base64(data) = &mut image.data {
662                    data.push_str(fragment);
663                    return Ok(());
664                }
665                return Err(ProviderError::Response(format!(
666                    "the reply wrote data to the image item {index}, which holds no base64 data"
667                )));
668            }
669            Body::Opaque { .. } => {
670                return Err(ProviderError::Response(format!(
671                    "the reply wrote text to the opaque item {index}"
672                )));
673            }
674        };
675        if !draft.started {
676            draft.started = true;
677            emit(
678                items,
679                StreamEvent::Start {
680                    part: draft.part,
681                    kind: draft.body.kind(),
682                    name: None,
683                },
684            );
685        }
686        emit(items, event);
687        Ok(())
688    }
689
690    /// The id a call closing now takes, and the provider item it keeps:
691    /// `id`, or, when the provider sent none or one an earlier call of the
692    /// reply holds, a fresh rig-issued id. A renamed call
693    /// keeps no item, since the item names the id it lost; a call sent with
694    /// no id keeps its item only on a wire whose items have an id slot,
695    /// which replay then fills with the id its result gets.
696    fn distinct_call_id(
697        &mut self,
698        id: Option<CallId>,
699        item: serde_json::Value,
700    ) -> (CallId, serde_json::Value) {
701        let (id, item) = match id {
702            Some(id) if !self.call_ids.contains(&id) => (id, item),
703            Some(id) => {
704                tracing::warn!(%id, "the provider named two tool calls with one id; renaming the second");
705                (CallId::Local(LocalCallId::new()), serde_json::Value::Null)
706            }
707            None if self.call_id_slot.is_some() => (CallId::Local(LocalCallId::new()), item),
708            None => (CallId::Local(LocalCallId::new()), serde_json::Value::Null),
709        };
710        self.call_ids.insert(id.clone());
711        (id, item)
712    }
713
714    /// Close the item at `index`: its block becomes visible, holding the
715    /// provider's item as its native when `closing` says it is complete.
716    fn close_item(
717        &mut self,
718        items: &mut Items,
719        index: usize,
720        closing: Closing,
721    ) -> Result<(), ProviderError> {
722        self.close_block(items, index, closing).map(drop)
723    }
724
725    /// [`Self::close_item`], saying the id of the call it closed, if it
726    /// closed one.
727    fn close_block(
728        &mut self,
729        items: &mut Items,
730        index: usize,
731        closing: Closing,
732    ) -> Result<Option<CallId>, ProviderError> {
733        let mut draft = self.open.remove(&index).ok_or_else(|| not_open(index))?;
734        // A call streams what it has not yet: its start, if it names a tool,
735        // and the argument text since its last fragment.
736        stream_call(items, &mut draft);
737        if self.run == Some(index) {
738            self.run = None;
739        }
740        self.ended.insert(index, draft.part.index());
741        let Draft {
742            part,
743            started,
744            item,
745            body,
746            citations,
747        } = draft;
748        let item = match (closing, &body) {
749            (Closing::Complete, _) | (Closing::Incomplete, Body::Opaque { .. }) => item,
750            (Closing::Incomplete, _) => {
751                self.cut_at(part.index());
752                serde_json::Value::Null
753            }
754        };
755        let content = match body {
756            Body::Text(text) => {
757                if text.is_empty() && item.is_null() {
758                    return Ok(None);
759                }
760                // Resolved at close: a provider may cite a text before its
761                // last fragment arrives.
762                let mut text = Text::new(text);
763                citation::attach(
764                    &mut text,
765                    Vec::new(),
766                    citations,
767                    &self.origin.provider,
768                    index,
769                );
770                self.with_item(AssistantContent::Text(text), item)
771            }
772            Body::Reasoning { text, redacted } => {
773                if text.is_empty() && !redacted && item.is_null() {
774                    return Ok(None);
775                }
776                let reasoning = Reasoning {
777                    text,
778                    redacted,
779                    native: None,
780                };
781                self.with_item(AssistantContent::Reasoning(reasoning), item)
782            }
783            Body::Image(image) => self.with_item(AssistantContent::Image(image), item),
784            Body::Opaque { replay } => AssistantContent::Opaque(Opaque {
785                item,
786                replay: replay && closing == Closing::Complete,
787            }),
788            Body::Call {
789                id,
790                name,
791                arguments,
792            } => {
793                let Ok(name) = ToolName::new(name) else {
794                    tracing::warn!(
795                        index,
796                        "the provider closed a tool call without a name; nothing can answer it"
797                    );
798                    return Ok(None);
799                };
800                let function = arguments.function(name);
801                let (id, item) = self.distinct_call_id(id, item);
802                // A call that streamed no text (none arrived, or only a
803                // placeholder `null`) states the arguments it ends with.
804                if arguments.streamed == 0 {
805                    let json = serde_json::Value::Object(function.arguments.clone()).to_string();
806                    emit(items, StreamEvent::Arguments { part, json });
807                }
808                let call = ToolCall::new(id.clone(), function);
809                let content = self.with_item(AssistantContent::ToolCall(call), item);
810                emit(items, StreamEvent::End { part, content });
811                return Ok(Some(id));
812            }
813        };
814        if !started {
815            emit(
816                items,
817                StreamEvent::Start {
818                    part,
819                    kind: kind_of(&content),
820                    name: None,
821                },
822            );
823        }
824        emit(items, StreamEvent::End { part, content });
825        Ok(None)
826    }
827
828    /// `block` holding `item` as its native, with where it keeps its blob,
829    /// unless there is no item.
830    fn with_item(&self, block: AssistantContent, item: serde_json::Value) -> AssistantContent {
831        if item.is_null() {
832            return block;
833        }
834        let mut block = block.with_native(item);
835        if let Some(Some(native)) = block.native_slot() {
836            native.find_blob(self.blob_slots);
837        }
838        block
839    }
840
841    /// Edit the item of the block the wire index `index` last closed, while
842    /// its end event waits to be taken or once the fold holds it.
843    fn edit_ended(
844        &mut self,
845        items: &mut Items,
846        index: usize,
847        edit: impl FnOnce(&mut serde_json::Value),
848    ) -> Result<(), ProviderError> {
849        let position = *self.ended.get(&index).ok_or_else(|| not_open(index))?;
850        let slots = self.blob_slots;
851        let queued = items.iter_mut().find_map(|item| match item {
852            Ok(Item::Event(StreamEvent::End { part, content })) if part.index() == position => {
853                Some(content)
854            }
855            _ => None,
856        });
857        let content = match queued {
858            Some(content) => Some(content),
859            None => self.choice.get_mut(position).and_then(Option::as_mut),
860        };
861        if let Some(content) = content {
862            if let Some(item) = item_of(content) {
863                edit(item);
864            }
865            if let Some(Some(native)) = content.native_slot() {
866                native.find_blob(slots);
867            }
868        }
869        Ok(())
870    }
871
872    /// Add `citations` to the text item at `index`, or with `replace` put
873    /// them in place of its own. An open item resolves them when it closes;
874    /// a closed one, against its final text now, while its end event waits
875    /// to be taken or once the fold holds it. Citations for any other item
876    /// are dropped with a warning.
877    fn cite_item(
878        &mut self,
879        items: &mut Items,
880        index: usize,
881        citations: Vec<WireCitation>,
882        replace: bool,
883    ) {
884        let provider = self.origin.provider.as_str();
885        if let Some(draft) = self.open.get_mut(&index) {
886            if !matches!(draft.body, Body::Text(_)) {
887                tracing::warn!(
888                    provider,
889                    index,
890                    "dropped citations of an item that is not text"
891                );
892            } else if replace {
893                draft.citations = citations;
894            } else {
895                draft.citations.extend(citations);
896            }
897            return;
898        }
899        let Some(position) = self.ended.get(&index).copied() else {
900            tracing::warn!(
901                provider,
902                index,
903                "dropped citations of an item the reply never opened"
904            );
905            return;
906        };
907        let queued = items.iter_mut().find_map(|item| match item {
908            Ok(Item::Event(StreamEvent::End { part, content })) if part.index() == position => {
909                Some(content)
910            }
911            _ => None,
912        });
913        let content = match queued {
914            Some(content) => Some(content),
915            None => self.choice.get_mut(position).and_then(Option::as_mut),
916        };
917        match content {
918            Some(AssistantContent::Text(text)) => {
919                let kept = if replace {
920                    Vec::new()
921                } else {
922                    text.citations().to_vec()
923                };
924                citation::attach(text, kept, citations, provider, index);
925            }
926            _ => tracing::warn!(
927                provider,
928                index,
929                "dropped citations of an item that is not text"
930            ),
931        }
932    }
933
934    pub(crate) fn run_item(
935        &mut self,
936        items: &mut Items,
937        block: Block,
938        fragment: &str,
939    ) -> Result<usize, ProviderError> {
940        let current = self.run.filter(|index| {
941            self.open.get(index).is_some_and(|draft| {
942                matches!(
943                    (&draft.body, &block),
944                    (Body::Text(_), Block::Text)
945                        | (
946                            Body::Reasoning {
947                                redacted: false,
948                                ..
949                            },
950                            Block::Reasoning { redacted: false }
951                        )
952                )
953            })
954        });
955        let index = match current {
956            Some(index) => index,
957            None => {
958                self.end_run(items)?;
959                let index = self.fresh_index();
960                self.open_item(items, index, block, serde_json::Value::Null)?;
961                self.run = Some(index);
962                index
963            }
964        };
965        self.push_item(items, index, fragment)?;
966        Ok(index)
967    }
968
969    pub(crate) fn end_run(&mut self, items: &mut Items) -> Result<(), ProviderError> {
970        match self.run.take() {
971            Some(index) => self.close_item(items, index, Closing::Complete),
972            None => Ok(()),
973        }
974    }
975
976    /// One whole block at the next position, its native kept as given.
977    pub(crate) fn write_content(
978        &mut self,
979        items: &mut Items,
980        content: AssistantContent,
981    ) -> Result<(), ProviderError> {
982        let content = match content {
983            AssistantContent::ToolCall(mut call) => {
984                let (id, _) = self.distinct_call_id(Some(call.id.clone()), serde_json::Value::Null);
985                if id != call.id {
986                    call.native = None;
987                }
988                call.id = id;
989                AssistantContent::ToolCall(call)
990            }
991            content => content,
992        };
993        let part = self.next();
994        let name = match &content {
995            AssistantContent::ToolCall(call) => Some(call.function.name.clone()),
996            _ => None,
997        };
998        emit(
999            items,
1000            StreamEvent::Start {
1001                part,
1002                kind: kind_of(&content),
1003                name,
1004            },
1005        );
1006        match &content {
1007            AssistantContent::Text(text) if !text.text.is_empty() => emit(
1008                items,
1009                StreamEvent::Text {
1010                    part,
1011                    text: text.text.clone(),
1012                },
1013            ),
1014            AssistantContent::Reasoning(reasoning) if !reasoning.text.is_empty() => emit(
1015                items,
1016                StreamEvent::Reasoning {
1017                    part,
1018                    text: reasoning.text.clone(),
1019                },
1020            ),
1021            AssistantContent::ToolCall(call) => emit(
1022                items,
1023                StreamEvent::Arguments {
1024                    part,
1025                    json: call.function.arguments_json(),
1026                },
1027            ),
1028            _ => {}
1029        }
1030        emit(items, StreamEvent::End { part, content });
1031        Ok(())
1032    }
1033
1034    /// Close every item still open when the provider ends the reply, in the
1035    /// order they opened. The block a boundary-less wire was streaming is
1036    /// complete; any other item was never stated complete, so its block has
1037    /// no native. Calls keep what their arguments state.
1038    pub(crate) fn close_open(&mut self, items: &mut Items) {
1039        let mut open: Vec<(Part, usize)> = self
1040            .open
1041            .iter()
1042            .map(|(index, draft)| (draft.part, *index))
1043            .collect();
1044        open.sort();
1045        for (_, index) in open {
1046            let closing = if self.run == Some(index) {
1047                Closing::Complete
1048            } else {
1049                Closing::Incomplete
1050            };
1051            match self.close_block(items, index, closing) {
1052                Ok(Some(call)) if closing == Closing::Incomplete => self.unfinished.push(call),
1053                Ok(_) => {}
1054                Err(error) => items.push_back(Err(error)),
1055            }
1056        }
1057    }
1058
1059    /// Record that the item at `position` never completed.
1060    fn cut_at(&mut self, position: usize) {
1061        self.first_incomplete = Some(
1062            self.first_incomplete
1063                .map_or(position, |first| first.min(position)),
1064        );
1065    }
1066
1067    /// Who the reply is from.
1068    pub fn origin(&self) -> &Origin {
1069        &self.origin
1070    }
1071
1072    /// What arrived so far as a response: every part that ended, the text
1073    /// the consumer already took of a text or reasoning part still open,
1074    /// and the provider's end when it arrived. A reply the provider did not
1075    /// end is a failed turn: `failure` is the error that ended it, and
1076    /// without one the caller stopped reading. Unless the provider ended the
1077    /// reply and the consumer took every part, no block keeps its provider
1078    /// item.
1079    pub(crate) fn partial(
1080        &self,
1081        end: Option<&Finish>,
1082        reply: &Reply,
1083        failure: Option<&ProviderError>,
1084    ) -> CompletionResponse {
1085        let mut response =
1086            self.response(end.cloned().unwrap_or_default(), reply.clone(), Vec::new());
1087        if end.is_none() {
1088            match failure {
1089                Some(error) => {
1090                    response.error.get_or_insert_with(|| error.to_string());
1091                }
1092                None => {
1093                    response.aborted = Some(
1094                        "the caller stopped reading before the provider ended the reply".to_owned(),
1095                    );
1096                }
1097            }
1098        }
1099        // A turn the consumer has not wholly taken replays only its canonical
1100        // fields, as pi drops an unfinished turn's provider data: no item of
1101        // it can then need a partner that never arrived.
1102        let unfinished =
1103            end.is_none() || !self.open.is_empty() || self.choice.iter().any(Option::is_none);
1104        let canonical_from = if unfinished {
1105            Some(0)
1106        } else {
1107            self.first_incomplete
1108        };
1109        let parts = self
1110            .choice
1111            .iter()
1112            .enumerate()
1113            .map(|(index, part)| {
1114                let part = part.clone().or_else(|| self.open_text.get(&index).cloned());
1115                (index, part)
1116            })
1117            .collect();
1118        response.choice = self.ordered(parts, canonical_from);
1119        response
1120    }
1121
1122    /// `parts` by position, in wire-index order on a wire that asked for
1123    /// it and after the lead block, each from position `canonical_from` on
1124    /// with no provider item.
1125    fn ordered(
1126        &self,
1127        parts: Vec<(usize, Option<AssistantContent>)>,
1128        canonical_from: Option<usize>,
1129    ) -> Vec<AssistantContent> {
1130        let mut parts: Vec<(usize, AssistantContent)> = parts
1131            .into_iter()
1132            .filter_map(|(position, part)| {
1133                let part = part?;
1134                Some(match canonical_from {
1135                    Some(first) if position >= first => (position, canonical(part)),
1136                    _ => (position, part),
1137                })
1138            })
1139            .collect();
1140        if let Some(by_index) = &self.by_index {
1141            parts.sort_by_key(|(position, _)| {
1142                (
1143                    by_index.get(position).copied().unwrap_or(usize::MAX),
1144                    *position,
1145                )
1146            });
1147        }
1148        if let Some(lead) = self.lead {
1149            parts.sort_by_key(|(position, _)| *position != lead);
1150        }
1151        parts.into_iter().map(|(_, part)| part).collect()
1152    }
1153
1154    /// The response of `choice`, the reply's parts, ended by `end`.
1155    fn response(
1156        &self,
1157        end: Finish,
1158        reply: Reply,
1159        choice: Vec<AssistantContent>,
1160    ) -> CompletionResponse {
1161        let Finish {
1162            usage,
1163            reason,
1164            response_id,
1165            model,
1166            error,
1167        } = end;
1168        use crate::provider_response::reported;
1169        let mut origin = self.origin.clone();
1170        let response_model = reported(model);
1171        named_by_reply(&mut origin, response_model.as_deref());
1172        origin.response_model = response_model;
1173        origin.response_id = reported(response_id);
1174        let usage = match usage.cost {
1175            Some(_) => usage,
1176            None => usage.cost(
1177                self.pricing(&origin)
1178                    .and_then(|pricing| pricing.cost(&usage)),
1179            ),
1180        };
1181        let error = error.or_else(|| {
1182            if !self.wire {
1183                return None;
1184            }
1185            match &reason {
1186                None => Some("the provider ended the reply without a finish reason".to_owned()),
1187                Some(FinishReason::Length) => None,
1188                Some(_) if !self.unfinished.is_empty() => Some(
1189                    "the provider ended the reply with a tool call it never finished".to_owned(),
1190                ),
1191                Some(_) => None,
1192            }
1193        });
1194        let mut response = CompletionResponse::new(choice, usage, origin, reply.raw)
1195            .with_optional_finish_reason(reason)
1196            .accept_unknown_finish_reasons(self.accept_unknown_finish)
1197            .with_unfinished_calls(self.unfinished.clone());
1198        response.error = error;
1199        response.provider_request_id = reported(reply.provider_request_id);
1200        response
1201    }
1202}
1203
1204/// `origin` with the model the reply names, `reported`, when the request
1205/// named none.
1206fn named_by_reply(origin: &mut Origin, reported: Option<&str>) {
1207    if origin.model.is_empty()
1208        && let Some(model) = reported
1209    {
1210        model.clone_into(&mut origin.model);
1211    }
1212}
1213
1214/// `block` with no provider item: an opaque item no longer replays.
1215pub(crate) fn canonical(block: AssistantContent) -> AssistantContent {
1216    match block {
1217        AssistantContent::Opaque(opaque) => AssistantContent::Opaque(Opaque {
1218            replay: false,
1219            ..opaque
1220        }),
1221        block => block.into_canonical(),
1222    }
1223}
1224
1225/// The provider item `content` holds, if any.
1226fn item_of(content: &mut AssistantContent) -> Option<&mut serde_json::Value> {
1227    let native = match content {
1228        AssistantContent::Text(text) => text.native.as_mut(),
1229        AssistantContent::ToolCall(call) => call.native.as_mut(),
1230        AssistantContent::Reasoning(reasoning) => reasoning.native.as_mut(),
1231        AssistantContent::Image(image) => image.native.as_mut(),
1232        AssistantContent::Opaque(opaque) => return Some(&mut opaque.item),
1233    };
1234    native.map(|native| &mut native.item)
1235}
1236
1237fn kind_of(content: &AssistantContent) -> PartKind {
1238    match content {
1239        AssistantContent::Text(_) => PartKind::Text,
1240        AssistantContent::Reasoning(_) => PartKind::Reasoning,
1241        AssistantContent::ToolCall(_) => PartKind::ToolCall,
1242        AssistantContent::Image(_) => PartKind::Image,
1243        AssistantContent::Opaque(_) => PartKind::Opaque,
1244    }
1245}
1246
1247pub(crate) type Items = std::collections::VecDeque<Result<Item<StreamEvent>, ProviderError>>;
1248
1249fn emit(items: &mut Items, event: StreamEvent) {
1250    items.push_back(Ok(Item::Event(event)));
1251}
1252
1253/// Merge `delta` into `item`: each string field but `type` appends to the
1254/// item's string of that key, an array extends the item's array, and any
1255/// other value replaces the key. A delta kind rig has never seen still
1256/// lands in the item.
1257pub fn merge(item: &mut serde_json::Value, delta: &serde_json::Map<String, serde_json::Value>) {
1258    use serde_json::Value;
1259    if !item.is_object() {
1260        *item = Value::Object(serde_json::Map::new());
1261    }
1262    let Value::Object(item) = item else {
1263        return;
1264    };
1265    for (key, value) in delta {
1266        if key == "type" {
1267            continue;
1268        }
1269        match (item.get_mut(key), value) {
1270            (Some(Value::String(existing)), Value::String(fragment)) => {
1271                existing.push_str(fragment);
1272            }
1273            (Some(Value::Array(existing)), Value::Array(more)) => {
1274                existing.extend(more.iter().cloned());
1275            }
1276            _ => {
1277                item.insert(key.clone(), value.clone());
1278            }
1279        }
1280    }
1281}
1282
1283impl Fold<Completion> for ReplyFold {
1284    fn absorb(&mut self, event: &StreamEvent) -> Result<(), ProviderError> {
1285        match event {
1286            StreamEvent::Start { part, .. } => {
1287                if self.choice.len() <= part.index() {
1288                    self.choice.resize(part.index() + 1, None);
1289                }
1290            }
1291            StreamEvent::End { part, content } => {
1292                if self.choice.len() <= part.index() {
1293                    self.choice.resize(part.index() + 1, None);
1294                }
1295                if let Some(slot) = self.choice.get_mut(part.index()) {
1296                    *slot = Some(content.clone());
1297                }
1298                self.open_text.remove(&part.index());
1299            }
1300            StreamEvent::Text { part, text } => {
1301                if let AssistantContent::Text(open) = self
1302                    .open_text
1303                    .entry(part.index())
1304                    .or_insert_with(|| AssistantContent::text(""))
1305                {
1306                    open.text.push_str(text);
1307                }
1308            }
1309            StreamEvent::Reasoning { part, text } => {
1310                if let AssistantContent::Reasoning(open) = self
1311                    .open_text
1312                    .entry(part.index())
1313                    .or_insert_with(|| AssistantContent::reasoning(""))
1314                {
1315                    open.text.push_str(text);
1316                }
1317            }
1318            StreamEvent::Arguments { .. } => {}
1319        }
1320        Ok(())
1321    }
1322
1323    fn finish(mut self, end: Finish, reply: Reply) -> Result<CompletionResponse, ProviderError> {
1324        if let Some(FinishReason::Other(reason)) = &end.reason {
1325            tracing::warn!(
1326                provider = %self.origin.provider,
1327                reason = %reason,
1328                accepted = self.accept_unknown_finish,
1329                "the provider ended the reply with an unknown finish reason"
1330            );
1331        }
1332        // The parts taken, in their position (in wire-index order on a wire
1333        // that asked for it, after the lead block); parts from the first
1334        // incomplete one on keep no provider item.
1335        let parts = std::mem::take(&mut self.choice).into_iter().enumerate();
1336        let choice = self.ordered(parts.collect(), self.first_incomplete);
1337        let response = self.response(end, reply, choice);
1338        self.span
1339            .record_response(response.response_id(), response.model(), &response.usage);
1340        Ok(response)
1341    }
1342}
1343
1344impl<'id> Out<'id, Completion> {
1345    /// The pricing in the wire's facts for the requested model, else for
1346    /// `reported`, the model the reply names, for a decoder that prices its
1347    /// reply itself: its provider bills other counters than the ones it
1348    /// reports as usage, or splits a price the catalog lists once.
1349    pub(crate) fn catalog_pricing(
1350        &self,
1351        reported: Option<&str>,
1352    ) -> Option<crate::catalog::Pricing> {
1353        let shared = self.lock();
1354        let mut origin = shared.fold.origin.clone();
1355        named_by_reply(&mut origin, reported);
1356        shared.fold.pricing(&origin)
1357    }
1358
1359    /// Open the provider item at wire `index` as `block`: it takes the next
1360    /// position in the reply. `item` is the item as its first event states
1361    /// it, the base its deltas merge into; `Null` when the block keeps no
1362    /// provider item. An index already open is an error.
1363    pub fn open(
1364        &mut self,
1365        index: usize,
1366        block: Block,
1367        item: serde_json::Value,
1368    ) -> Result<(), ProviderError> {
1369        let mut shared = self.lock();
1370        let Shared { fold, items, .. } = &mut *shared;
1371        fold.open_item(items, index, block, item)
1372    }
1373
1374    /// Append a fragment to the open item at `index`: text, reasoning, or a
1375    /// call's argument JSON, by the item's block.
1376    pub fn push(&mut self, index: usize, fragment: &str) -> Result<(), ProviderError> {
1377        let mut shared = self.lock();
1378        let Shared { fold, items, .. } = &mut *shared;
1379        fold.push_item(items, index, fragment)
1380    }
1381
1382    /// Edit the item at `index` in place: apply a delta with [`merge`], or
1383    /// replace the whole item a provider restates when it finishes.
1384    /// Until the reply ends, an index whose block already closed edits that
1385    /// block's item: the terminal backfill of a field the provider states
1386    /// only at its end.
1387    pub fn edit(
1388        &mut self,
1389        index: usize,
1390        edit: impl FnOnce(&mut serde_json::Value),
1391    ) -> Result<(), ProviderError> {
1392        let mut shared = self.lock();
1393        let Shared { fold, items, .. } = &mut *shared;
1394        match fold.open.get_mut(&index) {
1395            Some(draft) => {
1396                edit(&mut draft.item);
1397                Ok(())
1398            }
1399            None => fold.edit_ended(items, index, edit),
1400        }
1401    }
1402
1403    /// Close the item at `index` the provider never stated complete: its
1404    /// block becomes visible with no native, and an opaque item is kept but
1405    /// does not replay. Empty text and reasoning are dropped, and so is a
1406    /// call with no name. A call's arguments are read by
1407    /// [`ToolFunction::parse`], so malformed ones never fail the reply; a
1408    /// call that never got an id, or reuses one an earlier call took, gets
1409    /// one rig issues.
1410    pub fn close(&mut self, index: usize) -> Result<(), ProviderError> {
1411        let mut shared = self.lock();
1412        let Shared { fold, items, .. } = &mut *shared;
1413        fold.close_item(items, index, Closing::Incomplete)
1414    }
1415
1416    /// Close the item at `index` the provider stated complete: [`Self::close`],
1417    /// with the item as assembled becoming the block's native.
1418    pub fn finish(&mut self, index: usize) -> Result<(), ProviderError> {
1419        let mut shared = self.lock();
1420        let Shared { fold, items, .. } = &mut *shared;
1421        fold.close_item(items, index, Closing::Complete)
1422    }
1423
1424    /// End every item still open as stated complete, for a wire whose end
1425    /// of reply is the provider's statement that its items are done.
1426    pub fn finish_open(&mut self) -> Result<(), ProviderError> {
1427        let mut shared = self.lock();
1428        let Shared { fold, items, .. } = &mut *shared;
1429        let mut open: Vec<(Part, usize)> = fold
1430            .open
1431            .iter()
1432            .map(|(index, draft)| (draft.part, *index))
1433            .collect();
1434        open.sort();
1435        for (_, index) in open {
1436            fold.close_item(items, index, Closing::Complete)?;
1437        }
1438        Ok(())
1439    }
1440
1441    /// Open and finish the item at `index` in one step: a whole block a
1442    /// provider states in one piece, `item` its provider item and `text`
1443    /// its text, reasoning, argument JSON or base64 image data.
1444    pub fn whole(
1445        &mut self,
1446        index: usize,
1447        block: Block,
1448        item: serde_json::Value,
1449        text: &str,
1450    ) -> Result<(), ProviderError> {
1451        self.open(index, block, item)?;
1452        self.push(index, text)?;
1453        self.finish(index)
1454    }
1455
1456    /// Order the response's blocks by wire index rather than by when they
1457    /// opened, for a wire whose indices are the provider's item order and
1458    /// whose end may state items it never streamed. Call it before the
1459    /// first item opens; blocks with no wire index go last.
1460    pub fn order_by_index(&mut self) {
1461        let mut shared = self.lock();
1462        if shared.fold.by_index.is_none() {
1463            shared.fold.by_index = Some(BTreeMap::new());
1464        }
1465    }
1466
1467    /// Hold the open item at `index` first in the response whenever it
1468    /// opened, for a wire that states one block of the turn apart from its
1469    /// order. Its events keep their arrival order. A turn leads with the
1470    /// first item named.
1471    pub fn lead(&mut self, index: usize) -> Result<(), ProviderError> {
1472        let mut shared = self.lock();
1473        let position = shared.fold.draft(index)?.part.index();
1474        shared.fold.lead.get_or_insert(position);
1475        Ok(())
1476    }
1477
1478    /// Replace the text, reasoning or argument JSON of the open item at
1479    /// `index` with `text`, the whole of it as the provider restates it at
1480    /// its end. The fragments already streamed stand and nothing more
1481    /// streams for them; the block ends holding `text`.
1482    pub fn restate(&mut self, index: usize, text: &str) -> Result<(), ProviderError> {
1483        let mut shared = self.lock();
1484        let draft = shared.fold.draft(index)?;
1485        match &mut draft.body {
1486            Body::Text(body) | Body::Reasoning { text: body, .. } => text.clone_into(body),
1487            Body::Call { arguments, .. } => {
1488                text.clone_into(&mut arguments.text);
1489                // A call not yet started streams the restated text whole.
1490                arguments.streamed = if draft.started { text.len() } else { 0 };
1491            }
1492            Body::Image(_) | Body::Opaque { .. } => {}
1493        }
1494        Ok(())
1495    }
1496
1497    /// Write one fragment of a tool call the provider streams, opening the
1498    /// call at its first fragment. Its id and name may arrive in any
1499    /// fragment, and the first name it is given stays: the call starts
1500    /// streaming once it names a tool, carrying
1501    /// the argument text held until then, and its end states its id. Calls
1502    /// are told apart by `index` when the wire gives one. A new id under an
1503    /// index starts a new call once the held call's arguments are a complete
1504    /// object or when it names a tool: some providers send a fresh id, and
1505    /// no name, with every chunk of one call. Without an index (`None`, or
1506    /// the wire's `null`), a fragment with an unseen id opens a call, and one
1507    /// with no id continues the latest call while its arguments are
1508    /// incomplete.
1509    pub fn fragment(
1510        &mut self,
1511        index: Option<usize>,
1512        fragment: CallFragment<'_>,
1513    ) -> Result<(), ProviderError> {
1514        let mut shared = self.lock();
1515        let Shared {
1516            fold: turn, items, ..
1517        } = &mut *shared;
1518        let new_id = fragment.id.filter(|id| stated(id));
1519        let index = match index {
1520            Some(index) => {
1521                let names = fragment.name.is_some_and(|name| !name.is_empty());
1522                let held = turn.open.get(&index).and_then(|draft| match &draft.body {
1523                    Body::Call {
1524                        id: Some(id),
1525                        arguments,
1526                        ..
1527                    } if new_id.is_some_and(|new| id.wire() != new) => Some(arguments.complete()),
1528                    _ => None,
1529                });
1530                if let Some(complete) = held.filter(|complete| *complete || names) {
1531                    // Another call took over the index: the one it held ends.
1532                    let moved = turn.fresh_index();
1533                    if let Some(draft) = turn.open.remove(&index) {
1534                        turn.open.insert(moved, draft);
1535                    }
1536                    let closing = if complete {
1537                        Closing::Complete
1538                    } else {
1539                        Closing::Incomplete
1540                    };
1541                    turn.close_item(items, moved, closing)?;
1542                }
1543                index
1544            }
1545            None => {
1546                let owner = new_id.and_then(|new| {
1547                    turn.open
1548                        .iter()
1549                        .find_map(|(index, draft)| match &draft.body {
1550                            Body::Call { id: Some(id), .. } if id.wire() == new => Some(*index),
1551                            _ => None,
1552                        })
1553                });
1554                let continues = |last: &usize| {
1555                    turn.open.get(last).is_some_and(|draft| {
1556                        matches!(&draft.body, Body::Call { arguments, .. } if !arguments.complete())
1557                    })
1558                };
1559                match (owner, new_id, turn.last_call) {
1560                    (Some(index), _, _) => index,
1561                    (None, None, Some(last)) if continues(&last) => last,
1562                    _ => turn.fresh_index(),
1563                }
1564            }
1565        };
1566        turn.last_call = Some(index);
1567        if !turn.open.contains_key(&index) {
1568            let body = Body::Call {
1569                id: None,
1570                name: String::new(),
1571                arguments: Arguments::default(),
1572            };
1573            turn.insert(index, body, serde_json::Value::Null);
1574        }
1575        let Body::Call {
1576            id,
1577            name,
1578            arguments,
1579        } = &mut turn.draft(index)?.body
1580        else {
1581            return Err(ProviderError::Response(format!(
1582                "the reply sent a call fragment for item {index}, which is not a call"
1583            )));
1584        };
1585        if let Some(call_id) = fragment.id.filter(|id| stated(id)) {
1586            *id = Some(CallId::from_wire(call_id));
1587        }
1588        // A call keeps the first name it is given, so the tool its start
1589        // names is the tool its end names.
1590        if name.is_empty()
1591            && let Some(fragment) = fragment.name.filter(|name| !name.is_empty())
1592        {
1593            fragment.clone_into(name);
1594        }
1595        if let Some(fragment) = fragment.arguments {
1596            arguments.push(fragment, name);
1597        }
1598        stream_call(items, turn.draft(index)?);
1599        Ok(())
1600    }
1601
1602    /// Arguments the provider announced for the call at `index`, used only
1603    /// if no fragment arrives.
1604    pub fn announce(
1605        &mut self,
1606        index: usize,
1607        announced: serde_json::Value,
1608    ) -> Result<(), ProviderError> {
1609        if let Body::Call { arguments, .. } = &mut self.lock().fold.draft(index)?.body {
1610            arguments.announced = Some(announced);
1611        }
1612        Ok(())
1613    }
1614
1615    /// An index no provider item uses, for a wire that indexes nothing.
1616    pub fn fresh_index(&mut self) -> usize {
1617        self.lock().fold.fresh_index()
1618    }
1619
1620    /// A fragment on a wire that marks no item boundaries: it continues the
1621    /// block the last fragment went to while `block` is the same kind, and
1622    /// otherwise closes that block and opens a new one. Returns the block's
1623    /// index, for merging its provider item.
1624    pub fn run(&mut self, block: Block, fragment: &str) -> Result<usize, ProviderError> {
1625        let mut shared = self.lock();
1626        let Shared { fold, items, .. } = &mut *shared;
1627        fold.run_item(items, block, fragment)
1628    }
1629
1630    /// Close the block [`Self::run`] is extending, if any: output of
1631    /// another kind interleaved it.
1632    pub fn end_run(&mut self) -> Result<(), ProviderError> {
1633        let mut shared = self.lock();
1634        let Shared { fold, items, .. } = &mut *shared;
1635        fold.end_run(items)
1636    }
1637
1638    /// Add `citation` to the text item at wire `index`, open or ended. Its
1639    /// span is resolved to bytes of the text when the item closes, or at
1640    /// once for an item already closed. A span that does not resolve, or
1641    /// covers other text than its `quoted`, drops the citation with a
1642    /// warning, as do citations of an item that is not text: a citation
1643    /// never fails a reply.
1644    pub fn cite(&mut self, index: usize, citation: WireCitation) {
1645        let mut shared = self.lock();
1646        let Shared { fold, items, .. } = &mut *shared;
1647        fold.cite_item(items, index, vec![citation], false);
1648    }
1649
1650    /// Put `citations` in place of the text item's own, for a provider that
1651    /// restates them; [`Self::cite`] otherwise.
1652    pub fn set_citations(&mut self, index: usize, citations: Vec<WireCitation>) {
1653        let mut shared = self.lock();
1654        let Shared { fold, items, .. } = &mut *shared;
1655        fold.cite_item(items, index, citations, true);
1656    }
1657
1658    /// A whole block of an already assembled response, at the next
1659    /// position, its native kept as given.
1660    #[cfg(any(test, feature = "test-utils"))]
1661    pub(crate) fn content(&mut self, content: AssistantContent) -> Result<(), ProviderError> {
1662        let mut shared = self.lock();
1663        let Shared { fold, items, .. } = &mut *shared;
1664        fold.write_content(items, content)
1665    }
1666}
1667
1668/// The events a stream of `response` would have carried: each part whole,
1669/// in order.
1670pub(crate) fn events_of(
1671    response: &CompletionResponse,
1672) -> Result<Vec<Item<StreamEvent>>, ProviderError> {
1673    let mut turn = ReplyFold::new(response.origin.clone());
1674    let mut items = Items::new();
1675    for content in &response.choice {
1676        turn.write_content(&mut items, content.clone())?;
1677    }
1678    items.into_iter().collect()
1679}
1680
1681#[cfg(test)]
1682mod tests;