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

1use crate::{
2    transform::{
3        Report, TransformError, TransformErrorKind,
4        identity::{IdentityFlow, IdentityRole, KnownIdPrefix, SourceIdentity, TargetIdPolicy},
5    },
6    wire::gemini as g,
7};
8
9#[derive(Debug, Clone, Copy)]
10pub struct StreamLimits {
11    pub max_bytes: usize,
12    /// Retained payload/event bytes while waiting for facts or earlier blocks.
13    pub max_pending: usize,
14}
15
16impl Default for StreamLimits {
17    fn default() -> Self {
18        Self {
19            max_bytes: 16 * 1024 * 1024,
20            max_pending: 16 * 1024 * 1024,
21        }
22    }
23}
24
25#[derive(Debug)]
26pub struct StreamEnd<T> {
27    pub chunks: Vec<T>,
28    pub identities: IdentityFlow,
29    pub report: Report,
30}
31
32pub(super) fn invalid(field: &str, message: impl Into<String>) -> TransformError {
33    TransformError::invalid_result(field, message)
34}
35
36pub(super) fn limit() -> TransformError {
37    TransformError::new(
38        TransformErrorKind::Limit,
39        "claude_gemini.stream",
40        "stream limit exceeded",
41    )
42}
43
44pub(super) struct Budget {
45    limits: StreamLimits,
46
47    input_bytes: usize,
48
49    output_bytes: usize,
50}
51
52impl Budget {
53    pub fn new(limits: StreamLimits) -> Self {
54        Self {
55            limits,
56
57            input_bytes: 0,
58
59            output_bytes: 0,
60        }
61    }
62    pub fn input<T: serde::Serialize>(&mut self, value: &T) -> Result<(), TransformError> {
63        self.input_bytes += bound(
64            value,
65            self.limits.max_bytes.saturating_sub(self.input_bytes),
66        )?;
67        Ok(())
68    }
69    /// Charge exactly once, before cloning or retaining an output event.
70    pub fn output<T: serde::Serialize>(&mut self, value: &T) -> Result<usize, TransformError> {
71        let size = bound(
72            value,
73            self.limits.max_bytes.saturating_sub(self.output_bytes),
74        )?;
75        self.output_bytes += size;
76        Ok(size)
77    }
78}
79
80pub(super) fn bound<T: serde::Serialize>(value: &T, cap: usize) -> Result<usize, TransformError> {
81    let n = cap as u64;
82    crate::codec::encode_json(
83        value,
84        crate::codec::CodecLimits {
85            max_value_bytes: n,
86            max_body_bytes: n,
87            max_buffer_bytes: n,
88            max_line_bytes: n,
89            max_part_bytes: n,
90            max_parts: usize::MAX,
91        },
92    )
93    .map(|v| v.len())
94    .map_err(|e| {
95        if e.kind() == crate::codec::CodecErrorKind::Limit {
96            limit()
97        } else {
98            invalid("stream.event", e.to_string())
99        }
100    })
101}
102
103pub(super) fn claude_policy() -> TargetIdPolicy {
104    TargetIdPolicy::new(crate::Dialect::Claude)
105        .with_generated_prefix(IdentityRole::Response, KnownIdPrefix::Message)
106        .with_generated_prefix(IdentityRole::ToolCall, KnownIdPrefix::Tool)
107}
108
109pub(super) fn response_id(
110    flow: &mut IdentityFlow,
111    policy: &TargetIdPolicy,
112    dialect: crate::Dialect,
113    value: Option<String>,
114) -> Result<String, TransformError> {
115    flow.resolve_or_allocate(
116        IdentityRole::Response,
117        SourceIdentity::new(dialect, value, 0),
118        policy,
119    )
120    .map(|v| v.emitted_id)
121    .map_err(|e| invalid("stream.identity", e.to_string()))
122}
123
124/// Native Gemini appends parts. Only boundaries between adjacent plain text
125/// fragments with the exact same thought flag are immaterial to this pair.
126pub(super) fn normalize_gemini(body: &mut g::GenerateContentResponseBody) {
127    for candidate in body.candidates.iter_mut().flatten() {
128        if let Some(parts) = candidate.content.as_mut().and_then(|v| v.parts.as_mut()) {
129            let mut out: Vec<g::Part> = Vec::new();
130            for mut part in std::mem::take(parts) {
131                if plain(&part)
132                    && let Some(previous) = out.last_mut()
133                    && plain(previous)
134                    && previous.thought == part.thought
135                {
136                    previous
137                        .text
138                        .as_mut()
139                        .unwrap()
140                        .push_str(part.text.take().unwrap().as_str());
141                } else {
142                    out.push(part);
143                }
144            }
145            *parts = out;
146        }
147    }
148}
149
150fn plain(part: &g::Part) -> bool {
151    part.text.is_some()
152        && part.inline_data.is_none()
153        && part.file_data.is_none()
154        && part.function_call.is_none()
155        && part.function_response.is_none()
156        && part.executable_code.is_none()
157        && part.code_execution_result.is_none()
158        && part.tool_call.is_none()
159        && part.tool_response.is_none()
160        && part.video_metadata.is_none()
161        && part.media_processing.is_none()
162        && part.audio_transcription.is_none()
163        && part.speech_metadata.is_none()
164        && part.media_resolution.is_none()
165        && part.thought_signature.is_none()
166        && part.part_metadata.is_none()
167}
gproxy_protocol/transform/generate/claude_gemini/stream/common.rs

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