Skip to main content

llm/
llm_response.rs

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/// 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, 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    /// Whether cache token counts are separate from, rather than included in, `input_tokens`.
29    #[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    /// Build a `TokenUsage` with only the input/output token counts populated.
47    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}