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