1use serde::{de::DeserializeOwned, Deserialize, Serialize};
2use serde_json::Value;
3
4use crate::{LlmError, Result};
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
8#[serde(rename_all = "lowercase")]
9pub enum Role {
10 System,
11 User,
12 Assistant,
13 Tool,
16}
17
18#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
20#[serde(tag = "kind", rename_all = "snake_case")]
21pub enum ImageSource {
22 Url { url: String },
24 Base64 { media_type: String, data: String },
26}
27
28#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
30#[serde(tag = "type", rename_all = "snake_case")]
31#[non_exhaustive]
32pub enum Part {
33 Text {
34 text: String,
35 },
36 Image {
38 source: ImageSource,
39 },
40 Reasoning {
44 text: String,
45 #[serde(default)]
46 raw: Value,
47 },
48 ToolCall {
49 id: String,
50 name: String,
51 arguments: Value,
54 #[serde(default, skip_serializing_if = "Option::is_none")]
56 signature: Option<String>,
57 },
58 ToolResult {
59 call_id: String,
60 content: String,
61 is_error: bool,
62 },
63 Other {
66 raw: Value,
67 },
68}
69
70impl Part {
71 pub fn text(text: impl Into<String>) -> Self {
72 Part::Text { text: text.into() }
73 }
74
75 pub fn image_url(url: impl Into<String>) -> Self {
77 Part::Image {
78 source: ImageSource::Url { url: url.into() },
79 }
80 }
81
82 pub fn image_base64(media_type: impl Into<String>, data: impl Into<String>) -> Self {
84 Part::Image {
85 source: ImageSource::Base64 {
86 media_type: media_type.into(),
87 data: data.into(),
88 },
89 }
90 }
91
92 pub fn tool_result(
94 call_id: impl Into<String>,
95 content: impl Into<String>,
96 is_error: bool,
97 ) -> Self {
98 Part::ToolResult {
99 call_id: call_id.into(),
100 content: content.into(),
101 is_error,
102 }
103 }
104}
105
106#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
107pub struct Message {
108 pub role: Role,
109 pub parts: Vec<Part>,
110}
111
112impl Message {
113 pub fn new(role: Role, parts: Vec<Part>) -> Self {
114 Self { role, parts }
115 }
116
117 pub fn user(text: impl Into<String>) -> Self {
118 Self::text(Role::User, text)
119 }
120
121 pub fn user_with_image(text: impl Into<String>, image: Part) -> Self {
123 Self::new(Role::User, vec![Part::text(text), image])
124 }
125
126 pub fn assistant(text: impl Into<String>) -> Self {
127 Self::text(Role::Assistant, text)
128 }
129
130 pub fn text(role: Role, text: impl Into<String>) -> Self {
131 Self::new(role, vec![Part::text(text)])
132 }
133
134 pub fn tool_result(
137 call_id: impl Into<String>,
138 content: impl Into<String>,
139 is_error: bool,
140 ) -> Self {
141 Self::new(
142 Role::Tool,
143 vec![Part::tool_result(call_id, content, is_error)],
144 )
145 }
146
147 pub fn text_content(&self) -> String {
149 self.parts
150 .iter()
151 .filter_map(|p| match p {
152 Part::Text { text } => Some(text.as_str()),
153 _ => None,
154 })
155 .collect::<Vec<_>>()
156 .join("")
157 }
158
159 pub(crate) fn has_images(&self) -> bool {
160 self.parts.iter().any(|p| matches!(p, Part::Image { .. }))
161 }
162}
163
164#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
166pub struct Tool {
167 pub name: String,
168 pub description: String,
169 pub parameters: Value,
170}
171
172impl Tool {
173 pub fn new(name: impl Into<String>, description: impl Into<String>, parameters: Value) -> Self {
174 Self {
175 name: name.into(),
176 description: description.into(),
177 parameters,
178 }
179 }
180}
181
182#[derive(Debug, Clone, PartialEq)]
184pub struct ToolCall {
185 pub id: String,
186 pub name: String,
187 pub arguments: Value,
188}
189
190#[derive(Debug, Clone, PartialEq)]
192pub struct ResponseFormat {
193 pub name: String,
195 pub schema: Value,
197}
198
199#[derive(Debug, Clone, Copy, PartialEq, Eq)]
201pub enum ReasoningEffort {
202 Low,
203 Medium,
204 High,
205}
206
207impl ReasoningEffort {
208 pub(crate) fn as_str(self) -> &'static str {
209 match self {
210 ReasoningEffort::Low => "low",
211 ReasoningEffort::Medium => "medium",
212 ReasoningEffort::High => "high",
213 }
214 }
215}
216
217#[derive(Debug, Clone, Default)]
219pub struct ChatRequest {
220 pub model: String,
221 pub system: Option<String>,
224 pub messages: Vec<Message>,
225 pub tools: Vec<Tool>,
226 pub max_tokens: Option<u32>,
228 pub temperature: Option<f32>,
231 pub response_format: Option<ResponseFormat>,
233 pub reasoning: Option<ReasoningEffort>,
235}
236
237pub const DEFAULT_MAX_TOKENS: u32 = 4096;
239
240impl ChatRequest {
241 pub fn new(model: impl Into<String>) -> Self {
242 Self {
243 model: model.into(),
244 ..Default::default()
245 }
246 }
247
248 pub fn system(mut self, system: impl Into<String>) -> Self {
249 self.system = Some(system.into());
250 self
251 }
252
253 pub fn message(mut self, message: Message) -> Self {
254 self.messages.push(message);
255 self
256 }
257
258 pub fn tool(mut self, tool: Tool) -> Self {
259 self.tools.push(tool);
260 self
261 }
262
263 pub fn max_tokens(mut self, max_tokens: u32) -> Self {
264 self.max_tokens = Some(max_tokens);
265 self
266 }
267
268 pub fn temperature(mut self, temperature: f32) -> Self {
269 self.temperature = Some(temperature);
270 self
271 }
272
273 pub fn json_schema(mut self, name: impl Into<String>, schema: Value) -> Self {
275 self.response_format = Some(ResponseFormat {
276 name: name.into(),
277 schema,
278 });
279 self
280 }
281
282 pub fn reasoning(mut self, effort: ReasoningEffort) -> Self {
283 self.reasoning = Some(effort);
284 self
285 }
286}
287
288#[derive(Debug, Clone, PartialEq, Eq)]
290pub enum StopReason {
291 EndTurn,
293 ToolUse,
295 MaxTokens,
297 Paused,
299 Refusal,
301 Other(String),
303}
304
305#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
306pub struct Usage {
307 pub input_tokens: u32,
308 pub output_tokens: u32,
309}
310
311impl Usage {
312 pub fn cost(&self, input_per_million: f64, output_per_million: f64) -> f64 {
315 (f64::from(self.input_tokens) * input_per_million
316 + f64::from(self.output_tokens) * output_per_million)
317 / 1_000_000.0
318 }
319}
320
321#[derive(Debug, Clone, PartialEq)]
322pub struct ChatResponse {
323 pub model: String,
324 pub message: Message,
325 pub stop_reason: StopReason,
326 pub usage: Usage,
327}
328
329impl ChatResponse {
330 pub fn text(&self) -> String {
332 self.message.text_content()
333 }
334
335 pub fn reasoning(&self) -> String {
337 self.message
338 .parts
339 .iter()
340 .filter_map(|p| match p {
341 Part::Reasoning { text, .. } => Some(text.as_str()),
342 _ => None,
343 })
344 .collect::<Vec<_>>()
345 .join("")
346 }
347
348 pub fn tool_calls(&self) -> Vec<ToolCall> {
350 self.message
351 .parts
352 .iter()
353 .filter_map(|p| match p {
354 Part::ToolCall {
355 id,
356 name,
357 arguments,
358 ..
359 } => Some(ToolCall {
360 id: id.clone(),
361 name: name.clone(),
362 arguments: arguments.clone(),
363 }),
364 _ => None,
365 })
366 .collect()
367 }
368
369 pub fn json<T: DeserializeOwned>(&self) -> Result<T> {
371 let text = self.text();
372 serde_json::from_str(&text).map_err(|source| LlmError::Deserialize {
373 source,
374 content: crate::error::truncate(&text),
375 })
376 }
377}
378
379#[derive(Debug, Clone, PartialEq, Eq)]
381pub struct ModelInfo {
382 pub id: String,
383 pub display_name: Option<String>,
384}
385
386#[derive(Debug, Clone, PartialEq)]
388pub struct Embeddings {
389 pub vectors: Vec<Vec<f32>>,
390 pub usage: Usage,
391}