gproxy_protocol/transform/generate/stream/chat/
collector.rs

1use super::choice::ChoiceAccum;
2use crate::{
3    transform::{
4        Converted, Report, TransformError,
5        identity::{IdentityFlow, TargetIdPolicy},
6    },
7    wire::{
8        DeclaredFields,
9        openai::chat::{response as r, stream as s},
10    },
11};
12use std::collections::{BTreeMap, BTreeSet};
13
14#[derive(Debug, Clone, Copy)]
15pub struct ChatStreamLimits {
16    pub max_bytes: usize,
17}
18
19impl Default for ChatStreamLimits {
20    fn default() -> Self {
21        Self {
22            max_bytes: 16 * 1024 * 1024,
23        }
24    }
25}
26
27#[derive(Debug, Default, Clone)]
28pub struct ChatStreamContext {
29    pub created: Option<i64>,
30    pub model: Option<String>,
31    pub response_id: Option<String>,
32}
33
34/// Native Chat collector. The SSE decoder must call `push_done` for `[DONE]`;
35/// EOF after a finish chunk without the framing terminator remains incomplete.
36pub struct ChatStreamCollector {
37    context: ChatStreamContext,
38    choices: BTreeMap<i64, ChoiceAccum>,
39    usage: Option<s::ChunkUsage>,
40    service_tier: Option<Option<r::ResponseServiceTier>>,
41    fingerprint: Option<String>,
42    moderation: Option<Option<r::ModerationResponse>>,
43    done: bool,
44    failed: bool,
45    flow: IdentityFlow,
46    target: TargetIdPolicy,
47    limits: ChatStreamLimits,
48
49    bytes: usize,
50
51    report: Report,
52}
53
54impl ChatStreamCollector {
55    pub fn new(flow: IdentityFlow, target: TargetIdPolicy) -> Self {
56        Self::with_limits(flow, target, Default::default())
57    }
58    pub fn with_limits(
59        flow: IdentityFlow,
60        target: TargetIdPolicy,
61        limits: ChatStreamLimits,
62    ) -> Self {
63        Self {
64            context: Default::default(),
65            choices: BTreeMap::new(),
66            usage: None,
67            service_tier: None,
68            fingerprint: None,
69            moderation: None,
70            done: false,
71            failed: false,
72            flow,
73            target,
74            limits,
75
76            bytes: 0,
77
78            report: Default::default(),
79        }
80    }
81    pub fn push(&mut self, chunk: s::ChatCompletionChunk) -> Result<(), TransformError> {
82        if self.failed || self.done {
83            self.failed = true;
84            return Err(TransformError::invalid_result(
85                "stream",
86                "event after terminal or failed collector",
87            ));
88        }
89        let result = self.push_declared(chunk.into_declared());
90        if result.is_err() {
91            self.failed = true;
92        }
93        result
94    }
95    fn push_declared(&mut self, chunk: s::ChatCompletionChunk) -> Result<(), TransformError> {
96        if self.target.dialect != crate::Dialect::OpenAiChat {
97            return Err(TransformError::shape(
98                "identity.policy",
99                "native Chat collection requires OpenAiChat target policy",
100            ));
101        }
102
103        let remaining = self
104            .limits
105            .max_bytes
106            .checked_sub(self.bytes)
107            .ok_or_else(|| super::limit("bytes"))?;
108        self.bytes += super::encoded(&chunk, remaining)?;
109        if chunk.created < 0 {
110            return Err(TransformError::invalid_result(
111                "created",
112                "negative creation time",
113            ));
114        }
115        if self.context.created.is_some_and(|v| v != chunk.created) {
116            return Err(TransformError::invalid_result(
117                "created",
118                "creation time changed",
119            ));
120        }
121        self.context.created = Some(chunk.created);
122        identity(&mut self.context.response_id, chunk.id, "id")?;
123        identity(&mut self.context.model, chunk.model, "model")?;
124        if let Some(value) = chunk.service_tier {
125            if self
126                .service_tier
127                .flatten()
128                .zip(value)
129                .is_some_and(|(a, b)| a != b)
130            {
131                return Err(TransformError::invalid_result(
132                    "service_tier",
133                    "tier changed during response",
134                ));
135            }
136            if value.is_some() || self.service_tier.is_none() {
137                self.service_tier = Some(value);
138            }
139        }
140        if let Some(Some(value)) = chunk.system_fingerprint {
141            identity(&mut self.fingerprint, value, "system_fingerprint")?;
142        }
143        if let Some(value) = chunk.moderation {
144            self.moderation = Some(value);
145        }
146        if chunk.obfuscation.flatten().is_some() {
147            self.report
148                .omitted("obfuscation", "stream padding is not buffered content");
149        }
150        if let Some(Some(usage)) = chunk.usage {
151            if self.usage.as_ref().is_some_and(|v| v != &usage) {
152                return Err(TransformError::invalid_result(
153                    "usage",
154                    "conflicting final usage",
155                ));
156            }
157            self.usage = Some(usage);
158        }
159        let mut indexes = BTreeSet::new();
160        for choice in chunk.choices {
161            if choice.index < 0 || !indexes.insert(choice.index) {
162                return Err(TransformError::invalid_result(
163                    "choice.index",
164                    "negative or duplicate choice index",
165                ));
166            }
167
168            let state = self.choices.entry(choice.index).or_default();
169            state.push(choice)?;
170        }
171        Ok(())
172    }
173    pub fn push_done(&mut self) -> Result<(), TransformError> {
174        if self.failed || self.done {
175            self.failed = true;
176            return Err(TransformError::invalid_result(
177                "stream",
178                "duplicate DONE or failed collector",
179            ));
180        }
181        self.done = true;
182        if self.choices.is_empty() || self.choices.values().any(|c| !c.is_finished()) {
183            self.failed = true;
184            return Err(TransformError::invalid_result(
185                "stream",
186                "DONE before every choice finished",
187            ));
188        }
189        Ok(())
190    }
191    pub fn finish(mut self) -> Result<Converted<r::GenerateContentResponseBody>, TransformError> {
192        if self.failed || !self.done {
193            return Err(TransformError::invalid_result(
194                "stream",
195                "premature EOF or previous stream error",
196            ));
197        }
198        let id = self
199            .context
200            .response_id
201            .ok_or_else(|| TransformError::missing_metadata("id"))?;
202        let model = self
203            .context
204            .model
205            .ok_or_else(|| TransformError::missing_metadata("model"))?;
206        let created = self
207            .context
208            .created
209            .ok_or_else(|| TransformError::missing_metadata("created"))?;
210        let mut choices = Vec::new();
211        let mut tool_index = 0;
212        for (expected, (index, state)) in self.choices.into_iter().enumerate() {
213            if index != expected as i64 {
214                return Err(TransformError::invalid_result(
215                    "choice.index",
216                    "non-contiguous choices",
217                ));
218            }
219            choices.push(state.finish(index, &mut tool_index, &mut self.flow, &self.target)?);
220        }
221        let mut out = r::GenerateContentResponseBody::builder(
222            id,
223            choices,
224            created,
225            model,
226            r::CompletionObject::ChatCompletion,
227        )
228        .build();
229        out.usage = self
230            .usage
231            .map(|u| super::usage::collect(u, &mut self.report));
232        out.service_tier = self.service_tier;
233        out.system_fingerprint = self.fingerprint;
234        out.moderation = self.moderation;
235        Ok(Converted {
236            value: out,
237            report: self.report,
238        })
239    }
240}
241
242fn identity(old: &mut Option<String>, new: String, field: &str) -> Result<(), TransformError> {
243    if new.is_empty() {
244        return Ok(());
245    }
246    if old.as_ref().is_some_and(|v| v != &new) {
247        return Err(TransformError::invalid_result(
248            field,
249            "stream identity changed",
250        ));
251    }
252    *old = Some(new);
253    Ok(())
254}
gproxy_protocol/transform/generate/stream/chat/collector.rs

gproxy_protocol/transform/generate/stream/chat/
collector.rs
