af-llm 0.2.0

//! Strongly-typed OpenAI-compatible chat-completion request/response shapes.
//!
//! This is the Rust equivalent of the Pydantic models the Python stack relied
//! on: the wire format is validated at the type boundary via `serde`, so a
//! malformed provider response fails loudly at decode time rather than blowing
//! up three layers deep on a missing key.

use std::fmt;

use serde::{Deserialize, Deserializer, Serialize};

/// Role of a chat message.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    System,
    User,
    Assistant,
    Tool,
}

/// Provider-neutral structured assistant output. Provider adapters normalize
/// native annotations/content into this closed set at the response boundary.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)]
pub enum AssistantBlock {
    Text {
        text: String,
    },
    Resource {
        resource_id: String,
        media_type: String,
    },
    Data {
        slot: String,
        value: serde_json::Value,
    },
    Citation {
        resource_id: String,
        label: String,
        uri: String,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        excerpt: Option<String>,
    },
}

/// A single chat message, both for requests and for the assistant turn in a
/// response. `content` is optional because a tool-calling assistant turn may
/// carry only `tool_calls`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
    pub role: Role,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<String>,

    /// Present on assistant turns that invoke tools.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ToolCall>>,

    /// Set on a `role: tool` message to bind the result to its call.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,

    /// Optional name (tool name / participant name).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
}

impl ChatMessage {
    pub fn system(content: impl Into<String>) -> Self {
        Self::text(Role::System, content)
    }
    pub fn user(content: impl Into<String>) -> Self {
        Self::text(Role::User, content)
    }
    pub fn assistant(content: impl Into<String>) -> Self {
        Self::text(Role::Assistant, content)
    }

    fn text(role: Role, content: impl Into<String>) -> Self {
        Self {
            role,
            content: Some(content.into()),
            tool_calls: None,
            tool_call_id: None,
            name: None,
        }
    }
}

/// A tool the model is allowed to call. Only `function` tools exist today.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Tool {
    #[serde(rename = "type")]
    pub kind: String,
    pub function: FunctionDef,
}

