Skip to main content

liteforge/types/
chat.rs

1//! Chat completion types.
2
3use serde::{Deserialize, Serialize};
4
5use super::tools::ToolCall;
6
7/// A chat message.
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct Message {
10    /// The role of the message author (system, user, assistant, tool).
11    pub role: String,
12
13    /// The content of the message.
14    #[serde(skip_serializing_if = "Option::is_none")]
15    pub content: Option<String>,
16
17    /// Optional name for the participant.
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub name: Option<String>,
20
21    /// Tool calls made by the assistant.
22    #[serde(skip_serializing_if = "Option::is_none")]
23    pub tool_calls: Option<Vec<ToolCall>>,
24
25    /// Tool call ID for tool responses.
26    #[serde(skip_serializing_if = "Option::is_none")]
27    pub tool_call_id: Option<String>,
28}
29
30impl Message {
31    /// Create a new user message.
32    pub fn user(content: impl Into<String>) -> Self {
33        Self {
34            role: "user".to_string(),
35            content: Some(content.into()),
36            name: None,
37            tool_calls: None,
38            tool_call_id: None,
39        }
40    }
41
42    /// Create a new system message.
43    pub fn system(content: impl Into<String>) -> Self {
44        Self {
45            role: "system".to_string(),
46            content: Some(content.into()),
47            name: None,
48            tool_calls: None,
49            tool_call_id: None,
50        }
51    }
52
53    /// Create a new assistant message.
54    pub fn assistant(content: impl Into<String>) -> Self {
55        Self {
56            role: "assistant".to_string(),
57            content: Some(content.into()),
58            name: None,
59            tool_calls: None,
60            tool_call_id: None,
61        }
62    }
63
64    /// Create a tool response message.
65    pub fn tool(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
66        Self {
67            role: "tool".to_string(),
68            content: Some(content.into()),
69            name: None,
70            tool_calls: None,
71            tool_call_id: Some(tool_call_id.into()),
72        }
73    }
74}
75
76/// Token usage statistics.
77#[derive(Debug, Clone, Serialize, Deserialize)]
78pub struct Usage {
79    /// Number of tokens in the prompt.
80    pub prompt_tokens: u32,
81
82    /// Number of tokens in the completion.
83    pub completion_tokens: u32,
84
85    /// Total tokens used.
86    pub total_tokens: u32,
87}
88
89/// A completion choice.
90#[derive(Debug, Clone, Serialize, Deserialize)]
91pub struct Choice {
92    /// Index of this choice.
93    pub index: u32,
94
95    /// The generated message.
96    pub message: Message,
97
98    /// Reason the generation stopped.
99    #[serde(skip_serializing_if = "Option::is_none")]
100    pub finish_reason: Option<String>,
101}
102
103/// A chat completion response.
104#[derive(Debug, Clone, Serialize, Deserialize)]
105pub struct ChatCompletion {
106    /// Unique identifier for the completion.
107    pub id: String,
108
109    /// Object type (always "chat.completion").
110    pub object: String,
111
112    /// Unix timestamp of creation.
113    pub created: i64,
114
115    /// Model used for the completion.
116    pub model: String,
117
118    /// List of completion choices.
119    pub choices: Vec<Choice>,
120
121    /// Token usage statistics.
122    #[serde(skip_serializing_if = "Option::is_none")]
123    pub usage: Option<Usage>,
124}
125
126impl ChatCompletion {
127    /// Get the content of the first choice.
128    pub fn content(&self) -> Option<&str> {
129        self.choices
130            .first()
131            .and_then(|c| c.message.content.as_deref())
132    }
133
134    /// Get the first choice's message.
135    pub fn message(&self) -> Option<&Message> {
136        self.choices.first().map(|c| &c.message)
137    }
138}
139
140// --- Streaming types ---
141
142/// Delta content in a streaming chunk.
143#[derive(Debug, Clone, Serialize, Deserialize, Default)]
144pub struct ChoiceDelta {
145    /// The role (only in first chunk).
146    #[serde(skip_serializing_if = "Option::is_none")]
147    pub role: Option<String>,
148
149    /// The content delta.
150    #[serde(skip_serializing_if = "Option::is_none")]
151    pub content: Option<String>,
152
153    /// Tool calls delta.
154    #[serde(skip_serializing_if = "Option::is_none")]
155    pub tool_calls: Option<Vec<ToolCall>>,
156}
157
158/// A streaming choice.
159#[derive(Debug, Clone, Serialize, Deserialize)]
160pub struct StreamChoice {
161    /// Index of this choice.
162    pub index: u32,
163
164    /// The delta content.
165    pub delta: ChoiceDelta,
166
167    /// Reason the generation stopped.
168    #[serde(skip_serializing_if = "Option::is_none")]
169    pub finish_reason: Option<String>,
170}
171
172/// A streaming chunk from a chat completion.
173#[derive(Debug, Clone, Serialize, Deserialize)]
174pub struct ChatCompletionChunk {
175    /// Unique identifier for the completion.
176    pub id: String,
177
178    /// Object type (always "chat.completion.chunk").
179    pub object: String,
180
181    /// Unix timestamp of creation.
182    pub created: i64,
183
184    /// Model used for the completion.
185    pub model: String,
186
187    /// List of streaming choices.
188    pub choices: Vec<StreamChoice>,
189}
190
191impl ChatCompletionChunk {
192    /// Get the content delta of the first choice.
193    pub fn content(&self) -> Option<&str> {
194        self.choices
195            .first()
196            .and_then(|c| c.delta.content.as_deref())
197    }
198}
199
200// --- Request types ---
201
202/// Request body for chat completions.
203#[derive(Debug, Clone, Serialize)]
204pub struct ChatCompletionRequest {
205    /// The model to use.
206    pub model: String,
207
208    /// The messages to send.
209    pub messages: Vec<Message>,
210
211    /// Sampling temperature (0.0 to 2.0).
212    #[serde(skip_serializing_if = "Option::is_none")]
213    pub temperature: Option<f32>,
214
215    /// Maximum tokens to generate.
216    #[serde(skip_serializing_if = "Option::is_none")]
217    pub max_tokens: Option<u32>,
218
219    /// Whether to stream the response.
220    #[serde(skip_serializing_if = "Option::is_none")]
221    pub stream: Option<bool>,
222
223    /// Tools available to the model.
224    #[serde(skip_serializing_if = "Option::is_none")]
225    pub tools: Option<Vec<super::tools::ToolDefinition>>,
226
227    /// Top-p sampling.
228    #[serde(skip_serializing_if = "Option::is_none")]
229    pub top_p: Option<f32>,
230
231    /// Stop sequences.
232    #[serde(skip_serializing_if = "Option::is_none")]
233    pub stop: Option<Vec<String>>,
234
235    /// Presence penalty (-2.0 to 2.0).
236    #[serde(skip_serializing_if = "Option::is_none")]
237    pub presence_penalty: Option<f32>,
238
239    /// Frequency penalty (-2.0 to 2.0).
240    #[serde(skip_serializing_if = "Option::is_none")]
241    pub frequency_penalty: Option<f32>,
242
243    /// User identifier for tracking.
244    #[serde(skip_serializing_if = "Option::is_none")]
245    pub user: Option<String>,
246
247    /// Per-request metadata serialised as the top-level `metadata` field
248    /// in the request body. The LiteLLM gateway forwards this
249    /// dict as OTel span attributes, which is the supported channel for
250    /// tagging gateway spans with caller context (session_id, user_eid,
251    /// purpose, etc.), confirmed via endpoint probe (`extra_body` is
252    /// rejected by Bedrock-routed models).
253    ///
254    /// The transport layer additionally merges
255    /// [`ForgeConfig::default_metadata`](crate::ForgeConfig::default_metadata)
256    /// into the same field; per-request keys win on collision.
257    #[serde(skip_serializing_if = "Option::is_none")]
258    pub metadata: Option<std::collections::HashMap<String, serde_json::Value>>,
259}
260
261impl ChatCompletionRequest {
262    /// Create a new request with required fields.
263    pub fn new(model: impl Into<String>, messages: Vec<Message>) -> Self {
264        Self {
265            model: model.into(),
266            messages,
267            temperature: None,
268            max_tokens: None,
269            stream: None,
270            tools: None,
271            top_p: None,
272            stop: None,
273            presence_penalty: None,
274            frequency_penalty: None,
275            user: None,
276            metadata: None,
277        }
278    }
279
280    /// Set per-request metadata. Merged with
281    /// [`ForgeConfig::default_metadata`](crate::ForgeConfig::default_metadata)
282    /// at transport time.
283    pub fn metadata(
284        mut self,
285        metadata: std::collections::HashMap<String, serde_json::Value>,
286    ) -> Self {
287        self.metadata = Some(metadata);
288        self
289    }
290
291    /// Set the temperature.
292    pub fn temperature(mut self, temperature: f32) -> Self {
293        self.temperature = Some(temperature);
294        self
295    }
296
297    /// Set the max tokens.
298    pub fn max_tokens(mut self, max_tokens: u32) -> Self {
299        self.max_tokens = Some(max_tokens);
300        self
301    }
302
303    /// Enable streaming.
304    pub fn stream(mut self, stream: bool) -> Self {
305        self.stream = Some(stream);
306        self
307    }
308
309    /// Set tools.
310    pub fn tools(mut self, tools: Vec<super::tools::ToolDefinition>) -> Self {
311        self.tools = Some(tools);
312        self
313    }
314}