ic-rig 0.1.0

//! Completion model trait and associated request/response types.
//!
//! [`CompletionModel`] is the single interface that provider implementations
//! must satisfy. The [`Agent`](crate::agent::Agent) and any other high-level
//! construct in this crate are generic over it.
//!
//! # Implementing a provider
//!
//! ```rust,ignore
//! use irig::completion::{
//!     CompletionModel, CompletionRequest, CompletionResponse, CompletionError,
//! };
//!
//! pub struct MyModel<H> {
//!     client: H,
//!     model: String,
//!     api_key: String,
//! }
//!
//! impl<H: HttpClient> CompletionModel for MyModel<H> {
//!     type Error = CompletionError;
//!
//!     async fn complete(
//!         &self,
//!         request: CompletionRequest,
//!     ) -> Result<CompletionResponse, Self::Error> {
//!         // 1. Serialise `request` into the provider's API format.
//!         // 2. Call `self.client.post(...)`.
//!         // 3. Deserialise the response into `CompletionResponse`.
//!         todo!()
//!     }
//! }
//! ```

use crate::message::{Message, ToolCall};
use crate::tool::ToolDefinition;
use thiserror::Error;

// ── Request ───────────────────────────────────────────────────────────────────

/// Everything a completion model needs to generate a response.
#[derive(Debug, Clone)]
pub struct CompletionRequest {
    /// Conversation history, including the latest user turn.
    pub messages: Vec<Message>,
    /// Tools the model is allowed to call.
    pub tools: Vec<ToolDefinition>,
    /// Optional sampling temperature (provider-specific range).
    pub temperature: Option<f64>,
    /// Optional maximum number of tokens to generate.
    pub max_tokens: Option<u32>,
    /// Provider-specific extra parameters (passed through as-is).
    pub extra: Option<serde_json::Value>,
}

impl CompletionRequest {
    pub fn new(messages: Vec<Message>) -> Self {
        Self {
            messages,
            tools: Vec::new(),
            temperature: None,
            max_tokens: None,
            extra: None,
        }
    }
}

// ── Response ──────────────────────────────────────────────────────────────────

/// The model's choice of response for a single completion turn.
#[derive(Debug, Clone)]
pub enum ModelChoice {
    /// The model produced a text reply.
    Message(String),
    /// The model wants to call one or more tools.
    ToolCall(Vec<ToolCall>),
}

/// Token usage reported by the provider (optional — not all providers expose this).
#[derive(Debug, Clone, Default)]
pub struct Usage {
    pub prompt_tokens: u32,
    pub completion_tokens: u32,
}

/// The full response returned by a [`CompletionModel`].
#[derive(Debug, Clone)]
pub struct CompletionResponse {
    /// What the model decided to do this turn.
    pub choice: ModelChoice,
    /// Token usage, if the provider reported it.
    pub usage: Option<Usage>,
}

// ── Error ─────────────────────────────────────────────────────────────────────

/// Errors that can occur during a completion request.
#[derive(Debug, Error)]
pub enum CompletionError {
    /// The HTTP transport returned an error.
    #[error("HTTP error: {0}")]
    Http(String),

    /// The response body could not be deserialised.
    #[error("JSON error: {0}")]
    Json(#[from] serde_json::Error),

    /// The provider returned a non-2xx status with a message.
    #[error("Provider error ({status}): {message}")]
    Provider { status: u16, message: String },

    /// The response structure was unexpected or missing required fields.
    #[error("Response error: {0}")]
    Response(String),
}

// ── Trait ─────────────────────────────────────────────────────────────────────

/// The core abstraction for any text-generation model.
///
/// Implement this for each provider (OpenAI, Anthropic, …) or for your own
/// local model backend.
pub trait CompletionModel {
    type Error: std::error::Error + 'static;

    /// Send a completion request and return the model's response.
    fn complete(
        &self,
        request: CompletionRequest,
    ) -> impl std::future::Future<Output = Result<CompletionResponse, Self::Error>>;
}