1use schemars::JsonSchema;
2use serde::{Deserialize, Serialize};
3
4use super::{TokenUsage, 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#[doc = include_str!("docs/llm_response.md")]
19#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
20#[serde(tag = "type", rename_all = "camelCase")]
21pub enum LlmResponse {
22 Start {
23 message_id: String,
24 },
25 Text {
26 chunk: String,
27 },
28 Reasoning {
29 chunk: String,
30 },
31 EncryptedReasoning {
32 id: String,
33 content: String,
34 },
35 ToolRequestStart {
36 id: String,
37 name: String,
38 },
39 ToolRequestArg {
40 id: String,
41 chunk: String,
42 },
43 ToolRequestComplete {
44 tool_call: ToolCallRequest,
45 },
46 Done {
47 stop_reason: Option<StopReason>,
48 },
49 Error {
50 message: String,
51 },
52 Usage {
53 #[serde(flatten)]
54 tokens: TokenUsage,
55 },
56}
57
58impl LlmResponse {
59 pub fn start(message_id: &str) -> Self {
60 Self::Start { message_id: message_id.to_string() }
61 }
62
63 pub fn text(chunk: &str) -> Self {
64 Self::Text { chunk: chunk.to_string() }
65 }
66
67 pub fn reasoning(chunk: &str) -> Self {
68 Self::Reasoning { chunk: chunk.to_string() }
69 }
70
71 pub fn encrypted_reasoning(id: &str, encrypted: &str) -> Self {
72 Self::EncryptedReasoning { id: id.to_string(), content: encrypted.to_string() }
73 }
74
75 pub fn tool_request_start(id: &str, name: &str) -> Self {
76 Self::ToolRequestStart { id: id.to_string(), name: name.to_string() }
77 }
78
79 pub fn tool_request_arg(id: &str, chunk: &str) -> Self {
80 Self::ToolRequestArg { id: id.to_string(), chunk: chunk.to_string() }
81 }
82
83 pub fn tool_request_complete(id: &str, name: &str, arguments: &str) -> Self {
84 Self::ToolRequestComplete {
85 tool_call: ToolCallRequest { id: id.to_string(), name: name.to_string(), arguments: arguments.to_string() },
86 }
87 }
88
89 pub fn usage(input_tokens: u64, output_tokens: u64) -> Self {
90 Self::Usage { tokens: TokenUsage::new(input_tokens, output_tokens) }
91 }
92
93 pub fn done() -> Self {
94 Self::Done { stop_reason: None }
95 }
96
97 pub fn done_with_stop_reason(stop_reason: StopReason) -> Self {
98 Self::Done { stop_reason: Some(stop_reason) }
99 }
100}