a_agent/
model.rs

1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
4#[serde(rename_all = "snake_case")]
5pub enum Role {
6    System,
7    User,
8    Assistant,
9    Tool,
10}
11
12#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
13pub enum ContentBlock {
14    Text(String),
15    Reasoning(String),
16    ToolCall(ToolCall),
17    ToolResult(ToolResult),
18}
19
20#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
21pub struct ToolCall {
22    pub id: String,
23    pub name: String,
24    pub arguments: String,
25}
26
27impl ToolCall {
28    pub fn new(
29        id: impl Into<String>,
30        name: impl Into<String>,
31        arguments: impl Into<String>,
32    ) -> Self {
33        Self {
34            id: id.into(),
35            name: name.into(),
36            arguments: arguments.into(),
37        }
38    }
39}
40
41#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
42pub struct ToolResult {
43    pub call_id: String,
44    pub output: String,
45    pub is_error: bool,
46}
47
48impl ToolResult {
49    pub fn success(call_id: impl Into<String>, output: impl Into<String>) -> Self {
50        Self {
51            call_id: call_id.into(),
52            output: output.into(),
53            is_error: false,
54        }
55    }
56
57    pub fn error(call_id: impl Into<String>, output: impl Into<String>) -> Self {
58        Self {
59            call_id: call_id.into(),
60            output: output.into(),
61            is_error: true,
62        }
63    }
64}
65
66#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
67#[serde(default)]
68pub struct Usage {
69    pub input_tokens: Option<u64>,
70    pub output_tokens: Option<u64>,
71    pub cached_tokens: Option<u64>,
72    pub cache_write_tokens: Option<u64>,
73    pub total_tokens: Option<u64>,
74}
75
76impl Usage {
77    pub fn context_tokens(self) -> Option<u64> {
78        if let Some(total) = self.total_tokens.filter(|total| *total > 0) {
79            return Some(total);
80        }
81        [
82            self.input_tokens,
83            self.output_tokens,
84            self.cached_tokens,
85            self.cache_write_tokens,
86        ]
87        .iter()
88        .any(Option::is_some)
89        .then(|| {
90            self.input_tokens.unwrap_or(0)
91                + self.output_tokens.unwrap_or(0)
92                + self.cached_tokens.unwrap_or(0)
93                + self.cache_write_tokens.unwrap_or(0)
94        })
95    }
96}
97
98#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
99pub struct ModelMessage {
100    pub role: Role,
101    pub blocks: Vec<ContentBlock>,
102}
103
104#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
105pub struct ModelRequest {
106    pub system_prompt: String,
107    pub messages: Vec<ModelMessage>,
108    pub include_tools: bool,
109}
110
111#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
112pub struct ModelTurn {
113    pub blocks: Vec<ContentBlock>,
114    pub tool_calls: Vec<ToolCall>,
115    pub usage: Option<Usage>,
116    pub provider_state: Option<serde_json::Value>,
117}
118
119impl ModelTurn {
120    pub fn text(text: impl Into<String>) -> Self {
121        Self {
122            blocks: vec![ContentBlock::Text(text.into())],
123            ..Self::default()
124        }
125    }
126
127    pub fn with_tools(tool_calls: Vec<ToolCall>) -> Self {
128        Self {
129            blocks: tool_calls
130                .iter()
131                .cloned()
132                .map(ContentBlock::ToolCall)
133                .collect(),
134            tool_calls,
135            ..Self::default()
136        }
137    }
138
139    pub fn final_text(&self) -> Option<String> {
140        let text = self
141            .blocks
142            .iter()
143            .filter_map(|block| match block {
144                ContentBlock::Text(text) => Some(text.as_str()),
145                _ => None,
146            })
147            .collect::<String>();
148        (!text.is_empty()).then_some(text)
149    }
150}
151
152#[derive(Debug, Clone, PartialEq, Eq)]
153pub enum StreamEvent {
154    GenerationStart,
155    TextDelta { delta: String },
156    ReasoningDelta { delta: String },
157    ToolCallStart { id: String, name: String },
158    ToolCallArgsDelta { id: String, delta: String },
159    ToolCallEnd { id: String },
160    ToolExecutionStart { id: String },
161    ToolExecutionOutput { id: String, delta: String },
162    ToolExecutionEnd { id: String, result: ToolResult },
163    Usage(Usage),
164    Done,
165    Error { message: String },
166}
167
168#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
169pub struct ConversationItem {
170    pub id: String,
171    pub session_id: String,
172    pub parent_id: Option<String>,
173    pub role: Role,
174    pub blocks: Vec<ContentBlock>,
175    pub usage: Option<Usage>,
176    pub created_at: i64,
177}
a_agent/model.rs

a_agent/
model.rs
