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::MaxTokens => c::FinishReason::Length,
93        g::FinishReason::Safety
94        | g::FinishReason::Recitation
95        | g::FinishReason::Language
96        | g::FinishReason::Blocklist
97        | g::FinishReason::ProhibitedContent
98        | g::FinishReason::Spii
99        | g::FinishReason::ImageSafety
100        | g::FinishReason::ImageProhibitedContent
101        | g::FinishReason::ImageRecitation
102        | g::FinishReason::Escalation => c::FinishReason::ContentFilter,
103        g::FinishReason::Unspecified
104        | g::FinishReason::Other
105        | g::FinishReason::MalformedFunctionCall
106        | g::FinishReason::ImageOther
107        | g::FinishReason::NoImage
108        | g::FinishReason::UnexpectedToolCall
109        | g::FinishReason::TooManyToolCalls
110        | g::FinishReason::MissingThoughtSignature
111        | g::FinishReason::MalformedResponse
112        | g::FinishReason::PupLimitedDisabled => {
113            return Err(invalid(
114                "Gemini ended with an unrepresentable failure reason",
115            ));
116        }
117    })
118}
119
120pub(super) fn gemini_finish(v: c::FinishReason) -> g::FinishReason {
121    match v {
122        c::FinishReason::Stop | c::FinishReason::ToolCalls | c::FinishReason::FunctionCall => {
123            g::FinishReason::Stop
124        }
125        c::FinishReason::Length => g::FinishReason::MaxTokens,
126        c::FinishReason::ContentFilter => g::FinishReason::Safety,
127    }
128}
129
130pub(super) fn id(
131    flow: &mut IdentityFlow,
132    policy: &TargetIdPolicy,
133    role: IdentityRole,
134    dialect: crate::Dialect,
135    value: Option<String>,
136    index: u64,
137) -> Result<String, TransformError> {
138    flow.resolve_or_allocate(role, SourceIdentity::new(dialect, value, index), policy)
139        .map(|v| v.emitted_id)
140        .map_err(|e| {
141            TransformError::with_source(
142                TransformErrorKind::Conflict,
143                "stream.identity",
144                "identity association failed",
145                e,
146            )
147        })
148}
gproxy_protocol/transform/generate/gemini_chat/stream/common.rs

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