poe2-agent 0.2.1

//! ReAct-style tool-calling agent for build analysis.
//!
//! Uses OpenAI function calling to query PoB data on demand
//! rather than dumping everything into the system prompt upfront.

use std::sync::Arc;

use futures_core::Stream;

use crate::llm::{
    ChatGptClient, FunctionDefinition, LlmError, Message, ToolDefinition,
};
use crate::pob_parser::{PobParser, PobQuery};

const MAX_TOOL_ROUNDS: usize = 10;

const SYSTEM_PROMPT: &str = "\
You are a Path of Exile 2 build analysis assistant. The user has uploaded \
their Path of Building export.\n\
\n\
You have tools to inspect the build data. Use them to answer the user's \
questions accurately — do NOT guess at numbers.\n\
\n\
Start by calling get_build_stats to get an overview of the build's offense, \
defense, and resources. Then use get_skill_list or get_config if needed \
to answer the user's specific question.\n\
\n\
Use get_empty_slots to quickly scan all equipment slots and see which ones \
are empty and which have items equipped. This is useful for identifying \
obvious upgrade opportunities — empty slots mean free power. Call this \
before diving into individual items with get_item.\n\
\n\
Use get_item to inspect a specific equipment slot when the user asks about \
their gear, an item's mods, or how a particular slot could be upgraded. \
Do not call get_item unless the question is about specific equipment.\n\
\n\
Use get_passive_tree when the user asks about their passive tree, allocated \
nodes, keystones, notables, ascendancy choices, masteries, or jewel sockets. \
It returns all allocated nodes categorized by type.\n\
\n\
Use get_jewel to inspect a jewel socketed in a passive tree socket. First call \
get_passive_tree to get the jewel_sockets list with node IDs, then call \
get_jewel with the node_id to see the jewel's name, base, rarity, and mods.\n\
\n\
Use query_passive_stats to find how much of a specific stat comes from allocated \
passives and what's available nearby on the tree. Provide a stat pattern like \
\"fire damage\" or \"maximum life\". Optionally set radius (default 3) to control \
how far to search from current allocation.\n\
\n\
Use get_unallocated_ascendancy to see which ascendancy nodes the character has \
allocated and which are still available. Returns both primary and secondary \
ascendancy nodes with node names, types, and stats. Use this when recommending \
ascendancy choices or when the user asks what ascendancy nodes to take next.\n\
\n\
Be specific and reference actual numbers from the build data when relevant. \
If the data doesn't contain enough information to answer, say so.\n\
\n\
Path of Exile 2 differences from Path of Exile 1 — do NOT confuse these:\n\
- There are NO utility flasks. Players have 2 flask slots (life/mana style only).\n\
- Charms (3 slots) provide passive bonuses and trigger effects — they replace \
much of what utility flasks did in PoE1.\n\
- Spirit is a resource that reserves for persistent buffs, auras, and minions.\n\
- Gear does NOT have gem sockets. Skill gems are equipped independently in \
dedicated active-gem slots, each with support sockets.\n\
- Rune sockets on gear provide bonus stats (via socketed runes).\n\
- Do NOT reference PoE1-specific unique items, support gems, or league mechanics.\n\
- When recommending items, gems, or tree nodes, verify they exist using the \
available tools rather than relying on memory.";

/// A single turn from a prior conversation. Text only — no tool calls.
#[derive(Debug, Clone)]
pub struct ChatMessage {
    pub role: String,
    pub content: String,
}

/// Events yielded by the agent during a response.
pub enum AgentEvent {
    /// The agent is calling a tool (yields tool name for progress indication).
    ToolCall { name: String },
    /// A token of the final streamed response.
    Token(String),
}

/// Tool-calling build analysis agent.
///
/// Wraps an LLM client and a shared PoB parser. Each call to `respond`
/// runs a ReAct loop: the LLM decides which tools to call, the agent
/// executes them via the parser, and the results are fed back until
/// the LLM produces a final answer.
pub struct ToolAgent {
    llm: ChatGptClient,
    parser: Arc<PobParser>,
}

impl ToolAgent {
    pub fn new(llm: ChatGptClient, parser: Arc<PobParser>) -> Self {
        Self { llm, parser }
    }

