Skip to main content

af_llm/
types.rs

1//! Strongly-typed OpenAI-compatible chat-completion request/response shapes.
2//!
3//! This is the Rust equivalent of the Pydantic models the Python stack relied
4//! on: the wire format is validated at the type boundary via `serde`, so a
5//! malformed provider response fails loudly at decode time rather than blowing
6//! up three layers deep on a missing key.
7
8use af_context::ToolCallId;
9use std::fmt;
10
11use serde::{Deserialize, Deserializer, Serialize};
12
13/// Role of a chat message.
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
15#[serde(rename_all = "lowercase")]
16pub enum Role {
17    /// Instructions from the product.
18    System,
19    /// Input from the user.
20    User,
21    /// Model output.
22    Assistant,
23    /// Result of a tool call.
24    Tool,
25}
26
27/// Provider-neutral structured assistant output. Provider adapters normalize
28/// native annotations/content into this closed set at the response boundary.
29#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
30#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
31pub enum AssistantBlock {
32    /// Plain text.
33    Text {
34        /// The text.
35        text: String,
36    },
37    /// Reference to an attached resource resolved by the host.
38    Resource {
39        /// Host-resolvable resource identity.
40        resource_id: String,
41        /// MIME type of the resource.
42        media_type: String,
43    },
44    /// Product-typed structured data rendered by a registered UI slot.
45    Data {
46        /// Slot name the product registered for this data.
47        slot: String,
48        /// Slot payload.
49        value: serde_json::Value,
50    },
51    /// Citation of a retrieved resource.
52    Citation {
53        /// Cited resource identity.
54        resource_id: String,
55        /// Display label.
56        label: String,
57        /// Where the resource can be opened.
58        uri: String,
59        /// Quoted excerpt supporting the citation.
60        #[serde(default, skip_serializing_if = "Option::is_none")]
61        excerpt: Option<String>,
62    },
63}
64
65/// Durable image reference. Provider URLs are resolved only at the HTTP boundary.
66#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
67pub struct InputImage {
68    /// Stable asset identifier, never a URL or inline bytes.
69    pub asset_id: af_context::AssetId,
70    /// Declared image MIME type, verified by the host resource adapter.
71    pub media_type: String,
72}
73
74impl InputImage {
75    /// Reject inline data, URLs and unsupported media before persisting an image reference.
76    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/// A single chat message, both for requests and for the assistant turn in a
96/// response. `content` is optional because a tool-calling assistant turn may
97/// carry only `tool_calls`.
98#[derive(Debug, Clone, Serialize, Deserialize)]
99pub struct ChatMessage {
100    /// Stable images supplied to this message; empty for text-only input.
101    #[serde(default, skip_serializing_if = "Vec::is_empty")]
102    pub images: Vec<InputImage>,
103    /// Message role.
104    pub role: Role,
105
106    /// Content blocks carried by this record.
107    #[serde(skip_serializing_if = "Option::is_none")]
108    pub content: Option<String>,
109
110    /// Present on assistant turns that invoke tools.
111    #[serde(skip_serializing_if = "Option::is_none")]
112    pub tool_calls: Option<Vec<ToolCall>>,
113
114    /// Set on a `role: tool` message to bind the result to its call.
115    #[serde(skip_serializing_if = "Option::is_none")]
116    pub tool_call_id: Option<ToolCallId>,
117
118    /// Optional name (tool name / participant name).
119    #[serde(skip_serializing_if = "Option::is_none")]
120    pub name: Option<String>,
121}
122
123impl ChatMessage {
124    /// A `system` message.
125    pub fn system(content: impl Into<String>) -> Self {
126        Self::text(Role::System, content)
127    }
128    /// A `user` message.
129    pub fn user(content: impl Into<String>) -> Self {
130        Self::text(Role::User, content)
131    }
132    /// An `assistant` message without tool calls.
133    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/// A tool the model is allowed to call. Only `function` tools exist today.
150#[derive(Debug, Clone, Serialize, Deserialize)]
151pub struct Tool {
152    /// Discriminator naming the variant of this record.
153    #[serde(rename = "type")]
154    pub kind: String,
155    /// Function invoked by the tool call.
156    pub function: FunctionDef,
157}
158
159impl Tool {
160    /// Build a function tool. `parameters` is a JSON-Schema object.
161    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/// Function-calling definition advertised to the model.
178#[derive(Debug, Clone, Serialize, Deserialize)]
179pub struct FunctionDef {
180    /// Display name.
181    pub name: String,
182    /// Human-readable description.
183    #[serde(skip_serializing_if = "Option::is_none")]
184    pub description: Option<String>,
185    /// JSON-Schema for the arguments.
186    #[serde(skip_serializing_if = "Option::is_none")]
187    pub parameters: Option<serde_json::Value>,
188}
189
190/// A tool invocation emitted by the model.
191#[derive(Debug, Clone, Serialize, Deserialize)]
192pub struct ToolCall {
193    /// Stable identifier of this record.
194    pub id: ToolCallId,
195    /// Discriminator naming the variant of this record.
196    #[serde(rename = "type")]
197    pub kind: String,
198    /// Function invoked by the tool call.
199    pub function: FunctionCall,
200}
201
202/// Function name and JSON-encoded arguments of a tool call.
203#[derive(Debug, Clone, Serialize, Deserialize)]
204pub struct FunctionCall {
205    /// Display name.
206    pub name: String,
207    /// Raw JSON string of arguments — the provider does not pre-parse it.
208    pub arguments: String,
209}
210
211/// How the model should choose tools. Defaults to `auto` when tools are present.
212#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
213#[serde(rename_all = "lowercase")]
214pub enum ToolChoice {
215    /// The model decides whether to call a tool.
216    Auto,
217    /// The model must not call tools.
218    None,
219    /// The model must call at least one tool.
220    Required,
221}
222
223/// Provider-specific reasoning budget carried on the single request path.
224#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
225#[serde(rename_all = "lowercase")]
226pub enum ReasoningEffort {
227    /// Minimal reasoning.
228    Low,
229    /// Balanced reasoning.
230    Medium,
231    /// Maximum reasoning.
232    High,
233}
234
235/// A chat-completion request.
236#[derive(Debug, Clone, Serialize, Deserialize)]
237pub struct CompletionRequest {
238    /// Fresh caller identity for resource authorization; never serialized.
239    #[serde(skip)]
240    pub context: Option<af_context::RequestContext>,
241    /// Model identifier as registered in the model registry.
242    pub model: String,
243    /// Conversation messages in request order.
244    pub messages: Vec<ChatMessage>,
245
246    /// Tool names referenced by this record.
247    #[serde(skip_serializing_if = "Option::is_none")]
248    pub tools: Option<Vec<Tool>>,
249
250    /// How the model may use the advertised tools.
251    #[serde(skip_serializing_if = "Option::is_none")]
252    pub tool_choice: Option<ToolChoice>,
253
254    /// Sampling temperature.
255    pub temperature: f32,
256    /// Upper bound on prompt plus completion tokens.
257    pub max_tokens: u32,
258
259    /// Reasoning budget, when the provider supports it.
260    #[serde(skip_serializing_if = "Option::is_none")]
261    pub reasoning_effort: Option<ReasoningEffort>,
262
263    /// Durable runtime identity used for provider idempotency and reconciliation.
264    #[serde(skip)]
265    pub provider_attempt_id: Option<String>,
266
267    /// OpenAI-compatible streaming flag. Defaults to false and is omitted on
268    /// the wire in that case.
269    #[serde(skip_serializing_if = "std::ops::Not::not")]
270    pub stream: bool,
271
272    /// OpenAI-compatible streaming options. Providers only include the final
273    /// usage frame when `include_usage` is requested explicitly.
274    #[serde(skip_serializing_if = "Option::is_none")]
275    pub stream_options: Option<StreamOptions>,
276}
277
278/// Streaming options.
279#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
280pub struct StreamOptions {
281    /// Ask the provider to include a final usage chunk.
282    pub include_usage: bool,
283}
284
285impl CompletionRequest {
286    /// New request with the same defaults as the Python wrapper
287    /// (`temperature = 0.3`, `max_tokens = 4096`).
288    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    /// Enable or disable streaming.
305    pub fn stream(mut self, enabled: bool) -> Self {
306        self.stream = enabled;
307        self
308    }
309
310    /// Set the sampling temperature.
311    pub fn temperature(mut self, t: f32) -> Self {
312        self.temperature = t;
313        self
314    }
315
316    /// Set the maximum completion tokens.
317    pub fn max_tokens(mut self, n: u32) -> Self {
318        self.max_tokens = n;
319        self
320    }
321
322    /// Set the reasoning budget.
323    pub fn reasoning_effort(mut self, effort: ReasoningEffort) -> Self {
324        self.reasoning_effort = Some(effort);
325        self
326    }
327
328    /// Attach tools. Mirrors the Python default: if a caller adds tools without
329    /// an explicit choice, we set `tool_choice = auto`.
330    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    /// Set the tool-choice policy.
339    pub fn tool_choice(mut self, choice: ToolChoice) -> Self {
340        self.tool_choice = Some(choice);
341        self
342    }
343}
344
345/// A chat-completion response.
346#[derive(Debug, Clone, Deserialize)]
347pub struct CompletionResponse {
348    /// Stable identifier of this record.
349    #[serde(default)]
350    pub id: String,
351    /// Completion choices returned by the provider.
352    pub choices: Vec<Choice>,
353    /// Token usage reported by the provider.
354    #[serde(default)]
355    pub usage: Option<Usage>,
356}
357
358impl CompletionResponse {
359    /// Convenience: text content of the first choice, if any.
360    pub fn first_content(&self) -> Option<&str> {
361        self.choices
362            .first()
363            .and_then(|c| c.message.content.as_deref())
364    }
365
366    /// Convenience: tool calls of the first choice, if any.
367    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    /// Finish reason of the first choice, if the provider supplied one.
374    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/// Why the provider stopped generating a completion.
382#[derive(Debug, Clone, PartialEq, Eq)]
383pub enum FinishReason {
384    /// The model finished its answer.
385    Stop,
386    /// The model requested tool calls.
387    ToolCalls,
388    /// Output hit the token limit.
389    Length,
390    /// The provider filtered the output.
391    ContentFilter,
392    /// A reason this crate does not model.
393    Unknown(String),
394}
395
396impl FinishReason {
397    /// Provider wire name.
398    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/// One completion candidate.
443#[derive(Debug, Clone, Deserialize)]
444pub struct Choice {
445    /// Zero-based position.
446    #[serde(default)]
447    pub index: u32,
448    /// Human-readable message.
449    pub message: ChatMessage,
450    /// Why the provider stopped generating.
451    #[serde(default)]
452    pub finish_reason: Option<FinishReason>,
453    /// Structured assistant blocks when the provider returns them instead of prose.
454    #[serde(default, alias = "content_blocks")]
455    pub output_blocks: Vec<AssistantBlock>,
456}
457
458/// Token accounting reported by the provider.
459#[derive(Debug, Clone, Copy, Default, Deserialize)]
460pub struct Usage {
461    /// Prompt tokens consumed.
462    #[serde(default)]
463    pub prompt_tokens: u32,
464    /// Completion tokens produced.
465    #[serde(default)]
466    pub completion_tokens: u32,
467    /// Prompt plus completion tokens.
468    #[serde(default)]
469    pub total_tokens: u32,
470}