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    ToolUseStart {
73        index: usize,
74        id: String,
75        name: String,
76    },
77    ToolUseInputDelta {
78        index: usize,
79        partial_json: String,
80    },
81    ToolUseComplete {
82        index: usize,
83        tool_call: ToolCall,
84    },
85    Done {
86        stop_reason: String,
87    },
88    Usage(Usage),
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn stream_chunk_serializes_roundtrip() {
97        let chunk = StreamChunk::ToolUseStart {
98            index: 1,
99            id: "tool_1".to_string(),
100            name: "search".to_string(),
101        };
102        let serialized = serde_json::to_string(&chunk).unwrap();
103        let deserialized: StreamChunk = serde_json::from_str(&serialized).unwrap();
104        match deserialized {
105            StreamChunk::ToolUseStart { id, name, .. } => {
106                assert_eq!(id, "tool_1");
107                assert_eq!(name, "search");
108            }
109            _ => panic!("expected ToolUseStart"),
110        }
111    }
112
113    #[test]
114    fn tool_call_serializes_roundtrip() {
115        let call = ToolCall {
116            id: "call_1".to_string(),
117            call_type: "function".to_string(),
118            function: FunctionCall {
119                name: "lookup".to_string(),
120                arguments: "{\"q\":\"value\"}".to_string(),
121            },
122        };
123        let serialized = serde_json::to_string(&call).unwrap();
124        let deserialized: ToolCall = serde_json::from_str(&serialized).unwrap();
125        assert_eq!(deserialized, call);
126    }
127
128    #[test]
129    fn image_mime_serializes_roundtrip() {
130        let mime = ImageMime::PNG;
131        let serialized = serde_json::to_string(&mime).unwrap();
132        let deserialized: ImageMime = serde_json::from_str(&serialized).unwrap();
133        assert_eq!(deserialized, mime);
134    }
135
136    #[test]
137    fn image_mime_type_mapping() {
138        assert_eq!(ImageMime::JPEG.mime_type(), "image/jpeg");
139        assert_eq!(ImageMime::PNG.mime_type(), "image/png");
140        assert_eq!(ImageMime::GIF.mime_type(), "image/gif");
141        assert_eq!(ImageMime::WEBP.mime_type(), "image/webp");
142    }
143
144    #[test]
145    fn usage_serializes_with_details() {
146        let usage = Usage {
147            prompt_tokens: 1,
148            completion_tokens: 2,
149            total_tokens: 3,
150            completion_tokens_details: Some(CompletionTokensDetails {
151                reasoning_tokens: Some(1),
152                audio_tokens: None,
153            }),
154            prompt_tokens_details: Some(PromptTokensDetails {
155                cached_tokens: Some(2),
156                audio_tokens: None,
157            }),
158        };
159        let serialized = serde_json::to_value(&usage).unwrap();
160        assert!(serialized.get("completion_tokens_details").is_some());
161        assert!(serialized.get("prompt_tokens_details").is_some());
162    }
163}
autoagents_protocol/llm.rs

autoagents_protocol/
llm.rs
