Skip to main content

gproxy_protocol/transform/generate/chat_responses/stream/
responses_to_chat.rs

1mod content;
2use super::common::{Budget, StreamEnd, StreamLimits, declared, invalid, unsupported};
3use crate::{
4    Dialect,
5    transform::generate::stream::{
6        chat::{ChatStreamCollector, ChatStreamLimits},
7        responses::{ResponsesStreamCollector, ResponsesStreamLimits},
8    },
9    transform::{
10        Converted, Report, TransformError,
11        identity::{IdentityFlow, IdentityRole, SourceIdentity, TargetIdPolicy},
12    },
13    wire::openai::{
14        chat::{response as c, stream as cs},
15        responses::{input as i, response as r, stream as rs},
16    },
17};
18use std::collections::BTreeMap;
19
20pub(super) struct TextPart {
21    pub pending: String,
22    pub refusal: bool,
23    pub done: bool,
24    pub emitted: bool,
25}
26
27pub(super) enum ItemKind {
28    Excluded,
29    Message {
30        parts: BTreeMap<i64, TextPart>,
31        cursor: i64,
32        done: bool,
33    },
34    Function {
35        call_id: String,
36        name: String,
37        tool_index: i64,
38        sent: bool,
39    },
40    Reasoning,
41}
42
43pub(super) struct ItemState {
44    pub id: Option<String>,
45    pub kind: ItemKind,
46}
47
48/// Each push returns only newly emitted target chunks. Source item/part order is
49/// retained when concurrent Responses messages feed Chat's single content field.
50pub struct ResponsesToChatStream {
51    pub(super) source: Option<ResponsesStreamCollector>,
52    pub(super) target: Option<ChatStreamCollector>,
53    pub(super) flow: IdentityFlow,
54    pub(super) policy: TargetIdPolicy,
55    pub(super) budget: Budget,
56    pub(super) response_id: Option<String>,
57    pub(super) model: Option<String>,
58    pub(super) created: Option<i64>,
59    pub(super) items: BTreeMap<i64, ItemState>,
60    pub(super) content_cursor: i64,
61    pub(super) tools: usize,
62    pub(super) report: Report,
63    pub(super) expected: Option<c::GenerateContentResponseBody>,
64    pub(super) terminal: bool,
65    pub(super) failed: bool,
66}
67
68impl ResponsesToChatStream {
69    pub fn new(flow: IdentityFlow, limits: StreamLimits) -> Self {
70        Self::new_with_policy(
71            flow,
72            limits,
73            TargetIdPolicy::new(crate::Dialect::OpenAiChat),
74        )
75        .expect("default policy matches target dialect")
76    }
77    pub fn new_with_policy(
78        flow: IdentityFlow,
79        limits: StreamLimits,
80        policy: TargetIdPolicy,
81    ) -> Result<Self, TransformError> {
82        if policy.dialect != crate::Dialect::OpenAiChat {
83            return Err(TransformError::shape(
84                "stream.policy",
85                "target policy dialect mismatch",
86            ));
87        }
88
89        // These IDs were already allocated under the selected policy. The
90        // native target collector must preserve them while checking syntax;
91        // preserve_source_ids=false must not trigger a second allocation.
92        let mut collector_policy = policy.clone();
93        collector_policy.preserve_source_ids = true;
94        Ok(Self {
95            source: Some(ResponsesStreamCollector::new(ResponsesStreamLimits {
96                max_bytes: limits.max_bytes,
97
98                max_text_bytes: limits.max_bytes,
99                max_json_bytes: limits.max_bytes,
100            })),
101            target: Some(ChatStreamCollector::with_limits(
102                IdentityFlow::new(flow.namespace()),
103                collector_policy,
104                ChatStreamLimits {
105                    max_bytes: limits.max_bytes,
106                },
107            )),
108            flow,
109            policy,
110            budget: Budget::new(limits),
111            response_id: None,
112            model: None,
113            created: None,
114            items: BTreeMap::new(),
115            content_cursor: 0,
116            tools: 0,
117            report: Report::default(),
118            expected: None,
119            terminal: false,
120            failed: false,
121        })
122    }
123    pub(crate) fn reserve_external_ids(
124        &mut self,
125        role: crate::transform::identity::IdentityRole,
126        ids: &std::collections::BTreeSet<String>,
127    ) -> Result<(), TransformError> {
128        self.flow.reserve_external_ids(role, ids).map_err(|e| {
129            TransformError::new(
130                crate::transform::TransformErrorKind::Conflict,
131                "fanout.ids",
132                e.to_string(),
133            )
134        })
135    }
136    pub fn identities(&self) -> &IdentityFlow {
137        &self.flow
138    }
139    pub fn push(
140        &mut self,
141        event: rs::StreamEvent,
142    ) -> Result<Converted<Vec<cs::ChatCompletionChunk>>, TransformError> {
143        if self.failed || self.terminal {
144            self.failed = true;
145            return Err(invalid("event after terminal or failed stream"));
146        }
147        let result = self.push_inner(declared(event));
148        if result.is_err() {
149            self.failed = true;
150        }
151        result
152    }
153    fn push_inner(
154        &mut self,
155        event: rs::StreamEvent,
156    ) -> Result<Converted<Vec<cs::ChatCompletionChunk>>, TransformError> {
157        let event = crate::transform::generate::multi_agent::attribute_event(event)?;
158        self.budget.input(&event)?;
159        self.source
160            .as_mut()
161            .ok_or_else(|| invalid("source consumed"))?
162            .push(event.clone())?;
163        let mut out = Vec::new();
164        match event {
165            rs::StreamEvent::Created(v) => {
166                let id = self
167                    .flow
168                    .resolve_or_allocate(
169                        IdentityRole::Response,
170                        SourceIdentity::new(Dialect::OpenAi, Some(v.response.id), 0),
171                        &self.policy,
172                    )
173                    .map_err(|e| {
174                        TransformError::invalid_result("identity.response", e.to_string())
175                    })?
176                    .emitted_id;
177                self.response_id = Some(id);
178                self.model = Some(v.response.model);
179                self.created = Some(v.response.created_at);
180                self.emit_delta(
181                    cs::Delta::builder().role(cs::DeltaRole::Assistant).build(),
182                    &mut out,
183                )?;
184                for (index, item) in v.response.output.into_iter().enumerate() {
185                    self.add_item(index as i64, item, &mut out)?;
186                }
187            }
188            rs::StreamEvent::OutputItemAdded(v) => {
189                self.add_item(v.output_index, v.item, &mut out)?
190            }
191            rs::StreamEvent::ContentPartAdded(v) => {
192                self.add_part(v.output_index, v.content_index, v.part)?
193            }
194            rs::StreamEvent::OutputTextDelta(v) => {
195                self.append_text(v.output_index, v.content_index, &v.item_id, v.delta, false)?
196            }
197            rs::StreamEvent::RefusalDelta(v) => {
198                self.append_text(v.output_index, v.content_index, &v.item_id, v.delta, true)?
199            }
200            rs::StreamEvent::ContentPartDone(v) => {
201                if let Some(ItemState {
202                    kind: ItemKind::Message { parts, .. },
203                    ..
204                }) = self.items.get_mut(&v.output_index)
205                {
206                    parts
207                        .get_mut(&v.content_index)
208                        .ok_or_else(|| invalid("part done without part"))?
209                        .done = true;
210                }
211            }
212            rs::StreamEvent::OutputItemDone(v) => {
213                self.bind_item(v.output_index, item_id(&v.item))?;
214                if let Some(ItemState {
215                    kind: ItemKind::Message { parts, done, .. },
216                    ..
217                }) = self.items.get_mut(&v.output_index)
218                {
219                    *done = true;
220                    for p in parts.values_mut() {
221                        p.done = true;
222                    }
223                }
224            }
225            rs::StreamEvent::FunctionCallArgumentsDelta(v) => {
226                self.bind_item(v.output_index, Some(&v.item_id))?;
227                self.emit_tool(v.output_index, v.delta, &mut out)?;
228            }
229            rs::StreamEvent::FunctionCallArgumentsDone(_) => {}
230            rs::StreamEvent::Completed(_) | rs::StreamEvent::Incomplete(_) => {
231                self.finish_source(&mut out)?;
232            }
233            rs::StreamEvent::Queued(_)
234            | rs::StreamEvent::InProgress(_)
235            | rs::StreamEvent::OutputTextDone(_)
236            | rs::StreamEvent::RefusalDone(_) => {}
237            rs::StreamEvent::OutputTextAnnotationAdded(_) => self.report.omitted(
238                "response.annotations",
239                "Chat stream Delta has no annotation field",
240            ),
241            rs::StreamEvent::ReasoningTextDelta(_)
242            | rs::StreamEvent::ReasoningSummaryTextDelta(_) => {}
243            rs::StreamEvent::ReasoningTextDone(_)
244            | rs::StreamEvent::ReasoningSummaryTextDone(_)
245            | rs::StreamEvent::ReasoningSummaryPartAdded(_)
246            | rs::StreamEvent::ReasoningSummaryPartDone(_) => {}
247            rs::StreamEvent::Failed(_) | rs::StreamEvent::Error(_) => {
248                return Err(invalid("Responses stream failed"));
249            }
250            rs::StreamEvent::CustomToolInputDelta(_) | rs::StreamEvent::CustomToolInputDone(_) => {
251                return Err(unsupported(
252                    "custom_tool",
253                    "Chat streaming tool schema supports only function calls; buffered Chat has a separate custom-tool shape",
254                ));
255            }
256            rs::StreamEvent::AudioDelta(_)
257            | rs::StreamEvent::AudioDone(_)
258            | rs::StreamEvent::AudioTranscriptDelta(_)
259            | rs::StreamEvent::AudioTranscriptDone(_) => {
260                return Err(unsupported(
261                    "audio",
262                    "this native stream wire has no canonical output audio item",
263                ));
264            }
265            _ => {
266                return Err(unsupported(
267                    "response.execution",
268                    "native execution events require an invocation adapter",
269                ));
270            }
271        }
272        if !self.terminal {
273            self.flush_text(&mut out)?;
274        }
275        Ok(Converted {
276            value: out,
277            report: std::mem::take(&mut self.report),
278        })
279    }
280    fn finish_source(
281        &mut self,
282        out: &mut Vec<cs::ChatCompletionChunk>,
283    ) -> Result<(), TransformError> {
284        let actual = self
285            .source
286            .take()
287            .ok_or_else(|| invalid("source consumed"))?
288            .finish()?;
289        self.report.diagnostics.extend(actual.report.diagnostics);
290        let mapped =
291            super::super::responses_to_chat_response(actual.value, &mut self.flow, &self.policy)?;
292        self.report.diagnostics.extend(mapped.report.diagnostics);
293        for item in self.items.values_mut() {
294            if let ItemKind::Message { parts, done, .. } = &mut item.kind {
295                *done = true;
296                for p in parts.values_mut() {
297                    p.done = true;
298                }
299            }
300        }
301        self.flush_text(out)?;
302        let mut expected = mapped.value;
303        let message = &expected.choices[0].message;
304        if message.reasoning_details.is_some() || message.reasoning_content.is_some() {
305            let mut delta = cs::Delta::builder().build();
306            delta.reasoning_details = message.reasoning_details.clone();
307            delta.reasoning_content = message.reasoning_content.clone();
308            self.emit_delta(delta, out)?;
309        }
310        if expected.choices[0].message.annotations.take().is_some() {
311            self.report.omitted(
312                "message.annotations",
313                "Chat streaming Delta has no annotation field",
314            );
315        }
316        let logs = expected.choices[0]
317            .logprobs
318            .clone()
319            .map(crate::transform::generate::stream::chat::logs::synthesize)
320            .map(crate::transform::optional)
321            .transpose()?
322            .flatten();
323        let mut choice = cs::StreamChoice::builder(0, cs::Delta::builder().build())
324            .finish_reason(expected.choices[0].finish_reason)
325            .build();
326        choice.logprobs = logs.map(Some);
327        let mut terminal = self.chunk(vec![choice])?;
328        terminal.service_tier = expected.service_tier;
329        terminal.moderation = expected.moderation.clone();
330        self.emit(terminal, out)?;
331        if let Some(usage) = expected.usage.clone() {
332            let usage = crate::transform::generate::stream::chat::usage::synthesize(usage)?;
333            let mut chunk = self.chunk(Vec::new())?;
334            chunk.usage = Some(Some(usage));
335            self.emit(chunk, out)?;
336        }
337        self.expected = Some(expected);
338        self.terminal = true;
339        Ok(())
340    }
341    fn emit_delta(
342        &mut self,
343        delta: cs::Delta,
344        out: &mut Vec<cs::ChatCompletionChunk>,
345    ) -> Result<(), TransformError> {
346        let chunk = self.chunk(vec![cs::StreamChoice::builder(0, delta).build()])?;
347        self.emit(chunk, out)
348    }
349    fn chunk(
350        &self,
351        choices: Vec<cs::StreamChoice>,
352    ) -> Result<cs::ChatCompletionChunk, TransformError> {
353        Ok(cs::ChatCompletionChunk::builder(
354            self.response_id
355                .clone()
356                .ok_or_else(|| TransformError::missing_metadata("response.id"))?,
357            choices,
358            self.created
359                .ok_or_else(|| TransformError::missing_metadata("response.created_at"))?,
360            self.model
361                .clone()
362                .ok_or_else(|| TransformError::missing_metadata("response.model"))?,
363            cs::ChunkObject::ChatCompletionChunk,
364        )
365        .build())
366    }
367    fn emit(
368        &mut self,
369        chunk: cs::ChatCompletionChunk,
370        out: &mut Vec<cs::ChatCompletionChunk>,
371    ) -> Result<(), TransformError> {
372        self.budget.output(&chunk)?;
373        self.target
374            .as_mut()
375            .ok_or_else(|| invalid("target consumed"))?
376            .push(chunk.clone())?;
377        out.push(chunk);
378        Ok(())
379    }
380    pub fn finish(mut self) -> Result<StreamEnd<cs::ChatCompletionChunk>, TransformError> {
381        if self.failed || !self.terminal {
382            return Err(invalid("Responses EOF without terminal"));
383        }
384        let mut target = self
385            .target
386            .take()
387            .ok_or_else(|| invalid("target consumed"))?;
388        target.push_done()?;
389        let actual = target.finish()?;
390        if Some(&actual.value) != self.expected.as_ref() {
391            return Err(invalid(
392                "emitted Chat content/identity/metadata differs from native terminal response",
393            ));
394        }
395        self.report.diagnostics.extend(actual.report.diagnostics);
396        Ok(StreamEnd {
397            chunks: Vec::new(),
398            identities: self.flow,
399            report: self.report,
400        })
401    }
402}
403
404fn item_id(item: &r::ResponseOutputItem) -> Option<&str> {
405    match item {
406        r::ResponseOutputItem::Message(v) => Some(&v.id),
407        r::ResponseOutputItem::FunctionCall(v) => v.id.as_deref(),
408        r::ResponseOutputItem::Reasoning(v) => Some(&v.id),
409        _ => None,
410    }
411}