1use af_context::ToolCallId;
9use std::fmt;
10
11use serde::{Deserialize, Deserializer, Serialize};
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
15#[serde(rename_all = "lowercase")]
16pub enum Role {
17 System,
19 User,
21 Assistant,
23 Tool,
25}
26
27#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
30#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
31pub enum AssistantBlock {
32 Text {
34 text: String,
36 },
37 Resource {
39 resource_id: String,
41 media_type: String,
43 },
44 Data {
46 slot: String,
48 value: serde_json::Value,
50 },
51 Citation {
53 resource_id: String,
55 label: String,
57 uri: String,
59 #[serde(default, skip_serializing_if = "Option::is_none")]
61 excerpt: Option<String>,
62 },
63}
64
65#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
67pub struct InputImage {
68 pub asset_id: af_context::AssetId,
70 pub media_type: String,
72}
73
74impl InputImage {
75 pub fn validate(&self) -> crate::Result<()> {
77 let id = self.asset_id.as_str();
78 if id.len() > 256
79 || !id
80 .bytes()
81 .all(|c| c.is_ascii_alphanumeric() || b"-_.".contains(&c))
82 || !matches!(
83 self.media_type.as_str(),
84 "image/png" | "image/jpeg" | "image/webp"
85 )
86 {
87 return Err(crate::LlmError::InvalidInput(
88 "invalid image reference or MIME type".into(),
89 ));
90 }
91 Ok(())
92 }
93}
94
95#[derive(Debug, Clone, Serialize, Deserialize)]
99pub struct ChatMessage {
100 #[serde(default, skip_serializing_if = "Vec::is_empty")]
102 pub images: Vec<InputImage>,
103 pub role: Role,
105
106 #[serde(skip_serializing_if = "Option::is_none")]
108 pub content: Option<String>,
109
110 #[serde(skip_serializing_if = "Option::is_none")]
112 pub tool_calls: Option<Vec<ToolCall>>,
113
114 #[serde(skip_serializing_if = "Option::is_none")]
116 pub tool_call_id: Option<ToolCallId>,
117
118 #[serde(skip_serializing_if = "Option::is_none")]
120 pub name: Option<String>,
121}
122
123impl ChatMessage {
124 pub fn system(content: impl Into<String>) -> Self {
126 Self::text(Role::System, content)
127 }
128 pub fn user(content: impl Into<String>) -> Self {
130 Self::text(Role::User, content)
131 }
132 pub fn assistant(content: impl Into<String>) -> Self {
134 Self::text(Role::Assistant, content)
135 }
136
137 fn text(role: Role, content: impl Into<String>) -> Self {
138 Self {
139 images: Vec::new(),
140 role,
141 content: Some(content.into()),
142 tool_calls: None,
143 tool_call_id: None,
144 name: None,
145 }
146 }
147}
148
149#[derive(Debug, Clone, Serialize, Deserialize)]
151pub struct Tool {
152 #[serde(rename = "type")]
154 pub kind: String,
155 pub function: FunctionDef,
157}
158
159impl Tool {
160 pub fn function(
162 name: impl Into<String>,
163 description: impl Into<String>,
164 parameters: serde_json::Value,
165 ) -> Self {
166 Self {
167 kind: "function".to_string(),
168 function: FunctionDef {
169 name: name.into(),
170 description: Some(description.into()),
171 parameters: Some(parameters),
172 },
173 }
174 }
175}
176
177#[derive(Debug, Clone, Serialize, Deserialize)]
179pub struct FunctionDef {
180 pub name: String,
182 #[serde(skip_serializing_if = "Option::is_none")]
184 pub description: Option<String>,
185 #[serde(skip_serializing_if = "Option::is_none")]
187 pub parameters: Option<serde_json::Value>,
188}
189
190#[derive(Debug, Clone, Serialize, Deserialize)]
192pub struct ToolCall {
193 pub id: ToolCallId,
195 #[serde(rename = "type")]
197 pub kind: String,
198 pub function: FunctionCall,
200}
201
202#[derive(Debug, Clone, Serialize, Deserialize)]
204pub struct FunctionCall {
205 pub name: String,
207 pub arguments: String,
209}
210
211#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
213#[serde(rename_all = "lowercase")]
214pub enum ToolChoice {
215 Auto,
217 None,
219 Required,
221}
222
223#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
225#[serde(rename_all = "lowercase")]
226pub enum ReasoningEffort {
227 Low,
229 Medium,
231 High,
233}
234
235#[derive(Debug, Clone, Serialize, Deserialize)]
237pub struct CompletionRequest {
238 #[serde(skip)]
240 pub context: Option<af_context::RequestContext>,
241 pub model: String,
243 pub messages: Vec<ChatMessage>,
245
246 #[serde(skip_serializing_if = "Option::is_none")]
248 pub tools: Option<Vec<Tool>>,
249
250 #[serde(skip_serializing_if = "Option::is_none")]
252 pub tool_choice: Option<ToolChoice>,
253
254 pub temperature: f32,
256 pub max_tokens: u32,
258
259 #[serde(skip_serializing_if = "Option::is_none")]
261 pub reasoning_effort: Option<ReasoningEffort>,
262
263 #[serde(skip)]
265 pub provider_attempt_id: Option<String>,
266
267 #[serde(skip_serializing_if = "std::ops::Not::not")]
270 pub stream: bool,
271
272 #[serde(skip_serializing_if = "Option::is_none")]
275 pub stream_options: Option<StreamOptions>,
276}
277
278#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
280pub struct StreamOptions {
281 pub include_usage: bool,
283}
284
285impl CompletionRequest {
286 pub fn new(model: impl Into<String>, messages: Vec<ChatMessage>) -> Self {
289 Self {
290 context: None,
291 model: model.into(),
292 messages,
293 tools: None,
294 tool_choice: None,
295 temperature: 0.3,
296 max_tokens: 4096,
297 reasoning_effort: None,
298 provider_attempt_id: None,
299 stream: false,
300 stream_options: None,
301 }
302 }
303
304 pub fn stream(mut self, enabled: bool) -> Self {
306 self.stream = enabled;
307 self
308 }
309
310 pub fn temperature(mut self, t: f32) -> Self {
312 self.temperature = t;
313 self
314 }
315
316 pub fn max_tokens(mut self, n: u32) -> Self {
318 self.max_tokens = n;
319 self
320 }
321
322 pub fn reasoning_effort(mut self, effort: ReasoningEffort) -> Self {
324 self.reasoning_effort = Some(effort);
325 self
326 }
327
328 pub fn tools(mut self, tools: Vec<Tool>) -> Self {
331 if !tools.is_empty() && self.tool_choice.is_none() {
332 self.tool_choice = Some(ToolChoice::Auto);
333 }
334 self.tools = Some(tools);
335 self
336 }
337
338 pub fn tool_choice(mut self, choice: ToolChoice) -> Self {
340 self.tool_choice = Some(choice);
341 self
342 }
343}
344
345#[derive(Debug, Clone, Deserialize)]
347pub struct CompletionResponse {
348 #[serde(default)]
350 pub id: String,
351 pub choices: Vec<Choice>,
353 #[serde(default)]
355 pub usage: Option<Usage>,
356}
357
358impl CompletionResponse {
359 pub fn first_content(&self) -> Option<&str> {
361 self.choices
362 .first()
363 .and_then(|c| c.message.content.as_deref())
364 }
365
366 pub fn first_tool_calls(&self) -> Option<&[ToolCall]> {
368 self.choices
369 .first()
370 .and_then(|c| c.message.tool_calls.as_deref())
371 }
372
373 pub fn first_finish_reason(&self) -> Option<&FinishReason> {
375 self.choices
376 .first()
377 .and_then(|choice| choice.finish_reason.as_ref())
378 }
379}
380
381#[derive(Debug, Clone, PartialEq, Eq)]
383pub enum FinishReason {
384 Stop,
386 ToolCalls,
388 Length,
390 ContentFilter,
392 Unknown(String),
394}
395
396impl FinishReason {
397 pub fn as_str(&self) -> &str {
399 match self {
400 Self::Stop => "stop",
401 Self::ToolCalls => "tool_calls",
402 Self::Length => "length",
403 Self::ContentFilter => "content_filter",
404 Self::Unknown(reason) => reason,
405 }
406 }
407}
408
409impl fmt::Display for FinishReason {
410 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
411 formatter.write_str(self.as_str())
412 }
413}
414
415impl From<&str> for FinishReason {
416 fn from(reason: &str) -> Self {
417 match reason {
418 "stop" => Self::Stop,
419 "tool_calls" => Self::ToolCalls,
420 "length" => Self::Length,
421 "content_filter" => Self::ContentFilter,
422 unknown => Self::Unknown(unknown.to_string()),
423 }
424 }
425}
426
427impl From<String> for FinishReason {
428 fn from(reason: String) -> Self {
429 Self::from(reason.as_str())
430 }
431}
432
433impl<'de> Deserialize<'de> for FinishReason {
434 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
435 where
436 D: Deserializer<'de>,
437 {
438 String::deserialize(deserializer).map(Into::into)
439 }
440}
441
442#[derive(Debug, Clone, Deserialize)]
444pub struct Choice {
445 #[serde(default)]
447 pub index: u32,
448 pub message: ChatMessage,
450 #[serde(default)]
452 pub finish_reason: Option<FinishReason>,
453 #[serde(default, alias = "content_blocks")]
455 pub output_blocks: Vec<AssistantBlock>,
456}
457
458#[derive(Debug, Clone, Copy, Default, Deserialize)]
460pub struct Usage {
461 #[serde(default)]
463 pub prompt_tokens: u32,
464 #[serde(default)]
466 pub completion_tokens: u32,
467 #[serde(default)]
469 pub total_tokens: u32,
470}