Skip to main content

rpi_ai/
types.rs

1//! Mirrors `packages/ai/src/types.ts` — the core LLM type contract every other
2//! crate consumes.
3//!
4//! These types are the boundary between the provider-agnostic agent loop and the
5//! provider wire formats. Everything the loop talks in (`Content`, `Message`,
6//! `AssistantMessage`, `AssistantMessageEvent`) is defined here; provider code
7//! only produces/consumes these.
8
9use serde::{Deserialize, Serialize};
10use std::collections::BTreeMap;
11
12/// A newtype wrapping a JSON Schema value. Tools carry their parameters schema here
13/// (produced by `schemars::JsonSchema` derive, stored as a `serde_json::Value` so it
14/// can be sent over the wire to providers verbatim). Mirrors the TypeBox `TSchema`
15/// parameter shape on the TS `Tool`.
16#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
17#[serde(transparent)]
18pub struct Schema(pub serde_json::Value);
19
20impl Schema {
21    pub fn new(value: serde_json::Value) -> Self {
22        Self(value)
23    }
24
25    pub fn empty_object() -> Self {
26        Self(serde_json::Value::Object(serde_json::Map::new()))
27    }
28
29    pub fn as_value(&self) -> &serde_json::Value {
30        &self.0
31    }
32
33    pub fn into_value(self) -> serde_json::Value {
34        self.0
35    }
36
37    pub fn is_empty(&self) -> bool {
38        self.0.is_null() || (self.0.is_object() && self.0.as_object().map(|m| m.is_empty()).unwrap_or(true))
39    }
40}
41
42impl From<serde_json::Value> for Schema {
43    fn from(v: serde_json::Value) -> Self {
44        Self(v)
45    }
46}
47
48// ----------------------------------------------------------------------------
49// Content blocks — mirror TS `TextContent | ImageContent | ThinkingContent | ToolCall`
50// ----------------------------------------------------------------------------
51
52/// A text content block.
53#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
54#[serde(rename_all = "camelCase")]
55pub struct TextContent {
56    #[serde(rename = "type")]
57    pub kind: TextContentType,
58    pub text: String,
59    /// Provider message-metadata signature (OpenAI Responses legacy id string or
60    /// `TextSignatureV1` JSON). Optional.
61    #[serde(default, skip_serializing_if = "Option::is_none")]
62    pub text_signature: Option<String>,
63}
64
65/// Marker so the `type` field serializes as the literal `"text"`.
66#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
67pub struct TextContentType;
68impl Serialize for TextContentType {
69    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
70        s.serialize_str("text")
71    }
72}
73impl<'de> Deserialize<'de> for TextContentType {
74    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
75        let s = String::deserialize(d)?;
76        if s != "text" {
77            return Err(serde::de::Error::custom(format!("expected \"text\", got {s:?}")));
78        }
79        Ok(Self)
80    }
81}
82
83/// Thinking / reasoning content.
84#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
85#[serde(rename_all = "camelCase")]
86pub struct ThinkingContent {
87    #[serde(rename = "type")]
88    pub kind: ThinkingContentType,
89    pub thinking: String,
90    #[serde(default, skip_serializing_if = "Option::is_none")]
91    pub thinking_signature: Option<String>,
92    /// When true, safety filters redacted the thinking; the opaque payload is in
93    /// `thinking_signature`.
94    #[serde(default, skip_serializing_if = "is_false")]
95    pub redacted: bool,
96}
97
98#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
99pub struct ThinkingContentType;
100impl Serialize for ThinkingContentType {
101    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
102        s.serialize_str("thinking")
103    }
104}
105impl<'de> Deserialize<'de> for ThinkingContentType {
106    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
107        let s = String::deserialize(d)?;
108        if s != "thinking" {
109            return Err(serde::de::Error::custom(format!("expected \"thinking\", got {s:?}")));
110        }
111        Ok(Self)
112    }
113}
114
115fn is_false(b: &bool) -> bool {
116    !*b
117}
118
119/// An image content block (base64-encoded).
120#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
121#[serde(rename_all = "camelCase")]
122pub struct ImageContent {
123    #[serde(rename = "type")]
124    pub kind: ImageContentType,
125    pub data: String,
126    pub mime_type: String,
127}
128
129#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
130pub struct ImageContentType;
131impl Serialize for ImageContentType {
132    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
133        s.serialize_str("image")
134    }
135}
136impl<'de> Deserialize<'de> for ImageContentType {
137    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
138        let s = String::deserialize(d)?;
139        if s != "image" {
140            return Err(serde::de::Error::custom(format!("expected \"image\", got {s:?}")));
141        }
142        Ok(Self)
143    }
144}
145
146/// A tool call request from the model.
147#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
148#[serde(rename_all = "camelCase")]
149pub struct ToolCall {
150    #[serde(rename = "type")]
151    pub kind: ToolCallType,
152    pub id: String,
153    pub name: String,
154    /// LLM-supplied arguments. `Value::Object` in the happy path; we permit any
155    /// `Value` so partial-parse recovery has somewhere to park malformed input.
156    pub arguments: serde_json::Value,
157    /// Google-specific opaque signature for reusing thought context.
158    #[serde(default, skip_serializing_if = "Option::is_none")]
159    pub thought_signature: Option<String>,
160    /// OpenAI Responses namespace for dynamic/namespaced tools.
161    #[serde(default, skip_serializing_if = "Option::is_none")]
162    pub namespace: Option<String>,
163}
164
165#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
166pub struct ToolCallType;
167impl Serialize for ToolCallType {
168    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
169        s.serialize_str("toolCall")
170    }
171}
172impl<'de> Deserialize<'de> for ToolCallType {
173    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
174        let s = String::deserialize(d)?;
175        if s != "toolCall" {
176            return Err(serde::de::Error::custom(format!("expected \"toolCall\", got {s:?}")));
177        }
178        Ok(Self)
179    }
180}
181
182/// `TextContent | ImageContent | ThinkingContent | ToolCall` — assistant-side
183/// content blocks. User messages use only text/image; assistant messages may carry
184/// all four. Tagged via the `type` field so it round-trips the TS union exactly.
185#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
186#[serde(tag = "type", rename_all = "camelCase")]
187pub enum Content {
188    Text(TextContent),
189    Thinking(ThinkingContent),
190    Image(ImageContent),
191    ToolCall(ToolCall),
192}
193
194impl Content {
195    pub fn text<S: Into<String>>(s: S) -> Self {
196        Content::Text(TextContent {
197            kind: TextContentType,
198            text: s.into(),
199            text_signature: None,
200        })
201    }
202
203    pub fn thinking<S: Into<String>>(s: S) -> Self {
204        Content::Thinking(ThinkingContent {
205            kind: ThinkingContentType,
206            thinking: s.into(),
207            thinking_signature: None,
208            redacted: false,
209        })
210    }
211
212    pub fn tool_call<I: Into<String>, N: Into<String>>(id: I, name: N, arguments: serde_json::Value) -> Self {
213        Content::ToolCall(ToolCall {
214            kind: ToolCallType,
215            id: id.into(),
216            name: name.into(),
217            arguments,
218            thought_signature: None,
219            namespace: None,
220        })
221    }
222
223    /// Plain-text extraction (mirrors `contentText`): joins all `Text` blocks with
224    /// `\n`, dropping thinking/tool-call/image blocks.
225    pub fn text_only(content: &[Content], sep: &str) -> String {
226        content
227            .iter()
228            .filter_map(|c| match c {
229                Content::Text(t) => Some(t.text.as_str()),
230                _ => None,
231            })
232            .collect::<Vec<_>>()
233            .join(sep)
234    }
235}
236
237// ----------------------------------------------------------------------------
238// Usage + costs
239// ----------------------------------------------------------------------------
240
241#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
242#[serde(rename_all = "camelCase")]
243pub struct ModelCostRates {
244    pub input: f64,
245    pub output: f64,
246    pub cache_read: f64,
247    pub cache_write: f64,
248}
249
250#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
251#[serde(rename_all = "camelCase")]
252pub struct ModelCostTier {
253    #[serde(flatten)]
254    pub rates: ModelCostRates,
255    pub input_tokens_above: i64,
256}
257
258#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
259#[serde(rename_all = "camelCase")]
260pub struct ModelCost {
261    #[serde(flatten)]
262    pub rates: ModelCostRates,
263    #[serde(default, skip_serializing_if = "Vec::is_empty")]
264    pub tiers: Vec<ModelCostTier>,
265}
266
267impl Default for ModelCost {
268    fn default() -> Self {
269        Self {
270            rates: ModelCostRates::default(),
271            tiers: Vec::new(),
272        }
273    }
274}
275
276/// Per-request token usage + cost. Mirrors TS `Usage`.
277#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
278#[serde(rename_all = "camelCase")]
279pub struct Usage {
280    pub input: i64,
281    pub output: i64,
282    pub cache_read: i64,
283    pub cache_write: i64,
284    /// Subset of `cache_write` written with 1h retention (Anthropic only).
285    #[serde(default, skip_serializing_if = "Option::is_none")]
286    pub cache_write_1h: Option<i64>,
287    /// Reasoning tokens, when reported. Subset of `output`.
288    #[serde(default, skip_serializing_if = "Option::is_none")]
289    pub reasoning: Option<i64>,
290    pub total_tokens: i64,
291    pub cost: UsageCost,
292}
293
294#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
295#[serde(rename_all = "camelCase")]
296pub struct UsageCost {
297    pub input: f64,
298    pub output: f64,
299    pub cache_read: f64,
300    pub cache_write: f64,
301    pub total: f64,
302}
303
304impl Usage {
305    pub fn zero() -> Self {
306        Self {
307            input: 0,
308            output: 0,
309            cache_read: 0,
310            cache_write: 0,
311            cache_write_1h: None,
312            reasoning: None,
313            total_tokens: 0,
314            cost: UsageCost::default(),
315        }
316    }
317
318    pub fn add(&mut self, other: &Usage) {
319        self.input += other.input;
320        self.output += other.output;
321        self.cache_read += other.cache_read;
322        self.cache_write += other.cache_write;
323        if other.cache_write_1h.is_some() {
324            self.cache_write_1h = Some(self.cache_write_1h.unwrap_or(0) + other.cache_write_1h.unwrap());
325        }
326        if other.reasoning.is_some() {
327            self.reasoning = Some(self.reasoning.unwrap_or(0) + other.reasoning.unwrap());
328        }
329        self.total_tokens += other.total_tokens;
330        self.cost.input += other.cost.input;
331        self.cost.output += other.cost.output;
332        self.cost.cache_read += other.cost.cache_read;
333        self.cost.cache_write += other.cost.cache_write;
334        self.cost.total += other.cost.total;
335    }
336
337    /// `usage.totalTokens || (input + output + cacheRead + cacheWrite)` — mirrors
338    /// `calculateContextTokens`.
339    pub fn context_tokens(&self) -> i64 {
340        let sum = self.total_tokens;
341        if sum > 0 {
342            sum
343        } else {
344            self.input + self.output + self.cache_read + self.cache_write
345        }
346    }
347}
348
349// ----------------------------------------------------------------------------
350// Stop reason + provider/api identifiers
351// ----------------------------------------------------------------------------
352
353#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
354#[serde(rename_all = "lowercase")]
355pub enum StopReason {
356    Pending,
357    Stop,
358    Length,
359    ToolUse,
360    Error,
361    Aborted,
362    Deferred,
363}
364
365impl StopReason {
366    pub fn as_str(self) -> &'static str {
367        match self {
368            StopReason::Pending => "pending",
369            StopReason::Stop => "stop",
370            StopReason::Length => "length",
371            StopReason::ToolUse => "toolUse",
372            StopReason::Error => "error",
373            StopReason::Aborted => "aborted",
374            StopReason::Deferred => "deferred",
375        }
376    }
377}
378
379/// Pi known APIs. The string fallback covers custom APIs; we keep a typed enum for
380/// the common ones and expose `Api::Other(String)` for the rest.
381#[derive(Debug, Clone, PartialEq, Eq, Hash)]
382pub enum Api {
383    OpenaiCompletions,
384    MistralConversations,
385    OpenaiResponses,
386    AzureOpenaiResponses,
387    OpenaiCodexResponses,
388    AnthropicMessages,
389    BedrockConverseStream,
390    GoogleGenerativeAi,
391    GoogleVertex,
392    PiMessages,
393    Faux,
394    Other(String),
395}
396
397impl Api {
398    pub fn as_str(&self) -> &str {
399        match self {
400            Api::OpenaiCompletions => "openai-completions",
401            Api::MistralConversations => "mistral-conversations",
402            Api::OpenaiResponses => "openai-responses",
403            Api::AzureOpenaiResponses => "azure-openai-responses",
404            Api::OpenaiCodexResponses => "openai-codex-responses",
405            Api::AnthropicMessages => "anthropic-messages",
406            Api::BedrockConverseStream => "bedrock-converse-stream",
407            Api::GoogleGenerativeAi => "google-generative-ai",
408            Api::GoogleVertex => "google-vertex",
409            Api::PiMessages => "pi-messages",
410            Api::Faux => "faux",
411            Api::Other(s) => s.as_str(),
412        }
413    }
414}
415
416impl Serialize for Api {
417    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
418        s.serialize_str(self.as_str())
419    }
420}
421impl<'de> Deserialize<'de> for Api {
422    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
423        let s = String::deserialize(d)?;
424        Ok(match s.as_str() {
425            "openai-completions" => Api::OpenaiCompletions,
426            "mistral-conversations" => Api::MistralConversations,
427            "openai-responses" => Api::OpenaiResponses,
428            "azure-openai-responses" => Api::AzureOpenaiResponses,
429            "openai-codex-responses" => Api::OpenaiCodexResponses,
430            "anthropic-messages" => Api::AnthropicMessages,
431            "bedrock-converse-stream" => Api::BedrockConverseStream,
432            "google-generative-ai" => Api::GoogleGenerativeAi,
433            "google-vertex" => Api::GoogleVertex,
434            "pi-messages" => Api::PiMessages,
435            "faux" => Api::Faux,
436            other => Api::Other(other.to_string()),
437        })
438    }
439}
440
441/// Provider identifier (string; keeps the TS `KnownProvider | string` open enum).
442pub type ProviderId = String;
443
444/// Input modalities a model accepts.
445#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
446#[serde(rename_all = "lowercase")]
447pub enum InputModality {
448    Text,
449    Image,
450}
451
452/// Thinking levels, mirroring TS `ThinkingLevel | "off"`.
453#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
454#[serde(rename_all = "lowercase")]
455pub enum ThinkingLevel {
456    #[default]
457    Off,
458    Minimal,
459    Low,
460    Medium,
461    High,
462    Xhigh,
463    Max,
464}
465
466/// Maps pi thinking levels to provider-specific values; `None` marks a level as
467/// unsupported. Mirrors `ThinkingLevelMap = Partial<Record<...>>`.
468pub type ThinkingLevelMap = BTreeMap<ThinkingLevel, Option<String>>;
469
470/// Custom token budgets per thinking level for token-budgeted reasoning models.
471/// Mirrors TS `ThinkingBudgets`; all fields optional (defaults live in
472/// `simple-options.ts::adjustMaxTokensForThinking`).
473#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
474#[serde(rename_all = "camelCase")]
475pub struct ThinkingBudgets {
476    #[serde(default, skip_serializing_if = "Option::is_none")]
477    pub minimal: Option<u64>,
478    #[serde(default, skip_serializing_if = "Option::is_none")]
479    pub low: Option<u64>,
480    #[serde(default, skip_serializing_if = "Option::is_none")]
481    pub medium: Option<u64>,
482    #[serde(default, skip_serializing_if = "Option::is_none")]
483    pub high: Option<u64>,
484}
485
486// ----------------------------------------------------------------------------
487// Deferred responses
488// ----------------------------------------------------------------------------
489
490/// A durable handle to a deferred provider response (long-poll APIs).
491#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
492#[serde(rename_all = "camelCase")]
493pub struct DeferredHandle {
494    pub provider: String,
495    pub model_id: String,
496    pub api: String,
497    pub id: String,
498    #[serde(default, skip_serializing_if = "Option::is_none")]
499    pub expires_at: Option<i64>,
500    #[serde(default, skip_serializing_if = "Option::is_none")]
501    pub poll_after_ms: Option<i64>,
502    #[serde(default, skip_serializing_if = "Option::is_none")]
503    pub data: Option<serde_json::Value>,
504}
505
506// ----------------------------------------------------------------------------
507// Messages — mirror TS `Message = UserMessage | AssistantMessage | ToolResultMessage`
508// ----------------------------------------------------------------------------
509
510/// User content may be a plain string or a list of text/image blocks.
511#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
512#[serde(untagged)]
513pub enum UserContent {
514    Text(String),
515    Blocks(Vec<Content>),
516}
517
518impl UserContent {
519    pub fn as_text(&self) -> Option<&str> {
520        match self {
521            UserContent::Text(s) => Some(s),
522            UserContent::Blocks(_) => None,
523        }
524    }
525}
526
527impl From<String> for UserContent {
528    fn from(s: String) -> Self {
529        UserContent::Text(s)
530    }
531}
532
533impl From<&str> for UserContent {
534    fn from(s: &str) -> Self {
535        UserContent::Text(s.to_string())
536    }
537}
538
539#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
540#[serde(rename_all = "camelCase")]
541pub struct UserMessage {
542    #[serde(rename = "role")]
543    pub role: UserRole,
544    pub content: UserContent,
545    pub timestamp: i64,
546}
547
548#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
549pub struct UserRole;
550impl Serialize for UserRole {
551    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
552        s.serialize_str("user")
553    }
554}
555impl<'de> Deserialize<'de> for UserRole {
556    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
557        let s = String::deserialize(d)?;
558        if s != "user" {
559            return Err(serde::de::Error::custom(format!("expected \"user\", got {s:?}")));
560        }
561        Ok(Self)
562    }
563}
564
565impl UserMessage {
566    pub fn new(content: impl Into<UserContent>, timestamp: i64) -> Self {
567        Self {
568            role: UserRole,
569            content: content.into(),
570            timestamp,
571        }
572    }
573}
574
575#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
576#[serde(rename_all = "camelCase")]
577pub struct ToolResultMessage {
578    #[serde(rename = "role")]
579    pub role: ToolResultRole,
580    pub tool_call_id: String,
581    pub tool_name: String,
582    /// Result content (text/image blocks).
583    pub content: Vec<Content>,
584    #[serde(default, skip_serializing_if = "Option::is_none")]
585    pub details: Option<serde_json::Value>,
586    #[serde(default, skip_serializing_if = "Option::is_none")]
587    pub usage: Option<Usage>,
588    #[serde(default, skip_serializing_if = "Vec::is_empty")]
589    pub added_tool_names: Vec<String>,
590    pub is_error: bool,
591    pub timestamp: i64,
592}
593
594#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
595pub struct ToolResultRole;
596impl Serialize for ToolResultRole {
597    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
598        s.serialize_str("toolResult")
599    }
600}
601impl<'de> Deserialize<'de> for ToolResultRole {
602    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
603        let s = String::deserialize(d)?;
604        if s != "toolResult" {
605            return Err(serde::de::Error::custom(format!("expected \"toolResult\", got {s:?}")));
606        }
607        Ok(Self)
608    }
609}
610
611/// Concrete assistant message. Mirrors TS `AssistantMessage`.
612#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
613#[serde(rename_all = "camelCase")]
614pub struct AssistantMessage {
615    #[serde(rename = "role")]
616    pub role: AssistantRole,
617    pub content: Vec<Content>,
618    pub api: Api,
619    pub provider: ProviderId,
620    pub model: String,
621    #[serde(default, skip_serializing_if = "Option::is_none")]
622    pub response_model: Option<String>,
623    #[serde(default, skip_serializing_if = "Option::is_none")]
624    pub response_id: Option<String>,
625    pub usage: Usage,
626    pub stop_reason: StopReason,
627    #[serde(default, skip_serializing_if = "Option::is_none")]
628    pub deferred: Option<DeferredHandle>,
629    #[serde(default, skip_serializing_if = "Option::is_none")]
630    pub error_message: Option<String>,
631    #[serde(default, skip_serializing_if = "Option::is_none")]
632    pub raw_stop_reason: Option<String>,
633    #[serde(default, skip_serializing_if = "Option::is_none")]
634    pub end_turn: Option<bool>,
635    pub timestamp: i64,
636}
637
638#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
639pub struct AssistantRole;
640impl Serialize for AssistantRole {
641    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
642        s.serialize_str("assistant")
643    }
644}
645impl<'de> Deserialize<'de> for AssistantRole {
646    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
647        let s = String::deserialize(d)?;
648        if s != "assistant" {
649            return Err(serde::de::Error::custom(format!("expected \"assistant\", got {s:?}")));
650        }
651        Ok(Self)
652    }
653}
654
655impl AssistantMessage {
656    /// Construct a new assistant message with zeroed usage and `Pending` stop reason,
657    /// ready for the stream mapper to mutate per-event.
658    pub fn empty(api: Api, provider: impl Into<String>, model: impl Into<String>, timestamp: i64) -> Self {
659        Self {
660            role: AssistantRole,
661            content: Vec::new(),
662            api,
663            provider: provider.into(),
664            model: model.into(),
665            response_model: None,
666            response_id: None,
667            usage: Usage::zero(),
668            stop_reason: StopReason::Pending,
669            deferred: None,
670            error_message: None,
671            raw_stop_reason: None,
672            end_turn: None,
673            timestamp,
674        }
675    }
676
677    /// Convenience: an error/aborted terminal message (used by providers + faux).
678    pub fn terminal(
679        api: Api,
680        provider: impl Into<String>,
681        model: impl Into<String>,
682        stop_reason: StopReason,
683        error_message: impl Into<String>,
684        timestamp: i64,
685    ) -> Self {
686        Self {
687            role: AssistantRole,
688            content: Vec::new(),
689            api,
690            provider: provider.into(),
691            model: model.into(),
692            response_model: None,
693            response_id: None,
694            usage: Usage::zero(),
695            stop_reason,
696            deferred: None,
697            error_message: Some(error_message.into()),
698            raw_stop_reason: None,
699            end_turn: None,
700            timestamp,
701        }
702    }
703}
704
705/// `Message = UserMessage | AssistantMessage | ToolResultMessage`. The base
706/// provider-facing message type.
707#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
708#[serde(tag = "role", rename_all = "camelCase")]
709pub enum Message {
710    User(UserMessage),
711    Assistant(Box<AssistantMessage>),
712    ToolResult(Box<ToolResultMessage>),
713}
714
715impl From<UserMessage> for Message {
716    fn from(m: UserMessage) -> Self {
717        Message::User(m)
718    }
719}
720impl From<AssistantMessage> for Message {
721    fn from(m: AssistantMessage) -> Self {
722        Message::Assistant(Box::new(m))
723    }
724}
725impl From<ToolResultMessage> for Message {
726    fn from(m: ToolResultMessage) -> Self {
727        Message::ToolResult(Box::new(m))
728    }
729}
730
731// ----------------------------------------------------------------------------
732// Tool + Context
733// ----------------------------------------------------------------------------
734
735/// OpenAI grammar variants — mirror TS `GrammarFormat` (used by constrained sampling).
736#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
737pub enum GrammarFormat {
738    #[serde(rename = "openai_lark")]
739    OpenaiLark,
740    #[serde(rename = "openai_regex")]
741    OpenaiRegex,
742}
743
744pub type GrammarVariants = BTreeMap<GrammarFormat, String>;
745
746/// Provider-side constrained sampling config (json_schema strict, or grammar).
747#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
748#[serde(tag = "type", rename_all = "snake_case")]
749pub enum ConstrainedSamplingConfig {
750    JsonSchema {
751        #[serde(rename = "strict")]
752        strict: ConstrainedStrictness,
753    },
754    Grammar {
755        variants: GrammarVariants,
756    },
757}
758
759#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
760#[serde(rename_all = "lowercase")]
761pub enum ConstrainedStrictness {
762    Prefer,
763    Require,
764}
765
766/// A tool definition. Mirrors TS `Tool`.
767#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
768#[serde(rename_all = "camelCase")]
769pub struct Tool {
770    pub name: String,
771    pub description: String,
772    pub parameters: Schema,
773    #[serde(default, skip_serializing_if = "Option::is_none")]
774    pub constrained_sampling: Option<ConstrainedSamplingConfig>,
775}
776
777/// The context handed to a provider stream call. Mirrors TS `Context`.
778#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
779#[serde(rename_all = "camelCase")]
780pub struct Context {
781    #[serde(default, skip_serializing_if = "Option::is_none")]
782    pub system_prompt: Option<String>,
783    pub messages: Vec<Message>,
784    #[serde(default, skip_serializing_if = "Vec::is_empty")]
785    pub tools: Vec<Tool>,
786}
787
788impl Context {
789    pub fn new(messages: Vec<Message>) -> Self {
790        Self {
791            system_prompt: None,
792            messages,
793            tools: Vec::new(),
794        }
795    }
796}
797
798// ----------------------------------------------------------------------------
799// AssistantMessageEvent — the streaming protocol union
800// ----------------------------------------------------------------------------
801
802/// `AssistantMessageEvent` tagged union, mirrors TS exactly.
803///
804/// Providers emit events in this order: `Start` is emitted once before any per-block
805/// events, then a `{Start,Delta*,End}` triple per content block (text / thinking /
806/// toolcall), then a terminal `Done` (carrying the final successful message) or
807/// `Error` (carrying a terminal message with `StopReason::Error`/`Aborted`).
808///
809/// Every non-terminal event references the in-progress `AssistantMessage` via
810/// `partial` so subscribers can render a live view. We wrap the partial in `Arc`
811/// so the event is cheap to clone across broadcast subscribers.
812#[derive(Debug, Clone, PartialEq)]
813pub enum AssistantMessageEvent {
814    Start {
815        partial: std::sync::Arc<AssistantMessage>,
816    },
817    TextStart {
818        content_index: usize,
819        partial: std::sync::Arc<AssistantMessage>,
820    },
821    TextDelta {
822        content_index: usize,
823        delta: String,
824        partial: std::sync::Arc<AssistantMessage>,
825    },
826    TextEnd {
827        content_index: usize,
828        content: String,
829        partial: std::sync::Arc<AssistantMessage>,
830    },
831    ThinkingStart {
832        content_index: usize,
833        partial: std::sync::Arc<AssistantMessage>,
834    },
835    ThinkingDelta {
836        content_index: usize,
837        delta: String,
838        partial: std::sync::Arc<AssistantMessage>,
839    },
840    ThinkingEnd {
841        content_index: usize,
842        content: String,
843        partial: std::sync::Arc<AssistantMessage>,
844    },
845    ToolCallStart {
846        content_index: usize,
847        partial: std::sync::Arc<AssistantMessage>,
848    },
849    ToolCallDelta {
850        content_index: usize,
851        delta: String,
852        partial: std::sync::Arc<AssistantMessage>,
853    },
854    ToolCallEnd {
855        content_index: usize,
856        tool_call: ToolCall,
857        partial: std::sync::Arc<AssistantMessage>,
858    },
859    Done {
860        reason: DoneReason,
861        message: AssistantMessage,
862    },
863    Error {
864        reason: ErrorReason,
865        error: AssistantMessage,
866    },
867}
868
869/// The terminal reason on a `Done` event — the success subset of `StopReason`.
870#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
871pub enum DoneReason {
872    Stop,
873    Length,
874    ToolUse,
875    Deferred,
876}
877
878impl DoneReason {
879    pub fn as_str(self) -> &'static str {
880        match self {
881            DoneReason::Stop => "stop",
882            DoneReason::Length => "length",
883            DoneReason::ToolUse => "toolUse",
884            DoneReason::Deferred => "deferred",
885        }
886    }
887}
888
889impl From<DoneReason> for StopReason {
890    fn from(r: DoneReason) -> Self {
891        match r {
892            DoneReason::Stop => StopReason::Stop,
893            DoneReason::Length => StopReason::Length,
894            DoneReason::ToolUse => StopReason::ToolUse,
895            DoneReason::Deferred => StopReason::Deferred,
896        }
897    }
898}
899
900/// The terminal reason on an `Error` event — the failure subset of `StopReason`.
901#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
902pub enum ErrorReason {
903    Aborted,
904    Error,
905}
906
907impl ErrorReason {
908    pub fn as_str(self) -> &'static str {
909        match self {
910            ErrorReason::Aborted => "aborted",
911            ErrorReason::Error => "error",
912        }
913    }
914}
915
916impl From<ErrorReason> for StopReason {
917    fn from(r: ErrorReason) -> Self {
918        match r {
919            ErrorReason::Aborted => StopReason::Aborted,
920            ErrorReason::Error => StopReason::Error,
921        }
922    }
923}
924
925impl AssistantMessageEvent {
926    /// Returns the `type` tag, mirroring the TS discriminator.
927    pub fn type_tag(&self) -> &'static str {
928        match self {
929            AssistantMessageEvent::Start { .. } => "start",
930            AssistantMessageEvent::TextStart { .. } => "text_start",
931            AssistantMessageEvent::TextDelta { .. } => "text_delta",
932            AssistantMessageEvent::TextEnd { .. } => "text_end",
933            AssistantMessageEvent::ThinkingStart { .. } => "thinking_start",
934            AssistantMessageEvent::ThinkingDelta { .. } => "thinking_delta",
935            AssistantMessageEvent::ThinkingEnd { .. } => "thinking_end",
936            AssistantMessageEvent::ToolCallStart { .. } => "toolcall_start",
937            AssistantMessageEvent::ToolCallDelta { .. } => "toolcall_delta",
938            AssistantMessageEvent::ToolCallEnd { .. } => "toolcall_end",
939            AssistantMessageEvent::Done { .. } => "done",
940            AssistantMessageEvent::Error { .. } => "error",
941        }
942    }
943
944    /// True for the two terminal variants.
945    pub fn is_terminal(&self) -> bool {
946        matches!(self, AssistantMessageEvent::Done { .. } | AssistantMessageEvent::Error { .. })
947    }
948
949    /// Snapshot of the partial message for live rendering. Terminal events do not
950    /// carry a `partial`; returns the terminal message itself for `Done`/`Error`.
951    pub fn partial(&self) -> &AssistantMessage {
952        match self {
953            AssistantMessageEvent::Start { partial }
954            | AssistantMessageEvent::TextStart { partial, .. }
955            | AssistantMessageEvent::TextDelta { partial, .. }
956            | AssistantMessageEvent::TextEnd { partial, .. }
957            | AssistantMessageEvent::ThinkingStart { partial, .. }
958            | AssistantMessageEvent::ThinkingDelta { partial, .. }
959            | AssistantMessageEvent::ThinkingEnd { partial, .. }
960            | AssistantMessageEvent::ToolCallStart { partial, .. }
961            | AssistantMessageEvent::ToolCallDelta { partial, .. }
962            | AssistantMessageEvent::ToolCallEnd { partial, .. } => partial,
963            AssistantMessageEvent::Done { message, .. } => message,
964            AssistantMessageEvent::Error { error, .. } => error,
965        }
966    }
967}
968
969// serde for AssistantMessageEvent — tagged on "type" matching the wire protocol.
970impl Serialize for AssistantMessageEvent {
971    fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
972        use serde::ser::SerializeMap;
973        let tag = self.type_tag();
974        match self {
975            AssistantMessageEvent::Start { partial } => {
976                let mut m = s.serialize_map(Some(2))?;
977                m.serialize_entry("type", tag)?;
978                m.serialize_entry("partial", &**partial)?;
979                m.end()
980            }
981            AssistantMessageEvent::TextStart { content_index, partial }
982            | AssistantMessageEvent::ThinkingStart { content_index, partial }
983            | AssistantMessageEvent::ToolCallStart { content_index, partial } => {
984                let mut m = s.serialize_map(Some(3))?;
985                m.serialize_entry("type", tag)?;
986                m.serialize_entry("contentIndex", content_index)?;
987                m.serialize_entry("partial", &**partial)?;
988                m.end()
989            }
990            AssistantMessageEvent::TextDelta { content_index, delta, partial }
991            | AssistantMessageEvent::ThinkingDelta { content_index, delta, partial } => {
992                let mut m = s.serialize_map(Some(4))?;
993                m.serialize_entry("type", tag)?;
994                m.serialize_entry("contentIndex", content_index)?;
995                m.serialize_entry("delta", delta)?;
996                m.serialize_entry("partial", &**partial)?;
997                m.end()
998            }
999            AssistantMessageEvent::ToolCallDelta { content_index, delta, partial } => {
1000                let mut m = s.serialize_map(Some(4))?;
1001                m.serialize_entry("type", tag)?;
1002                m.serialize_entry("contentIndex", content_index)?;
1003                m.serialize_entry("delta", delta)?;
1004                m.serialize_entry("partial", &**partial)?;
1005                m.end()
1006            }
1007            AssistantMessageEvent::TextEnd { content_index, content, partial }
1008            | AssistantMessageEvent::ThinkingEnd { content_index, content, partial } => {
1009                let mut m = s.serialize_map(Some(4))?;
1010                m.serialize_entry("type", tag)?;
1011                m.serialize_entry("contentIndex", content_index)?;
1012                m.serialize_entry("content", content)?;
1013                m.serialize_entry("partial", &**partial)?;
1014                m.end()
1015            }
1016            AssistantMessageEvent::ToolCallEnd { content_index, tool_call, partial } => {
1017                let mut m = s.serialize_map(Some(4))?;
1018                m.serialize_entry("type", tag)?;
1019                m.serialize_entry("contentIndex", content_index)?;
1020                m.serialize_entry("toolCall", tool_call)?;
1021                m.serialize_entry("partial", &**partial)?;
1022                m.end()
1023            }
1024            AssistantMessageEvent::Done { reason, message } => {
1025                let mut m = s.serialize_map(Some(3))?;
1026                m.serialize_entry("type", tag)?;
1027                m.serialize_entry("reason", reason.as_str())?;
1028                m.serialize_entry("message", message)?;
1029                m.end()
1030            }
1031            AssistantMessageEvent::Error { reason, error } => {
1032                let mut m = s.serialize_map(Some(3))?;
1033                m.serialize_entry("type", tag)?;
1034                m.serialize_entry("reason", reason.as_str())?;
1035                m.serialize_entry("error", error)?;
1036                m.end()
1037            }
1038        }
1039    }
1040}