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 self.failed || self.terminal {
116            self.failed = true;
117            return Err(invalid("event after terminal or failure"));
118        }
119        let result = self.inner(declared(event));
120        if result.is_err() {
121            self.failed = true;
122        }
123        result
124    }
125    fn inner(
126        &mut self,
127        event: s::StreamEvent,
128    ) -> Result<Converted<Vec<g::GenerateContentResponseBody>>, TransformError> {
129        let event = crate::transform::generate::multi_agent::attribute_event(event)?;
130        self.budget.input(&event)?;
131        self.source
132            .as_mut()
133            .ok_or_else(|| invalid("source consumed"))?
134            .push(event.clone())?;
135        let mut out = Vec::new();
136        match event {
137            s::StreamEvent::Keepalive => {}
138            s::StreamEvent::Created(v) => {
139                self.model = Some(v.response.model.clone());
140                let id = identity::response_id(
141                    &mut self.flow,
142                    &self.policy,
143                    Dialect::OpenAi,
144                    Some(v.response.id.clone()),
145                )?;
146                self.response_id = Some(id.clone());
147                self.observe(&v.response)?;
148                self.emit(
149                    g::GenerateContentResponseBody::builder()
150                        .response_id(id)
151                        .model_version(v.response.model)
152                        .build(),
153                    &mut out,
154                )?;
155                for (index, item) in v.response.output.into_iter().enumerate() {
156                    self.add_item(index as i64, item)?;
157                }
158            }
159            s::StreamEvent::InProgress(v) => self.observe(&v.response)?,
160            s::StreamEvent::Queued(v) => self.observe(&v.response)?,
161            s::StreamEvent::OutputItemAdded(v) => self.add_item(v.output_index, v.item)?,
162            s::StreamEvent::ContentPartAdded(v) => {
163                self.bind(v.output_index, &v.item_id)?;
164                self.add_part(v.output_index, v.content_index, v.part)?;
165            }
166            s::StreamEvent::OutputTextDelta(v) => {
167                self.bind(v.output_index, &v.item_id)?;
168                self.append_text(v.output_index, v.content_index, v.delta, false)?;
169            }
170            s::StreamEvent::RefusalDelta(v) => {
171                self.bind(v.output_index, &v.item_id)?;
172                self.append_text(v.output_index, v.content_index, v.delta, false)?;
173            }
174            s::StreamEvent::ReasoningTextDelta(v) => {
175                self.bind(v.output_index, &v.item_id)?;
176                self.append_text(v.output_index, v.content_index, v.delta, true)?;
177            }
178            s::StreamEvent::ContentPartDone(v) => {
179                self.part_done(v.output_index, v.content_index)?
180            }
181            s::StreamEvent::FunctionCallArgumentsDelta(v) => {
182                self.bind(v.output_index, &v.item_id)?;
183                self.reserve_arguments(v.output_index, v.delta.len())?;
184            }
185            s::StreamEvent::FunctionCallArgumentsDone(v) => {
186                self.bind(v.output_index, &v.item_id)?;
187                self.arguments_done(v.output_index, v.arguments)?;
188            }
189            s::StreamEvent::OutputItemDone(v) => self.item_done(v.output_index, v.item)?,
190            s::StreamEvent::ReasoningSummaryPartAdded(v) => {
191                self.count_part()?;
192                self.bind(v.output_index, &v.item_id)?;
193                self.summary(v.output_index, v.summary_index, v.part.text, true)?;
194            }
195            s::StreamEvent::ReasoningSummaryTextDelta(v) => {
196                self.bind(v.output_index, &v.item_id)?;
197                self.summary(v.output_index, v.summary_index, v.delta, false)?;
198            }
199            s::StreamEvent::ReasoningSummaryPartDone(v) => {
200                self.bind(v.output_index, &v.item_id)?;
201                self.summary_done(v.output_index, v.summary_index)?;
202            }
203            s::StreamEvent::Completed(v) => {
204                self.observe(&v.response)?;
205                self.terminal = true;
206            }
207            s::StreamEvent::Incomplete(v) => {
208                self.observe(&v.response)?;
209                self.terminal = true;
210            }
211            s::StreamEvent::OutputTextDone(_)
212            | s::StreamEvent::RefusalDone(_)
213            | s::StreamEvent::ReasoningTextDone(_)
214            | s::StreamEvent::OutputTextAnnotationAdded(_)
215            | s::StreamEvent::ReasoningSummaryTextDone(_) => {}
216            s::StreamEvent::ImageCall(_)
217            | s::StreamEvent::ImageGenerating(_)
218            | s::StreamEvent::ImageCompleted(_)
219            | s::StreamEvent::ImagePartial(_) => {
220                self.image_progress = true;
221            }
222            s::StreamEvent::Failed(_) | s::StreamEvent::Error(_) => {
223                return Err(invalid("native source failure"));
224            }
225            _ if self.image_only => {}
226            _ => {}
227        }
228        self.flush(&mut out)?;
229        Ok(Converted {
230            value: out,
231            report: Report::default(),
232        })
233    }
234    fn observe(&mut self, body: &r::GenerateContentResponseBody) -> Result<(), TransformError> {
235        if let Some(usage) = body.usage.as_ref().and_then(Option::as_ref) {
236            super::super::usage::to_gemini(usage.clone(), &mut Report::default())?;
237            if let Some(old) = &self.observed_usage
238                && (usage.input_tokens < old.input_tokens
239                    || usage.output_tokens < old.output_tokens
240                    || usage.total_tokens < old.total_tokens
241                    || usage.input_tokens_details.cached_tokens
242                        < old.input_tokens_details.cached_tokens
243                    || usage.input_tokens_details.cache_write_tokens
244                        < old.input_tokens_details.cache_write_tokens
245                    || usage.output_tokens_details.reasoning_tokens
246                        < old.output_tokens_details.reasoning_tokens)
247            {
248                return Err(invalid("native cumulative usage decreased"));
249            }
250            self.observed_usage = Some(usage.clone());
251        }
252        Ok(())
253    }
254    pub(super) fn reserve(&mut self, bytes: usize) -> Result<(), TransformError> {
255        self.held = self
256            .held
257            .checked_add(bytes)
258            .filter(|n| *n <= self.limits.max_pending)
259            .ok_or_else(limit)?;
260        Ok(())
261    }
262    pub(super) fn release(&mut self, bytes: usize) {
263        self.held -= bytes;
264    }
265    pub(super) fn emit(
266        &mut self,
267        chunk: g::GenerateContentResponseBody,
268        out: &mut Vec<g::GenerateContentResponseBody>,
269    ) -> Result<(), TransformError> {
270        self.budget.output(&chunk)?;
271        self.target
272            .as_mut()
273            .ok_or_else(|| invalid("target consumed"))?
274            .push(chunk.clone())?;
275        out.push(chunk);
276        Ok(())
277    }
278    pub(super) fn emit_part(
279        &mut self,
280        part: g::Part,
281        out: &mut Vec<g::GenerateContentResponseBody>,
282    ) -> Result<(), TransformError> {
283        self.emit(
284            g::GenerateContentResponseBody::builder()
285                .candidates(vec![
286                    g::Candidate::builder()
287                        .index(0)
288                        .content(
289                            g::Content::builder()
290                                .role("model")
291                                .parts(vec![part])
292                                .build(),
293                        )
294                        .build(),
295                ])
296                .build(),
297            out,
298        )
299    }
300    pub fn finish(mut self) -> Result<StreamEnd<g::GenerateContentResponseBody>, TransformError> {
301        if self.failed || !self.terminal {
302            return Err(invalid("failed stream or premature EOF"));
303        }
304        let mut source = self
305            .source
306            .take()
307            .ok_or_else(|| invalid("source consumed"))?
308            .finish()?
309            .value;
310        // Once visible summaries have been streamed, keep that representation
311        // for this item rather than switching to later raw content or replaying
312        // both. Validate against the same selected view at completion.
313        let mut summary_report = Report::default();
314        for (index, item) in source.output.iter_mut().enumerate() {
315            if self.items.get(&(index as i64)).is_some_and(|state| {
316                matches!(
317                    state.kind,
318                    super::response_items::Kind::Reasoning {
319                        summary_emitted: true,
320                        ..
321                    }
322                )
323            }) && let r::ResponseOutputItem::Reasoning(reasoning) = item
324                && reasoning.content.take().is_some()
325            {
326                summary_report.omitted(
327                    format!("output[{index}].content"),
328                    "visible reasoning summary was already streamed; late raw content is not replayed",
329                );
330            }
331        }
332        let only_image = [g::Modality::Image];
333        let mut converted = super::super::responses_to_gemini_response_with_modalities(
334            source.clone(),
335            Default::default(),
336            self.image_only.then_some(only_image.as_slice()),
337        )?;
338        converted
339            .report
340            .diagnostics
341            .extend(summary_report.diagnostics);
342        if self.image_progress {
343            converted.report.omitted("image.progress", "Gemini inline images expose the final image; preview/progress events are not additional generated images");
344        }
345        let mut check_flow = self.flow.clone();
346        converted.value.response_id = Some(identity::response_id(
347            &mut check_flow,
348            &self.policy,
349            Dialect::OpenAi,
350            Some(source.id.clone()),
351        )?);
352        let originals: Vec<_> = source
353            .output
354            .iter()
355            .enumerate()
356            .filter_map(|(index, item)| {
357                if self.image_only {
358                    return None;
359                }
360                if let r::ResponseOutputItem::FunctionCall(call) = item {
361                    Some((index as i64, call))
362                } else {
363                    None
364                }
365            })
366            .collect();
367        let projected: Vec<_> = converted
368            .value
369            .candidates
370            .iter_mut()
371            .flatten()
372            .filter_map(|c| c.content.as_mut())
373            .flat_map(|c| c.parts.iter_mut().flatten())
374            .filter(|p| p.function_call.is_some())
375            .collect();
376        if originals.len() != projected.len() {
377            return Err(invalid("canonical function count changed"));
378        }
379        let mut used = BTreeSet::new();
380        for ((index, call), part) in originals.into_iter().zip(projected) {
381            identity::call(
382                part,
383                call,
384                index,
385                (&mut check_flow, &self.policy),
386                &mut used,
387            )?;
388        }
389        measure(&converted.value, self.limits.max_bytes)?;
390        let mut tail = converted.value.clone();
391        for candidate in tail.candidates.iter_mut().flatten() {
392            if let Some(content) = &mut candidate.content {
393                content.parts = Some(Vec::new());
394            }
395        }
396        let mut out = Vec::new();
397        self.emit(tail, &mut out)?;
398        let mut actual = self
399            .target
400            .take()
401            .ok_or_else(|| invalid("target consumed"))?
402            .finish()?
403            .value;
404        identity::normalize_text(&mut actual);
405        identity::normalize_text(&mut converted.value);
406        if actual != converted.value {
407            return Err(invalid(
408                "native Gemini stream differs from completed Responses conversion",
409            ));
410        }
411        Ok(StreamEnd {
412            chunks: out,
413            identities: self.flow,
414            report: converted.report,
415        })
416    }
417}
gproxy_protocol/transform/generate/gemini_responses/stream/responses_to_gemini.rs

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