1use crate::TokenUsage;
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
5#[serde(rename_all = "lowercase")]
6pub enum Role {
7 User,
8 Assistant,
9}
10
11#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
12#[serde(rename_all = "lowercase")]
13pub enum CacheType {
14 Ephemeral,
15}
16
17#[derive(Debug, Clone, Serialize, Deserialize)]
18#[serde(rename_all = "snake_case")]
19pub struct Request {
20 pub model: String,
21 pub messages: Vec<Message>,
22 pub max_tokens: u32,
23 #[serde(skip_serializing_if = "Option::is_none")]
24 pub temperature: Option<f32>,
25 #[serde(skip_serializing_if = "Option::is_none")]
26 pub top_p: Option<f32>,
27 #[serde(skip_serializing_if = "Option::is_none")]
28 pub system: Option<SystemContent>,
29 #[serde(skip_serializing_if = "Option::is_none")]
30 pub tools: Option<Vec<Tool>>,
31 pub stream: bool,
32 #[serde(skip_serializing_if = "Option::is_none")]
33 pub thinking: Option<Thinking>,
34 #[serde(skip_serializing_if = "Option::is_none")]
35 pub cache_control: Option<CacheControl>,
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize)]
39pub struct Thinking {
40 #[serde(rename = "type")]
41 pub thinking_type: String,
42 pub budget_tokens: u32,
43}
44
45impl Thinking {
46 pub fn new(budget_tokens: u32) -> Self {
47 Self { thinking_type: "enabled".to_string(), budget_tokens }
48 }
49}
50
51impl Request {
52 pub fn new(model: String, messages: Vec<Message>) -> Self {
53 Self {
54 model,
55 messages,
56 max_tokens: 4096,
57 temperature: None,
58 top_p: None,
59 system: None,
60 tools: None,
61 stream: false,
62 thinking: None,
63 cache_control: None,
64 }
65 }
66
67 pub fn with_auto_caching(mut self) -> Self {
68 self.cache_control = Some(CacheControl::ephemeral());
69 self
70 }
71
72 pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
73 self.max_tokens = max_tokens;
74 self
75 }
76
77 pub fn with_temperature(mut self, temperature: f32) -> Self {
78 self.temperature = Some(temperature);
79 self
80 }
81
82 pub fn with_top_p(mut self, top_p: f32) -> Self {
83 self.top_p = Some(top_p);
84 self
85 }
86
87 pub fn with_system_cached(mut self, system: String) -> Self {
88 self.system = Some(SystemContent::Blocks(vec![SystemContentBlock::Text {
89 text: system,
90 cache_control: Some(CacheControl::ephemeral()),
91 }]));
92 self
93 }
94
95 pub fn with_tools(mut self, tools: Vec<Tool>) -> Self {
96 self.tools = Some(tools);
97 self
98 }
99
100 pub fn with_stream(mut self, stream: bool) -> Self {
101 self.stream = stream;
102 self
103 }
104
105 pub fn with_thinking(mut self, thinking: Thinking) -> Self {
106 self.thinking = Some(thinking);
107 self
108 }
109}
110
111#[derive(Debug, Clone, Serialize, Deserialize)]
112#[serde(rename_all = "snake_case")]
113pub struct Message {
114 pub role: Role,
115 pub content: Content,
116 #[serde(skip_serializing_if = "Option::is_none")]
117 pub cache_control: Option<CacheControl>,
118}
119
120#[derive(Debug, Clone, Serialize, Deserialize)]
121#[serde(untagged)]
122pub enum Content {
123 Text(String),
124 Blocks(Vec<ContentBlock>),
125}
126
127#[derive(Debug, Clone, Serialize, Deserialize)]
128#[serde(untagged)]
129pub enum SystemContent {
130 Text(String),
131 Blocks(Vec<SystemContentBlock>),
132}
133
134#[derive(Debug, Clone, Serialize, Deserialize)]
135#[serde(tag = "type")]
136pub enum SystemContentBlock {
137 #[serde(rename = "text")]
138 Text {
139 text: String,
140 #[serde(skip_serializing_if = "Option::is_none")]
141 cache_control: Option<CacheControl>,
142 },
143}
144
145#[derive(Debug, Clone, Serialize, Deserialize)]
146#[serde(tag = "type")]
147pub enum ContentBlock {
148 #[serde(rename = "text")]
149 Text {
150 text: String,
151 #[serde(skip_serializing_if = "Option::is_none")]
152 cache_control: Option<CacheControl>,
153 },
154 #[serde(rename = "image")]
155 Image { source: ImageSource },
156 #[serde(rename = "tool_use")]
157 ToolUse { id: String, name: String, input: serde_json::Value },
158 #[serde(rename = "tool_result")]
159 ToolResult {
160 tool_use_id: String,
161 content: String,
162 #[serde(skip_serializing_if = "Option::is_none")]
163 is_error: Option<bool>,
164 },
165}
166
167#[derive(Debug, Clone, Serialize, Deserialize)]
168pub struct ImageSource {
169 #[serde(rename = "type")]
170 pub source_type: String,
171 pub media_type: String,
172 pub data: String,
173}
174
175#[derive(Debug, Clone, Serialize, Deserialize)]
176pub struct CacheControl {
177 #[serde(rename = "type")]
178 pub cache_type: CacheType,
179}
180
181impl CacheControl {
182 pub fn ephemeral() -> Self {
183 Self { cache_type: CacheType::Ephemeral }
184 }
185}
186
187#[derive(Debug, Clone, Serialize, Deserialize)]
188#[serde(rename_all = "snake_case")]
189pub struct Tool {
190 pub name: String,
191 pub description: String,
192 pub input_schema: serde_json::Value,
193 #[serde(skip_serializing_if = "Option::is_none")]
194 pub cache_control: Option<CacheControl>,
195}
196
197#[derive(Debug, Clone, Deserialize)]
198#[serde(tag = "type", rename_all = "snake_case")]
199pub enum StreamEvent {
200 MessageStart {
201 #[serde(flatten)]
202 data: MessageStart,
203 },
204 ContentBlockStart {
205 #[serde(flatten)]
206 data: ContentBlockStart,
207 },
208 ContentBlockDelta {
209 #[serde(flatten)]
210 data: ContentBlockDelta,
211 },
212 ContentBlockStop {
213 #[serde(flatten)]
214 data: ContentBlockStop,
215 },
216 MessageDelta {
217 #[serde(flatten)]
218 data: MessageDelta,
219 },
220 MessageStop {
221 #[serde(flatten)]
222 data: MessageStop,
223 },
224 #[serde(rename = "error")]
225 Error {
226 #[serde(flatten)]
227 data: ErrorEvent,
228 },
229 #[serde(rename = "ping")]
230 Ping,
231}
232
233#[derive(Debug, Clone, Deserialize)]
234pub struct MessageStart {
235 pub message: ResponseMessage,
236}
237
238#[derive(Debug, Clone, Deserialize)]
239#[serde(rename_all = "snake_case")]
240pub struct ResponseMessage {
241 pub id: String,
242 #[serde(rename = "type")]
243 pub message_type: String,
244 pub role: Role,
245 pub content: Vec<serde_json::Value>,
246 pub model: String,
247 pub stop_reason: Option<String>,
248 pub stop_sequence: Option<String>,
249 pub usage: Usage,
250}
251
252#[derive(Debug, Clone, Deserialize)]
253#[serde(rename_all = "snake_case")]
254pub struct Usage {
255 pub input_tokens: u32,
256 pub output_tokens: u32,
257 #[serde(skip_serializing_if = "Option::is_none")]
258 pub cache_creation_input_tokens: Option<u32>,
259 #[serde(skip_serializing_if = "Option::is_none")]
260 pub cache_read_input_tokens: Option<u32>,
261}
262
263impl From<&Usage> for TokenUsage {
264 fn from(usage: &Usage) -> Self {
265 TokenUsage {
266 input_tokens: usage.input_tokens,
267 output_tokens: usage.output_tokens,
268 cache_read_tokens: usage.cache_read_input_tokens,
269 cache_creation_tokens: usage.cache_creation_input_tokens,
270 cache_reporting_exclusive: Some(true),
271 ..TokenUsage::default()
272 }
273 }
274}
275
276#[derive(Debug, Clone, Deserialize)]
277#[serde(rename_all = "snake_case")]
278pub struct ContentBlockStart {
279 pub index: u32,
280 pub content_block: ContentBlockStartData,
281}
282
283#[derive(Debug, Clone, Deserialize)]
284#[serde(tag = "type")]
285pub enum ContentBlockStartData {
286 #[serde(rename = "text")]
287 Text { text: String },
288 #[serde(rename = "tool_use")]
289 ToolUse { id: String, name: String },
290 #[serde(rename = "thinking")]
291 Thinking { thinking: String },
292}
293
294#[derive(Debug, Clone, Deserialize)]
295#[serde(rename_all = "snake_case")]
296pub struct ContentBlockDelta {
297 pub index: u32,
298 pub delta: ContentBlockDeltaData,
299}
300
301#[derive(Debug, Clone, Deserialize)]
302#[serde(tag = "type", rename_all = "snake_case")]
303pub enum ContentBlockDeltaData {
304 TextDelta { text: String },
305 InputJsonDelta { partial_json: String },
306 ThinkingDelta { thinking: String },
307}
308
309#[derive(Debug, Clone, Deserialize)]
310pub struct ContentBlockStop {
311 pub index: u32,
312}
313
314#[derive(Debug, Clone, Deserialize)]
315#[serde(rename_all = "snake_case")]
316pub struct MessageDelta {
317 pub delta: MessageDeltaData,
318 #[serde(default)]
319 pub usage: Option<Usage>,
320}
321
322#[derive(Debug, Clone, Deserialize)]
323#[serde(rename_all = "snake_case")]
324pub struct MessageDeltaData {
325 #[serde(skip_serializing_if = "Option::is_none")]
326 pub stop_reason: Option<String>,
327 #[serde(skip_serializing_if = "Option::is_none")]
328 pub stop_sequence: Option<String>,
329}
330
331#[derive(Debug, Clone, Deserialize)]
332pub struct MessageStop {}
333
334#[derive(Debug, Clone, Deserialize)]
335#[serde(rename_all = "snake_case")]
336pub struct ErrorEvent {
337 pub error: Error,
338}
339
340#[derive(Debug, Clone, Deserialize)]
341#[serde(rename_all = "snake_case")]
342pub struct Error {
343 #[serde(rename = "type")]
344 pub error_type: String,
345 pub message: String,
346}