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 = Turn;
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 Turn {
85 span,
86 wire: replay.is_some_and(|target| target.states_finish_reason()),
87 call_id_slot: replay.and_then(|target| target.call_id_slot()),
88 accept_unknown_finish: request.accept_unknown_finish_reasons,
89 ..Turn::new(origin)
90 }
91 }
92
93 fn prepare(
102 mut request: Self::Request,
103 wire: &Descriptor<'_>,
104 ) -> Result<Self::Request, ProviderError> {
105 let Some(target) = wire.replay else {
106 return Err(ProviderError::request(format!(
107 "completion wire `{}` names no replay target",
108 wire.name
109 )));
110 };
111 request.validate_message_content()?;
114 request.model = request
117 .model
118 .take()
119 .filter(|model| !model.is_empty())
120 .or_else(|| Some(target.model().to_owned()).filter(|model| !model.is_empty()));
121 let target = target.route(&request).unwrap_or(target);
122 crate::completion::options::check(target, &mut request)?;
125 if !target.takes_documents() {
128 request.chat_history = request.chat_history_with_documents();
129 request.documents.clear();
130 }
131 let stored = target.continues_stored(&request);
132 let shape = crate::completion::history::Request {
133 model: request.model.as_deref(),
134 stored,
135 tools: stored || target.declares_tools(&request),
138 context: (!target.drops_unbound_items(&request)).then(|| {
139 crate::completion::history::context_of(
140 &request,
141 target,
142 request.model.as_deref().unwrap_or(target.model()),
143 )
144 }),
145 };
146 request.chat_history =
147 crate::completion::history::adapt_for(&request.chat_history, target, &shape);
148 request.validate_message_content()?;
149 Ok(request)
150 }
151}
152
153impl crate::wire::reply::Closing<Completion> for Assembled {
154 fn close(shared: &mut Shared<Completion>) {
155 let Shared { fold, items, .. } = shared;
156 fold.close_open(items);
157 }
158}
159
160impl Emit<Completion> for Assembled {}
161
162#[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)]
164pub struct Finish {
165 pub usage: Usage,
171 #[serde(default, skip_serializing_if = "Option::is_none")]
173 pub reason: Option<FinishReason>,
174 pub response_id: Option<String>,
176 pub model: Option<String>,
178 #[serde(default, skip_serializing_if = "Option::is_none")]
181 pub error: Option<String>,
182}
183
184#[derive(Debug, Clone, PartialEq)]
186pub enum Block {
187 Text,
189 Reasoning {
191 redacted: bool,
193 },
194 Image(Image),
197 Call {
199 id: CallId,
201 name: ToolName,
203 },
204 Opaque {
206 replay: bool,
208 },
209}
210
211pub struct Turn {
218 span: tracing::Span,
219 origin: Origin,
221 open: BTreeMap<usize, Draft>,
223 ended: HashMap<usize, usize>,
226 next_part: u32,
227 call_ids: HashSet<CallId>,
228 wire: bool,
231 call_id_slot: Option<&'static str>,
235 accept_unknown_finish: bool,
240 first_incomplete: Option<usize>,
244 unfinished_call: bool,
246 last_call: Option<usize>,
248 run: Option<usize>,
250 next_auto: usize,
251 by_index: Option<BTreeMap<usize, usize>>,
254 lead: Option<usize>,
256 choice: Vec<Option<AssistantContent>>,
258 open_text: BTreeMap<usize, AssistantContent>,
261}
262
263struct Draft {
265 part: Part,
266 started: bool,
268 item: serde_json::Value,
270 body: Body,
271 citations: Vec<WireCitation>,
273}
274
275#[derive(Clone, Copy, PartialEq, Eq)]
277enum Closing {
278 Complete,
280 Incomplete,
283}
284
285enum Body {
286 Text(String),
287 Image(Image),
288 Reasoning {
289 text: String,
290 redacted: bool,
291 },
292 Call {
293 id: Option<CallId>,
294 name: String,
295 arguments: Arguments,
296 },
297 Opaque {
298 replay: bool,
299 },
300}
301
302impl Body {
303 fn kind(&self) -> PartKind {
304 match self {
305 Self::Text(_) => PartKind::Text,
306 Self::Reasoning { .. } => PartKind::Reasoning,
307 Self::Image(_) => PartKind::Image,
308 Self::Call { .. } => PartKind::ToolCall,
309 Self::Opaque { .. } => PartKind::Opaque,
310 }
311 }
312}
313
314#[derive(Default)]
316struct Arguments {
317 text: String,
318 streamed: usize,
320 overflowed: bool,
321 substantive: bool,
323 announced: Option<serde_json::Value>,
326}
327
328const MAX_TOOL_INPUT_BYTES: usize = 32 * 1024 * 1024;
330
331impl Arguments {
332 fn push(&mut self, fragment: &str, name: &str) {
333 self.substantive |= !fragment.trim().is_empty();
334 if self.text.trim() == "null" && !fragment.trim().is_empty() {
337 self.text.clear();
338 }
339 if self.text.len().saturating_add(fragment.len()) > MAX_TOOL_INPUT_BYTES {
340 if !self.overflowed {
341 self.overflowed = true;
342 tracing::warn!(
343 tool = name,
344 "streamed tool-call input exceeded the accumulation bound; truncating"
345 );
346 }
347 } else {
348 self.text.push_str(fragment);
349 }
350 }
351
352 fn function(&self, name: ToolName) -> ToolFunction {
355 if self.text.is_empty()
356 && let Some(announced) = &self.announced
357 {
358 return ToolFunction::new(name, announced.clone());
359 }
360 let mut function = ToolFunction::parse(name, &self.text);
361 if self.overflowed {
362 function.invalid_arguments = Some(self.text.clone());
363 }
364 function
365 }
366
367 fn unstreamed(&mut self) -> Option<String> {
371 if "null".starts_with(self.text.trim()) {
372 return None;
373 }
374 let fragment = self.text.get(self.streamed..)?;
375 if fragment.is_empty() {
376 return None;
377 }
378 let fragment = fragment.to_owned();
379 self.streamed = self.text.len();
380 Some(fragment)
381 }
382
383 fn complete(&self) -> bool {
385 !self.overflowed
386 && matches!(
387 crate::json_utils::parse_tool_arguments(&self.text),
388 Ok(serde_json::Value::Object(_))
389 )
390 }
391}
392
393fn stated(id: &str) -> bool {
396 !id.is_empty() && id != "null"
397}
398
399#[derive(Debug, Clone, Copy, Default)]
401pub struct CallFragment<'a> {
402 pub id: Option<&'a str>,
404 pub name: Option<&'a str>,
406 pub arguments: Option<&'a str>,
408}
409
410pub const AUTO_INDEX: usize = 1 << 30;
413
414fn not_open(index: usize) -> ProviderError {
415 ProviderError::Response(format!(
416 "the reply wrote to item {index}, which is not open"
417 ))
418}
419
420fn stream_call(items: &mut Items, draft: &mut Draft) {
424 let Body::Call {
425 name, arguments, ..
426 } = &mut draft.body
427 else {
428 return;
429 };
430 if !draft.started {
431 let Ok(name) = ToolName::new(name.clone()) else {
432 return;
433 };
434 draft.started = true;
435 emit(
436 items,
437 StreamEvent::Start {
438 part: draft.part,
439 kind: PartKind::ToolCall,
440 name: Some(name),
441 },
442 );
443 }
444 if let Some(json) = arguments.unstreamed() {
445 emit(
446 items,
447 StreamEvent::Arguments {
448 part: draft.part,
449 json,
450 },
451 );
452 }
453}
454
455impl Turn {
456 pub(crate) fn new(origin: Origin) -> Self {
458 Self {
459 span: tracing::Span::none(),
460 origin,
461 open: BTreeMap::new(),
462 ended: HashMap::new(),
463 next_part: 0,
464 call_ids: HashSet::new(),
465 wire: false,
466 call_id_slot: None,
467 accept_unknown_finish: false,
468 first_incomplete: None,
469 unfinished_call: false,
470 last_call: None,
471 run: None,
472 next_auto: AUTO_INDEX,
473 by_index: None,
474 lead: None,
475 choice: Vec::new(),
476 open_text: BTreeMap::new(),
477 }
478 }
479
480 pub(crate) fn relayed(label: impl Into<String>) -> Self {
484 let label = label.into();
485 Self::new(Origin::new(label.clone(), label.clone(), label))
486 }
487
488 pub(crate) fn set_origin(&mut self, origin: Origin) {
490 self.origin = origin;
491 }
492
493 fn next(&mut self) -> Part {
495 let part = Part::new(self.next_part);
496 self.next_part += 1;
497 part
498 }
499
500 fn fresh_index(&mut self) -> usize {
501 let index = self.next_auto;
502 self.next_auto += 1;
503 index
504 }
505
506 fn draft(&mut self, index: usize) -> Result<&mut Draft, ProviderError> {
507 self.open.get_mut(&index).ok_or_else(|| not_open(index))
508 }
509
510 pub(crate) fn open_item(
511 &mut self,
512 items: &mut Items,
513 index: usize,
514 block: Block,
515 item: serde_json::Value,
516 ) -> Result<(), ProviderError> {
517 if self.open.contains_key(&index) {
518 return Err(ProviderError::Response(format!(
519 "the reply opened item {index} twice"
520 )));
521 }
522 let body = match block {
523 Block::Text => Body::Text(String::new()),
524 Block::Image(image) => Body::Image(image),
525 Block::Reasoning { redacted } => Body::Reasoning {
526 text: String::new(),
527 redacted,
528 },
529 Block::Call { id, name } => Body::Call {
530 id: Some(id),
531 name: name.into(),
532 arguments: Arguments::default(),
533 },
534 Block::Opaque { replay } => Body::Opaque { replay },
535 };
536 self.insert(index, body, item);
537 stream_call(items, self.draft(index)?);
538 Ok(())
539 }
540
541 fn insert(&mut self, index: usize, body: Body, item: serde_json::Value) {
542 self.ended.remove(&index);
543 let part = self.next();
544 if let Some(by_index) = &mut self.by_index {
545 by_index.insert(part.index(), index);
546 }
547 let started = false;
548 self.open.insert(
549 index,
550 Draft {
551 part,
552 started,
553 item,
554 body,
555 citations: Vec::new(),
556 },
557 );
558 }
559
560 pub(crate) fn push_item(
561 &mut self,
562 items: &mut Items,
563 index: usize,
564 fragment: &str,
565 ) -> Result<(), ProviderError> {
566 if fragment.is_empty() {
567 return Ok(());
568 }
569 let draft = self.draft(index)?;
570 let event = match &mut draft.body {
571 Body::Text(text) => {
572 text.push_str(fragment);
573 StreamEvent::Text {
574 part: draft.part,
575 text: fragment.to_owned(),
576 }
577 }
578 Body::Reasoning { text, .. } => {
579 text.push_str(fragment);
580 StreamEvent::Reasoning {
581 part: draft.part,
582 text: fragment.to_owned(),
583 }
584 }
585 Body::Call {
586 name, arguments, ..
587 } => {
588 arguments.push(fragment, name);
589 stream_call(items, draft);
590 return Ok(());
591 }
592 Body::Image(image) => {
593 if let crate::message::DocumentSourceKind::Base64(data) = &mut image.data {
594 data.push_str(fragment);
595 return Ok(());
596 }
597 return Err(ProviderError::Response(format!(
598 "the reply wrote data to the image item {index}, which holds no base64 data"
599 )));
600 }
601 Body::Opaque { .. } => {
602 return Err(ProviderError::Response(format!(
603 "the reply wrote text to the opaque item {index}"
604 )));
605 }
606 };
607 if !draft.started {
608 draft.started = true;
609 emit(
610 items,
611 StreamEvent::Start {
612 part: draft.part,
613 kind: draft.body.kind(),
614 name: None,
615 },
616 );
617 }
618 emit(items, event);
619 Ok(())
620 }
621
622 fn distinct_call_id(
629 &mut self,
630 id: Option<CallId>,
631 item: serde_json::Value,
632 ) -> (CallId, serde_json::Value) {
633 let (id, item) = match id {
634 Some(id) if !self.call_ids.contains(&id) => (id, item),
635 Some(id) => {
636 tracing::warn!(%id, "the provider named two tool calls with one id; renaming the second");
637 (CallId::Local(LocalCallId::new()), serde_json::Value::Null)
638 }
639 None if self.call_id_slot.is_some() => (CallId::Local(LocalCallId::new()), item),
640 None => (CallId::Local(LocalCallId::new()), serde_json::Value::Null),
641 };
642 self.call_ids.insert(id.clone());
643 (id, item)
644 }
645
646 fn close_item(
649 &mut self,
650 items: &mut Items,
651 index: usize,
652 closing: Closing,
653 ) -> Result<(), ProviderError> {
654 let mut draft = self.open.remove(&index).ok_or_else(|| not_open(index))?;
655 stream_call(items, &mut draft);
658 if self.run == Some(index) {
659 self.run = None;
660 }
661 self.ended.insert(index, draft.part.index());
662 let Draft {
663 part,
664 started,
665 item,
666 body,
667 citations,
668 } = draft;
669 let item = match (closing, &body) {
670 (Closing::Complete, _) | (Closing::Incomplete, Body::Opaque { .. }) => item,
671 (Closing::Incomplete, _) => {
672 self.cut_at(part.index());
673 serde_json::Value::Null
674 }
675 };
676 let content = match body {
677 Body::Text(text) => {
678 if text.is_empty() && item.is_null() {
679 return Ok(());
680 }
681 let mut text = Text::new(text);
684 citation::attach(
685 &mut text,
686 Vec::new(),
687 citations,
688 &self.origin.provider,
689 index,
690 );
691 with_item(AssistantContent::Text(text), item)
692 }
693 Body::Reasoning { text, redacted } => {
694 if text.is_empty() && !redacted && item.is_null() {
695 return Ok(());
696 }
697 let reasoning = Reasoning {
698 text,
699 redacted,
700 native: None,
701 };
702 with_item(AssistantContent::Reasoning(reasoning), item)
703 }
704 Body::Image(image) => with_item(AssistantContent::Image(image), item),
705 Body::Opaque { replay } => AssistantContent::Opaque(Opaque {
706 item,
707 replay: replay && closing == Closing::Complete,
708 }),
709 Body::Call {
710 id,
711 name,
712 arguments,
713 } => {
714 let Ok(name) = ToolName::new(name) else {
715 tracing::warn!(
716 index,
717 "the provider closed a tool call without a name; nothing can answer it"
718 );
719 return Ok(());
720 };
721 let function = arguments.function(name);
722 let (id, item) = self.distinct_call_id(id, item);
723 if arguments.streamed == 0 {
726 let json = serde_json::Value::Object(function.arguments.clone()).to_string();
727 emit(items, StreamEvent::Arguments { part, json });
728 }
729 let call = ToolCall::new(id, function);
730 let content = with_item(AssistantContent::ToolCall(call), item);
731 emit(items, StreamEvent::End { part, content });
732 return Ok(());
733 }
734 };
735 if !started {
736 emit(
737 items,
738 StreamEvent::Start {
739 part,
740 kind: kind_of(&content),
741 name: None,
742 },
743 );
744 }
745 emit(items, StreamEvent::End { part, content });
746 Ok(())
747 }
748
749 fn edit_ended(
752 &mut self,
753 items: &mut Items,
754 index: usize,
755 edit: impl FnOnce(&mut serde_json::Value),
756 ) -> Result<(), ProviderError> {
757 let position = *self.ended.get(&index).ok_or_else(|| not_open(index))?;
758 let queued = items.iter_mut().find_map(|item| match item {
759 Ok(Item::Event(StreamEvent::End { part, content })) if part.index() == position => {
760 Some(content)
761 }
762 _ => None,
763 });
764 let content = match queued {
765 Some(content) => Some(content),
766 None => self.choice.get_mut(position).and_then(Option::as_mut),
767 };
768 if let Some(item) = content.and_then(item_of) {
769 edit(item);
770 }
771 Ok(())
772 }
773
774 fn cite_item(
780 &mut self,
781 items: &mut Items,
782 index: usize,
783 citations: Vec<WireCitation>,
784 replace: bool,
785 ) {
786 let provider = self.origin.provider.as_str();
787 if let Some(draft) = self.open.get_mut(&index) {
788 if !matches!(draft.body, Body::Text(_)) {
789 tracing::warn!(
790 provider,
791 index,
792 "dropped citations of an item that is not text"
793 );
794 } else if replace {
795 draft.citations = citations;
796 } else {
797 draft.citations.extend(citations);
798 }
799 return;
800 }
801 let Some(position) = self.ended.get(&index).copied() else {
802 tracing::warn!(
803 provider,
804 index,
805 "dropped citations of an item the reply never opened"
806 );
807 return;
808 };
809 let queued = items.iter_mut().find_map(|item| match item {
810 Ok(Item::Event(StreamEvent::End { part, content })) if part.index() == position => {
811 Some(content)
812 }
813 _ => None,
814 });
815 let content = match queued {
816 Some(content) => Some(content),
817 None => self.choice.get_mut(position).and_then(Option::as_mut),
818 };
819 match content {
820 Some(AssistantContent::Text(text)) => {
821 let kept = if replace {
822 Vec::new()
823 } else {
824 text.citations().to_vec()
825 };
826 citation::attach(text, kept, citations, provider, index);
827 }
828 _ => tracing::warn!(
829 provider,
830 index,
831 "dropped citations of an item that is not text"
832 ),
833 }
834 }
835
836 pub(crate) fn run_item(
837 &mut self,
838 items: &mut Items,
839 block: Block,
840 fragment: &str,
841 ) -> Result<usize, ProviderError> {
842 let current = self.run.filter(|index| {
843 self.open.get(index).is_some_and(|draft| {
844 matches!(
845 (&draft.body, &block),
846 (Body::Text(_), Block::Text)
847 | (
848 Body::Reasoning {
849 redacted: false,
850 ..
851 },
852 Block::Reasoning { redacted: false }
853 )
854 )
855 })
856 });
857 let index = match current {
858 Some(index) => index,
859 None => {
860 self.end_run(items)?;
861 let index = self.fresh_index();
862 self.open_item(items, index, block, serde_json::Value::Null)?;
863 self.run = Some(index);
864 index
865 }
866 };
867 self.push_item(items, index, fragment)?;
868 Ok(index)
869 }
870
871 pub(crate) fn end_run(&mut self, items: &mut Items) -> Result<(), ProviderError> {
872 match self.run.take() {
873 Some(index) => self.close_item(items, index, Closing::Complete),
874 None => Ok(()),
875 }
876 }
877
878 pub(crate) fn write_content(
880 &mut self,
881 items: &mut Items,
882 content: AssistantContent,
883 ) -> Result<(), ProviderError> {
884 let content = match content {
885 AssistantContent::ToolCall(mut call) => {
886 let item = call
887 .native
888 .take()
889 .map_or(serde_json::Value::Null, |native| native.item);
890 let (id, item) = self.distinct_call_id(Some(call.id), item);
891 call.id = id;
892 with_item(AssistantContent::ToolCall(call), item)
893 }
894 content => content,
895 };
896 let part = self.next();
897 let name = match &content {
898 AssistantContent::ToolCall(call) => Some(call.function.name.clone()),
899 _ => None,
900 };
901 emit(
902 items,
903 StreamEvent::Start {
904 part,
905 kind: kind_of(&content),
906 name,
907 },
908 );
909 match &content {
910 AssistantContent::Text(text) if !text.text.is_empty() => emit(
911 items,
912 StreamEvent::Text {
913 part,
914 text: text.text.clone(),
915 },
916 ),
917 AssistantContent::Reasoning(reasoning) if !reasoning.text.is_empty() => emit(
918 items,
919 StreamEvent::Reasoning {
920 part,
921 text: reasoning.text.clone(),
922 },
923 ),
924 AssistantContent::ToolCall(call) => emit(
925 items,
926 StreamEvent::Arguments {
927 part,
928 json: call.function.arguments_value().to_string(),
929 },
930 ),
931 _ => {}
932 }
933 emit(items, StreamEvent::End { part, content });
934 Ok(())
935 }
936
937 pub(crate) fn close_open(&mut self, items: &mut Items) {
942 let mut open: Vec<(Part, usize)> = self
943 .open
944 .iter()
945 .map(|(index, draft)| (draft.part, *index))
946 .collect();
947 open.sort();
948 for (_, index) in open {
949 let closing = if self.run == Some(index) {
950 Closing::Complete
951 } else {
952 Closing::Incomplete
953 };
954 if self
955 .open
956 .get(&index)
957 .is_some_and(|draft| matches!(draft.body, Body::Call { .. }))
958 {
959 self.unfinished_call = true;
960 }
961 if let Err(error) = self.close_item(items, index, closing) {
962 items.push_back(Err(error));
963 }
964 }
965 }
966
967 fn cut_at(&mut self, position: usize) {
969 self.first_incomplete = Some(
970 self.first_incomplete
971 .map_or(position, |first| first.min(position)),
972 );
973 }
974
975 pub fn snapshot(&self) -> Vec<AssistantContent> {
980 self.ordered(
981 self.choice.iter().cloned().enumerate().collect(),
982 self.first_incomplete,
983 )
984 }
985
986 pub fn origin(&self) -> &Origin {
988 &self.origin
989 }
990
991 pub(crate) fn partial(
999 &self,
1000 end: Option<&Finish>,
1001 reply: &Reply,
1002 failure: Option<&ProviderError>,
1003 ) -> CompletionResponse {
1004 let mut response = self.response(end.cloned().unwrap_or_default(), reply.clone());
1005 if end.is_none() {
1006 match failure {
1007 Some(error) => {
1008 response.error.get_or_insert_with(|| error.to_string());
1009 }
1010 None => {
1011 response.aborted = Some(
1012 "the caller stopped reading before the provider ended the reply".to_owned(),
1013 );
1014 }
1015 }
1016 }
1017 let unfinished =
1021 end.is_none() || !self.open.is_empty() || self.choice.iter().any(Option::is_none);
1022 let canonical_from = if unfinished {
1023 Some(0)
1024 } else {
1025 self.first_incomplete
1026 };
1027 let parts = self
1028 .choice
1029 .iter()
1030 .enumerate()
1031 .map(|(index, part)| {
1032 let part = part.clone().or_else(|| self.open_text.get(&index).cloned());
1033 (index, part)
1034 })
1035 .collect();
1036 response.choice = self.ordered(parts, canonical_from);
1037 response
1038 }
1039
1040 fn ordered(
1044 &self,
1045 parts: Vec<(usize, Option<AssistantContent>)>,
1046 canonical_from: Option<usize>,
1047 ) -> Vec<AssistantContent> {
1048 let mut parts: Vec<(usize, AssistantContent)> = parts
1049 .into_iter()
1050 .filter_map(|(position, part)| {
1051 let part = part?;
1052 Some(match canonical_from {
1053 Some(first) if position >= first => (position, canonical(part)),
1054 _ => (position, part),
1055 })
1056 })
1057 .collect();
1058 if let Some(by_index) = &self.by_index {
1059 parts.sort_by_key(|(position, _)| {
1060 (
1061 by_index.get(position).copied().unwrap_or(usize::MAX),
1062 *position,
1063 )
1064 });
1065 }
1066 if let Some(lead) = self.lead {
1067 parts.sort_by_key(|(position, _)| *position != lead);
1068 }
1069 parts.into_iter().map(|(_, part)| part).collect()
1070 }
1071
1072 fn response(&self, end: Finish, reply: Reply) -> CompletionResponse {
1073 let Finish {
1074 usage,
1075 reason,
1076 response_id,
1077 model,
1078 error,
1079 } = end;
1080 use crate::provider_response::reported;
1081 let mut origin = self.origin.clone();
1082 origin.response_model = reported(model);
1083 if origin.model.is_empty()
1086 && let Some(model) = &origin.response_model
1087 {
1088 origin.model.clone_from(model);
1089 }
1090 origin.response_id = reported(response_id);
1091 let usage = priced(usage, &origin);
1092 let error = error.or_else(|| {
1093 if !self.wire {
1094 return None;
1095 }
1096 match &reason {
1097 None => Some("the provider ended the reply without a finish reason".to_owned()),
1098 Some(FinishReason::Length) => None,
1099 Some(_) if self.unfinished_call => Some(
1100 "the provider ended the reply with a tool call it never finished".to_owned(),
1101 ),
1102 Some(_) => None,
1103 }
1104 });
1105 let mut response = CompletionResponse::new(self.snapshot(), usage, origin, reply.raw)
1106 .with_optional_finish_reason(reason)
1107 .accept_unknown_finish_reasons(self.accept_unknown_finish);
1108 response.error = error;
1109 response.provider_request_id = reported(reply.provider_request_id);
1110 response
1111 }
1112}
1113
1114fn priced(usage: Usage, origin: &Origin) -> Usage {
1117 if usage.cost.is_some() {
1118 return usage;
1119 }
1120 let cost = catalog_cost(origin, &usage);
1121 usage.cost(cost)
1122}
1123
1124const UNPRICED: [&str; 3] = [
1130 crate::providers::chatgpt::PROVIDER_NAME,
1131 crate::providers::copilot::PROVIDER_NAME,
1132 crate::providers::ollama::PROVIDER_NAME,
1133];
1134
1135fn catalog_cost(origin: &Origin, usage: &Usage) -> Option<crate::completion::Cost> {
1139 if UNPRICED.contains(&origin.provider.as_str()) {
1140 return None;
1141 }
1142 crate::catalog::lookup(&origin.provider, &origin.model)
1143 .or_else(|| crate::catalog::lookup_snapshot(&origin.provider, &origin.model))
1144 .and_then(|spec| spec.pricing.as_ref())
1145 .and_then(|pricing| pricing.cost(usage))
1146}
1147
1148pub(crate) fn canonical(block: AssistantContent) -> AssistantContent {
1150 match block {
1151 AssistantContent::Opaque(opaque) => AssistantContent::Opaque(Opaque {
1152 replay: false,
1153 ..opaque
1154 }),
1155 block => block.canonical(),
1156 }
1157}
1158
1159fn item_of(content: &mut AssistantContent) -> Option<&mut serde_json::Value> {
1161 let native = match content {
1162 AssistantContent::Text(text) => text.native.as_mut(),
1163 AssistantContent::ToolCall(call) => call.native.as_mut(),
1164 AssistantContent::Reasoning(reasoning) => reasoning.native.as_mut(),
1165 AssistantContent::Image(image) => image.native.as_mut(),
1166 AssistantContent::Opaque(opaque) => return Some(&mut opaque.item),
1167 };
1168 native.map(|native| &mut native.item)
1169}
1170
1171fn with_item(block: AssistantContent, item: serde_json::Value) -> AssistantContent {
1173 if item.is_null() {
1174 block
1175 } else {
1176 block.with_native(item)
1177 }
1178}
1179
1180fn kind_of(content: &AssistantContent) -> PartKind {
1181 match content {
1182 AssistantContent::Text(_) => PartKind::Text,
1183 AssistantContent::Reasoning(_) => PartKind::Reasoning,
1184 AssistantContent::ToolCall(_) => PartKind::ToolCall,
1185 AssistantContent::Image(_) => PartKind::Image,
1186 AssistantContent::Opaque(_) => PartKind::Opaque,
1187 }
1188}
1189
1190pub(crate) type Items = std::collections::VecDeque<Result<Item<StreamEvent>, ProviderError>>;
1191
1192fn emit(items: &mut Items, event: StreamEvent) {
1193 items.push_back(Ok(Item::Event(event)));
1194}
1195
1196pub fn merge(item: &mut serde_json::Value, delta: &serde_json::Map<String, serde_json::Value>) {
1201 use serde_json::Value;
1202 if !item.is_object() {
1203 *item = Value::Object(serde_json::Map::new());
1204 }
1205 let Value::Object(item) = item else {
1206 return;
1207 };
1208 for (key, value) in delta {
1209 if key == "type" {
1210 continue;
1211 }
1212 match (item.get_mut(key), value) {
1213 (Some(Value::String(existing)), Value::String(fragment)) => {
1214 existing.push_str(fragment);
1215 }
1216 (Some(Value::Array(existing)), Value::Array(more)) => {
1217 existing.extend(more.iter().cloned());
1218 }
1219 _ => {
1220 item.insert(key.clone(), value.clone());
1221 }
1222 }
1223 }
1224}
1225
1226impl Fold<Completion> for Turn {
1227 fn absorb(&mut self, event: &StreamEvent) -> Result<(), ProviderError> {
1228 match event {
1229 StreamEvent::Start { part, .. } => {
1230 if self.choice.len() <= part.index() {
1231 self.choice.resize(part.index() + 1, None);
1232 }
1233 }
1234 StreamEvent::End { part, content } => {
1235 if self.choice.len() <= part.index() {
1236 self.choice.resize(part.index() + 1, None);
1237 }
1238 if let Some(slot) = self.choice.get_mut(part.index()) {
1239 *slot = Some(content.clone());
1240 }
1241 self.open_text.remove(&part.index());
1242 }
1243 StreamEvent::Text { part, text } => {
1244 if let AssistantContent::Text(open) = self
1245 .open_text
1246 .entry(part.index())
1247 .or_insert_with(|| AssistantContent::text(""))
1248 {
1249 open.text.push_str(text);
1250 }
1251 }
1252 StreamEvent::Reasoning { part, text } => {
1253 if let AssistantContent::Reasoning(open) = self
1254 .open_text
1255 .entry(part.index())
1256 .or_insert_with(|| AssistantContent::reasoning(""))
1257 {
1258 open.text.push_str(text);
1259 }
1260 }
1261 StreamEvent::Arguments { .. } => {}
1262 }
1263 Ok(())
1264 }
1265
1266 fn finish(self, end: Finish, reply: Reply) -> Result<CompletionResponse, ProviderError> {
1267 if let Some(FinishReason::Other(reason)) = &end.reason {
1268 tracing::warn!(
1269 provider = %self.origin.provider,
1270 reason = %reason,
1271 accepted = self.accept_unknown_finish,
1272 "the provider ended the reply with an unknown finish reason"
1273 );
1274 }
1275 let response = self.response(end, reply);
1276 self.span
1277 .record_response(response.response_id(), response.model(), &response.usage);
1278 Ok(response)
1279 }
1280}
1281
1282impl<'id> Out<'id, Completion> {
1283 pub(crate) fn catalog_cost(&self, usage: &Usage) -> Option<crate::completion::Cost> {
1287 catalog_cost(&self.lock().fold.origin, usage)
1288 }
1289
1290 pub fn open(
1295 &mut self,
1296 index: usize,
1297 block: Block,
1298 item: serde_json::Value,
1299 ) -> Result<(), ProviderError> {
1300 let mut shared = self.lock();
1301 let Shared { fold, items, .. } = &mut *shared;
1302 fold.open_item(items, index, block, item)
1303 }
1304
1305 pub fn push(&mut self, index: usize, fragment: &str) -> Result<(), ProviderError> {
1308 let mut shared = self.lock();
1309 let Shared { fold, items, .. } = &mut *shared;
1310 fold.push_item(items, index, fragment)
1311 }
1312
1313 pub fn edit(
1319 &mut self,
1320 index: usize,
1321 edit: impl FnOnce(&mut serde_json::Value),
1322 ) -> Result<(), ProviderError> {
1323 let mut shared = self.lock();
1324 let Shared { fold, items, .. } = &mut *shared;
1325 match fold.open.get_mut(&index) {
1326 Some(draft) => {
1327 edit(&mut draft.item);
1328 Ok(())
1329 }
1330 None => fold.edit_ended(items, index, edit),
1331 }
1332 }
1333
1334 pub fn close(&mut self, index: usize) -> Result<(), ProviderError> {
1342 let mut shared = self.lock();
1343 let Shared { fold, items, .. } = &mut *shared;
1344 fold.close_item(items, index, Closing::Incomplete)
1345 }
1346
1347 pub fn finish(&mut self, index: usize) -> Result<(), ProviderError> {
1350 let mut shared = self.lock();
1351 let Shared { fold, items, .. } = &mut *shared;
1352 fold.close_item(items, index, Closing::Complete)
1353 }
1354
1355 pub fn finish_open(&mut self) -> Result<(), ProviderError> {
1358 let mut shared = self.lock();
1359 let Shared { fold, items, .. } = &mut *shared;
1360 let mut open: Vec<(Part, usize)> = fold
1361 .open
1362 .iter()
1363 .map(|(index, draft)| (draft.part, *index))
1364 .collect();
1365 open.sort();
1366 for (_, index) in open {
1367 fold.close_item(items, index, Closing::Complete)?;
1368 }
1369 Ok(())
1370 }
1371
1372 pub fn whole(
1376 &mut self,
1377 index: usize,
1378 block: Block,
1379 item: serde_json::Value,
1380 text: &str,
1381 ) -> Result<(), ProviderError> {
1382 self.open(index, block, item)?;
1383 self.push(index, text)?;
1384 self.finish(index)
1385 }
1386
1387 pub fn order_by_index(&mut self) {
1392 let mut shared = self.lock();
1393 if shared.fold.by_index.is_none() {
1394 shared.fold.by_index = Some(BTreeMap::new());
1395 }
1396 }
1397
1398 pub fn lead(&mut self, index: usize) -> Result<(), ProviderError> {
1403 let mut shared = self.lock();
1404 let position = shared.fold.draft(index)?.part.index();
1405 shared.fold.lead.get_or_insert(position);
1406 Ok(())
1407 }
1408
1409 pub fn restate(&mut self, index: usize, text: &str) -> Result<(), ProviderError> {
1414 let mut shared = self.lock();
1415 let draft = shared.fold.draft(index)?;
1416 match &mut draft.body {
1417 Body::Text(body) | Body::Reasoning { text: body, .. } => text.clone_into(body),
1418 Body::Call { arguments, .. } => {
1419 text.clone_into(&mut arguments.text);
1420 arguments.streamed = if draft.started { text.len() } else { 0 };
1422 }
1423 Body::Image(_) | Body::Opaque { .. } => {}
1424 }
1425 Ok(())
1426 }
1427
1428 pub fn fragment(
1441 &mut self,
1442 index: Option<usize>,
1443 fragment: CallFragment<'_>,
1444 ) -> Result<(), ProviderError> {
1445 let mut shared = self.lock();
1446 let Shared {
1447 fold: turn, items, ..
1448 } = &mut *shared;
1449 let new_id = fragment.id.filter(|id| stated(id));
1450 let index = match index {
1451 Some(index) => {
1452 let names = fragment.name.is_some_and(|name| !name.is_empty());
1453 let held = turn.open.get(&index).and_then(|draft| match &draft.body {
1454 Body::Call {
1455 id: Some(id),
1456 arguments,
1457 ..
1458 } if new_id.is_some_and(|new| id.wire() != new) => Some(arguments.complete()),
1459 _ => None,
1460 });
1461 if let Some(complete) = held.filter(|complete| *complete || names) {
1462 let moved = turn.fresh_index();
1464 if let Some(draft) = turn.open.remove(&index) {
1465 turn.open.insert(moved, draft);
1466 }
1467 let closing = if complete {
1468 Closing::Complete
1469 } else {
1470 Closing::Incomplete
1471 };
1472 turn.close_item(items, moved, closing)?;
1473 }
1474 index
1475 }
1476 None => {
1477 let owner = new_id.and_then(|new| {
1478 turn.open
1479 .iter()
1480 .find_map(|(index, draft)| match &draft.body {
1481 Body::Call { id: Some(id), .. } if id.wire() == new => Some(*index),
1482 _ => None,
1483 })
1484 });
1485 let continues = |last: &usize| {
1486 turn.open.get(last).is_some_and(|draft| {
1487 matches!(&draft.body, Body::Call { arguments, .. } if !arguments.complete())
1488 })
1489 };
1490 match (owner, new_id, turn.last_call) {
1491 (Some(index), _, _) => index,
1492 (None, None, Some(last)) if continues(&last) => last,
1493 _ => turn.fresh_index(),
1494 }
1495 }
1496 };
1497 turn.last_call = Some(index);
1498 if !turn.open.contains_key(&index) {
1499 let body = Body::Call {
1500 id: None,
1501 name: String::new(),
1502 arguments: Arguments::default(),
1503 };
1504 turn.insert(index, body, serde_json::Value::Null);
1505 }
1506 let Body::Call {
1507 id,
1508 name,
1509 arguments,
1510 } = &mut turn.draft(index)?.body
1511 else {
1512 return Err(ProviderError::Response(format!(
1513 "the reply sent a call fragment for item {index}, which is not a call"
1514 )));
1515 };
1516 if let Some(call_id) = fragment.id.filter(|id| stated(id)) {
1517 *id = Some(CallId::from_wire(call_id));
1518 }
1519 if name.is_empty()
1522 && let Some(fragment) = fragment.name.filter(|name| !name.is_empty())
1523 {
1524 fragment.clone_into(name);
1525 }
1526 if let Some(fragment) = fragment.arguments {
1527 arguments.push(fragment, name);
1528 }
1529 stream_call(items, turn.draft(index)?);
1530 Ok(())
1531 }
1532
1533 pub fn announce(
1536 &mut self,
1537 index: usize,
1538 announced: serde_json::Value,
1539 ) -> Result<(), ProviderError> {
1540 if let Body::Call { arguments, .. } = &mut self.lock().fold.draft(index)?.body {
1541 arguments.announced = Some(announced);
1542 }
1543 Ok(())
1544 }
1545
1546 pub fn fresh_index(&mut self) -> usize {
1548 self.lock().fold.fresh_index()
1549 }
1550
1551 pub fn run(&mut self, block: Block, fragment: &str) -> Result<usize, ProviderError> {
1556 let mut shared = self.lock();
1557 let Shared { fold, items, .. } = &mut *shared;
1558 fold.run_item(items, block, fragment)
1559 }
1560
1561 pub fn end_run(&mut self) -> Result<(), ProviderError> {
1564 let mut shared = self.lock();
1565 let Shared { fold, items, .. } = &mut *shared;
1566 fold.end_run(items)
1567 }
1568
1569 pub fn cite(&mut self, index: usize, citation: WireCitation) {
1576 let mut shared = self.lock();
1577 let Shared { fold, items, .. } = &mut *shared;
1578 fold.cite_item(items, index, vec![citation], false);
1579 }
1580
1581 pub fn set_citations(&mut self, index: usize, citations: Vec<WireCitation>) {
1584 let mut shared = self.lock();
1585 let Shared { fold, items, .. } = &mut *shared;
1586 fold.cite_item(items, index, citations, true);
1587 }
1588
1589 #[cfg(any(test, feature = "test-utils"))]
1592 pub(crate) fn content(&mut self, content: AssistantContent) -> Result<(), ProviderError> {
1593 let mut shared = self.lock();
1594 let Shared { fold, items, .. } = &mut *shared;
1595 fold.write_content(items, content)
1596 }
1597}
1598
1599pub(crate) fn events_of(
1602 response: &CompletionResponse,
1603) -> Result<Vec<Item<StreamEvent>>, ProviderError> {
1604 let mut turn = Turn::new(response.origin.clone());
1605 let mut items = Items::new();
1606 for content in &response.choice {
1607 turn.write_content(&mut items, content.clone())?;
1608 }
1609 items.into_iter().collect()
1610}
1611
1612#[cfg(test)]
1613mod tests;