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    Error,
15    Unknown(String),
16}
17
18/// Token usage reported by a single LLM API response. Providers fill in only
19/// the dimensions they expose; the rest stay `None`.
20#[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    /// Build a `TokenUsage` with only the input/output token counts populated.
44    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}