Skip to main content

gproxy_protocol/transform/generate/claude_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
16pub struct StreamEnd<T> {
17    pub chunks: Vec<T>,
18    pub identities: IdentityFlow,
19    pub report: Report,
20}
21
22pub(super) struct Budget {
23    limits: StreamLimits,
24
25    input_bytes: usize,
26
27    output_bytes: usize,
28}
29
30impl Budget {
31    pub fn new(limits: StreamLimits) -> Self {
32        Self {
33            limits,
34
35            input_bytes: 0,
36
37            output_bytes: 0,
38        }
39    }
40    pub fn input<T: serde::Serialize>(&mut self, v: &T) -> Result<(), TransformError> {
41        self.input_bytes += bounded(v, self.limits.max_bytes.saturating_sub(self.input_bytes))?;
42        Ok(())
43    }
44    pub fn emit<T: serde::Serialize>(
45        &mut self,
46        out: &mut Vec<T>,
47        v: T,
48    ) -> Result<(), TransformError> {
49        self.output_bytes += bounded(&v, self.limits.max_bytes.saturating_sub(self.output_bytes))?;
50        out.push(v);
51        Ok(())
52    }
53}
54
55pub(super) fn bounded<T: serde::Serialize>(value: &T, cap: usize) -> Result<usize, TransformError> {
56    let n = cap as u64;
57    crate::codec::encode_json(
58        value,
59        crate::codec::CodecLimits {
60            max_body_bytes: n,
61            max_value_bytes: n,
62            max_buffer_bytes: n,
63            max_line_bytes: n,
64            max_part_bytes: n,
65            max_parts: usize::MAX,
66        },
67    )
68    .map(|v| v.len())
69    .map_err(|e| {
70        if e.kind() == crate::codec::CodecErrorKind::Limit {
71            limit()
72        } else {
73            invalid("invalid declared event")
74        }
75    })
76}
77
78pub(super) fn id(
79    flow: &mut IdentityFlow,
80    policy: &TargetIdPolicy,
81    role: IdentityRole,
82    source: crate::Dialect,
83    value: Option<String>,
84    index: u64,
85) -> Result<String, TransformError> {
86    flow.resolve_or_allocate(role, SourceIdentity::new(source, value, index), policy)
87        .map(|v| v.emitted_id)
88        .map_err(|e| {
89            TransformError::with_source(
90                TransformErrorKind::Conflict,
91                "stream.identity",
92                "identity association failed",
93                e,
94            )
95        })
96}
97
98/// A legacy Chat `function_call` has no ID; its alias records that it is one.
99pub(super) fn legacy_id(
100    flow: &mut IdentityFlow,
101    policy: &TargetIdPolicy,
102    index: u64,
103) -> Result<String, TransformError> {
104    flow.resolve_legacy_chat_call(
105        SourceIdentity::new(crate::Dialect::OpenAiChat, None, index),
106        policy,
107    )
108    .map(|v| v.emitted_id)
109    .map_err(|e| {
110        TransformError::with_source(
111            TransformErrorKind::Conflict,
112            "stream.identity",
113            "identity association failed",
114            e,
115        )
116    })
117}
118
119pub(super) fn claude_policy() -> TargetIdPolicy {
120    TargetIdPolicy::new(crate::Dialect::Claude)
121        .with_generated_prefix(IdentityRole::Response, KnownIdPrefix::Message)
122        .with_generated_prefix(IdentityRole::ToolCall, KnownIdPrefix::Tool)
123}
124
125pub(super) fn chat_finish(reason: c::StopReason) -> Result<o::FinishReason, TransformError> {
126    Ok(match reason {
127        c::StopReason::EndTurn | c::StopReason::StopSequence => o::FinishReason::Stop,
128        c::StopReason::MaxTokens | c::StopReason::ModelContextWindowExceeded => {
129            o::FinishReason::Length
130        }
131        c::StopReason::ToolUse => o::FinishReason::ToolCalls,
132        c::StopReason::Refusal => o::FinishReason::ContentFilter,
133        c::StopReason::PauseTurn | c::StopReason::Compaction => o::FinishReason::Stop,
134    })
135}
136
137pub(super) fn claude_finish(reason: o::FinishReason, refusal: bool) -> c::StopReason {
138    match reason {
139        o::FinishReason::Stop if refusal => c::StopReason::Refusal,
140        o::FinishReason::Stop => c::StopReason::EndTurn,
141        o::FinishReason::Length => c::StopReason::MaxTokens,
142        o::FinishReason::ToolCalls | o::FinishReason::FunctionCall => c::StopReason::ToolUse,
143        o::FinishReason::ContentFilter => c::StopReason::Refusal,
144    }
145}