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

1use super::*;
2
3#[derive(Debug, Clone, Copy)]
4pub struct StreamLimits {
5    pub max_bytes: usize,
6}
7
8impl Default for StreamLimits {
9    fn default() -> Self {
10        Self {
11            max_bytes: 16 * 1024 * 1024,
12        }
13    }
14}
15
16/// Final chunks and the invocation's identity associations. A successful Chat
17/// target finish is followed by the SSE `[DONE]` framing token.
18pub struct StreamEnd<T> {
19    pub chunks: Vec<T>,
20    pub identities: IdentityFlow,
21    pub report: Report,
22}
23
24pub(super) struct Budget {
25    limits: StreamLimits,
26
27    bytes: usize,
28
29    output_bytes: usize,
30}
31
32impl Budget {
33    pub fn new(limits: StreamLimits) -> Self {
34        Self {
35            limits,
36
37            bytes: 0,
38
39            output_bytes: 0,
40        }
41    }
42    pub fn input<T: serde::Serialize>(&mut self, value: &T) -> Result<(), TransformError> {
43        self.bytes += bounded(value, self.limits.max_bytes.saturating_sub(self.bytes))?;
44        Ok(())
45    }
46    pub fn output<T: serde::Serialize>(&mut self, values: &[T]) -> Result<(), TransformError> {
47        for value in values {
48            self.output_bytes += bounded(
49                value,
50                self.limits.max_bytes.saturating_sub(self.output_bytes),
51            )?;
52        }
53        Ok(())
54    }
55}
56
57pub(super) fn bounded<T: serde::Serialize>(value: &T, cap: usize) -> Result<usize, TransformError> {
58    let cap = cap as u64;
59    crate::codec::encode_json(
60        value,
61        crate::codec::CodecLimits {
62            max_body_bytes: cap,
63            max_value_bytes: cap,
64            max_buffer_bytes: cap,
65            max_line_bytes: cap,
66            max_part_bytes: cap,
67            max_parts: usize::MAX,
68        },
69    )
70    .map(|v| v.len())
71    .map_err(|e| {
72        if e.kind() == crate::codec::CodecErrorKind::Limit {
73            limit()
74        } else {
75            invalid("invalid declared stream value")
76        }
77    })
78}
79
80pub(super) fn chat_finish(
81    v: g::FinishReason,
82    has_tools: bool,
83) -> Result<c::FinishReason, TransformError> {
84    Ok(match v {
85        g::FinishReason::Stop => {
86            if has_tools {
87                c::FinishReason::ToolCalls
88            } else {
89                c::FinishReason::Stop
90            }
91        }
92        g::FinishReason::Continuation => {
93            return Err(crate::transform::generate::gemini_controls::continuation_error());
94        }
95        g::FinishReason::MaxTokens => c::FinishReason::Length,
96        g::FinishReason::Safety
97        | g::FinishReason::Recitation
98        | g::FinishReason::Language
99        | g::FinishReason::Blocklist
100        | g::FinishReason::ProhibitedContent
101        | g::FinishReason::Spii
102        | g::FinishReason::ImageSafety
103        | g::FinishReason::ImageProhibitedContent
104        | g::FinishReason::ImageRecitation
105        | g::FinishReason::Escalation => c::FinishReason::ContentFilter,
106        g::FinishReason::Unspecified
107        | g::FinishReason::Other
108        | g::FinishReason::MalformedFunctionCall
109        | g::FinishReason::ImageOther
110        | g::FinishReason::NoImage
111        | g::FinishReason::UnexpectedToolCall
112        | g::FinishReason::TooManyToolCalls
113        | g::FinishReason::MissingThoughtSignature
114        | g::FinishReason::MalformedResponse
115        | g::FinishReason::PupLimitedDisabled => {
116            return Err(invalid(
117                "Gemini ended with an unrepresentable failure reason",
118            ));
119        }
120    })
121}
122
123pub(super) fn gemini_finish(v: c::FinishReason) -> g::FinishReason {
124    match v {
125        c::FinishReason::Stop | c::FinishReason::ToolCalls | c::FinishReason::FunctionCall => {
126            g::FinishReason::Stop
127        }
128        c::FinishReason::Length => g::FinishReason::MaxTokens,
129        c::FinishReason::ContentFilter => g::FinishReason::Safety,
130    }
131}
132
133pub(super) fn id(
134    flow: &mut IdentityFlow,
135    policy: &TargetIdPolicy,
136    role: IdentityRole,
137    dialect: crate::Dialect,
138    value: Option<String>,
139    index: u64,
140) -> Result<String, TransformError> {
141    flow.resolve_or_allocate(role, SourceIdentity::new(dialect, value, index), policy)
142        .map(|v| v.emitted_id)
143        .map_err(|e| {
144            TransformError::with_source(
145                TransformErrorKind::Conflict,
146                "stream.identity",
147                "identity association failed",
148                e,
149            )
150        })
151}
gproxy_protocol/transform/generate/gemini_chat/stream/common.rs

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