symbi-runtime 1.21.1

//! Unified inference provider trait
//!
//! Defines the `InferenceProvider` trait that abstracts over cloud LLM APIs
//! and local SLM runners, adding tool calling and structured output support
//! on top of the existing `LlmClient` and `SlmRunner`.

use crate::reasoning::conversation::Conversation;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// A tool definition that can be provided to an inference call.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDefinition {
    /// Tool name (must match the name the LLM will use to call it).
    pub name: String,
    /// Human-readable description of what the tool does.
    pub description: String,
    /// JSON Schema describing the tool's parameters.
    pub parameters: serde_json::Value,
}

/// A tool call request returned by the model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCallRequest {
    /// Unique identifier for this tool call.
    pub id: String,
    /// Name of the tool to invoke.
    pub name: String,
    /// JSON-encoded arguments for the tool.
    pub arguments: String,
}

/// The reason the model stopped generating.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum FinishReason {
    /// Model produced a complete response.
    Stop,
    /// Model wants to call one or more tools.
    ToolCalls,
    /// Generation was truncated due to max_tokens.
    MaxTokens,
    /// Generation was truncated due to content filter.
    ContentFilter,
    /// The model declined the request (e.g. Anthropic `stop_reason: "refusal"`).
    /// Distinct from [`FinishReason::Stop`] so callers can retry or fail over
    /// instead of reading a refused turn as an empty, successful completion.
    Refusal,
}

/// Desired response format from the model.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum ResponseFormat {
    /// Free-form text response.
    #[serde(rename = "text")]
    Text,
    /// JSON object response (model is instructed to return valid JSON).
    #[serde(rename = "json_object")]
    JsonObject,
    /// JSON response conforming to a specific schema.
    #[serde(rename = "json_schema")]
    JsonSchema {
        /// The JSON schema the response must conform to.
        schema: serde_json::Value,
        /// Optional name for the schema (used in API calls).
        #[serde(default, skip_serializing_if = "Option::is_none")]
        name: Option<String>,
    },
}

/// Token usage information.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Usage {
    /// Tokens in the prompt/input.
    pub prompt_tokens: u32,
    /// Tokens in the completion/output.
    pub completion_tokens: u32,
    /// Total tokens used.
    pub total_tokens: u32,
}

/// Controls whether the model is required to call a tool on this turn.
///
/// Default is `Auto`, which preserves the historic behavior of letting
/// the model decide whether to respond with a tool_use block or with
/// plain text. Callers that drive an iterate-until-done workflow
/// (research agents, code-execution agents, multi-step planners) should
/// set `Any` so the model is forced into tool_use mode and the loop
/// terminates on a sentinel tool rather than on a stray text response.
///
/// `Tool { name }` constrains the model to a specific named tool; useful
/// when the next action is unambiguous (e.g., always start by calling
/// `read_threat_model`).
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ToolChoice {
    /// Model decides whether to call a tool (provider default).
    #[default]
    Auto,
    /// Model MUST call some tool — any one of the available definitions.
    Any,
    /// Model MUST call this specific named tool.
    Tool { name: String },
}

/// Options for an inference call.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InferenceOptions {
    /// Maximum tokens to generate.
    #[serde(default = "default_max_tokens")]
    pub max_tokens: u32,
    /// Sampling temperature (0.0 = deterministic, 1.0 = creative).
    #[serde(default = "default_temperature")]
    pub temperature: f32,
    /// Tool definitions available for this call.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub tool_definitions: Vec<ToolDefinition>,
    /// Constraint on whether the model must call a tool this turn.
    /// `None` defers to the provider default (Anthropic + OpenAI both
    /// default to "auto"). Set to `Some(Any)` to force tool_use.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tool_choice: Option<ToolChoice>,
    /// Desired response format.
    #[serde(default = "default_response_format")]
    pub response_format: ResponseFormat,
    /// Optional model override (provider decides default otherwise).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub model: Option<String>,
    /// Additional provider-specific parameters.
    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
    pub extra: HashMap<String, serde_json::Value>,
}

fn default_max_tokens() -> u32 {
    4096
}

fn default_temperature() -> f32 {
    0.3
}

fn default_response_format() -> ResponseFormat {
    ResponseFormat::Text
}

impl Default for InferenceOptions {
    fn default() -> Self {
        Self {
            max_tokens: default_max_tokens(),
            temperature: default_temperature(),
            tool_definitions: Vec::new(),
            tool_choice: None,
            response_format: ResponseFormat::Text,
            model: None,
            extra: HashMap::new(),
        }
    }
}

/// Response from an inference call.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InferenceResponse {
    /// Text content of the response.
    pub content: String,
    /// Tool calls requested by the model (empty if none).
    pub tool_calls: Vec<ToolCallRequest>,
    /// Why the model stopped generating.
    pub finish_reason: FinishReason,
    /// Token usage statistics.
    pub usage: Usage,
    /// The model ID that actually served the request.
    pub model: String,
}

impl InferenceResponse {
    /// Returns true if the model requested tool calls.
    pub fn has_tool_calls(&self) -> bool {
        !self.tool_calls.is_empty()
    }
}

