Skip to main content

poe2_agent/tools/
mod.rs

1//! Trait-based tool dispatch for the agent.
2//!
3//! Each tool implements [`Tool`] — providing its LLM definition and execution
4//! logic in one place. [`ToolRegistry`] collects tools and handles dispatch.
5
6mod mutation;
7mod pob;
8mod trade;
9
10use std::collections::HashMap;
11
12use async_trait::async_trait;
13
14use crate::llm::ToolDefinition;
15use crate::pob_parser::PobParser;
16use crate::trade::TradeClient;
17
18/// Shared context available to every tool during execution.
19pub struct ToolContext<'a> {
20    pub parser: &'a PobParser,
21    pub build_xml: &'a [u8],
22    pub trade: Option<&'a TradeClient>,
23}
24
25/// A build mutation produced by a tool (updated XML + human-readable label).
26pub struct BuildMutation {
27    /// Complete build XML with the mutation applied.
28    pub xml: String,
29    /// Human-readable description of the change (e.g. "Equipped item in Helmet").
30    pub label: String,
31}
32
33/// Result of executing a tool — response for the LLM plus an optional build mutation.
34pub struct ToolResult {
35    /// JSON value sent back to the LLM as the tool's response.
36    pub response: serde_json::Value,
37    /// If the tool modified the build, the mutation to apply.
38    pub mutation: Option<BuildMutation>,
39}
40
41/// A single agent tool — definition + execution in one place.
42#[async_trait]
43pub trait Tool: Send + Sync {
44    /// LLM tool definition (name, description, JSON Schema parameters).
45    fn definition(&self) -> ToolDefinition;
46
47    /// Execute the tool with the given JSON arguments string.
48    async fn execute(&self, ctx: &ToolContext<'_>, args: &str) -> Result<ToolResult, String>;
49}
50
51/// Collects tools and provides dispatch by name.
52pub struct ToolRegistry {
53    tools: Vec<Box<dyn Tool>>,
54    index: HashMap<String, usize>,
55}
56
57impl ToolRegistry {
58    /// Build a registry with all available tools.
59    ///
60    /// When `has_trade` is true, trade tools are included.
61    pub fn new(has_trade: bool) -> Self {
62        let mut tools: Vec<Box<dyn Tool>> = Vec::new();
63
64        // PoB query tools.
65        pob::register(&mut tools);
66
67        // Mutation tools.
68        mutation::register(&mut tools);
69
70        // Trade tools (conditional).
71        if has_trade {
72            trade::register(&mut tools);
73        }
74
75        let index = tools
76            .iter()
77            .enumerate()
78            .map(|(i, t)| (t.definition().name.clone(), i))
79            .collect();
80
81        Self { tools, index }
82    }
83
84    /// Tool definitions for the LLM (Responses API format).
85    pub fn definitions(&self) -> Vec<ToolDefinition> {
86        self.tools.iter().map(|t| t.definition()).collect()
87    }
88
89    /// Execute a tool by name. Returns `Err` for unknown tools.
90    pub async fn execute(
91        &self,
92        ctx: &ToolContext<'_>,
93        tool_name: &str,
94        args: &str,
95    ) -> Result<ToolResult, String> {
96        let idx = self
97            .index
98            .get(tool_name)
99            .ok_or_else(|| format!("unknown tool: {tool_name}"))?;
100        self.tools[*idx].execute(ctx, args).await
101    }
102}
103
104// -- Helpers shared by tool implementations -----------------------------------
105
106/// Parse a JSON arguments string.
107fn parse_args(args: &str) -> Result<serde_json::Value, String> {
108    serde_json::from_str(args).map_err(|e| format!("invalid arguments: {e}"))
109}
110
111/// Execute a PobQuery through the parser, returning a query-only `ToolResult`.
112async fn pob_query(
113    ctx: &ToolContext<'_>,
114    query: crate::pob_parser::PobQuery,
115) -> Result<ToolResult, String> {
116    ctx.parser
117        .query(ctx.build_xml, query)
118        .await
119        .map(|v| ToolResult {
120            response: v,
121            mutation: None,
122        })
123        .map_err(|e| e.to_string())
124}