llm-universal 0.2.0

use serde::{Deserialize, Serialize};
use serde_json::Value;

/// Author of a message.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    System,
    User,
    Assistant,
    /// Tool results. Sent as role `tool` to OpenAI and as `tool_result` blocks in a
    /// `user` turn to Anthropic.
    Tool,
}

/// One piece of a message.
///
/// Thinking blocks are not represented: if you enable extended thinking on Anthropic,
/// keep the raw response yourself so the signed blocks can be sent back unchanged.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum Part {
    Text {
        text: String,
    },
    ToolCall {
        id: String,
        name: String,
        /// Parsed arguments. If a provider returns arguments that are not valid JSON
        /// (for example cut off by `max_tokens`), the raw string is kept as `Value::String`.
        arguments: Value,
    },
    ToolResult {
        call_id: String,
        content: String,
        is_error: bool,
    },
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Message {
    pub role: Role,
    pub parts: Vec<Part>,
}

impl Message {
    pub fn new(role: Role, parts: Vec<Part>) -> Self {
        Self { role, parts }
    }

    pub fn user(text: impl Into<String>) -> Self {
        Self::text(Role::User, text)
    }

    pub fn assistant(text: impl Into<String>) -> Self {
        Self::text(Role::Assistant, text)
    }

    pub fn text(role: Role, text: impl Into<String>) -> Self {
        Self::new(role, vec![Part::Text { text: text.into() }])
    }

    /// A single tool result. To return several results from one step, put them in one
    /// message with `Message::new(Role::Tool, parts)`.
    pub fn tool_result(
        call_id: impl Into<String>,
        content: impl Into<String>,
        is_error: bool,
    ) -> Self {
        Self::new(
            Role::Tool,
            vec![Part::ToolResult {
                call_id: call_id.into(),
                content: content.into(),
                is_error,
            }],
        )
    }

    /// All text parts joined together.
    pub fn text_content(&self) -> String {
        self.parts
            .iter()
            .filter_map(|p| match p {
                Part::Text { text } => Some(text.as_str()),
                _ => None,
            })
            .collect::<Vec<_>>()
            .join("")
    }
}

/// A function the model may call. `parameters` is a JSON Schema with `"type": "object"`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Tool {
    pub name: String,
    pub description: String,
    pub parameters: Value,
}

impl Tool {
    pub fn new(name: impl Into<String>, description: impl Into<String>, parameters: Value) -> Self {
        Self {
            name: name.into(),
            description: description.into(),
            parameters,
        }
    }
}

/// Provider-neutral chat request.
#[derive(Debug, Clone, Default)]
pub struct ChatRequest {
    pub model: String,
    /// System prompt. Sent as a `system` message to OpenAI and as the top-level
    /// `system` field to Anthropic.
    pub system: Option<String>,
    pub messages: Vec<Message>,
    pub tools: Vec<Tool>,
    /// Required by Anthropic; [`DEFAULT_MAX_TOKENS`] is used when unset.
    pub max_tokens: Option<u32>,
    /// Anthropic accepts 0.0–1.0, and newer Anthropic models accept only 1.0.
    /// Leave unset unless the model supports it.
    pub temperature: Option<f32>,
}

/// `max_tokens` sent to Anthropic when the request does not set one.
pub const DEFAULT_MAX_TOKENS: u32 = 4096;

impl ChatRequest {
    pub fn new(model: impl Into<String>) -> Self {
        Self {
            model: model.into(),
            ..Default::default()
        }
    }

    pub fn system(mut self, system: impl Into<String>) -> Self {
        self.system = Some(system.into());
        self
    }

    pub fn message(mut self, message: Message) -> Self {
        self.messages.push(message);
        self
    }

    pub fn tool(mut self, tool: Tool) -> Self {
        self.tools.push(tool);
        self
    }

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

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

/// Why generation stopped, normalized across providers.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StopReason {
    /// Natural end of the turn (`stop` / `end_turn` / `stop_sequence`).
    EndTurn,
    /// The model wants tools to be called (`tool_calls` / `tool_use`).
    ToolUse,
    /// Token limit reached (`length` / `max_tokens` / `model_context_window_exceeded`).
    MaxTokens,
    /// A long-running turn was paused (`pause_turn`). Send the response back as-is to continue.
    Paused,
    /// The model declined for policy reasons (`refusal` / `content_filter`).
    Refusal,
    /// Any other value returned by the provider.
    Other(String),
}

#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Usage {
    pub input_tokens: u32,
    pub output_tokens: u32,
}

#[derive(Debug, Clone, PartialEq)]
pub struct ChatResponse {
    pub model: String,
    pub message: Message,
    pub stop_reason: StopReason,
    pub usage: Usage,
}

impl ChatResponse {
    /// All text of the response.
    pub fn text(&self) -> String {
        self.message.text_content()
    }

    /// Tool calls as `(id, name, arguments)`.
    pub fn tool_calls(&self) -> Vec<(String, String, Value)> {
        self.message
            .parts
            .iter()
            .filter_map(|p| match p {
                Part::ToolCall {
                    id,
                    name,
                    arguments,
                } => Some((id.clone(), name.clone(), arguments.clone())),
                _ => None,
            })
            .collect()
    }
}