ic_rig/

agent.rs

//! High-level agent that drives the completion ↔ tool-call loop.
//!
//! An [`Agent`] wraps a [`CompletionModel`], a system prompt, an optional set
//! of [`Tool`]s, and common sampling parameters. Call [`Agent::prompt`] for a
//! one-shot exchange or [`Agent::chat`] to continue an existing history.
//!
//! # Agentic loop
//!
//! When the model returns [`ModelChoice::ToolCall`], the agent:
//! 1. Executes each requested tool.
//! 2. Appends the results as a user-turn message.
//! 3. Sends a new completion request.
//!
//! This repeats until the model returns [`ModelChoice::Message`] or the
//! `max_iterations` limit is reached.
//!
//! # Example
//!
//! ```rust,ignore
//! let agent = Agent::builder(model)
//!     .preamble("You are a helpful assistant.")
//!     .tool(calculator)
//!     .max_tokens(1024)
//!     .build();
//!
//! let reply = agent.prompt("What is 3 * 7?").await?;
//! println!("{reply}");
//! ```

use crate::{
    completion::{CompletionModel, CompletionRequest, ModelChoice},
    message::{AssistantContent, Message, ToolCall, ToolResult, UserContent},
    tool::{ToolError, ToolSet},
};
use thiserror::Error;

// ── AgentError ────────────────────────────────────────────────────────────────

/// Errors that can occur while running the agentic loop.
///
/// Not using `#[from]` on both variants because when `M = ToolError` the
/// blanket impls would conflict. Instead, conversions are explicit at the
/// call sites inside [`Agent`].
#[derive(Debug, Error)]
pub enum AgentError<M: std::error::Error + 'static> {
    #[error("Completion error: {0}")]
    Completion(M),

    #[error("Tool error: {0}")]
    Tool(ToolError),

    /// The model kept issuing tool calls beyond the iteration cap.
    #[error("Exceeded maximum iterations ({0}) without a final text reply")]
    MaxIterations(u32),
}

// ── Agent ─────────────────────────────────────────────────────────────────────

/// A configured LLM agent.
///
/// Build one with [`Agent::builder`] then call [`prompt`](Agent::prompt)
/// or [`chat`](Agent::chat).
pub struct Agent<M> {
    model: M,
    preamble: Option<String>,
    tools: ToolSet,
    temperature: Option<f64>,
    max_tokens: Option<u32>,
    /// Hard cap on tool-call iterations to prevent infinite loops.
    max_iterations: u32,
    /// Static context documents prepended before the first user message.
    context: Vec<String>,
}

impl<M: CompletionModel> Agent<M> {
    /// Create a builder for this model.
    pub fn builder(model: M) -> AgentBuilder<M> {
        AgentBuilder::new(model)
    }

    /// Send a single prompt and return the model's text reply.
    ///
    /// Any tool calls are resolved internally; the caller receives the final
    /// text response.
    pub async fn prompt(&self, prompt: &str) -> Result<String, AgentError<M::Error>> {
        self.chat(prompt, vec![]).await
    }

    /// Continue a conversation with `history` and a new `prompt`.
    ///
    /// Returns the model's final text reply after resolving any tool calls.
    pub async fn chat(
        &self,
        prompt: &str,
        history: Vec<Message>,
    ) -> Result<String, AgentError<M::Error>> {
        let mut messages = self.build_messages(prompt, history);

        for _ in 0..self.max_iterations {
            let request = self.build_request(messages.clone());
            let response = self.model.complete(request).await.map_err(AgentError::Completion)?;

            match response.choice {
                ModelChoice::Message(text) => return Ok(text),
                ModelChoice::ToolCall(calls) => {
                    // Append the assistant's tool-call turn.
                    messages.push(Message::Assistant {
                        content: calls
                            .iter()
                            .map(|c| AssistantContent::ToolCall(c.clone()))
                            .collect(),
                    });

                    // Execute all requested tool calls and collect results.
                    let mut results: Vec<UserContent> = Vec::with_capacity(calls.len());
                    for call in &calls {
                        let result = self.dispatch_tool(call).await;
                        results.push(UserContent::ToolResult(result));
                    }

                    // Feed results back as a user turn.
                    messages.push(Message::User { content: results });
                }
            }
        }

        Err(AgentError::MaxIterations(self.max_iterations))
    }

    // ── Private helpers ───────────────────────────────────────────────────────

    fn build_messages(&self, prompt: &str, mut history: Vec<Message>) -> Vec<Message> {
        let mut messages: Vec<Message> = Vec::new();

        if let Some(preamble) = &self.preamble {
            messages.push(Message::system(preamble));
        }

        // Static context documents become a single user message prepended
        // before the history so the model has them in its window.
        if !self.context.is_empty() {
            let combined = self.context.join("\n\n");
            messages.push(Message::user(combined));
        }

        messages.append(&mut history);
        messages.push(Message::user(prompt));
        messages
    }

    fn build_request(&self, messages: Vec<Message>) -> CompletionRequest {
        let mut req = CompletionRequest::new(messages);
        req.tools = self.tools.definitions();
        req.temperature = self.temperature;
        req.max_tokens = self.max_tokens;
        req
    }

    async fn dispatch_tool(&self, call: &ToolCall) -> ToolResult {
        let result = self.tools.call(&call.name, call.arguments.clone()).await;
        let content = match result {
            Ok(v) => v.to_string(),
            // Surface tool errors as plain-text so the model can react.
            Err(e) => format!("Error: {e}"),
        };
        ToolResult {
            call_id: call.id.clone(),
            name: call.name.clone(),
            content,
        }
    }
}

// ── AgentBuilder ──────────────────────────────────────────────────────────────

/// Fluent builder for [`Agent`].
pub struct AgentBuilder<M> {
    model: M,
    preamble: Option<String>,
    tools: ToolSet,
    temperature: Option<f64>,
    max_tokens: Option<u32>,
    max_iterations: u32,
    context: Vec<String>,
}

impl<M: CompletionModel> AgentBuilder<M> {
    fn new(model: M) -> Self {
        Self {
            model,
            preamble: None,
            tools: ToolSet::new(),
            temperature: None,
            max_tokens: None,
            max_iterations: 10,
            context: Vec::new(),
        }
    }

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

    /// Register a tool.
    pub fn tool<T: crate::tool::Tool + 'static>(mut self, tool: T) -> Self {
        self.tools.add(tool);
        self
    }

    /// Set the sampling temperature.
    pub fn temperature(mut self, temperature: f64) -> Self {
        self.temperature = Some(temperature);
        self
    }

    /// Set the maximum tokens to generate.
    pub fn max_tokens(mut self, max_tokens: u32) -> Self {
        self.max_tokens = Some(max_tokens);
        self
    }

    /// Override the default iteration cap (10) for the tool-call loop.
    pub fn max_iterations(mut self, n: u32) -> Self {
        self.max_iterations = n;
        self
    }

    /// Append a static context document (e.g. retrieved RAG chunks).
    pub fn context(mut self, doc: impl Into<String>) -> Self {
        self.context.push(doc.into());
        self
    }

    /// Consume the builder and produce an [`Agent`].
    pub fn build(self) -> Agent<M> {
        Agent {
            model: self.model,
            preamble: self.preamble,
            tools: self.tools,
            temperature: self.temperature,
            max_tokens: self.max_tokens,
            max_iterations: self.max_iterations,
            context: self.context,
        }
    }
}