gproxy_protocol/adapt/generate/stream/event/
chat.rs

1use super::*;
2use crate::{
3    transform::generate::stream::chat::{ChatStreamCollector, ChatStreamLimits},
4    wire::openai::chat::{self as c, stream as s},
5};
6use std::collections::BTreeMap;
7
8#[derive(Default)]
9struct ChoiceIds {
10    legacy: bool,
11    tools: BTreeMap<i64, Option<String>>,
12}
13
14pub struct ChatCollector {
15    native: ChatStreamCollector,
16    ids: BTreeMap<i64, ChoiceIds>,
17}
18
19impl sealed::Event for s::ChatCompletionChunk {}
20impl NativeEvent for s::ChatCompletionChunk {
21    type Full = c::GenerateContentResponseBody;
22    type Collector = ChatCollector;
23    const DIALECT: Dialect = Dialect::OpenAiChat;
24    const DONE: bool = true;
25    fn collector(
26        flow: IdentityFlow,
27        mut policy: TargetIdPolicy,
28        limits: EventLimits,
29    ) -> Self::Collector {
30        policy.preserve_source_ids = true;
31        ChatCollector {
32            native: ChatStreamCollector::with_limits(
33                flow,
34                policy,
35                ChatStreamLimits {
36                    max_bytes: limits.max_bytes,
37                },
38            ),
39            ids: BTreeMap::new(),
40        }
41    }
42    fn collect(collector: &mut Self::Collector, event: Self) -> Result<(), TransformError> {
43        collector.native.push(event.clone())?;
44        for choice in event.choices {
45            let observed = collector.ids.entry(choice.index).or_default();
46            if choice.delta.function_call.flatten().is_some() {
47                observed.legacy = true;
48            }
49            for tool in choice.delta.tool_calls.flatten().into_iter().flatten() {
50                let original = observed.tools.entry(tool.index).or_default();
51                if let Some(id) = tool.id.flatten() {
52                    *original = Some(id);
53                }
54            }
55        }
56        Ok(())
57    }
58    fn collect_done(collector: &mut Self::Collector) -> Result<(), TransformError> {
59        collector.native.push_done()
60    }
61    fn collected(
62        collector: Self::Collector,
63    ) -> Result<Converted<Collected<Self::Full>>, TransformError> {
64        let value = collector.native.finish()?;
65        let observed: Vec<_> = collector
66            .ids
67            .into_values()
68            .flat_map(|choice| {
69                if choice.legacy {
70                    vec![(None, super::super::super::ChatCallForm::LegacyFunction)]
71                } else {
72                    choice
73                        .tools
74                        .into_values()
75                        .map(|id| (id, super::super::super::ChatCallForm::Modern))
76                        .collect()
77                }
78            })
79            .collect();
80        let tools = value.value.tools();
81        if observed.len() != tools.len()
82            || observed
83                .iter()
84                .zip(&tools)
85                .any(|((_, form), tool)| tool.chat_form != Some(*form))
86        {
87            return Err(invalid(
88                "native Chat tool observation cardinality or original wire form changed",
89            ));
90        }
91        let ids = observed.into_iter().map(|(id, _)| id).collect();
92        Ok(Converted {
93            value: Collected {
94                value: value.value,
95                original_tool_ids: Some(ids),
96            },
97            report: value.report,
98        })
99    }
100    fn event_name(&self) -> Option<&'static str> {
101        None
102    }
103    fn is_terminal(&self) -> bool {
104        self.choices
105            .iter()
106            .any(|c| c.finish_reason.is_some_and(|v| v.is_some()))
107    }
108    fn is_usage_only(&self) -> bool {
109        self.choices.is_empty() && self.usage.as_ref().and_then(Option::as_ref).is_some()
110    }
111    fn tool_declarations(
112        &self,
113        complete_tool_names: bool,
114    ) -> Vec<(String, ToolCallKind, Option<String>)> {
115        self.choices
116            .iter()
117            .flat_map(|c| {
118                c.delta
119                    .tool_calls
120                    .as_ref()
121                    .and_then(Option::as_ref)
122                    .into_iter()
123                    .flatten()
124            })
125            .filter_map(|tool| {
126                tool.id.as_ref().and_then(Option::as_ref).map(|id| {
127                    (
128                        id.clone(),
129                        ToolCallKind::Function,
130                        if complete_tool_names {
131                            tool.function
132                                .as_ref()
133                                .and_then(Option::as_ref)
134                                .and_then(|v| v.name.as_ref().and_then(Option::as_ref))
135                                .filter(|v| !v.is_empty())
136                                .cloned()
137                        } else {
138                            None
139                        },
140                    )
141                })
142            })
143            .collect()
144    }
145}
gproxy_protocol/adapt/generate/stream/event/chat.rs

gproxy_protocol/adapt/generate/stream/event/
chat.rs
