gproxy_protocol/transform/generate/gemini_responses/stream/
gemini_to_responses.rs1use super::{
2 common::{Budget, StreamEnd, StreamLimits, declared, invalid, limit, measure},
3 context::{GeminiToResponsesContext, clean_context, clone_response},
4 identity,
5 response_events::ResponseEvents,
6 usage::GeminiUsageProgress,
7};
8use crate::{
9 Dialect,
10 transform::generate::stream::gemini::{GeminiStreamCollector, GeminiStreamLimits},
11 transform::{
12 Converted, Report, TransformError,
13 identity::{IdentityFlow, TargetIdPolicy},
14 },
15 wire::{
16 gemini as g,
17 openai::responses::{response as r, stream as s},
18 },
19};
20
21pub struct GeminiToResponsesStream {
22 source: Option<GeminiStreamCollector>,
23 pub(super) target: ResponseEvents,
24 pub(super) flow: IdentityFlow,
25 pub(super) policy: TargetIdPolicy,
26 pub(super) context: GeminiToResponsesContext,
27 pub(super) limits: StreamLimits,
28 pub(super) budget: Budget,
29 model: Option<String>,
30 source_id: Option<String>,
31 emitted_id: Option<String>,
32 started: bool,
33 failed: bool,
34 pending: Vec<g::Part>,
35 pending_bytes: usize,
36 usage: GeminiUsageProgress,
37 usage_after_content: bool,
38 pub(super) part_index: usize,
39 pub(super) item_count: usize,
40
41 pub(super) call_ids: std::collections::BTreeSet<String>,
42 pub(super) client_tools: crate::transform::generate::client_tools::Bindings,
43}
44
45impl GeminiToResponsesStream {
46 pub fn new(
47 context: GeminiToResponsesContext,
48 flow: IdentityFlow,
49 limits: StreamLimits,
50 ) -> Result<Self, TransformError> {
51 Self::new_with_policy(context, flow, TargetIdPolicy::new(Dialect::OpenAi), limits)
52 }
53 pub fn new_with_policy(
54 context: GeminiToResponsesContext,
55 flow: IdentityFlow,
56 policy: TargetIdPolicy,
57 limits: StreamLimits,
58 ) -> Result<Self, TransformError> {
59 if policy.dialect != Dialect::OpenAi {
60 return Err(invalid("Responses target policy required"));
61 }
62 let context = clean_context(context, limits)?;
63 let client_tools = crate::transform::generate::client_tools::Bindings::for_target(
64 &context.response.request,
65 Dialect::Gemini,
66 )?;
67 Ok(Self {
68 source: Some(GeminiStreamCollector::new(GeminiStreamLimits {
69 max_bytes: limits.max_bytes,
70 })),
71 target: ResponseEvents::new(limits),
72 flow,
73 policy,
74 model: context.actual_model.clone(),
75 context,
76 limits,
77 budget: Budget::new(limits),
78 source_id: None,
79 emitted_id: None,
80 started: false,
81 failed: false,
82 pending: Vec::new(),
83 pending_bytes: 0,
84 usage: GeminiUsageProgress::default(),
85 usage_after_content: false,
86 part_index: 0,
87 item_count: 0,
88
89 call_ids: Default::default(),
90 client_tools,
91 })
92 }
93 pub fn identities(&self) -> &IdentityFlow {
94 &self.flow
95 }
96 pub fn push(
97 &mut self,
98 chunk: g::GenerateContentResponseBody,
99 ) -> Result<Converted<Vec<s::StreamEvent>>, TransformError> {
100 if self.failed {
101 return Err(invalid("stream failed"));
102 }
103 let result = self.inner(declared(chunk));
104 if result.is_err() {
105 self.failed = true;
106 }
107 result
108 }
109 fn inner(
110 &mut self,
111 chunk: g::GenerateContentResponseBody,
112 ) -> Result<Converted<Vec<s::StreamEvent>>, TransformError> {
113 self.budget.input(&chunk)?;
114 self.source
115 .as_mut()
116 .ok_or_else(|| invalid("source consumed"))?
117 .push(chunk.clone())?;
118 if let Some(model) = &chunk.model_version {
119 if self.model.as_ref().is_some_and(|old| old != model) {
120 return Err(invalid("native model conflicts with actual model fact"));
121 }
122 self.model = Some(model.clone());
123 }
124 if let Some(id) = &chunk.response_id {
125 self.source_id = Some(id.clone());
126 if let Some(previous) = &self.emitted_id {
127 let bound = identity::response_id(
128 &mut self.flow,
129 &self.policy,
130 Dialect::Gemini,
131 Some(id.clone()),
132 )?;
133 if &bound != previous {
134 return Err(invalid("late source response ID changed emitted ID"));
135 }
136 }
137 }
138 let usage_present = chunk.usage_metadata.is_some();
139 if let Some(usage) = &chunk.usage_metadata {
140 self.usage.observe(usage);
141 self.usage_after_content = true;
142 }
143 let mut out = Vec::new();
144 let mut report = Report::default();
145 self.start(&mut out)?;
146 for candidate in chunk.candidates.into_iter().flatten() {
147 if candidate.index.is_some_and(|n| n != 0) {
148 continue;
149 }
150 if let Some(content) = candidate.content {
151 if content.role.as_deref().is_some_and(|r| r != "model") {
152 return Err(invalid("generated content must have model role"));
153 }
154 for part in content.parts.into_iter().flatten() {
155 crate::transform::generate::gemini_controls::omitted(&part, &mut report);
156 if !usage_present {
157 self.usage_after_content = false;
158 }
159 if self.started {
160 self.part(part, &mut out)?;
161 } else {
162 let size = measure(
163 &part,
164 self.limits.max_pending.saturating_sub(self.pending_bytes),
165 )?;
166 self.pending_bytes =
167 self.pending_bytes.checked_add(size).ok_or_else(limit)?;
168 if self.pending_bytes > self.limits.max_pending {
169 return Err(limit());
170 }
171 self.pending.push(part);
172 }
173 }
174 }
175 }
176 Ok(Converted { value: out, report })
177 }
178 fn start(&mut self, out: &mut Vec<s::StreamEvent>) -> Result<(), TransformError> {
179 if self.started {
180 return Ok(());
181 }
182 let Some(model) = self.model.clone() else {
183 return Ok(());
184 };
185 let id = identity::response_id(
186 &mut self.flow,
187 &self.policy,
188 Dialect::Gemini,
189 self.source_id.clone(),
190 )?;
191 let mut body = clone_response(&self.context.response).into_response(
192 id.clone(),
193 self.context.response.created_at,
194 model,
195 )?;
196 body.status = Some(r::ResponseStatus::InProgress);
197 self.target.emit(&mut self.budget, out, |sequence_number| {
198 s::StreamEvent::Created(s::ResponseCreated {
199 sequence_number,
200 response: body,
201 rest: Default::default(),
202 })
203 })?;
204 self.emitted_id = Some(id);
205 self.started = true;
206 self.pending_bytes = 0;
207 for part in std::mem::take(&mut self.pending) {
208 self.part(part, out)?;
209 }
210 Ok(())
211 }
212 pub fn finish(mut self) -> Result<StreamEnd<s::StreamEvent>, TransformError> {
213 if self.failed {
214 return Err(invalid("stream failed"));
215 }
216 let mut source = self
217 .source
218 .take()
219 .ok_or_else(|| invalid("source consumed"))?
220 .finish()?
221 .value;
222 if !self.usage_after_content {
223 return Err(TransformError::missing_metadata(
224 "final Gemini usage after last content",
225 ));
226 }
227 let usage = source
228 .usage_metadata
229 .as_mut()
230 .ok_or_else(|| TransformError::missing_metadata("final Gemini usage"))?;
231 let changed = self
232 .usage
233 .effective(usage, self.context.final_thinking_tokens)?;
234 let model = self
235 .model
236 .clone()
237 .ok_or_else(|| TransformError::missing_metadata("actual model for Responses header"))?;
238 source.model_version = Some(model);
239 let mut out = Vec::new();
240 self.start(&mut out)?;
241 let converted = super::super::gemini_to_responses_response(
242 source,
243 self.context.response,
244 &mut self.flow.clone(),
245 &self.policy,
246 )?;
247 let expected = converted.value;
248 measure(&expected, self.limits.max_bytes)?;
249 if expected.output.len() != self.item_count {
250 return Err(invalid("native Gemini output item count changed"));
251 }
252 for (index, item) in expected.output.iter().enumerate() {
253 if matches!(item, r::ResponseOutputItem::ImageGenerationCall(_)) {
256 continue;
257 }
258 if matches!(item, r::ResponseOutputItem::FunctionCall(_)) {
259 self.target
260 .arguments_done(&mut self.budget, &mut out, index as i64, item)?;
261 }
262 self.target
263 .item_done(&mut self.budget, &mut out, index as i64, item.clone())?;
264 }
265 self.target
266 .emit(&mut self.budget, &mut out, |sequence_number| {
267 if expected.status == Some(r::ResponseStatus::Incomplete) {
268 s::StreamEvent::Incomplete(s::ResponseIncomplete {
269 sequence_number,
270 response: expected.clone(),
271 rest: Default::default(),
272 })
273 } else {
274 s::StreamEvent::Completed(s::ResponseCompleted {
275 sequence_number,
276 response: expected.clone(),
277 rest: Default::default(),
278 })
279 }
280 })?;
281 let actual = self.target.finish()?;
282 if actual != expected {
283 return Err(invalid(
284 "native Responses result differs from completed Gemini conversion",
285 ));
286 }
287 let mut report = converted.report;
288 if changed {
289 report.changed("usage","final split resolved from current counters/explicit final thinking; stale prefix components are lower bounds");
290 }
291 Ok(StreamEnd {
292 chunks: out,
293 identities: self.flow,
294 report,
295 })
296 }
297}