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.summary(v.output_index, v.part.text.len())?;
193 }
194 s::StreamEvent::ReasoningSummaryTextDelta(v) => {
195 self.summary(v.output_index, v.delta.len())?
196 }
197 s::StreamEvent::Completed(v) => {
198 self.observe(&v.response)?;
199 self.terminal = true;
200 }
201 s::StreamEvent::Incomplete(v) => {
202 self.observe(&v.response)?;
203 self.terminal = true;
204 }
205 s::StreamEvent::OutputTextDone(_)
206 | s::StreamEvent::RefusalDone(_)
207 | s::StreamEvent::ReasoningTextDone(_)
208 | s::StreamEvent::OutputTextAnnotationAdded(_)
209 | s::StreamEvent::ReasoningSummaryTextDone(_)
210 | s::StreamEvent::ReasoningSummaryPartDone(_) => {}
211 s::StreamEvent::ImageCall(_)
212 | s::StreamEvent::ImageGenerating(_)
213 | s::StreamEvent::ImageCompleted(_)
214 | s::StreamEvent::ImagePartial(_) => {
215 self.image_progress = true;
216 }
217 s::StreamEvent::Failed(_) | s::StreamEvent::Error(_) => {
218 return Err(invalid("native source failure"));
219 }
220 _ if self.image_only => {}
221 _ => {}
222 }
223 self.flush(&mut out)?;
224 Ok(Converted {
225 value: out,
226 report: Report::default(),
227 })
228 }
229 fn observe(&mut self, body: &r::GenerateContentResponseBody) -> Result<(), TransformError> {
230 if let Some(usage) = body.usage.as_ref().and_then(Option::as_ref) {
231 super::super::usage::to_gemini(usage.clone(), &mut Report::default())?;
232 if let Some(old) = &self.observed_usage
233 && (usage.input_tokens < old.input_tokens
234 || usage.output_tokens < old.output_tokens
235 || usage.total_tokens < old.total_tokens
236 || usage.input_tokens_details.cached_tokens
237 < old.input_tokens_details.cached_tokens
238 || usage.input_tokens_details.cache_write_tokens
239 < old.input_tokens_details.cache_write_tokens
240 || usage.output_tokens_details.reasoning_tokens
241 < old.output_tokens_details.reasoning_tokens)
242 {
243 return Err(invalid("native cumulative usage decreased"));
244 }
245 self.observed_usage = Some(usage.clone());
246 }
247 Ok(())
248 }
249 pub(super) fn reserve(&mut self, bytes: usize) -> Result<(), TransformError> {
250 self.held = self
251 .held
252 .checked_add(bytes)
253 .filter(|n| *n <= self.limits.max_pending)
254 .ok_or_else(limit)?;
255 Ok(())
256 }
257 pub(super) fn release(&mut self, bytes: usize) {
258 self.held -= bytes;
259 }
260 pub(super) fn emit(
261 &mut self,
262 chunk: g::GenerateContentResponseBody,
263 out: &mut Vec<g::GenerateContentResponseBody>,
264 ) -> Result<(), TransformError> {
265 self.budget.output(&chunk)?;
266 self.target
267 .as_mut()
268 .ok_or_else(|| invalid("target consumed"))?
269 .push(chunk.clone())?;
270 out.push(chunk);
271 Ok(())
272 }
273 pub(super) fn emit_part(
274 &mut self,
275 part: g::Part,
276 out: &mut Vec<g::GenerateContentResponseBody>,
277 ) -> Result<(), TransformError> {
278 self.emit(
279 g::GenerateContentResponseBody::builder()
280 .candidates(vec![
281 g::Candidate::builder()
282 .index(0)
283 .content(
284 g::Content::builder()
285 .role("model")
286 .parts(vec![part])
287 .build(),
288 )
289 .build(),
290 ])
291 .build(),
292 out,
293 )
294 }
295 pub fn finish(mut self) -> Result<StreamEnd<g::GenerateContentResponseBody>, TransformError> {
296 if self.failed || !self.terminal {
297 return Err(invalid("failed stream or premature EOF"));
298 }
299 let source = self
300 .source
301 .take()
302 .ok_or_else(|| invalid("source consumed"))?
303 .finish()?
304 .value;
305 let only_image = [g::Modality::Image];
306 let mut converted = super::super::responses_to_gemini_response_with_modalities(
307 source.clone(),
308 Default::default(),
309 self.image_only.then_some(only_image.as_slice()),
310 )?;
311 if self.image_progress {
312 converted.report.omitted("image.progress", "Gemini inline images expose the final image; preview/progress events are not additional generated images");
313 }
314 let mut check_flow = self.flow.clone();
315 converted.value.response_id = Some(identity::response_id(
316 &mut check_flow,
317 &self.policy,
318 Dialect::OpenAi,
319 Some(source.id.clone()),
320 )?);
321 let originals: Vec<_> = source
322 .output
323 .iter()
324 .enumerate()
325 .filter_map(|(index, item)| {
326 if self.image_only {
327 return None;
328 }
329 if let r::ResponseOutputItem::FunctionCall(call) = item {
330 Some((index as i64, call))
331 } else {
332 None
333 }
334 })
335 .collect();
336 let projected: Vec<_> = converted
337 .value
338 .candidates
339 .iter_mut()
340 .flatten()
341 .filter_map(|c| c.content.as_mut())
342 .flat_map(|c| c.parts.iter_mut().flatten())
343 .filter(|p| p.function_call.is_some())
344 .collect();
345 if originals.len() != projected.len() {
346 return Err(invalid("canonical function count changed"));
347 }
348 let mut used = BTreeSet::new();
349 for ((index, call), part) in originals.into_iter().zip(projected) {
350 identity::call(
351 part,
352 call,
353 index,
354 (&mut check_flow, &self.policy),
355 &mut used,
356 )?;
357 }
358 measure(&converted.value, self.limits.max_bytes)?;
359 let mut tail = converted.value.clone();
360 for candidate in tail.candidates.iter_mut().flatten() {
361 if let Some(content) = &mut candidate.content {
362 content.parts = Some(Vec::new());
363 }
364 }
365 let mut out = Vec::new();
366 self.emit(tail, &mut out)?;
367 let mut actual = self
368 .target
369 .take()
370 .ok_or_else(|| invalid("target consumed"))?
371 .finish()?
372 .value;
373 identity::normalize_text(&mut actual);
374 identity::normalize_text(&mut converted.value);
375 if actual != converted.value {
376 return Err(invalid(
377 "native Gemini stream differs from completed Responses conversion",
378 ));
379 }
380 Ok(StreamEnd {
381 chunks: out,
382 identities: self.flow,
383 report: converted.report,
384 })
385 }
386}