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 decode(mut value: serde_json::Value) -> Result<Self, TransformError> {
26        if let Some(error) = value.get_mut("error").filter(|error| !error.is_null()) {
27            let error: s::StreamError = super::decode_event(error.take())?;
28            return Err(TransformError::with_source(
29                crate::transform::TransformErrorKind::InvalidResult,
30                "chat.stream.error",
31                error.message.clone(),
32                crate::transform::UpstreamError(error),
33            ));
34        }
35        super::decode_event(value)
36    }
37    fn collector(
38        flow: IdentityFlow,
39        mut policy: TargetIdPolicy,
40        limits: EventLimits,
41    ) -> Self::Collector {
42        policy.preserve_source_ids = true;
43        ChatCollector {
44            native: ChatStreamCollector::with_limits(
45                flow,
46                policy,
47                ChatStreamLimits {
48                    max_bytes: limits.max_bytes,
49                },
50            ),
51            ids: BTreeMap::new(),
52        }
53    }
54    fn collect(collector: &mut Self::Collector, event: Self) -> Result<(), TransformError> {
55        collector.native.push(event.clone())?;
56        for choice in event.choices {
57            let observed = collector.ids.entry(choice.index).or_default();
58            if choice.delta.function_call.flatten().is_some() {
59                observed.legacy = true;
60            }
61            for tool in choice.delta.tool_calls.flatten().into_iter().flatten() {
62                let original = observed.tools.entry(tool.index).or_default();
63                if let Some(id) = tool.id.flatten() {
64                    *original = Some(id);
65                }
66            }
67        }
68        Ok(())
69    }
70    fn collect_done(collector: &mut Self::Collector) -> Result<(), TransformError> {
71        collector.native.push_done()
72    }
73    fn collected(
74        collector: Self::Collector,
75    ) -> Result<Converted<Collected<Self::Full>>, TransformError> {
76        let value = collector.native.finish()?;
77        let observed: Vec<_> = collector
78            .ids
79            .into_values()
80            .flat_map(|choice| {
81                if choice.legacy {
82                    vec![(None, super::super::super::ChatCallForm::LegacyFunction)]
83                } else {
84                    choice
85                        .tools
86                        .into_values()
87                        .map(|id| (id, super::super::super::ChatCallForm::Modern))
88                        .collect()
89                }
90            })
91            .collect();
92        let tools = value.value.tools();
93        if observed.len() != tools.len()
94            || observed
95                .iter()
96                .zip(&tools)
97                .any(|((_, form), tool)| tool.chat_form != Some(*form))
98        {
99            return Err(invalid(
100                "native Chat tool observation cardinality or original wire form changed",
101            ));
102        }
103        let ids = observed.into_iter().map(|(id, _)| id).collect();
104        Ok(Converted {
105            value: Collected {
106                value: value.value,
107                original_tool_ids: Some(ids),
108            },
109            report: value.report,
110        })
111    }
112    fn event_name(&self) -> Option<&'static str> {
113        None
114    }
115    fn is_terminal(&self) -> bool {
116        self.choices
117            .iter()
118            .any(|c| c.finish_reason.is_some_and(|v| v.is_some()))
119    }
120    fn is_usage_only(&self) -> bool {
121        self.choices.is_empty() && self.usage.as_ref().and_then(Option::as_ref).is_some()
122    }
123    fn tool_declarations(
124        &self,
125        complete_tool_names: bool,
126    ) -> Vec<(String, ToolCallKind, Option<String>)> {
127        self.choices
128            .iter()
129            .flat_map(|c| {
130                c.delta
131                    .tool_calls
132                    .as_ref()
133                    .and_then(Option::as_ref)
134                    .into_iter()
135                    .flatten()
136            })
137            .filter_map(|tool| {
138                tool.id.as_ref().and_then(Option::as_ref).map(|id| {
139                    (
140                        id.clone(),
141                        ToolCallKind::Function,
142                        if complete_tool_names {
143                            tool.function
144                                .as_ref()
145                                .and_then(Option::as_ref)
146                                .and_then(|v| v.name.as_ref().and_then(Option::as_ref))
147                                .filter(|v| !v.is_empty())
148                                .cloned()
149                        } else {
150                            None
151                        },
152                    )
153                })
154            })
155            .collect()
156    }
157}
gproxy_protocol/adapt/generate/stream/event/chat.rs

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