    /// Stream a response to a user question about the given build.
    ///
    /// `build_xml` is the raw PoB XML export. The agent loads it into PoB
    /// on each tool call so queries always reflect the full build.
    pub fn respond(
        &self,
        build_xml: &[u8],
        message: &str,
        history: Vec<ChatMessage>,
    ) -> impl Stream<Item = Result<AgentEvent, LlmError>> + Send {
        let llm = self.llm.clone();
        let parser = Arc::clone(&self.parser);
        let build_xml = build_xml.to_vec();
        let message = message.to_owned();

        async_stream::try_stream! {
            let tools = tool_definitions();
            let mut messages = vec![Message::system(SYSTEM_PROMPT)];
            for msg in history {
                match msg.role.as_str() {
                    "user" => messages.push(Message::user(&msg.content)),
                    "assistant" => messages.push(Message::assistant(&msg.content)),
                    _ => {}
                }
            }
            messages.push(Message::user(message));

            // Tool-calling loop: let the LLM call tools until it has
            // enough data, then break to stream the final answer.
            let mut tools_were_called = false;

            for _ in 0..MAX_TOOL_ROUNDS {
                let (assistant_msg, finish_reason) = llm
                    .chat_with_tools(messages.clone(), Some(&tools))
                    .await?;

                let reason = finish_reason.as_deref().unwrap_or("stop");

                if reason != "tool_calls" {
                    if !tools_were_called {
                        // LLM answered directly without tools — yield its text.
                        if let Some(text) = assistant_msg.content {
                            yield AgentEvent::Token(text);
                        }
                        return;
                    }
                    // Tools were used in a prior round; break to stream.
                    break;
                }

                if let Some(ref tool_calls) = assistant_msg.tool_calls {
                    tools_were_called = true;

                    for tc in tool_calls {
                        yield AgentEvent::ToolCall {
                            name: tc.function.name.clone(),
                        };
                    }

                    messages.push(assistant_msg.clone());

                    for tc in tool_calls {
                        let result = execute_tool(&parser, &build_xml, &tc.function.name, &tc.function.arguments).await;
                        let content = match result {
                            Ok(val) => val.to_string(),
                            Err(e) => format!("{{\"error\": \"{e}\"}}"),
                        };
                        messages.push(Message::tool_result(&tc.id, content));
                    }
                }
            }

            // Stream the final answer so tokens arrive progressively.
            let stream = llm.chat_stream(messages);
            tokio::pin!(stream);
            while let Some(token_result) = futures_lite::StreamExt::next(&mut stream).await {
                yield AgentEvent::Token(token_result?);
            }
        }
    }
}

