gproxy_protocol/transform/generate/gemini_chat/stream/
chat_to_gemini.rs

1use super::{
2    common::{Budget, StreamEnd, gemini_finish},
3    tools::{ChatChoice, ChatTool},
4    *,
5};
6use crate::transform::generate::stream::chat::{ChatStreamCollector, ChatStreamLimits};
7
8/// Text is emitted immediately. Function arguments remain bounded strings until
9/// their choice finishes because Gemini requires one complete JSON object.
10pub struct ChatToGeminiStream {
11    source: Option<ChatStreamCollector>,
12    policy: TargetIdPolicy,
13    flow: IdentityFlow,
14    budget: Budget,
15    choices: BTreeMap<i64, ChatChoice>,
16    next_tool: u64,
17    failed: bool,
18    done: bool,
19    response_id: Option<String>,
20    model: Option<String>,
21    report: Report,
22}
23
24impl ChatToGeminiStream {
25    pub fn new(flow: IdentityFlow, limits: StreamLimits) -> Self {
26        Self::new_with_policy(flow, limits, TargetIdPolicy::new(crate::Dialect::Gemini))
27            .expect("default policy matches target dialect")
28    }
29    pub fn new_with_policy(
30        flow: IdentityFlow,
31        limits: StreamLimits,
32        policy: TargetIdPolicy,
33    ) -> Result<Self, TransformError> {
34        if policy.dialect != crate::Dialect::Gemini {
35            return Err(TransformError::shape(
36                "stream.policy",
37                "target policy dialect mismatch",
38            ));
39        }
40
41        let source = ChatStreamCollector::with_limits(
42            IdentityFlow::new(flow.namespace()),
43            TargetIdPolicy::new(crate::Dialect::OpenAiChat),
44            ChatStreamLimits {
45                max_bytes: limits.max_bytes,
46            },
47        );
48        Ok(Self {
49            source: Some(source),
50            policy,
51            flow,
52            budget: Budget::new(limits),
53            choices: BTreeMap::new(),
54            next_tool: 0,
55            failed: false,
56            done: false,
57            response_id: None,
58            model: None,
59            report: Default::default(),
60        })
61    }
62    pub fn identities(&self) -> &IdentityFlow {
63        &self.flow
64    }
65    pub fn push(
66        &mut self,
67        chunk: s::ChatCompletionChunk,
68    ) -> Result<Converted<Vec<g::GenerateContentResponseBody>>, TransformError> {
69        if self.failed || self.done {
70            self.failed = true;
71            return Err(invalid("event after DONE or failed stream"));
72        }
73        let result = self.push_inner(chunk.into_declared());
74        if result.is_err() {
75            self.failed = true;
76        }
77        result
78    }
79    fn push_inner(
80        &mut self,
81        chunk: s::ChatCompletionChunk,
82    ) -> Result<Converted<Vec<g::GenerateContentResponseBody>>, TransformError> {
83        self.budget.input(&chunk)?;
84        self.source
85            .as_mut()
86            .ok_or_else(|| invalid("source consumed"))?
87            .push(chunk.clone())?;
88        if !chunk.id.is_empty() {
89            self.response_id = Some(super::common::id(
90                &mut self.flow,
91                &self.policy,
92                IdentityRole::Response,
93                crate::Dialect::OpenAiChat,
94                Some(chunk.id.clone()),
95                0,
96            )?);
97        }
98        if !chunk.model.is_empty() {
99            self.model = Some(chunk.model.clone());
100        }
101        let mut output = Vec::new();
102        for choice in chunk.choices {
103            let index = choice.index;
104            let d = choice.delta;
105            if let Some(text) = crate::wire::openai::chat::visible_reasoning(
106                &d.reasoning_content,
107                &d.reasoning,
108                &d.reasoning_details,
109            ) {
110                output.push(candidate(
111                    index,
112                    vec![g::Part::builder().text(text).thought(true).build()],
113                ));
114            }
115
116            let state = self.choices.entry(index).or_default();
117            if let Some(text) = d.content.flatten() {
118                output.push(candidate(
119                    index,
120                    vec![g::Part::builder().text(text).build()],
121                ));
122            }
123            if let Some(refusal) = d.refusal.flatten() {
124                state.refusal.get_or_insert_default().push_str(&refusal);
125            }
126            for call in d.tool_calls.flatten().into_iter().flatten() {
127                let tool = match state.tools.entry(call.index) {
128                    std::collections::btree_map::Entry::Occupied(entry) => entry.into_mut(),
129                    std::collections::btree_map::Entry::Vacant(entry) => {
130                        let ordinal = self.next_tool;
131                        self.next_tool += 1;
132                        entry.insert(ChatTool {
133                            ordinal,
134                            ..Default::default()
135                        })
136                    }
137                };
138                tool.append(call.id.flatten(), call.function.flatten())?;
139            }
140            if let Some(function) = d.function_call.flatten() {
141                let legacy = match &mut state.legacy {
142                    Some(tool) => tool,
143                    slot => {
144                        let ordinal = self.next_tool;
145                        self.next_tool += 1;
146                        slot.insert(ChatTool {
147                            ordinal,
148                            ..Default::default()
149                        })
150                    }
151                };
152                legacy.append(None, Some(function))?;
153            }
154            if choice.finish_reason.flatten().is_some() {
155                let mut parts = Vec::new();
156                if let Some(refusal) = state.refusal.take() {
157                    parts.push(g::Part::builder().text(refusal).build());
158                    self.report.changed("refusal","Gemini represents refusal as text; the source finish reason controls the terminal reason");
159                }
160                if let Some(tool) = state.legacy.take() {
161                    parts.push(tool.part(&mut self.flow, true, &self.policy)?);
162                }
163                for (_, tool) in std::mem::take(&mut state.tools) {
164                    parts.push(tool.part(&mut self.flow, false, &self.policy)?);
165                }
166                if !parts.is_empty() {
167                    output.push(candidate(index, parts));
168                }
169            }
170        }
171        let mut metered = Vec::new();
172        for mut body in output {
173            body.response_id = self.response_id.clone();
174            body.model_version = self.model.clone();
175            // Retain only a bounded batch; metadata may be much larger than
176            // the small text/tool fragment associated with an individual choice.
177            self.budget.output(std::slice::from_ref(&body))?;
178            metered.push(body);
179        }
180        Ok(Converted {
181            value: metered,
182            report: std::mem::take(&mut self.report),
183        })
184    }
185    pub fn push_done(&mut self) -> Result<(), TransformError> {
186        if self.failed || self.done {
187            self.failed = true;
188            return Err(invalid("duplicate DONE or failed stream"));
189        }
190        let result = self
191            .source
192            .as_mut()
193            .ok_or_else(|| invalid("source consumed"))?
194            .push_done();
195        if result.is_err() {
196            self.failed = true;
197        } else {
198            self.done = true;
199        }
200        result
201    }
202    pub fn finish(mut self) -> Result<StreamEnd<g::GenerateContentResponseBody>, TransformError> {
203        if self.failed || !self.done {
204            return Err(invalid("Chat stream ended without valid SSE DONE"));
205        }
206        let source = self
207            .source
208            .take()
209            .ok_or_else(|| invalid("source consumed"))?
210            .finish()?;
211        self.report.diagnostics.extend(source.report.diagnostics);
212        let source = source.value;
213        let mut candidates = Vec::new();
214        for choice in source.choices {
215            let mut candidate = g::Candidate::builder()
216                .index(choice.index)
217                // A completed Chat choice always contains an assistant message,
218                // including an explicitly empty one. Empty parts add no text.
219                .content(
220                    g::Content::builder()
221                        .role("model")
222                        .parts(Vec::new())
223                        .build(),
224                )
225                .finish_reason(gemini_finish(choice.finish_reason))
226                .build();
227            candidate.logprobs_result = choice
228                .logprobs
229                .as_ref()
230                .map(|v| super::super::logs::to_gemini(v, &mut self.report))
231                .map(crate::transform::optional)
232                .transpose()?
233                .flatten();
234            candidates.push(candidate);
235        }
236        let usage = source
237            .usage
238            .as_ref()
239            .map(|u| super::super::usage::to_gemini(u, &mut self.report))
240            .map(crate::transform::optional)
241            .transpose()?
242            .flatten();
243        let mut tail = g::GenerateContentResponseBody::builder()
244            .candidates(candidates)
245            .response_id(super::common::id(
246                &mut self.flow,
247                &self.policy,
248                IdentityRole::Response,
249                crate::Dialect::OpenAiChat,
250                Some(source.id),
251                0,
252            )?)
253            .model_version(source.model)
254            .build();
255        tail.usage_metadata = usage;
256        if source.service_tier.is_some() {
257            self.report.omitted(
258                "service_tier",
259                "Gemini response has no equivalent Chat service tier",
260            );
261        }
262        if source.system_fingerprint.is_some() {
263            self.report.omitted(
264                "system_fingerprint",
265                "Gemini has no corresponding fingerprint field",
266            );
267        }
268        if source.moderation.is_some() {
269            self.report.omitted(
270                "moderation",
271                "Chat moderation has no equivalent Gemini metadata schema",
272            );
273        }
274        self.report
275            .omitted("created", "Gemini response has no creation timestamp field");
276        let chunks = vec![tail];
277        self.budget.output(&chunks)?;
278        Ok(StreamEnd {
279            chunks,
280            identities: self.flow,
281            report: self.report,
282        })
283    }
284}
285
286fn candidate(index: i64, parts: Vec<g::Part>) -> g::GenerateContentResponseBody {
287    g::GenerateContentResponseBody::builder()
288        .candidates(vec![
289            g::Candidate::builder()
290                .index(index)
291                .content(g::Content::builder().role("model").parts(parts).build())
292                .build(),
293        ])
294        .build()
295}
gproxy_protocol/transform/generate/gemini_chat/stream/chat_to_gemini.rs

gproxy_protocol/transform/generate/gemini_chat/stream/
chat_to_gemini.rs
