gproxy_protocol/transform/generate/gemini_responses/stream/
responses_to_gemini.rs

1use super::{
2    common::{Budget, StreamEnd, StreamLimits, declared, invalid, limit, measure},
3    context::ResponsesToGeminiContext,
4    identity,
5    response_items::Item,
6};
7use crate::{
8    Dialect,
9    transform::generate::stream::{
10        gemini::{GeminiStreamCollector, GeminiStreamLimits},
11        responses::{ResponsesStreamCollector, ResponsesStreamLimits},
12    },
13    transform::{
14        Converted, Report, TransformError,
15        identity::{IdentityFlow, TargetIdPolicy},
16    },
17    wire::{
18        gemini as g,
19        openai::responses::{response as r, stream as s},
20    },
21};
22use std::collections::{BTreeMap, BTreeSet};
23
24pub struct ResponsesToGeminiStream {
25    source: Option<ResponsesStreamCollector>,
26    target: Option<GeminiStreamCollector>,
27    pub(super) flow: IdentityFlow,
28    pub(super) policy: TargetIdPolicy,
29    pub(super) limits: StreamLimits,
30    budget: Budget,
31    pub(super) used: BTreeSet<String>,
32    pub(super) native_calls: BTreeSet<String>,
33    pub(super) model: Option<String>,
34    response_id: Option<String>,
35    observed_usage: Option<r::ResponseUsage>,
36    failed: bool,
37    pub(super) terminal: bool,
38    image_progress: bool,
39    pub(super) image_only: bool,
40    pub(super) jpeg_only: bool,
41    pub(super) items: BTreeMap<i64, Item>,
42    pub(super) cursor: i64,
43    pub(super) held: usize,
44}
45
46impl ResponsesToGeminiStream {
47    pub fn new(
48        context: ResponsesToGeminiContext,
49        flow: IdentityFlow,
50        limits: StreamLimits,
51    ) -> Result<Self, TransformError> {
52        Self::new_with_policy(context, flow, TargetIdPolicy::new(Dialect::Gemini), limits)
53    }
54    pub fn new_with_policy(
55        context: ResponsesToGeminiContext,
56        flow: IdentityFlow,
57        policy: TargetIdPolicy,
58        limits: StreamLimits,
59    ) -> Result<Self, TransformError> {
60        if policy.dialect != Dialect::Gemini {
61            return Err(invalid("Gemini target policy required"));
62        }
63        Ok(Self {
64            source: Some(ResponsesStreamCollector::new(ResponsesStreamLimits {
65                max_bytes: limits.max_bytes,
66
67                max_text_bytes: limits.max_bytes,
68                max_json_bytes: limits.max_bytes,
69            })),
70            target: Some(GeminiStreamCollector::new(GeminiStreamLimits {
71                max_bytes: limits.max_bytes,
72            })),
73            flow,
74            policy,
75            limits,
76            budget: Budget::new(limits),
77            used: Default::default(),
78            native_calls: Default::default(),
79            model: None,
80            response_id: None,
81            observed_usage: None,
82            failed: false,
83            terminal: false,
84            image_progress: false,
85            jpeg_only: context.image_mime == Some(g::ImageMimeType::ImageJpeg),
86            image_only: context
87                .response_modalities
88                .as_ref()
89                .is_some_and(|v| !v.is_empty() && v.iter().all(|v| *v == g::Modality::Image)),
90            items: Default::default(),
91            cursor: 0,
92            held: 0,
93        })
94    }
95    pub(crate) fn reserve_external_ids(
96        &mut self,
97        role: crate::transform::identity::IdentityRole,
98        ids: &std::collections::BTreeSet<String>,
99    ) -> Result<(), TransformError> {
100        self.flow.reserve_external_ids(role, ids).map_err(|e| {
101            TransformError::new(
102                crate::transform::TransformErrorKind::Conflict,
103                "fanout.ids",
104                e.to_string(),
105            )
106        })
107    }
108    pub fn identities(&self) -> &IdentityFlow {
109        &self.flow
110    }
111    pub fn push(
112        &mut self,
113        event: s::StreamEvent,
114    ) -> Result<Converted<Vec<g::GenerateContentResponseBody>>, TransformError> {
115        if matches!(event, s::StreamEvent::Unknown) {
116            return Ok(Converted {
117                value: Vec::new(),
118                report: Report::default(),
119            });
120        }
121        if self.failed || self.terminal {
122            self.failed = true;
123            return Err(invalid("event after terminal or failure"));
124        }
125        let result = self.inner(declared(event));
126        if result.is_err() {
127            self.failed = true;
128        }
129        result
130    }
131    fn inner(
132        &mut self,
133        event: s::StreamEvent,
134    ) -> Result<Converted<Vec<g::GenerateContentResponseBody>>, TransformError> {
135        let event = crate::transform::generate::multi_agent::attribute_event(event)?;
136        self.budget.input(&event)?;
137        self.source
138            .as_mut()
139            .ok_or_else(|| invalid("source consumed"))?
140            .push(event.clone())?;
141        let mut out = Vec::new();
142        match event {
143            s::StreamEvent::Keepalive | s::StreamEvent::Metadata(_) | s::StreamEvent::Unknown => {}
144            s::StreamEvent::Created(v) => {
145                self.model = Some(v.response.model.clone());
146                let id = identity::response_id(
147                    &mut self.flow,
148                    &self.policy,
149                    Dialect::OpenAi,
150                    Some(v.response.id.clone()),
151                )?;
152                self.response_id = Some(id.clone());
153                self.observe(&v.response)?;
154                self.emit(
155                    g::GenerateContentResponseBody::builder()
156                        .response_id(id)
157                        .model_version(v.response.model)
158                        .build(),
159                    &mut out,
160                )?;
161                for (index, item) in v.response.output.into_iter().enumerate() {
162                    self.add_item(index as i64, item)?;
163                }
164            }
165            s::StreamEvent::InProgress(v) => self.observe(&v.response)?,
166            s::StreamEvent::Queued(v) => self.observe(&v.response)?,
167            s::StreamEvent::OutputItemAdded(v) => self.add_item(v.output_index, v.item)?,
168            s::StreamEvent::ContentPartAdded(v) => {
169                self.bind(v.output_index, &v.item_id)?;
170                self.add_part(v.output_index, v.content_index, v.part)?;
171            }
172            s::StreamEvent::OutputTextDelta(v) => {
173                self.bind(v.output_index, &v.item_id)?;
174                self.append_text(v.output_index, v.content_index, v.delta, false)?;
175            }
176            s::StreamEvent::RefusalDelta(v) => {
177                self.bind(v.output_index, &v.item_id)?;
178                self.append_text(v.output_index, v.content_index, v.delta, false)?;
179            }
180            s::StreamEvent::ReasoningTextDelta(v) => {
181                self.bind(v.output_index, &v.item_id)?;
182                self.append_text(v.output_index, v.content_index, v.delta, true)?;
183            }
184            s::StreamEvent::ContentPartDone(v) => {
185                self.part_done(v.output_index, v.content_index)?
186            }
187            s::StreamEvent::FunctionCallArgumentsDelta(v) => {
188                self.bind(v.output_index, &v.item_id)?;
189                self.reserve_arguments(v.output_index, v.delta.len())?;
190            }
191            s::StreamEvent::FunctionCallArgumentsDone(v) => {
192                self.bind(v.output_index, &v.item_id)?;
193                self.arguments_done(v.output_index, v.arguments)?;
194            }
195            s::StreamEvent::OutputItemDone(v) => self.item_done(v.output_index, v.item)?,
196            s::StreamEvent::ReasoningSummaryPartAdded(v) => {
197                self.count_part()?;
198                self.bind(v.output_index, &v.item_id)?;
199                self.summary(v.output_index, v.summary_index, v.part.text, true)?;
200            }
201            s::StreamEvent::ReasoningSummaryTextDelta(v) => {
202                self.bind(v.output_index, &v.item_id)?;
203                self.summary(v.output_index, v.summary_index, v.delta, false)?;
204            }
205            s::StreamEvent::ReasoningSummaryPartDone(v) => {
206                self.bind(v.output_index, &v.item_id)?;
207                self.summary_done(v.output_index, v.summary_index)?;
208            }
209            s::StreamEvent::Completed(v) => {
210                self.observe(&v.response)?;
211                self.terminal = true;
212            }
213            s::StreamEvent::Incomplete(v) => {
214                self.observe(&v.response)?;
215                self.terminal = true;
216            }
217            s::StreamEvent::OutputTextDone(_)
218            | s::StreamEvent::RefusalDone(_)
219            | s::StreamEvent::ReasoningTextDone(_)
220            | s::StreamEvent::OutputTextAnnotationAdded(_)
221            | s::StreamEvent::ReasoningSummaryTextDone(_) => {}
222            s::StreamEvent::ImageCall(_)
223            | s::StreamEvent::ImageGenerating(_)
224            | s::StreamEvent::ImageCompleted(_)
225            | s::StreamEvent::ImagePartial(_) => {
226                self.image_progress = true;
227            }
228            s::StreamEvent::Failed(_) | s::StreamEvent::Error(_) => {
229                return Err(invalid("native source failure"));
230            }
231            _ if self.image_only => {}
232            _ => {}
233        }
234        self.flush(&mut out)?;
235        Ok(Converted {
236            value: out,
237            report: Report::default(),
238        })
239    }
240    fn observe(&mut self, body: &r::GenerateContentResponseBody) -> Result<(), TransformError> {
241        if let Some(usage) = body.usage.as_ref().and_then(Option::as_ref) {
242            super::super::usage::to_gemini(usage.clone(), &mut Report::default())?;
243            if let Some(old) = &self.observed_usage
244                && (usage.input_tokens < old.input_tokens
245                    || usage.output_tokens < old.output_tokens
246                    || usage.total_tokens < old.total_tokens
247                    || usage.input_tokens_details.cached_tokens
248                        < old.input_tokens_details.cached_tokens
249                    || usage.input_tokens_details.cache_write_tokens
250                        < old.input_tokens_details.cache_write_tokens
251                    || usage.output_tokens_details.reasoning_tokens
252                        < old.output_tokens_details.reasoning_tokens)
253            {
254                return Err(invalid("native cumulative usage decreased"));
255            }
256            self.observed_usage = Some(usage.clone());
257        }
258        Ok(())
259    }
260    pub(super) fn reserve(&mut self, bytes: usize) -> Result<(), TransformError> {
261        self.held = self
262            .held
263            .checked_add(bytes)
264            .filter(|n| *n <= self.limits.max_pending)
265            .ok_or_else(limit)?;
266        Ok(())
267    }
268    pub(super) fn release(&mut self, bytes: usize) {
269        self.held -= bytes;
270    }
271    pub(super) fn emit(
272        &mut self,
273        chunk: g::GenerateContentResponseBody,
274        out: &mut Vec<g::GenerateContentResponseBody>,
275    ) -> Result<(), TransformError> {
276        self.budget.output(&chunk)?;
277        self.target
278            .as_mut()
279            .ok_or_else(|| invalid("target consumed"))?
280            .push(chunk.clone())?;
281        out.push(chunk);
282        Ok(())
283    }
284    pub(super) fn emit_part(
285        &mut self,
286        part: g::Part,
287        out: &mut Vec<g::GenerateContentResponseBody>,
288    ) -> Result<(), TransformError> {
289        self.emit(
290            g::GenerateContentResponseBody::builder()
291                .candidates(vec![
292                    g::Candidate::builder()
293                        .index(0)
294                        .content(
295                            g::Content::builder()
296                                .role("model")
297                                .parts(vec![part])
298                                .build(),
299                        )
300                        .build(),
301                ])
302                .build(),
303            out,
304        )
305    }
306    pub fn finish(mut self) -> Result<StreamEnd<g::GenerateContentResponseBody>, TransformError> {
307        if self.failed || !self.terminal {
308            return Err(invalid("failed stream or premature EOF"));
309        }
310        let mut source = self
311            .source
312            .take()
313            .ok_or_else(|| invalid("source consumed"))?
314            .finish()?
315            .value;
316        // Once visible summaries have been streamed, keep that representation
317        // for this item rather than switching to later raw content or replaying
318        // both. Validate against the same selected view at completion.
319        let mut summary_report = Report::default();
320        for (index, item) in source.output.iter_mut().enumerate() {
321            if self.items.get(&(index as i64)).is_some_and(|state| {
322                matches!(
323                    state.kind,
324                    super::response_items::Kind::Reasoning {
325                        summary_emitted: true,
326                        ..
327                    }
328                )
329            }) && let r::ResponseOutputItem::Reasoning(reasoning) = item
330                && reasoning.content.take().is_some()
331            {
332                summary_report.omitted(
333                    format!("output[{index}].content"),
334                    "visible reasoning summary was already streamed; late raw content is not replayed",
335                );
336            }
337        }
338        let only_image = [g::Modality::Image];
339        let mut converted = super::super::responses_to_gemini_response_with_modalities(
340            source.clone(),
341            Default::default(),
342            self.image_only.then_some(only_image.as_slice()),
343        )?;
344        converted
345            .report
346            .diagnostics
347            .extend(summary_report.diagnostics);
348        if self.image_progress {
349            converted.report.omitted("image.progress", "Gemini inline images expose the final image; preview/progress events are not additional generated images");
350        }
351        let mut check_flow = self.flow.clone();
352        converted.value.response_id = Some(identity::response_id(
353            &mut check_flow,
354            &self.policy,
355            Dialect::OpenAi,
356            Some(source.id.clone()),
357        )?);
358        let originals: Vec<_> = source
359            .output
360            .iter()
361            .enumerate()
362            .filter_map(|(index, item)| {
363                if self.image_only {
364                    return None;
365                }
366                if let r::ResponseOutputItem::FunctionCall(call) = item {
367                    Some((index as i64, call))
368                } else {
369                    None
370                }
371            })
372            .collect();
373        let projected: Vec<_> = converted
374            .value
375            .candidates
376            .iter_mut()
377            .flatten()
378            .filter_map(|c| c.content.as_mut())
379            .flat_map(|c| c.parts.iter_mut().flatten())
380            .filter(|p| p.function_call.is_some())
381            .collect();
382        if originals.len() != projected.len() {
383            return Err(invalid("canonical function count changed"));
384        }
385        let mut used = BTreeSet::new();
386        for ((index, call), part) in originals.into_iter().zip(projected) {
387            identity::call(
388                part,
389                call,
390                index,
391                (&mut check_flow, &self.policy),
392                &mut used,
393            )?;
394        }
395        measure(&converted.value, self.limits.max_bytes)?;
396        let mut tail = converted.value.clone();
397        for candidate in tail.candidates.iter_mut().flatten() {
398            if let Some(content) = &mut candidate.content {
399                content.parts = Some(Vec::new());
400            }
401        }
402        let mut out = Vec::new();
403        self.emit(tail, &mut out)?;
404        let mut actual = self
405            .target
406            .take()
407            .ok_or_else(|| invalid("target consumed"))?
408            .finish()?
409            .value;
410        identity::normalize_text(&mut actual);
411        identity::normalize_text(&mut converted.value);
412        if actual != converted.value {
413            return Err(invalid(
414                "native Gemini stream differs from completed Responses conversion",
415            ));
416        }
417        Ok(StreamEnd {
418            chunks: out,
419            identities: self.flow,
420            report: converted.report,
421        })
422    }
423}
gproxy_protocol/transform/generate/gemini_responses/stream/responses_to_gemini.rs

gproxy_protocol/transform/generate/gemini_responses/stream/
responses_to_gemini.rs
