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/// A single chat message, both for requests and for the assistant turn in a
66/// response. `content` is optional because a tool-calling assistant turn may
67/// carry only `tool_calls`.
68#[derive(Debug, Clone, Serialize, Deserialize)]
69pub struct ChatMessage {
70    /// Message role.
71    pub role: Role,
72
73    /// Content blocks carried by this record.
74    #[serde(skip_serializing_if = "Option::is_none")]
75    pub content: Option<String>,
76
77    /// Present on assistant turns that invoke tools.
78    #[serde(skip_serializing_if = "Option::is_none")]
79    pub tool_calls: Option<Vec<ToolCall>>,
80
81    /// Set on a `role: tool` message to bind the result to its call.
82    #[serde(skip_serializing_if = "Option::is_none")]
83    pub tool_call_id: Option<ToolCallId>,
84
85    /// Optional name (tool name / participant name).
86    #[serde(skip_serializing_if = "Option::is_none")]
87    pub name: Option<String>,
88}
89
90impl ChatMessage {
91    /// A `system` message.
92    pub fn system(content: impl Into<String>) -> Self {
93        Self::text(Role::System, content)
94    }
95    /// A `user` message.
96    pub fn user(content: impl Into<String>) -> Self {
97        Self::text(Role::User, content)
98    }
99    /// An `assistant` message without tool calls.
100    pub fn assistant(content: impl Into<String>) -> Self {
101        Self::text(Role::Assistant, content)
102    }
103
104    fn text(role: Role, content: impl Into<String>) -> Self {
105        Self {
106            role,
107            content: Some(content.into()),
108            tool_calls: None,
109            tool_call_id: None,
110            name: None,
111        }
112    }
113}
114
115/// A tool the model is allowed to call. Only `function` tools exist today.
116#[derive(Debug, Clone, Serialize, Deserialize)]
117pub struct Tool {
118    /// Discriminator naming the variant of this record.
119    #[serde(rename = "type")]
120    pub kind: String,
121    /// Function invoked by the tool call.
122    pub function: FunctionDef,
123}
124
125impl Tool {
126    /// Build a function tool. `parameters` is a JSON-Schema object.
127    pub fn function(
128        name: impl Into<String>,
129        description: impl Into<String>,
130        parameters: serde_json::Value,
131    ) -> Self {
132        Self {
133            kind: "function".to_string(),
134            function: FunctionDef {
135                name: name.into(),
136                description: Some(description.into()),
137                parameters: Some(parameters),
138            },
139        }
140    }
141}
142
143/// Function-calling definition advertised to the model.
144#[derive(Debug, Clone, Serialize, Deserialize)]
145pub struct FunctionDef {
146    /// Display name.
147    pub name: String,
148    /// Human-readable description.
149    #[serde(skip_serializing_if = "Option::is_none")]
150    pub description: Option<String>,
151    /// JSON-Schema for the arguments.
152    #[serde(skip_serializing_if = "Option::is_none")]
153    pub parameters: Option<serde_json::Value>,
154}
155
156/// A tool invocation emitted by the model.
157#[derive(Debug, Clone, Serialize, Deserialize)]
158pub struct ToolCall {
159    /// Stable identifier of this record.
160    pub id: ToolCallId,
161    /// Discriminator naming the variant of this record.
162    #[serde(rename = "type")]
163    pub kind: String,
164    /// Function invoked by the tool call.
165    pub function: FunctionCall,
166}
167
168/// Function name and JSON-encoded arguments of a tool call.
169#[derive(Debug, Clone, Serialize, Deserialize)]
170pub struct FunctionCall {
171    /// Display name.
172    pub name: String,
173    /// Raw JSON string of arguments — the provider does not pre-parse it.
174    pub arguments: String,
175}
176
177/// How the model should choose tools. Defaults to `auto` when tools are present.
178#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
179#[serde(rename_all = "lowercase")]
180pub enum ToolChoice {
181    /// The model decides whether to call a tool.
182    Auto,
183    /// The model must not call tools.
184    None,
185    /// The model must call at least one tool.
186    Required,
187}
188
189/// Provider-specific reasoning budget carried on the single request path.
190#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
191#[serde(rename_all = "lowercase")]
192pub enum ReasoningEffort {
193    /// Minimal reasoning.
194    Low,
195    /// Balanced reasoning.
196    Medium,
197    /// Maximum reasoning.
198    High,
199}
200
201/// A chat-completion request.
202#[derive(Debug, Clone, Serialize, Deserialize)]
203pub struct CompletionRequest {
204    /// Model identifier as registered in the model registry.
205    pub model: String,
206    /// Conversation messages in request order.
207    pub messages: Vec<ChatMessage>,
208
209    /// Tool names referenced by this record.
210    #[serde(skip_serializing_if = "Option::is_none")]
211    pub tools: Option<Vec<Tool>>,
212
213    /// How the model may use the advertised tools.
214    #[serde(skip_serializing_if = "Option::is_none")]
215    pub tool_choice: Option<ToolChoice>,
216
217    /// Sampling temperature.
218    pub temperature: f32,
219    /// Upper bound on prompt plus completion tokens.
220    pub max_tokens: u32,
221
222    /// Reasoning budget, when the provider supports it.
223    #[serde(skip_serializing_if = "Option::is_none")]
224    pub reasoning_effort: Option<ReasoningEffort>,
225
226    /// Durable runtime identity used for provider idempotency and reconciliation.
227    #[serde(skip)]
228    pub provider_attempt_id: Option<String>,
229
230    /// OpenAI-compatible streaming flag. Defaults to false and is omitted on
231    /// the wire in that case.
232    #[serde(skip_serializing_if = "std::ops::Not::not")]
233    pub stream: bool,
234
235    /// OpenAI-compatible streaming options. Providers only include the final
236    /// usage frame when `include_usage` is requested explicitly.
237    #[serde(skip_serializing_if = "Option::is_none")]
238    pub stream_options: Option<StreamOptions>,
239}
240
241/// Streaming options.
242#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
243pub struct StreamOptions {
244    /// Ask the provider to include a final usage chunk.
245    pub include_usage: bool,
246}
247
248impl CompletionRequest {
249    /// New request with the same defaults as the Python wrapper
250    /// (`temperature = 0.3`, `max_tokens = 4096`).
251    pub fn new(model: impl Into<String>, messages: Vec<ChatMessage>) -> Self {
252        Self {
253            model: model.into(),
254            messages,
255            tools: None,
256            tool_choice: None,
257            temperature: 0.3,
258            max_tokens: 4096,
259            reasoning_effort: None,
260            provider_attempt_id: None,
261            stream: false,
262            stream_options: None,
263        }
264    }
265
266    /// Enable or disable streaming.
267    pub fn stream(mut self, enabled: bool) -> Self {
268        self.stream = enabled;
269        self
270    }
271
272    /// Set the sampling temperature.
273    pub fn temperature(mut self, t: f32) -> Self {
274        self.temperature = t;
275        self
276    }
277
278    /// Set the maximum completion tokens.
279    pub fn max_tokens(mut self, n: u32) -> Self {
280        self.max_tokens = n;
281        self
282    }
283
284    /// Set the reasoning budget.
285    pub fn reasoning_effort(mut self, effort: ReasoningEffort) -> Self {
286        self.reasoning_effort = Some(effort);
287        self
288    }
289
290    /// Attach tools. Mirrors the Python default: if a caller adds tools without
291    /// an explicit choice, we set `tool_choice = auto`.
292    pub fn tools(mut self, tools: Vec<Tool>) -> Self {
293        if !tools.is_empty() && self.tool_choice.is_none() {
294            self.tool_choice = Some(ToolChoice::Auto);
295        }
296        self.tools = Some(tools);
297        self
298    }
299
300    /// Set the tool-choice policy.
301    pub fn tool_choice(mut self, choice: ToolChoice) -> Self {
302        self.tool_choice = Some(choice);
303        self
304    }
305}
306
307/// A chat-completion response.
308#[derive(Debug, Clone, Deserialize)]
309pub struct CompletionResponse {
310    /// Stable identifier of this record.
311    #[serde(default)]
312    pub id: String,
313    /// Completion choices returned by the provider.
314    pub choices: Vec<Choice>,
315    /// Token usage reported by the provider.
316    #[serde(default)]
317    pub usage: Option<Usage>,
318}
319
320impl CompletionResponse {
321    /// Convenience: text content of the first choice, if any.
322    pub fn first_content(&self) -> Option<&str> {
323        self.choices
324            .first()
325            .and_then(|c| c.message.content.as_deref())
326    }
327
328    /// Convenience: tool calls of the first choice, if any.
329    pub fn first_tool_calls(&self) -> Option<&[ToolCall]> {
330        self.choices
331            .first()
332            .and_then(|c| c.message.tool_calls.as_deref())
333    }
334
335    /// Finish reason of the first choice, if the provider supplied one.
336    pub fn first_finish_reason(&self) -> Option<&FinishReason> {
337        self.choices
338            .first()
339            .and_then(|choice| choice.finish_reason.as_ref())
340    }
341}
342
343/// Why the provider stopped generating a completion.
344#[derive(Debug, Clone, PartialEq, Eq)]
345pub enum FinishReason {
346    /// The model finished its answer.
347    Stop,
348    /// The model requested tool calls.
349    ToolCalls,
350    /// Output hit the token limit.
351    Length,
352    /// The provider filtered the output.
353    ContentFilter,
354    /// A reason this crate does not model.
355    Unknown(String),
356}
357
358impl FinishReason {
359    /// Provider wire name.
360    pub fn as_str(&self) -> &str {
361        match self {
362            Self::Stop => "stop",
363            Self::ToolCalls => "tool_calls",
364            Self::Length => "length",
365            Self::ContentFilter => "content_filter",
366            Self::Unknown(reason) => reason,
367        }
368    }
369}
370
371impl fmt::Display for FinishReason {
372    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
373        formatter.write_str(self.as_str())
374    }
375}
376
377impl From<&str> for FinishReason {
378    fn from(reason: &str) -> Self {
379        match reason {
380            "stop" => Self::Stop,
381            "tool_calls" => Self::ToolCalls,
382            "length" => Self::Length,
383            "content_filter" => Self::ContentFilter,
384            unknown => Self::Unknown(unknown.to_string()),
385        }
386    }
387}
388
389impl From<String> for FinishReason {
390    fn from(reason: String) -> Self {
391        Self::from(reason.as_str())
392    }
393}
394
395impl<'de> Deserialize<'de> for FinishReason {
396    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
397    where
398        D: Deserializer<'de>,
399    {
400        String::deserialize(deserializer).map(Into::into)
401    }
402}
403
404/// One completion candidate.
405#[derive(Debug, Clone, Deserialize)]
406pub struct Choice {
407    /// Zero-based position.
408    #[serde(default)]
409    pub index: u32,
410    /// Human-readable message.
411    pub message: ChatMessage,
412    /// Why the provider stopped generating.
413    #[serde(default)]
414    pub finish_reason: Option<FinishReason>,
415    /// Structured assistant blocks when the provider returns them instead of prose.
416    #[serde(default, alias = "content_blocks")]
417    pub output_blocks: Vec<AssistantBlock>,
418}
419
420/// Token accounting reported by the provider.
421#[derive(Debug, Clone, Copy, Default, Deserialize)]
422pub struct Usage {
423    /// Prompt tokens consumed.
424    #[serde(default)]
425    pub prompt_tokens: u32,
426    /// Completion tokens produced.
427    #[serde(default)]
428    pub completion_tokens: u32,
429    /// Prompt plus completion tokens.
430    #[serde(default)]
431    pub total_tokens: u32,
432}
af_llm/types.rs

af_llm/
types.rs
