1use serde::{Deserialize, Serialize};
2
3use super::ToolCallRequest;
4
5#[doc = include_str!("docs/stop_reason.md")]
6#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
7#[serde(rename_all = "snake_case")]
8pub enum StopReason {
9 EndTurn,
10 Length,
11 ToolCalls,
12 ContentFilter,
13 FunctionCall,
14 Error,
15 Unknown(String),
16}
17
18#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
21pub struct TokenUsage {
22 pub input_tokens: u32,
23 pub output_tokens: u32,
24 #[serde(default)]
25 pub cache_read_tokens: Option<u32>,
26 #[serde(default)]
27 pub cache_creation_tokens: Option<u32>,
28 #[serde(default)]
29 pub input_audio_tokens: Option<u32>,
30 #[serde(default)]
31 pub input_video_tokens: Option<u32>,
32 #[serde(default)]
33 pub reasoning_tokens: Option<u32>,
34 #[serde(default)]
35 pub output_audio_tokens: Option<u32>,
36 #[serde(default)]
37 pub accepted_prediction_tokens: Option<u32>,
38 #[serde(default)]
39 pub rejected_prediction_tokens: Option<u32>,
40}
41
42impl TokenUsage {
43 pub fn new(input_tokens: u32, output_tokens: u32) -> Self {
45 Self { input_tokens, output_tokens, ..Self::default() }
46 }
47}
48
49#[doc = include_str!("docs/llm_response.md")]
50#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
51#[serde(tag = "type", rename_all = "camelCase")]
52pub enum LlmResponse {
53 Start {
54 message_id: String,
55 },
56 Text {
57 chunk: String,
58 },
59 Reasoning {
60 chunk: String,
61 },
62 EncryptedReasoning {
63 id: String,
64 content: String,
65 },
66 ToolRequestStart {
67 id: String,
68 name: String,
69 },
70 ToolRequestArg {
71 id: String,
72 chunk: String,
73 },
74 ToolRequestComplete {
75 tool_call: ToolCallRequest,
76 },
77 Done {
78 stop_reason: Option<StopReason>,
79 },
80 Error {
81 message: String,
82 },
83 Usage {
84 #[serde(flatten)]
85 tokens: TokenUsage,
86 },
87}
88
89impl LlmResponse {
90 pub fn start(message_id: &str) -> Self {
91 Self::Start { message_id: message_id.to_string() }
92 }
93
94 pub fn text(chunk: &str) -> Self {
95 Self::Text { chunk: chunk.to_string() }
96 }
97
98 pub fn reasoning(chunk: &str) -> Self {
99 Self::Reasoning { chunk: chunk.to_string() }
100 }
101
102 pub fn encrypted_reasoning(id: &str, encrypted: &str) -> Self {
103 Self::EncryptedReasoning { id: id.to_string(), content: encrypted.to_string() }
104 }
105
106 pub fn tool_request_start(id: &str, name: &str) -> Self {
107 Self::ToolRequestStart { id: id.to_string(), name: name.to_string() }
108 }
109
110 pub fn tool_request_arg(id: &str, chunk: &str) -> Self {
111 Self::ToolRequestArg { id: id.to_string(), chunk: chunk.to_string() }
112 }
113
114 pub fn tool_request_complete(id: &str, name: &str, arguments: &str) -> Self {
115 Self::ToolRequestComplete {
116 tool_call: ToolCallRequest { id: id.to_string(), name: name.to_string(), arguments: arguments.to_string() },
117 }
118 }
119
120 pub fn usage(input_tokens: u32, output_tokens: u32) -> Self {
121 Self::Usage { tokens: TokenUsage::new(input_tokens, output_tokens) }
122 }
123
124 pub fn done() -> Self {
125 Self::Done { stop_reason: None }
126 }
127
128 pub fn done_with_stop_reason(stop_reason: StopReason) -> Self {
129 Self::Done { stop_reason: Some(stop_reason) }
130 }
131}