Skip to main content

llm/
llm_response.rs

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/// Token usage reported by a single LLM API response. Providers fill in only
18/// the dimensions they expose; the rest stay `None`.
19#[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    /// Build a `TokenUsage` with only the input/output token counts populated.
43    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}