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