pi_ai/
types.rs

1use serde::{Deserialize, Serialize};
2use serde_json::Value;
3
4#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
5#[serde(rename_all = "lowercase")]
6pub enum ThinkingLevel {
7    #[default]
8    Off,
9    Minimal,
10    Low,
11    Medium,
12    High,
13    Xhigh,
14}
15
16#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
17#[serde(tag = "type", rename_all = "lowercase")]
18pub enum Content {
19    #[serde(rename = "text")]
20    Text { text: String },
21    #[serde(rename = "thinking")]
22    Thinking {
23        thinking: String,
24        #[serde(skip_serializing_if = "Option::is_none", default)]
25        thinking_signature: Option<String>,
26    },
27    #[serde(rename = "image")]
28    Image { data: String, mime_type: String },
29    #[serde(rename = "toolCall")]
30    ToolCall {
31        id: String,
32        name: String,
33        #[serde(default)]
34        arguments: Value,
35    },
36}
37
38impl Content {
39    pub fn text(s: impl Into<String>) -> Self {
40        Content::Text { text: s.into() }
41    }
42    pub fn as_text(&self) -> Option<&str> {
43        if let Content::Text { text } = self {
44            Some(text)
45        } else {
46            None
47        }
48    }
49}
50
51#[derive(Debug, Clone, Default, Serialize, Deserialize)]
52pub struct Usage {
53    #[serde(default)]
54    pub input: u64,
55    #[serde(default)]
56    pub output: u64,
57    #[serde(default)]
58    pub cache_read: u64,
59    #[serde(default)]
60    pub cache_write: u64,
61    #[serde(default)]
62    pub total_tokens: u64,
63}
64
65#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
66#[serde(rename_all = "camelCase")]
67pub enum StopReason {
68    Stop,
69    Length,
70    ToolUse,
71    Error,
72    Aborted,
73}
74
75#[derive(Debug, Clone, Serialize, Deserialize)]
76#[serde(tag = "role", rename_all = "camelCase")]
77pub enum Message {
78    #[serde(rename = "user")]
79    User {
80        content: Vec<Content>,
81        #[serde(default = "now_ms")]
82        timestamp: i64,
83    },
84    #[serde(rename = "assistant")]
85    Assistant(AssistantMessage),
86    #[serde(rename = "toolResult")]
87    ToolResult(ToolResultMessage),
88}
89
90#[derive(Debug, Clone, Serialize, Deserialize)]
91pub struct AssistantMessage {
92    pub content: Vec<Content>,
93    pub api: String,
94    pub provider: String,
95    pub model: String,
96    #[serde(default)]
97    pub usage: Usage,
98    pub stop_reason: StopReason,
99    #[serde(skip_serializing_if = "Option::is_none", default)]
100    pub error_message: Option<String>,
101    #[serde(default = "now_ms")]
102    pub timestamp: i64,
103}
104
105#[derive(Debug, Clone, Serialize, Deserialize)]
106pub struct ToolResultMessage {
107    pub tool_call_id: String,
108    pub tool_name: String,
109    pub content: Vec<Content>,
110    pub is_error: bool,
111    #[serde(default = "now_ms")]
112    pub timestamp: i64,
113}
114
115impl Message {
116    pub fn user_text(s: impl Into<String>) -> Self {
117        Message::User {
118            content: vec![Content::text(s)],
119            timestamp: now_ms(),
120        }
121    }
122}
123
124pub fn now_ms() -> i64 {
125    chrono::Utc::now().timestamp_millis()
126}
127
128/// Tool schema definition matching the unified pi-ai Tool interface.
129#[derive(Debug, Clone, Serialize, Deserialize)]
130pub struct Tool {
131    pub name: String,
132    pub description: String,
133    /// JSON Schema for the tool parameters (object schema).
134    pub parameters: Value,
135}
136
137#[derive(Debug, Clone, Default)]
138pub struct Context {
139    pub system_prompt: Option<String>,
140    pub messages: Vec<Message>,
141    pub tools: Vec<Tool>,
142}
143
144#[derive(Debug, Clone, Default)]
145pub struct StreamOptions {
146    pub temperature: Option<f32>,
147    pub max_tokens: Option<u32>,
148    pub api_key: Option<String>,
149    pub reasoning: Option<ThinkingLevel>,
150    /// Cancellation token honored by the provider HTTP client and SSE loop.
151    pub cancel: Option<tokio_util::sync::CancellationToken>,
152    /// Override base URL (used by OpenAI-compatible passthrough providers).
153    pub base_url: Option<String>,
154    /// Optional custom request headers.
155    pub headers: std::collections::BTreeMap<String, String>,
156}
157
158/// Model descriptor — analogous to the Model<TApi> interface in pi-ai.
159#[derive(Debug, Clone)]
160pub struct Model {
161    pub id: String,
162    pub name: String,
163    pub api: String,
164    pub provider: String,
165    pub base_url: String,
166    pub reasoning: bool,
167    pub context_window: u32,
168    pub max_tokens: u32,
169}
170
171impl Model {
172    /// Anthropic Claude Sonnet 4.6.
173    pub fn anthropic_claude_sonnet_4_6() -> Self {
174        Self {
175            id: "claude-sonnet-4-6".into(),
176            name: "Claude Sonnet 4.6".into(),
177            api: "anthropic-messages".into(),
178            provider: "anthropic".into(),
179            base_url: "https://api.anthropic.com".into(),
180            reasoning: true,
181            context_window: 200_000,
182            max_tokens: 8_192,
183        }
184    }
185
186    pub fn anthropic_claude_opus_4_7() -> Self {
187        Self {
188            id: "claude-opus-4-7".into(),
189            name: "Claude Opus 4.7".into(),
190            api: "anthropic-messages".into(),
191            provider: "anthropic".into(),
192            base_url: "https://api.anthropic.com".into(),
193            reasoning: true,
194            context_window: 200_000,
195            max_tokens: 8_192,
196        }
197    }
198
199    pub fn openai_gpt_4o_mini() -> Self {
200        Self {
201            id: "gpt-4o-mini".into(),
202            name: "GPT-4o mini".into(),
203            api: "openai-completions".into(),
204            provider: "openai".into(),
205            base_url: "https://api.openai.com/v1".into(),
206            reasoning: false,
207            context_window: 128_000,
208            max_tokens: 16_384,
209        }
210    }
211
212    pub fn openai_gpt_4o() -> Self {
213        Self {
214            id: "gpt-4o".into(),
215            name: "GPT-4o".into(),
216            api: "openai-completions".into(),
217            provider: "openai".into(),
218            base_url: "https://api.openai.com/v1".into(),
219            reasoning: false,
220            context_window: 128_000,
221            max_tokens: 16_384,
222        }
223    }
224
225    pub fn gemini_2_0_flash() -> Self {
226        Self {
227            id: "gemini-2.0-flash".into(),
228            name: "Gemini 2.0 Flash".into(),
229            api: "google-generative-ai".into(),
230            provider: "google".into(),
231            base_url: "https://generativelanguage.googleapis.com".into(),
232            reasoning: false,
233            context_window: 1_000_000,
234            max_tokens: 8_192,
235        }
236    }
237
238    /// OpenAI-compatible passthrough: any provider whose API matches OpenAI Chat
239    /// Completions (OpenRouter, Together, Groq, Cerebras, DeepSeek, Fireworks,
240    /// xAI, etc.). Construct via `openai_compat("groq", "...", "...", ...)`
241    /// or by overriding `base_url`.
242    pub fn openai_compat(
243        provider: impl Into<String>,
244        id: impl Into<String>,
245        base_url: impl Into<String>,
246        context_window: u32,
247        max_tokens: u32,
248    ) -> Self {
249        let id = id.into();
250        Self {
251            name: id.clone(),
252            id,
253            api: "openai-completions".into(),
254            provider: provider.into(),
255            base_url: base_url.into(),
256            reasoning: false,
257            context_window,
258            max_tokens,
259        }
260    }
261}
262
263/// Events emitted by streaming completions, matching `AssistantMessageEvent` in pi-ai.
264#[derive(Debug, Clone)]
265pub enum AssistantMessageEvent {
266    Start,
267    TextStart {
268        content_index: usize,
269    },
270    TextDelta {
271        content_index: usize,
272        delta: String,
273    },
274    TextEnd {
275        content_index: usize,
276        content: String,
277    },
278    ThinkingStart {
279        content_index: usize,
280    },
281    ThinkingDelta {
282        content_index: usize,
283        delta: String,
284    },
285    ThinkingEnd {
286        content_index: usize,
287        content: String,
288    },
289    ToolCallStart {
290        content_index: usize,
291        id: String,
292        name: String,
293    },
294    ToolCallDelta {
295        content_index: usize,
296        delta: String,
297    },
298    ToolCallEnd {
299        content_index: usize,
300        id: String,
301        name: String,
302        arguments: Value,
303    },
304    Done {
305        reason: StopReason,
306        message: AssistantMessage,
307    },
308    Error {
309        reason: StopReason,
310        error: AssistantMessage,
311    },
312}
pi_ai/types.rs

pi_ai/
types.rs