/// Errors that can occur during inference.
#[derive(Debug, thiserror::Error)]
pub enum InferenceError {
    #[error("Provider error: {0}")]
    Provider(String),

    #[error("Rate limited, retry after {retry_after_ms}ms")]
    RateLimited { retry_after_ms: u64 },

    #[error("Context window exceeded: {0} tokens requested, {1} available")]
    ContextOverflow(usize, usize),

    #[error("Model not available: {0}")]
    ModelUnavailable(String),

    #[error("Invalid request: {0}")]
    InvalidRequest(String),

    #[error("Timeout after {0:?}")]
    Timeout(std::time::Duration),

    #[error("Response parse error: {0}")]
    ParseError(String),
}

/// Unified trait for inference providers (cloud LLMs and local SLMs).
///
/// Wraps existing `LlmClient` and `SlmRunner` to add:
/// - Multi-turn conversation support
/// - Tool calling
/// - Structured output (response_format)
/// - Token usage tracking
#[async_trait]
pub trait InferenceProvider: Send + Sync {
    /// Stable, non-secret configuration identity for opt-in evaluated releases.
    /// Unimplemented providers remain usable for ordinary runs.
    fn configuration_identity(&self) -> Option<String> {
        None
    }
    /// Input allowance reserved before dispatch, including tools and framing.
    /// The default conservatively budgets serialized UTF-8 bytes plus framing
    /// overhead for text providers. It is not an exact tokenizer or billing
    /// guarantee: providers with additional hidden input must override it.
    /// Reported usage above the reservation closes the shared budget before
    /// response-driven actions. Implementations may supply an exact count.
    fn input_token_reservation(
        &self,
        conversation: &Conversation,
        options: &InferenceOptions,
    ) -> Result<u32, InferenceError> {
        let bytes = serde_json::to_vec(&(conversation, options))
            .map_err(|error| InferenceError::InvalidRequest(error.to_string()))?;
        input_reservation_from_bytes(bytes.len())
    }

    /// Run inference on a conversation with the given options.
    async fn complete(
        &self,
        conversation: &Conversation,
        options: &InferenceOptions,
    ) -> Result<InferenceResponse, InferenceError>;

    /// Get the provider's name for logging and routing.
    fn provider_name(&self) -> &str;

    /// Get the default model ID for this provider.
    fn default_model(&self) -> &str;

    /// Check if this provider supports tool calling natively.
    fn supports_native_tools(&self) -> bool;

    /// Check if this provider supports structured output natively.
    fn supports_structured_output(&self) -> bool;
}

pub(crate) fn input_reservation_from_bytes(bytes: usize) -> Result<u32, InferenceError> {
    bytes
        .checked_add(1024)
        .and_then(|value| u32::try_from(value).ok())
        .ok_or_else(|| {
            InferenceError::InvalidRequest("input reservation exceeds token range".into())
        })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_inference_options_default() {
        let opts = InferenceOptions::default();
        assert_eq!(opts.max_tokens, 4096);
        assert!((opts.temperature - 0.3).abs() < f32::EPSILON);
        assert!(opts.tool_definitions.is_empty());
        assert!(matches!(opts.response_format, ResponseFormat::Text));
    }

    #[test]
    fn test_tool_definition_serde() {
        let tool = ToolDefinition {
            name: "web_search".into(),
            description: "Search the web".into(),
            parameters: serde_json::json!({
                "type": "object",
                "properties": {
                    "query": { "type": "string" }
                },
                "required": ["query"]
            }),
        };
        let json = serde_json::to_string(&tool).unwrap();
        let restored: ToolDefinition = serde_json::from_str(&json).unwrap();
        assert_eq!(restored.name, "web_search");
    }

    #[test]
    fn test_response_format_serde() {
        let text = ResponseFormat::Text;
        let json = serde_json::to_string(&text).unwrap();
        assert!(json.contains("text"));

        let schema = ResponseFormat::JsonSchema {
            schema: serde_json::json!({"type": "object"}),
            name: Some("MySchema".into()),
        };
        let json = serde_json::to_string(&schema).unwrap();
        assert!(json.contains("json_schema"));
        assert!(json.contains("MySchema"));
    }

    #[test]
    fn test_inference_response_has_tool_calls() {
        let resp = InferenceResponse {
            content: String::new(),
            tool_calls: vec![ToolCallRequest {
                id: "tc_1".into(),
                name: "search".into(),
                arguments: "{}".into(),
            }],
            finish_reason: FinishReason::ToolCalls,
            usage: Usage::default(),
            model: "test".into(),
        };
        assert!(resp.has_tool_calls());

        let resp_no_tools = InferenceResponse {
            content: "Hello".into(),
            tool_calls: vec![],
            finish_reason: FinishReason::Stop,
            usage: Usage::default(),
            model: "test".into(),
        };
        assert!(!resp_no_tools.has_tool_calls());
    }

    #[test]
    fn test_finish_reason_serde() {
        let json = serde_json::to_string(&FinishReason::ToolCalls).unwrap();
        assert_eq!(json, "\"tool_calls\"");
        let restored: FinishReason = serde_json::from_str(&json).unwrap();
        assert_eq!(restored, FinishReason::ToolCalls);
    }
}