agent_sdk/llm/
types.rs

1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone)]
4pub struct ChatRequest {
5    pub system: String,
6    pub messages: Vec<Message>,
7    pub tools: Option<Vec<Tool>>,
8    pub max_tokens: u32,
9}
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct Message {
13    pub role: Role,
14    pub content: Content,
15}
16
17impl Message {
18    #[must_use]
19    pub fn user(text: impl Into<String>) -> Self {
20        Self {
21            role: Role::User,
22            content: Content::Text(text.into()),
23        }
24    }
25
26    #[must_use]
27    pub fn assistant(text: impl Into<String>) -> Self {
28        Self {
29            role: Role::Assistant,
30            content: Content::Text(text.into()),
31        }
32    }
33
34    #[must_use]
35    pub fn assistant_with_tool_use(
36        text: Option<String>,
37        id: impl Into<String>,
38        name: impl Into<String>,
39        input: serde_json::Value,
40    ) -> Self {
41        let mut blocks = Vec::new();
42        if let Some(t) = text {
43            blocks.push(ContentBlock::Text { text: t });
44        }
45        blocks.push(ContentBlock::ToolUse {
46            id: id.into(),
47            name: name.into(),
48            input,
49        });
50        Self {
51            role: Role::Assistant,
52            content: Content::Blocks(blocks),
53        }
54    }
55
56    #[must_use]
57    pub fn tool_result(
58        tool_use_id: impl Into<String>,
59        content: impl Into<String>,
60        is_error: bool,
61    ) -> Self {
62        Self {
63            role: Role::User,
64            content: Content::Blocks(vec![ContentBlock::ToolResult {
65                tool_use_id: tool_use_id.into(),
66                content: content.into(),
67                is_error: if is_error { Some(true) } else { None },
68            }]),
69        }
70    }
71}
72
73#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
74#[serde(rename_all = "lowercase")]
75pub enum Role {
76    User,
77    Assistant,
78}
79
80#[derive(Debug, Clone, Serialize, Deserialize)]
81#[serde(untagged)]
82pub enum Content {
83    Text(String),
84    Blocks(Vec<ContentBlock>),
85}
86
87impl Content {
88    #[must_use]
89    pub fn first_text(&self) -> Option<&str> {
90        match self {
91            Self::Text(s) => Some(s),
92            Self::Blocks(blocks) => blocks.iter().find_map(|b| match b {
93                ContentBlock::Text { text } => Some(text.as_str()),
94                _ => None,
95            }),
96        }
97    }
98}
99
100#[derive(Debug, Clone, Serialize, Deserialize)]
101#[serde(tag = "type")]
102pub enum ContentBlock {
103    #[serde(rename = "text")]
104    Text { text: String },
105
106    #[serde(rename = "tool_use")]
107    ToolUse {
108        id: String,
109        name: String,
110        input: serde_json::Value,
111    },
112
113    #[serde(rename = "tool_result")]
114    ToolResult {
115        tool_use_id: String,
116        content: String,
117        #[serde(skip_serializing_if = "Option::is_none")]
118        is_error: Option<bool>,
119    },
120}
121
122#[derive(Debug, Clone, Serialize)]
123pub struct Tool {
124    pub name: String,
125    pub description: String,
126    pub input_schema: serde_json::Value,
127}
128
129#[derive(Debug, Clone)]
130pub struct ChatResponse {
131    pub id: String,
132    pub content: Vec<ContentBlock>,
133    pub model: String,
134    pub stop_reason: Option<StopReason>,
135    pub usage: Usage,
136}
137
138impl ChatResponse {
139    #[must_use]
140    pub fn first_text(&self) -> Option<&str> {
141        self.content.iter().find_map(|b| match b {
142            ContentBlock::Text { text } => Some(text.as_str()),
143            _ => None,
144        })
145    }
146
147    pub fn tool_uses(&self) -> impl Iterator<Item = (&str, &str, &serde_json::Value)> {
148        self.content.iter().filter_map(|b| match b {
149            ContentBlock::ToolUse { id, name, input } => Some((id.as_str(), name.as_str(), input)),
150            _ => None,
151        })
152    }
153
154    #[must_use]
155    pub fn has_tool_use(&self) -> bool {
156        self.content
157            .iter()
158            .any(|b| matches!(b, ContentBlock::ToolUse { .. }))
159    }
160}
161
162#[derive(Debug, Clone, Copy, Deserialize, PartialEq, Eq)]
163#[serde(rename_all = "snake_case")]
164pub enum StopReason {
165    EndTurn,
166    ToolUse,
167    MaxTokens,
168    StopSequence,
169}
170
171#[derive(Debug, Clone, Deserialize)]
172pub struct Usage {
173    pub input_tokens: u32,
174    pub output_tokens: u32,
175}
176
177#[derive(Debug)]
178pub enum ChatOutcome {
179    Success(ChatResponse),
180    RateLimited,
181    InvalidRequest(String),
182    ServerError(String),
183}