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

rpi_ai/
types.rs
