1use 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#[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 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 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 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 fn observed_usage(finish: &Finish) -> Option<&crate::completion::Usage> {
98 Some(&finish.usage)
99 }
100
101 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 request.validate_message_content()?;
122 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 crate::completion::options::check(target, &mut request)?;
133 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 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 tools: stored || target.declares_tools(&request),
159 context: binds.then(|| crate::completion::history::context_of(&request, target, model)),
160 };
161 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 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#[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)]
204pub struct Finish {
205 pub usage: Usage,
211 #[serde(default, skip_serializing_if = "Option::is_none")]
213 pub reason: Option<FinishReason>,
214 pub response_id: Option<String>,
216 pub model: Option<String>,
218 #[serde(default, skip_serializing_if = "Option::is_none")]
221 pub error: Option<String>,
222}
223
224#[derive(Debug, Clone, PartialEq)]
226pub enum Block {
227 Text,
229 Reasoning {
231 redacted: bool,
233 },
234 Image(Image),
237 Call {
239 id: CallId,
241 name: ToolName,
243 },
244 Opaque {
246 replay: bool,
248 },
249}
250
251pub struct ReplyFold {
258 span: tracing::Span,
259 origin: Origin,
261 facts: crate::catalog::ModelFacts,
263 open: BTreeMap<usize, Draft>,
265 ended: HashMap<usize, usize>,
268 next_part: u32,
269 call_ids: HashSet<CallId>,
270 wire: bool,
273 call_id_slot: Option<&'static str>,
277 catalog_priced: Option<bool>,
283 blob_slots: &'static [&'static str],
287 accept_unknown_finish: bool,
292 first_incomplete: Option<usize>,
296 unfinished: Vec<CallId>,
299 last_call: Option<usize>,
301 run: Option<usize>,
303 next_auto: usize,
304 by_index: Option<BTreeMap<usize, usize>>,
307 lead: Option<usize>,
309 choice: Vec<Option<AssistantContent>>,
311 open_text: BTreeMap<usize, AssistantContent>,
314}
315
316struct Draft {
318 part: Part,
319 started: bool,
321 item: serde_json::Value,
323 body: Body,
324 citations: Vec<WireCitation>,
326}
327
328#[derive(Clone, Copy, PartialEq, Eq)]
330enum Closing {
331 Complete,
333 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#[derive(Default)]
369struct Arguments {
370 text: String,
371 streamed: usize,
373 overflowed: bool,
374 substantive: bool,
376 announced: Option<serde_json::Value>,
379}
380
381const 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 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 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 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 fn complete(&self) -> bool {
438 !self.overflowed && crate::json_utils::is_whole_object(&self.text)
439 }
440}
441
442fn stated(id: &str) -> bool {
445 !id.is_empty() && id != "null"
446}
447
448#[derive(Debug, Clone, Copy, Default)]
450pub struct CallFragment<'a> {
451 pub id: Option<&'a str>,
453 pub name: Option<&'a str>,
455 pub arguments: Option<&'a str>,
457}
458
459pub 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
469fn 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 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 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 pub(crate) fn set_origin(&mut self, origin: Origin) {
542 self.origin = origin;
543 }
544
545 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 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 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 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 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 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 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 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 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 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 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 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 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 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 pub fn origin(&self) -> &Origin {
1069 &self.origin
1070 }
1071
1072 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 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 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 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
1204fn 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
1214pub(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
1225fn 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
1253pub 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 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 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 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 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 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 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 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 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 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 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 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 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 arguments.streamed = if draft.started { text.len() } else { 0 };
1491 }
1492 Body::Image(_) | Body::Opaque { .. } => {}
1493 }
1494 Ok(())
1495 }
1496
1497 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 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 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 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 pub fn fresh_index(&mut self) -> usize {
1617 self.lock().fold.fresh_index()
1618 }
1619
1620 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 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 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 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 #[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
1668pub(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;