Skip to main content

llm/
llm_response.rs

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}