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