/// Build the tool definitions for the agent.
fn tool_definitions() -> Vec<ToolDefinition> {
    vec![
        ToolDefinition {
            tool_type: "function".to_owned(),
            function: FunctionDefinition {
                name: "get_build_stats".to_owned(),
                description: "Get extended build statistics including offense, defense, \
                    resources, speed, and charges. Returns ~40 fields grouped by category."
                    .to_owned(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {},
                    "required": [],
                    "additionalProperties": false
                }),
            },
        },
        ToolDefinition {
            tool_type: "function".to_owned(),
            function: FunctionDefinition {
                name: "get_skill_list".to_owned(),
                description: "Get the list of skills with their DPS values, trigger info, \
                    and gem links (socket groups with gems, levels, and quality)."
                    .to_owned(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {},
                    "required": [],
                    "additionalProperties": false
                }),
            },
        },
        ToolDefinition {
            tool_type: "function".to_owned(),
            function: FunctionDefinition {
                name: "get_config".to_owned(),
                description: "Get the build's configuration flags (enemy settings, \
                    charge generation, conditions, etc.)."
                    .to_owned(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {},
                    "required": [],
                    "additionalProperties": false
                }),
            },
        },
        ToolDefinition {
            tool_type: "function".to_owned(),
            function: FunctionDefinition {
                name: "get_item".to_owned(),
                description: "Retrieve the item equipped in a specific gear slot, including \
                    its name, base type, rarity, and all mod lines (implicit, explicit, \
                    enchant, rune)."
                    .to_owned(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {
                        "slot": {
                            "type": "string",
                            "enum": [
                                "Weapon 1", "Weapon 2", "Helmet", "Body Armour",
                                "Gloves", "Boots", "Amulet", "Ring 1", "Ring 2", "Ring 3",
                                "Belt", "Charm 1", "Charm 2", "Charm 3",
                                "Flask 1", "Flask 2"
                            ],
                            "description": "The equipment slot to inspect"
                        }
                    },
                    "required": ["slot"],
                    "additionalProperties": false
                }),
            },
        },
        ToolDefinition {
            tool_type: "function".to_owned(),
            function: FunctionDefinition {
                name: "get_empty_slots".to_owned(),
                description: "Scan all equipment slots and return which are empty and which \
                    have items equipped. Returns item name and rarity for filled slots. \
                    Useful for quickly identifying missing gear without calling get_item \
                    for every slot."
                    .to_owned(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {},
                    "required": [],
                    "additionalProperties": false
                }),
            },
        },
        ToolDefinition {
            tool_type: "function".to_owned(),
            function: FunctionDefinition {
                name: "get_jewel".to_owned(),
                description: "Retrieve a jewel socketed in a passive tree socket, including \
                    its name, base type, rarity, and all mod lines. Use socket node IDs \
                    from get_passive_tree's jewel_sockets array."
                    .to_owned(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {
                        "node_id": {
                            "type": "integer",
                            "description": "The passive tree socket node ID (from get_passive_tree jewel_sockets)"
                        }
                    },
                    "required": ["node_id"],
                    "additionalProperties": false
                }),
            },
        },
        ToolDefinition {
            tool_type: "function".to_owned(),
            function: FunctionDefinition {
                name: "get_passive_tree".to_owned(),
                description: "Get the allocated passive tree nodes, grouped by type: \
                    keystones, notables, ascendancy nodes, masteries, and jewel sockets. \
                    Also returns class, ascendancy, and total allocated node count."
                    .to_owned(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {},
                    "required": [],
                    "additionalProperties": false
                }),
            },
        },
        ToolDefinition {
            tool_type: "function".to_owned(),
            function: FunctionDefinition {
                name: "query_passive_stats".to_owned(),
                description: "Query how much of a specific stat comes from allocated passive \
                    tree nodes, and how much more is available on nearby unallocated nodes. \
                    Uses case-insensitive pattern matching on stat descriptions."
                    .to_owned(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {
                        "stat": {
                            "type": "string",
                            "description": "Stat pattern to search for (e.g. \"fire damage\", \"maximum life\", \"critical strike\")"
                        },
                        "radius": {
                            "type": "integer",
                            "description": "How many hops from allocated nodes to search for nearby stats (default: 3)"
                        }
                    },
                    "required": ["stat"],
                    "additionalProperties": false
                }),
            },
        },
        ToolDefinition {
            tool_type: "function".to_owned(),
            function: FunctionDefinition {
                name: "get_unallocated_ascendancy".to_owned(),
                description: "Get the character's ascendancy nodes — both allocated and \
                    available — for primary and secondary ascendancies. Returns node names, \
                    types, stats, and points spent. Use this to recommend which ascendancy \
                    nodes to take next."
                    .to_owned(),
                parameters: serde_json::json!({
                    "type": "object",
                    "properties": {},
                    "required": [],
                    "additionalProperties": false
                }),
            },
        },
    ]
}

/// Execute a single tool call via the PoB parser.
async fn execute_tool(
    parser: &PobParser,
    build_xml: &[u8],
    tool_name: &str,
    tool_args: &str,
) -> Result<serde_json::Value, String> {
    let query = match tool_name {
        "get_build_stats" => PobQuery::BuildStats,
        "get_skill_list" => PobQuery::SkillList,
        "get_config" => PobQuery::Config,
        "get_item" => {
            let args: serde_json::Value =
                serde_json::from_str(tool_args).map_err(|e| format!("invalid arguments: {e}"))?;
            let slot = args["slot"]
                .as_str()
                .ok_or("missing required parameter: slot")?
                .to_owned();
            PobQuery::Item(slot)
        }
        "get_empty_slots" => PobQuery::EmptySlots,
        "get_jewel" => {
            let args: serde_json::Value =
                serde_json::from_str(tool_args).map_err(|e| format!("invalid arguments: {e}"))?;
            let node_id = args["node_id"]
                .as_i64()
                .ok_or("missing required parameter: node_id")?;
            PobQuery::Jewel(node_id)
        }
        "get_passive_tree" => PobQuery::PassiveTree,
        "query_passive_stats" => {
            let args: serde_json::Value =
                serde_json::from_str(tool_args).map_err(|e| format!("invalid arguments: {e}"))?;
            let stat = args["stat"]
                .as_str()
                .ok_or("missing required parameter: stat")?
                .to_owned();
            let radius = args["radius"].as_u64().unwrap_or(3) as u32;
            PobQuery::PassiveStats { stat, radius }
        }
        "get_unallocated_ascendancy" => PobQuery::UnallocatedAscendancy,
        other => return Err(format!("unknown tool: {other}")),
    };

    parser
        .query(build_xml, query)
        .await
        .map_err(|e| e.to_string())
}