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