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 std::fmt;
9
10use serde::{Deserialize, Deserializer, Serialize};
11
12/// Role of a chat message.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(rename_all = "lowercase")]
15pub enum Role {
16    System,
17    User,
18    Assistant,
19    Tool,
20}
21
22/// Provider-neutral structured assistant output. Provider adapters normalize
23/// native annotations/content into this closed set at the response boundary.
24#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
25#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
26pub enum AssistantBlock {
27    Text {
28        text: String,
29    },
30    Resource {
31        resource_id: String,
32        media_type: String,
33    },
34    Data {
35        slot: String,
36        value: serde_json::Value,
37    },
38    Citation {
39        resource_id: String,
40        label: String,
41        uri: String,
42        #[serde(default, skip_serializing_if = "Option::is_none")]
43        excerpt: Option<String>,
44    },
45}
46
47/// A single chat message, both for requests and for the assistant turn in a
48/// response. `content` is optional because a tool-calling assistant turn may
49/// carry only `tool_calls`.
50#[derive(Debug, Clone, Serialize, Deserialize)]
51pub struct ChatMessage {
52    pub role: Role,
53
54    #[serde(skip_serializing_if = "Option::is_none")]
55    pub content: Option<String>,
56
57    /// Present on assistant turns that invoke tools.
58    #[serde(skip_serializing_if = "Option::is_none")]
59    pub tool_calls: Option<Vec<ToolCall>>,
60
61    /// Set on a `role: tool` message to bind the result to its call.
62    #[serde(skip_serializing_if = "Option::is_none")]
63    pub tool_call_id: Option<String>,
64
65    /// Optional name (tool name / participant name).
66    #[serde(skip_serializing_if = "Option::is_none")]
67    pub name: Option<String>,
68}
69
70impl ChatMessage {
71    pub fn system(content: impl Into<String>) -> Self {
72        Self::text(Role::System, content)
73    }
74    pub fn user(content: impl Into<String>) -> Self {
75        Self::text(Role::User, content)
76    }
77    pub fn assistant(content: impl Into<String>) -> Self {
78        Self::text(Role::Assistant, content)
79    }
80
81    fn text(role: Role, content: impl Into<String>) -> Self {
82        Self {
83            role,
84            content: Some(content.into()),
85            tool_calls: None,
86            tool_call_id: None,
87            name: None,
88        }
89    }
90}
91
92/// A tool the model is allowed to call. Only `function` tools exist today.
93#[derive(Debug, Clone, Serialize, Deserialize)]
94pub struct Tool {
95    #[serde(rename = "type")]
96    pub kind: String,
97    pub function: FunctionDef,
98}
99
100impl Tool {
101    /// Build a function tool. `parameters` is a JSON-Schema object.
102    pub fn function(
103        name: impl Into<String>,
104        description: impl Into<String>,
105        parameters: serde_json::Value,
106    ) -> Self {
107        Self {
108            kind: "function".to_string(),
109            function: FunctionDef {
110                name: name.into(),
111                description: Some(description.into()),
112                parameters: Some(parameters),
113            },
114        }
115    }
116}
117
118#[derive(Debug, Clone, Serialize, Deserialize)]
119pub struct FunctionDef {
120    pub name: String,
121    #[serde(skip_serializing_if = "Option::is_none")]
122    pub description: Option<String>,
123    /// JSON-Schema for the arguments.
124    #[serde(skip_serializing_if = "Option::is_none")]
125    pub parameters: Option<serde_json::Value>,
126}
127
128/// A tool invocation emitted by the model.
129#[derive(Debug, Clone, Serialize, Deserialize)]
130pub struct ToolCall {
131    pub id: String,
132    #[serde(rename = "type")]
133    pub kind: String,
134    pub function: FunctionCall,
135}
136
137#[derive(Debug, Clone, Serialize, Deserialize)]
138pub struct FunctionCall {
139    pub name: String,
140    /// Raw JSON string of arguments — the provider does not pre-parse it.
141    pub arguments: String,
142}
143
144/// How the model should choose tools. Defaults to `auto` when tools are present.
145#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
146#[serde(rename_all = "lowercase")]
147pub enum ToolChoice {
148    Auto,
149    None,
150    Required,
151}
152
153#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
154#[serde(rename_all = "lowercase")]
155pub enum ReasoningEffort {
156    Low,
157    Medium,
158    High,
159}
160
161/// A chat-completion request.
162#[derive(Debug, Clone, Serialize, Deserialize)]
163pub struct CompletionRequest {
164    pub model: String,
165    pub messages: Vec<ChatMessage>,
166
167    #[serde(skip_serializing_if = "Option::is_none")]
168    pub tools: Option<Vec<Tool>>,
169
170    #[serde(skip_serializing_if = "Option::is_none")]
171    pub tool_choice: Option<ToolChoice>,
172
173    pub temperature: f32,
174    pub max_tokens: u32,
175
176    #[serde(skip_serializing_if = "Option::is_none")]
177    pub reasoning_effort: Option<ReasoningEffort>,
178
179    /// Durable runtime identity used for provider idempotency and reconciliation.
180    #[serde(skip)]
181    pub provider_attempt_id: Option<String>,
182
183    /// OpenAI-compatible streaming flag. Defaults to false and is omitted on
184    /// the wire in that case.
185    #[serde(skip_serializing_if = "std::ops::Not::not")]
186    pub stream: bool,
187
188    /// OpenAI-compatible streaming options. Providers only include the final
189    /// usage frame when `include_usage` is requested explicitly.
190    #[serde(skip_serializing_if = "Option::is_none")]
191    pub stream_options: Option<StreamOptions>,
192}
193
194#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
195pub struct StreamOptions {
196    pub include_usage: bool,
197}
198
199impl CompletionRequest {
200    /// New request with the same defaults as the Python wrapper
201    /// (`temperature = 0.3`, `max_tokens = 4096`).
202    pub fn new(model: impl Into<String>, messages: Vec<ChatMessage>) -> Self {
203        Self {
204            model: model.into(),
205            messages,
206            tools: None,
207            tool_choice: None,
208            temperature: 0.3,
209            max_tokens: 4096,
210            reasoning_effort: None,
211            provider_attempt_id: None,
212            stream: false,
213            stream_options: None,
214        }
215    }
216
217    pub fn stream(mut self, enabled: bool) -> Self {
218        self.stream = enabled;
219        self
220    }
221
222    pub fn temperature(mut self, t: f32) -> Self {
223        self.temperature = t;
224        self
225    }
226
227    pub fn max_tokens(mut self, n: u32) -> Self {
228        self.max_tokens = n;
229        self
230    }
231
232    pub fn reasoning_effort(mut self, effort: ReasoningEffort) -> Self {
233        self.reasoning_effort = Some(effort);
234        self
235    }
236
237    /// Attach tools. Mirrors the Python default: if a caller adds tools without
238    /// an explicit choice, we set `tool_choice = auto`.
239    pub fn tools(mut self, tools: Vec<Tool>) -> Self {
240        if !tools.is_empty() && self.tool_choice.is_none() {
241            self.tool_choice = Some(ToolChoice::Auto);
242        }
243        self.tools = Some(tools);
244        self
245    }
246
247    pub fn tool_choice(mut self, choice: ToolChoice) -> Self {
248        self.tool_choice = Some(choice);
249        self
250    }
251}
252
253/// A chat-completion response.
254#[derive(Debug, Clone, Deserialize)]
255pub struct CompletionResponse {
256    #[serde(default)]
257    pub id: String,
258    pub choices: Vec<Choice>,
259    #[serde(default)]
260    pub usage: Option<Usage>,
261}
262
263impl CompletionResponse {
264    /// Convenience: text content of the first choice, if any.
265    pub fn first_content(&self) -> Option<&str> {
266        self.choices
267            .first()
268            .and_then(|c| c.message.content.as_deref())
269    }
270
271    /// Convenience: tool calls of the first choice, if any.
272    pub fn first_tool_calls(&self) -> Option<&[ToolCall]> {
273        self.choices
274            .first()
275            .and_then(|c| c.message.tool_calls.as_deref())
276    }
277
278    /// Finish reason of the first choice, if the provider supplied one.
279    pub fn first_finish_reason(&self) -> Option<&FinishReason> {
280        self.choices
281            .first()
282            .and_then(|choice| choice.finish_reason.as_ref())
283    }
284}
285
286/// Why the provider stopped generating a completion.
287#[derive(Debug, Clone, PartialEq, Eq)]
288pub enum FinishReason {
289    Stop,
290    ToolCalls,
291    Length,
292    ContentFilter,
293    Unknown(String),
294}
295
296impl FinishReason {
297    pub fn as_str(&self) -> &str {
298        match self {
299            Self::Stop => "stop",
300            Self::ToolCalls => "tool_calls",
301            Self::Length => "length",
302            Self::ContentFilter => "content_filter",
303            Self::Unknown(reason) => reason,
304        }
305    }
306}
307
308impl fmt::Display for FinishReason {
309    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
310        formatter.write_str(self.as_str())
311    }
312}
313
314impl From<&str> for FinishReason {
315    fn from(reason: &str) -> Self {
316        match reason {
317            "stop" => Self::Stop,
318            "tool_calls" => Self::ToolCalls,
319            "length" => Self::Length,
320            "content_filter" => Self::ContentFilter,
321            unknown => Self::Unknown(unknown.to_string()),
322        }
323    }
324}
325
326impl From<String> for FinishReason {
327    fn from(reason: String) -> Self {
328        Self::from(reason.as_str())
329    }
330}
331
332impl<'de> Deserialize<'de> for FinishReason {
333    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
334    where
335        D: Deserializer<'de>,
336    {
337        String::deserialize(deserializer).map(Into::into)
338    }
339}
340
341#[derive(Debug, Clone, Deserialize)]
342pub struct Choice {
343    #[serde(default)]
344    pub index: u32,
345    pub message: ChatMessage,
346    #[serde(default)]
347    pub finish_reason: Option<FinishReason>,
348    #[serde(default, alias = "content_blocks")]
349    pub output_blocks: Vec<AssistantBlock>,
350}
351
352#[derive(Debug, Clone, Copy, Default, Deserialize)]
353pub struct Usage {
354    #[serde(default)]
355    pub prompt_tokens: u32,
356    #[serde(default)]
357    pub completion_tokens: u32,
358    #[serde(default)]
359    pub total_tokens: u32,
360}
af_llm/types.rs

af_llm/
types.rs
