Skip to main content

llm_universal/
types.rs

1use serde::{de::DeserializeOwned, Deserialize, Serialize};
2use serde_json::Value;
3
4use crate::{LlmError, Result};
5
6/// Author of a message.
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
8#[serde(rename_all = "lowercase")]
9pub enum Role {
10    System,
11    User,
12    Assistant,
13    /// Tool results. Sent as role `tool` to OpenAI and Cohere, as `tool_result` blocks in a
14    /// `user` turn to Anthropic, and as `functionResponse` parts to Gemini.
15    Tool,
16}
17
18/// Where an image comes from.
19#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
20#[serde(tag = "kind", rename_all = "snake_case")]
21pub enum ImageSource {
22    /// A public `http(s)://` URL. Not supported by Gemini; use [`ImageSource::Base64`] there.
23    Url { url: String },
24    /// Inline image data. `media_type` is `image/png`, `image/jpeg`, `image/gif` or `image/webp`.
25    Base64 { media_type: String, data: String },
26}
27
28/// One piece of a message.
29#[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    /// An image in a user message.
37    Image {
38        source: ImageSource,
39    },
40    /// Model reasoning ("thinking"). `raw` is the provider's original block; it is sent back
41    /// unchanged on the next turn (Anthropic and Gemini require that) and skipped where the
42    /// provider does not accept it back. `text` may be empty when the provider hides reasoning.
43    Reasoning {
44        text: String,
45        #[serde(default)]
46        raw: Value,
47    },
48    ToolCall {
49        id: String,
50        name: String,
51        /// Parsed arguments. If a provider returns arguments that are not valid JSON
52        /// (for example cut off by `max_tokens`), the raw string is kept as `Value::String`.
53        arguments: Value,
54        /// Gemini's `thoughtSignature` for this call; it must be sent back with the call.
55        #[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    /// A provider block this crate does not model (e.g. Anthropic server-tool blocks or Gemini
64    /// code-execution parts). Kept as received and sent back unchanged to the same provider.
65    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    /// An image from a public URL.
76    pub fn image_url(url: impl Into<String>) -> Self {
77        Part::Image {
78            source: ImageSource::Url { url: url.into() },
79        }
80    }
81
82    /// An inline image. `data` is base64 without the `data:` prefix.
83    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    /// A tool result to send back for `call_id`.
93    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    /// A user message with text and one image (see [`Part::image_url`] and [`Part::image_base64`]).
122    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    /// A single tool result. To return several results from one step, put them in one
135    /// message with `Message::new(Role::Tool, parts)`.
136    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    /// All text parts joined together (reasoning is not included).
148    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/// A function the model may call. `parameters` is a JSON Schema with `"type": "object"`.
165#[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/// A tool call taken out of a response.
183#[derive(Debug, Clone, PartialEq)]
184pub struct ToolCall {
185    pub id: String,
186    pub name: String,
187    pub arguments: Value,
188}
189
190/// Asks the model to answer with JSON matching `schema`.
191#[derive(Debug, Clone, PartialEq)]
192pub struct ResponseFormat {
193    /// Schema name (required by OpenAI; letters, digits, `_` and `-`).
194    pub name: String,
195    /// JSON Schema of the answer.
196    pub schema: Value,
197}
198
199/// How hard the model should think, on models that support reasoning.
200#[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/// Provider-neutral chat request.
218#[derive(Debug, Clone, Default)]
219pub struct ChatRequest {
220    pub model: String,
221    /// System prompt. Sent as a `system` message to OpenAI/Cohere, as the top-level
222    /// `system` field to Anthropic and as `systemInstruction` to Gemini.
223    pub system: Option<String>,
224    pub messages: Vec<Message>,
225    pub tools: Vec<Tool>,
226    /// Required by Anthropic; [`DEFAULT_MAX_TOKENS`] is used when unset.
227    pub max_tokens: Option<u32>,
228    /// Anthropic accepts 0.0–1.0, and newer Anthropic models accept only 1.0.
229    /// Leave unset unless the model supports it.
230    pub temperature: Option<f32>,
231    /// Structured output: the answer is JSON matching this schema. Read it with [`ChatResponse::json`].
232    pub response_format: Option<ResponseFormat>,
233    /// Reasoning effort / thinking. Reasoning comes back as [`Part::Reasoning`].
234    pub reasoning: Option<ReasoningEffort>,
235}
236
237/// `max_tokens` sent to Anthropic when the request does not set one.
238pub 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    /// Asks for JSON matching `schema`.
274    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/// Why generation stopped, normalized across providers.
289#[derive(Debug, Clone, PartialEq, Eq)]
290pub enum StopReason {
291    /// Natural end of the turn (`stop` / `end_turn` / `stop_sequence` / `STOP` / `COMPLETE`).
292    EndTurn,
293    /// The model wants tools to be called.
294    ToolUse,
295    /// Token limit reached.
296    MaxTokens,
297    /// A long-running turn was paused (`pause_turn`). Send the response back as-is to continue.
298    Paused,
299    /// The model declined or the output was blocked for policy reasons.
300    Refusal,
301    /// Any other value returned by the provider.
302    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    /// Cost in your currency, given prices per million input and output tokens.
313    /// Prices change and differ per model, so they are passed in rather than built in.
314    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    /// All text of the response.
331    pub fn text(&self) -> String {
332        self.message.text_content()
333    }
334
335    /// All reasoning text of the response.
336    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    /// Tool calls the model asked for.
349    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    /// Parses the text as JSON into `T` (use with [`ChatRequest::json_schema`]).
370    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/// A model available to the key.
380#[derive(Debug, Clone, PartialEq, Eq)]
381pub struct ModelInfo {
382    pub id: String,
383    pub display_name: Option<String>,
384}
385
386/// Result of an embeddings call: one vector per input, in input order.
387#[derive(Debug, Clone, PartialEq)]
388pub struct Embeddings {
389    pub vectors: Vec<Vec<f32>>,
390    pub usage: Usage,
391}