gproxy_protocol/transform/claude/stream_generate_content/
utils.rs

1use crate::claude::count_tokens::types::{BetaThinkingBlockType, BetaToolUseBlockType};
2use crate::claude::create_message::stream::{
3    BetaMessageDeltaUsage, BetaRawContentBlockDelta, BetaRawMessageDelta, ClaudeStreamEvent,
4};
5use crate::claude::create_message::types::{
6    BetaContentBlock, BetaMessage, BetaMessageRole, BetaMessageType, BetaServiceTier,
7    BetaStopReason, BetaTextBlock, BetaTextBlockType, BetaThinkingBlock, BetaToolUseBlock, Model,
8};
9use crate::claude::types::{BetaApiError, BetaApiErrorType, BetaError};
10use crate::transform::claude::generate_content::utils::beta_usage_from_counts;
11
12pub fn message_start_event(
13    id: String,
14    model: String,
15    service_tier: BetaServiceTier,
16    input_tokens: u64,
17    cached_input_tokens: u64,
18) -> ClaudeStreamEvent {
19    ClaudeStreamEvent::MessageStart {
20        message: BetaMessage {
21            id,
22            container: None,
23            content: Vec::new(),
24            context_management: None,
25            model: Model::Custom(model),
26            role: BetaMessageRole::Assistant,
27            stop_reason: None,
28            stop_sequence: None,
29            type_: BetaMessageType::Message,
30            usage: beta_usage_from_counts(input_tokens, cached_input_tokens, 0, service_tier),
31        },
32    }
33}
34
35pub fn start_text_block_event(index: u64) -> ClaudeStreamEvent {
36    ClaudeStreamEvent::ContentBlockStart {
37        content_block: BetaContentBlock::Text(BetaTextBlock {
38            citations: None,
39            text: String::new(),
40            type_: BetaTextBlockType::Text,
41        }),
42        index,
43    }
44}
45
46pub fn start_thinking_block_event(index: u64, signature: String) -> ClaudeStreamEvent {
47    ClaudeStreamEvent::ContentBlockStart {
48        content_block: BetaContentBlock::Thinking(BetaThinkingBlock {
49            signature,
50            thinking: String::new(),
51            type_: BetaThinkingBlockType::Thinking,
52        }),
53        index,
54    }
55}
56
57pub fn start_tool_use_block_event(index: u64, id: String, name: String) -> ClaudeStreamEvent {
58    ClaudeStreamEvent::ContentBlockStart {
59        content_block: BetaContentBlock::ToolUse(BetaToolUseBlock {
60            id,
61            input: Default::default(),
62            name,
63            type_: BetaToolUseBlockType::ToolUse,
64            cache_control: None,
65            caller: None,
66        }),
67        index,
68    }
69}
70
71pub fn text_delta_event(index: u64, text: String) -> ClaudeStreamEvent {
72    ClaudeStreamEvent::ContentBlockDelta {
73        delta: BetaRawContentBlockDelta::Text { text },
74        index,
75    }
76}
77
78pub fn thinking_delta_event(index: u64, thinking: String) -> ClaudeStreamEvent {
79    ClaudeStreamEvent::ContentBlockDelta {
80        delta: BetaRawContentBlockDelta::Thinking { thinking },
81        index,
82    }
83}
84
85pub fn input_json_delta_event(index: u64, partial_json: String) -> ClaudeStreamEvent {
86    ClaudeStreamEvent::ContentBlockDelta {
87        delta: BetaRawContentBlockDelta::InputJson { partial_json },
88        index,
89    }
90}
91
92pub fn stop_block_event(index: u64) -> ClaudeStreamEvent {
93    ClaudeStreamEvent::ContentBlockStop { index }
94}
95
96pub fn message_delta_event(
97    stop_reason: Option<BetaStopReason>,
98    input_tokens: u64,
99    cached_input_tokens: u64,
100    output_tokens: u64,
101) -> ClaudeStreamEvent {
102    ClaudeStreamEvent::MessageDelta {
103        context_management: None,
104        delta: BetaRawMessageDelta {
105            container: None,
106            stop_reason,
107            stop_sequence: None,
108        },
109        usage: BetaMessageDeltaUsage {
110            cache_creation_input_tokens: Some(0),
111            cache_read_input_tokens: Some(cached_input_tokens),
112            input_tokens: Some(input_tokens),
113            iterations: None,
114            output_tokens,
115            server_tool_use: None,
116        },
117    }
118}
119
120pub fn message_stop_event() -> ClaudeStreamEvent {
121    ClaudeStreamEvent::MessageStop {}
122}
123
124pub fn stream_error_event(message: String) -> ClaudeStreamEvent {
125    ClaudeStreamEvent::Error {
126        error: BetaError::Api(BetaApiError {
127            message,
128            type_: BetaApiErrorType::ApiError,
129        }),
130    }
131}
132
133pub fn push_text_block(
134    out: &mut Vec<ClaudeStreamEvent>,
135    next_block_index: &mut u64,
136    text: String,
137) -> bool {
138    if text.is_empty() {
139        return false;
140    }
141    let block_index = *next_block_index;
142    *next_block_index = next_block_index.saturating_add(1);
143    out.push(start_text_block_event(block_index));
144    out.push(text_delta_event(block_index, text));
145    out.push(stop_block_event(block_index));
146    true
147}
148
149pub fn push_thinking_block(
150    out: &mut Vec<ClaudeStreamEvent>,
151    next_block_index: &mut u64,
152    signature: String,
153    thinking: String,
154) -> bool {
155    if thinking.is_empty() {
156        return false;
157    }
158    let block_index = *next_block_index;
159    *next_block_index = next_block_index.saturating_add(1);
160    out.push(start_thinking_block_event(block_index, signature));
161    out.push(thinking_delta_event(block_index, thinking));
162    out.push(stop_block_event(block_index));
163    true
164}
165
166pub fn push_tool_use_block(
167    out: &mut Vec<ClaudeStreamEvent>,
168    next_block_index: &mut u64,
169    id: String,
170    name: String,
171    input_json: Option<String>,
172) -> u64 {
173    let block_index = *next_block_index;
174    *next_block_index = next_block_index.saturating_add(1);
175    out.push(start_tool_use_block_event(block_index, id, name));
176    if let Some(input_json) = input_json
177        && !input_json.is_empty()
178    {
179        out.push(input_json_delta_event(block_index, input_json));
180    }
181    out.push(stop_block_event(block_index));
182    block_index
183}
gproxy_protocol/transform/claude/stream_generate_content/utils.rs

gproxy_protocol/transform/claude/stream_generate_content/
utils.rs
