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

1use crate::{
2    transform::{Report, TransformError, TransformErrorKind},
3    wire::DeclaredFields,
4};
5
6#[derive(Debug, Clone, Copy)]
7pub struct StreamLimits {
8    pub max_bytes: usize,
9}
10
11impl Default for StreamLimits {
12    fn default() -> Self {
13        Self {
14            max_bytes: 16 * 1024 * 1024,
15        }
16    }
17}
18
19pub struct StreamEnd<T> {
20    pub chunks: Vec<T>,
21    pub identities: crate::transform::identity::IdentityFlow,
22    pub report: Report,
23}
24
25pub(super) fn invalid(message: &'static str) -> TransformError {
26    TransformError::invalid_result("chat_responses.stream", message)
27}
28
29pub(super) fn unsupported(field: &'static str, message: &'static str) -> TransformError {
30    TransformError::unsupported(field, message)
31}
32
33pub(super) fn limit() -> TransformError {
34    TransformError::new(
35        TransformErrorKind::Limit,
36        "chat_responses.stream",
37        "stream limit exceeded",
38    )
39}
40
41pub(super) struct Budget {
42    limits: StreamLimits,
43
44    input_bytes: usize,
45
46    output_bytes: usize,
47}
48
49impl Budget {
50    pub(super) fn new(limits: StreamLimits) -> Self {
51        Self {
52            limits,
53
54            input_bytes: 0,
55
56            output_bytes: 0,
57        }
58    }
59
60    pub(super) fn input<T: serde::Serialize>(&mut self, value: &T) -> Result<(), TransformError> {
61        self.input_bytes = self
62            .input_bytes
63            .checked_add(measure(
64                value,
65                self.limits.max_bytes.saturating_sub(self.input_bytes),
66            )?)
67            .ok_or_else(limit)?;
68        Ok(())
69    }
70
71    pub(super) fn output<T: serde::Serialize>(&mut self, value: &T) -> Result<(), TransformError> {
72        self.output_bytes = self
73            .output_bytes
74            .checked_add(measure(
75                value,
76                self.limits.max_bytes.saturating_sub(self.output_bytes),
77            )?)
78            .ok_or_else(limit)?;
79        Ok(())
80    }
81}
82
83pub(super) fn measure<T: serde::Serialize>(
84    value: &T,
85    remaining: usize,
86) -> Result<usize, TransformError> {
87    crate::codec::encode_json(
88        value,
89        crate::codec::CodecLimits {
90            max_value_bytes: remaining as u64,
91            max_body_bytes: remaining as u64,
92            max_buffer_bytes: remaining as u64,
93            max_line_bytes: remaining as u64,
94            max_part_bytes: remaining as u64,
95            max_parts: usize::MAX,
96        },
97    )
98    .map(|v| v.len())
99    .map_err(|e| {
100        if e.kind() == crate::codec::CodecErrorKind::Limit {
101            limit()
102        } else {
103            invalid("invalid declared stream value")
104        }
105    })
106}
107
108pub(super) fn declared<T: DeclaredFields>(value: T) -> T {
109    value.into_declared()
110}
111
112pub(super) fn stream_logs(
113    logs: &[crate::wire::openai::responses::input::OutputLogprob],
114) -> Vec<crate::wire::openai::responses::stream::StreamLogprob> {
115    use crate::wire::openai::responses::stream as s;
116    logs.iter()
117        .map(|v| s::StreamLogprob {
118            token: v.token.clone(),
119            logprob: v.logprob.clone(),
120            top_logprobs: Some(
121                v.top_logprobs
122                    .iter()
123                    .map(|t| s::StreamTopLogprob {
124                        token: Some(t.token.clone()),
125                        logprob: Some(t.logprob.clone()),
126                        rest: Default::default(),
127                    })
128                    .collect(),
129            ),
130            rest: Default::default(),
131        })
132        .collect()
133}
gproxy_protocol/transform/generate/chat_responses/stream/common.rs

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