impl Tool {
    /// Build a function tool. `parameters` is a JSON-Schema object.
    pub fn function(
        name: impl Into<String>,
        description: impl Into<String>,
        parameters: serde_json::Value,
    ) -> Self {
        Self {
            kind: "function".to_string(),
            function: FunctionDef {
                name: name.into(),
                description: Some(description.into()),
                parameters: Some(parameters),
            },
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionDef {
    pub name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    /// JSON-Schema for the arguments.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub parameters: Option<serde_json::Value>,
}

/// A tool invocation emitted by the model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    pub id: String,
    #[serde(rename = "type")]
    pub kind: String,
    pub function: FunctionCall,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FunctionCall {
    pub name: String,
    /// Raw JSON string of arguments — the provider does not pre-parse it.
    pub arguments: String,
}

/// How the model should choose tools. Defaults to `auto` when tools are present.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ToolChoice {
    Auto,
    None,
    Required,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ReasoningEffort {
    Low,
    Medium,
    High,
}

/// A chat-completion request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompletionRequest {
    pub model: String,
    pub messages: Vec<ChatMessage>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub tools: Option<Vec<Tool>>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_choice: Option<ToolChoice>,

    pub temperature: f32,
    pub max_tokens: u32,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub reasoning_effort: Option<ReasoningEffort>,

    /// Durable runtime identity used for provider idempotency and reconciliation.
    #[serde(skip)]
    pub provider_attempt_id: Option<String>,

    /// OpenAI-compatible streaming flag. Defaults to false and is omitted on
    /// the wire in that case.
    #[serde(skip_serializing_if = "std::ops::Not::not")]
    pub stream: bool,

    /// OpenAI-compatible streaming options. Providers only include the final
    /// usage frame when `include_usage` is requested explicitly.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream_options: Option<StreamOptions>,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct StreamOptions {
    pub include_usage: bool,
}

impl CompletionRequest {
    /// New request with the same defaults as the Python wrapper
    /// (`temperature = 0.3`, `max_tokens = 4096`).
    pub fn new(model: impl Into<String>, messages: Vec<ChatMessage>) -> Self {
        Self {
            model: model.into(),
            messages,
            tools: None,
            tool_choice: None,
            temperature: 0.3,
            max_tokens: 4096,
            reasoning_effort: None,
            provider_attempt_id: None,
            stream: false,
            stream_options: None,
        }
    }

    pub fn stream(mut self, enabled: bool) -> Self {
        self.stream = enabled;
        self
    }

    pub fn temperature(mut self, t: f32) -> Self {
        self.temperature = t;
        self
    }

    pub fn max_tokens(mut self, n: u32) -> Self {
        self.max_tokens = n;
        self
    }

    pub fn reasoning_effort(mut self, effort: ReasoningEffort) -> Self {
        self.reasoning_effort = Some(effort);
        self
    }

    /// Attach tools. Mirrors the Python default: if a caller adds tools without
    /// an explicit choice, we set `tool_choice = auto`.
    pub fn tools(mut self, tools: Vec<Tool>) -> Self {
        if !tools.is_empty() && self.tool_choice.is_none() {
            self.tool_choice = Some(ToolChoice::Auto);
        }
        self.tools = Some(tools);
        self
    }

    pub fn tool_choice(mut self, choice: ToolChoice) -> Self {
        self.tool_choice = Some(choice);
        self
    }
}

/// A chat-completion response.
#[derive(Debug, Clone, Deserialize)]
pub struct CompletionResponse {
    #[serde(default)]
    pub id: String,
    pub choices: Vec<Choice>,
    #[serde(default)]
    pub usage: Option<Usage>,
}

impl CompletionResponse {
    /// Convenience: text content of the first choice, if any.
    pub fn first_content(&self) -> Option<&str> {
        self.choices
            .first()
            .and_then(|c| c.message.content.as_deref())
    }

    /// Convenience: tool calls of the first choice, if any.
    pub fn first_tool_calls(&self) -> Option<&[ToolCall]> {
        self.choices
            .first()
            .and_then(|c| c.message.tool_calls.as_deref())
    }

    /// Finish reason of the first choice, if the provider supplied one.
    pub fn first_finish_reason(&self) -> Option<&FinishReason> {
        self.choices
            .first()
            .and_then(|choice| choice.finish_reason.as_ref())
    }
}

/// Why the provider stopped generating a completion.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum FinishReason {
    Stop,
    ToolCalls,
    Length,
    ContentFilter,
    Unknown(String),
}

impl FinishReason {
    pub fn as_str(&self) -> &str {
        match self {
            Self::Stop => "stop",
            Self::ToolCalls => "tool_calls",
            Self::Length => "length",
            Self::ContentFilter => "content_filter",
            Self::Unknown(reason) => reason,
        }
    }
}

impl fmt::Display for FinishReason {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(self.as_str())
    }
}

impl From<&str> for FinishReason {
    fn from(reason: &str) -> Self {
        match reason {
            "stop" => Self::Stop,
            "tool_calls" => Self::ToolCalls,
            "length" => Self::Length,
            "content_filter" => Self::ContentFilter,
            unknown => Self::Unknown(unknown.to_string()),
        }
    }
}

impl From<String> for FinishReason {
    fn from(reason: String) -> Self {
        Self::from(reason.as_str())
    }
}

impl<'de> Deserialize<'de> for FinishReason {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        String::deserialize(deserializer).map(Into::into)
    }
}

#[derive(Debug, Clone, Deserialize)]
pub struct Choice {
    #[serde(default)]
    pub index: u32,
    pub message: ChatMessage,
    #[serde(default)]
    pub finish_reason: Option<FinishReason>,
    #[serde(default, alias = "content_blocks")]
    pub output_blocks: Vec<AssistantBlock>,
}

#[derive(Debug, Clone, Copy, Default, Deserialize)]
pub struct Usage {
    #[serde(default)]
    pub prompt_tokens: u32,
    #[serde(default)]
    pub completion_tokens: u32,
    #[serde(default)]
    pub total_tokens: u32,
}