Skip to main content

rtb_ai/
message.rs

1//! Provider-agnostic chat-message types.
2
3use serde::{Deserialize, Serialize};
4
5/// Who said what.
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
7#[serde(rename_all = "lowercase")]
8pub enum Role {
9    /// System / instruction prompt.
10    System,
11    /// User input.
12    User,
13    /// Assistant reply.
14    Assistant,
15}
16
17/// One message in a chat exchange. The body is a list of content
18/// blocks; most callers pass a single [`ContentBlock::Text`].
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct Message {
21    /// Who's speaking.
22    pub role: Role,
23    /// Body — can be one or more text blocks. Multi-block mode is
24    /// useful for prompt-caching (Anthropic-direct path) where each
25    /// block can have its own `cache_control`.
26    pub content: Vec<ContentBlock>,
27}
28
29impl Message {
30    /// Convenience: a `Message::user("…")` with a single text block.
31    #[must_use]
32    pub fn user(text: impl Into<String>) -> Self {
33        Self { role: Role::User, content: vec![ContentBlock::Text(text.into())] }
34    }
35
36    /// Convenience: a `Message::system("…")` with a single text block.
37    #[must_use]
38    pub fn system(text: impl Into<String>) -> Self {
39        Self { role: Role::System, content: vec![ContentBlock::Text(text.into())] }
40    }
41
42    /// Convenience: an `Message::assistant("…")` with a single text
43    /// block.
44    #[must_use]
45    pub fn assistant(text: impl Into<String>) -> Self {
46        Self { role: Role::Assistant, content: vec![ContentBlock::Text(text.into())] }
47    }
48}
49
50/// One block of message content. Today: just text; future: image /
51/// tool-use (Anthropic) / function-call (`OpenAI`).
52#[derive(Debug, Clone, Serialize, Deserialize)]
53#[serde(tag = "type", rename_all = "snake_case")]
54#[non_exhaustive]
55pub enum ContentBlock {
56    /// Plain text.
57    Text(String),
58}
59
60impl ContentBlock {
61    /// Borrow the inner text. `None` for non-text blocks (none today,
62    /// future variants may add image / tool-use).
63    #[must_use]
64    pub fn as_text(&self) -> Option<&str> {
65        let Self::Text(t) = self;
66        Some(t)
67    }
68}
69
70/// Token usage reported by the provider on a non-streaming response
71/// (or as the final event of a stream). Fields default to `0` on
72/// providers that don't surface that breakdown.
73#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
74pub struct Usage {
75    /// Tokens in the prompt (system + history + user input).
76    pub input_tokens: u32,
77    /// Tokens in the assistant's reply.
78    pub output_tokens: u32,
79    /// Anthropic-only — tokens written to the prompt cache.
80    pub cache_creation_input_tokens: u32,
81    /// Anthropic-only — tokens served from the prompt cache.
82    pub cache_read_input_tokens: u32,
83}
84
85/// Source citation produced by the assistant. Populated only on the
86/// Anthropic-direct path when the model emits citations. Other
87/// providers return an empty vector.
88#[derive(Debug, Clone, Serialize, Deserialize)]
89pub struct Citation {
90    /// The cited text snippet.
91    pub cited_text: String,
92    /// Provider-specific source identifier (file path, URL, doc ID).
93    pub source: String,
94    /// Character offset in the source where the cited span starts,
95    /// when the provider supplies it.
96    pub start_index: Option<u32>,
97    /// Character offset where the cited span ends.
98    pub end_index: Option<u32>,
99}