poe2-agent 0.5.0

//! Trait-based tool dispatch for the agent.
//!
//! Each tool implements [`Tool`] — providing its LLM definition and execution
//! logic in one place. [`ToolRegistry`] collects tools and handles dispatch.

mod mutation;
mod pob;
mod trade;

use std::collections::HashMap;

use async_trait::async_trait;

use crate::llm::ToolDefinition;
use crate::pob_parser::PobParser;
use crate::trade::TradeClient;

/// Shared context available to every tool during execution.
pub struct ToolContext<'a> {
    pub parser: &'a PobParser,
    pub build_xml: &'a [u8],
    pub trade: Option<&'a TradeClient>,
}

/// A build mutation produced by a tool (updated XML + human-readable label).
pub struct BuildMutation {
    /// Complete build XML with the mutation applied.
    pub xml: String,
    /// Human-readable description of the change (e.g. "Equipped item in Helmet").
    pub label: String,
}

/// Result of executing a tool — response for the LLM plus an optional build mutation.
pub struct ToolResult {
    /// JSON value sent back to the LLM as the tool's response.
    pub response: serde_json::Value,
    /// If the tool modified the build, the mutation to apply.
    pub mutation: Option<BuildMutation>,
}

/// A single agent tool — definition + execution in one place.
#[async_trait]
pub trait Tool: Send + Sync {
    /// LLM tool definition (name, description, JSON Schema parameters).
    fn definition(&self) -> ToolDefinition;

    /// Execute the tool with the given JSON arguments string.
    async fn execute(&self, ctx: &ToolContext<'_>, args: &str) -> Result<ToolResult, String>;
}

/// Collects tools and provides dispatch by name.
pub struct ToolRegistry {
    tools: Vec<Box<dyn Tool>>,
    index: HashMap<String, usize>,
}

impl ToolRegistry {
    /// Build a registry with all available tools.
    ///
    /// When `has_trade` is true, trade tools are included.
    pub fn new(has_trade: bool) -> Self {
        let mut tools: Vec<Box<dyn Tool>> = Vec::new();

        // PoB query tools.
        pob::register(&mut tools);

        // Mutation tools.
        mutation::register(&mut tools);

        // Trade tools (conditional).
        if has_trade {
            trade::register(&mut tools);
        }

        let index = tools
            .iter()
            .enumerate()
            .map(|(i, t)| (t.definition().name.clone(), i))
            .collect();

        Self { tools, index }
    }

    /// Tool definitions for the LLM (Responses API format).
    pub fn definitions(&self) -> Vec<ToolDefinition> {
        self.tools.iter().map(|t| t.definition()).collect()
    }

    /// Execute a tool by name. Returns `Err` for unknown tools.
    pub async fn execute(
        &self,
        ctx: &ToolContext<'_>,
        tool_name: &str,
        args: &str,
    ) -> Result<ToolResult, String> {
        let idx = self
            .index
            .get(tool_name)
            .ok_or_else(|| format!("unknown tool: {tool_name}"))?;
        self.tools[*idx].execute(ctx, args).await
    }
}

// -- Helpers shared by tool implementations -----------------------------------

/// Parse a JSON arguments string.
fn parse_args(args: &str) -> Result<serde_json::Value, String> {
    serde_json::from_str(args).map_err(|e| format!("invalid arguments: {e}"))
}

/// Execute a PobQuery through the parser, returning a query-only `ToolResult`.
async fn pob_query(
    ctx: &ToolContext<'_>,
    query: crate::pob_parser::PobQuery,
) -> Result<ToolResult, String> {
    ctx.parser
        .query(ctx.build_xml, query)
        .await
        .map(|v| ToolResult {
            response: v,
            mutation: None,
        })
        .map_err(|e| e.to_string())
}