autoagents_protocol/
llm.rs

1use serde::{Deserialize, Serialize};
2
3/// Usage metadata for a chat response.
4#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
5pub struct Usage {
6    pub prompt_tokens: u32,
7    pub completion_tokens: u32,
8    pub total_tokens: u32,
9    #[serde(skip_serializing_if = "Option::is_none")]
10    pub completion_tokens_details: Option<CompletionTokensDetails>,
11    #[serde(skip_serializing_if = "Option::is_none")]
12    pub prompt_tokens_details: Option<PromptTokensDetails>,
13}
14
15#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
16pub struct CompletionTokensDetails {
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub reasoning_tokens: Option<u32>,
19    #[serde(skip_serializing_if = "Option::is_none")]
20    pub audio_tokens: Option<u32>,
21}
22
23#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
24pub struct PromptTokensDetails {
25    #[serde(skip_serializing_if = "Option::is_none")]
26    pub cached_tokens: Option<u32>,
27    #[serde(skip_serializing_if = "Option::is_none")]
28    pub audio_tokens: Option<u32>,
29}
30
31/// The supported MIME type of an image.
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
33#[non_exhaustive]
34pub enum ImageMime {
35    JPEG,
36    PNG,
37    GIF,
38    WEBP,
39}
40
41impl ImageMime {
42    pub fn mime_type(&self) -> &'static str {
43        match self {
44            ImageMime::JPEG => "image/jpeg",
45            ImageMime::PNG => "image/png",
46            ImageMime::GIF => "image/gif",
47            ImageMime::WEBP => "image/webp",
48        }
49    }
50}
51
52/// Tool call represents a function call that an LLM wants to make.
53#[derive(Debug, Deserialize, Serialize, Clone, Eq, PartialEq)]
54pub struct ToolCall {
55    pub id: String,
56    #[serde(rename = "type")]
57    pub call_type: String,
58    pub function: FunctionCall,
59}
60
61/// FunctionCall contains details about which function to call and with what arguments.
62#[derive(Debug, Deserialize, Serialize, Clone, Eq, PartialEq)]
63pub struct FunctionCall {
64    pub name: String,
65    pub arguments: String,
66}
67
68/// A streaming chunk that can be either text or a tool call event.
69#[derive(Debug, Clone, Serialize, Deserialize)]
70pub enum StreamChunk {
71    Text(String),
72    ReasoningContent(String),
73    ToolUseStart {
74        index: usize,
75        id: String,
76        name: String,
77    },
78    ToolUseInputDelta {
79        index: usize,
80        partial_json: String,
81    },
82    ToolUseComplete {
83        index: usize,
84        tool_call: ToolCall,
85    },
86    Done {
87        stop_reason: String,
88    },
89    Usage(Usage),
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[test]
97    fn stream_chunk_serializes_roundtrip() {
98        let chunk = StreamChunk::ToolUseStart {
99            index: 1,
100            id: "tool_1".to_string(),
101            name: "search".to_string(),
102        };
103        let serialized = serde_json::to_string(&chunk).unwrap();
104        let deserialized: StreamChunk = serde_json::from_str(&serialized).unwrap();
105        match deserialized {
106            StreamChunk::ToolUseStart { id, name, .. } => {
107                assert_eq!(id, "tool_1");
108                assert_eq!(name, "search");
109            }
110            _ => panic!("expected ToolUseStart"),
111        }
112    }
113
114    #[test]
115    fn tool_call_serializes_roundtrip() {
116        let call = ToolCall {
117            id: "call_1".to_string(),
118            call_type: "function".to_string(),
119            function: FunctionCall {
120                name: "lookup".to_string(),
121                arguments: "{\"q\":\"value\"}".to_string(),
122            },
123        };
124        let serialized = serde_json::to_string(&call).unwrap();
125        let deserialized: ToolCall = serde_json::from_str(&serialized).unwrap();
126        assert_eq!(deserialized, call);
127    }
128
129    #[test]
130    fn image_mime_serializes_roundtrip() {
131        let mime = ImageMime::PNG;
132        let serialized = serde_json::to_string(&mime).unwrap();
133        let deserialized: ImageMime = serde_json::from_str(&serialized).unwrap();
134        assert_eq!(deserialized, mime);
135    }
136
137    #[test]
138    fn image_mime_type_mapping() {
139        assert_eq!(ImageMime::JPEG.mime_type(), "image/jpeg");
140        assert_eq!(ImageMime::PNG.mime_type(), "image/png");
141        assert_eq!(ImageMime::GIF.mime_type(), "image/gif");
142        assert_eq!(ImageMime::WEBP.mime_type(), "image/webp");
143    }
144
145    #[test]
146    fn usage_serializes_with_details() {
147        let usage = Usage {
148            prompt_tokens: 1,
149            completion_tokens: 2,
150            total_tokens: 3,
151            completion_tokens_details: Some(CompletionTokensDetails {
152                reasoning_tokens: Some(1),
153                audio_tokens: None,
154            }),
155            prompt_tokens_details: Some(PromptTokensDetails {
156                cached_tokens: Some(2),
157                audio_tokens: None,
158            }),
159        };
160        let serialized = serde_json::to_value(&usage).unwrap();
161        assert!(serialized.get("completion_tokens_details").is_some());
162        assert!(serialized.get("prompt_tokens_details").is_some());
163    }
164}
autoagents_protocol/llm.rs

autoagents_protocol/
llm.rs
