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

rpi_ai/
types.rs
