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 ..TokenUsage::default()
271 }
272 }
273}
274
275#[derive(Debug, Clone, Deserialize)]
276#[serde(rename_all = "snake_case")]
277pub struct ContentBlockStart {
278 pub index: u32,
279 pub content_block: ContentBlockStartData,
280}
281
282#[derive(Debug, Clone, Deserialize)]
283#[serde(tag = "type")]
284pub enum ContentBlockStartData {
285 #[serde(rename = "text")]
286 Text { text: String },
287 #[serde(rename = "tool_use")]
288 ToolUse { id: String, name: String },
289 #[serde(rename = "thinking")]
290 Thinking { thinking: String },
291}
292
293#[derive(Debug, Clone, Deserialize)]
294#[serde(rename_all = "snake_case")]
295pub struct ContentBlockDelta {
296 pub index: u32,
297 pub delta: ContentBlockDeltaData,
298}
299
300#[derive(Debug, Clone, Deserialize)]
301#[serde(tag = "type", rename_all = "snake_case")]
302pub enum ContentBlockDeltaData {
303 TextDelta { text: String },
304 InputJsonDelta { partial_json: String },
305 ThinkingDelta { thinking: String },
306}
307
308#[derive(Debug, Clone, Deserialize)]
309pub struct ContentBlockStop {
310 pub index: u32,
311}
312
313#[derive(Debug, Clone, Deserialize)]
314#[serde(rename_all = "snake_case")]
315pub struct MessageDelta {
316 pub delta: MessageDeltaData,
317 #[serde(default)]
318 pub usage: Option<Usage>,
319}
320
321#[derive(Debug, Clone, Deserialize)]
322#[serde(rename_all = "snake_case")]
323pub struct MessageDeltaData {
324 #[serde(skip_serializing_if = "Option::is_none")]
325 pub stop_reason: Option<String>,
326 #[serde(skip_serializing_if = "Option::is_none")]
327 pub stop_sequence: Option<String>,
328}
329
330#[derive(Debug, Clone, Deserialize)]
331pub struct MessageStop {}
332
333#[derive(Debug, Clone, Deserialize)]
334#[serde(rename_all = "snake_case")]
335pub struct ErrorEvent {
336 pub error: Error,
337}
338
339#[derive(Debug, Clone, Deserialize)]
340#[serde(rename_all = "snake_case")]
341pub struct Error {
342 #[serde(rename = "type")]
343 pub error_type: String,
344 pub message: String,
345}