poe2-agent 0.5.0

//! Profiling tests for the agent pipeline.
//!
//! Measures wall-clock time for each phase of a chat request to identify
//! bottlenecks. Run with:
//!
//!   cargo test --test profiling -- --nocapture
//!
//! For the full agent round-trip test, set OPENAI_API_KEY in the environment
//! or in a `.env` file.

use std::path::Path;
use std::sync::Arc;
use std::time::Instant;

use poe2_agent::{ChatGptClient, PobParser, PobQuery, ToolAgent};

const POB_PATH: &str = "vendor/PathOfBuilding-PoE2";

fn fixture_xml() -> Vec<u8> {
    std::fs::read(Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/ranger-with-gear.xml"))
        .expect("fixture XML missing — run from repo root")
}

// ---------------------------------------------------------------------------
// PoB layer profiling
// ---------------------------------------------------------------------------

#[tokio::test]
async fn profile_pob_init() {
    let t = Instant::now();
    let _parser = PobParser::new(Path::new(POB_PATH))
        .await
        .expect("PobParser::new failed");
    let elapsed = t.elapsed();

    eprintln!("\n=== PobParser::new (Lua VM + PoB bootstrap) ===");
    eprintln!("  {elapsed:?}");
}

#[tokio::test]
async fn profile_pob_queries() {
    let parser = PobParser::new(Path::new(POB_PATH))
        .await
        .expect("PobParser::new failed");
    let xml = fixture_xml();

    let queries: &[(&str, PobQuery)] = &[
        ("BuildStats", PobQuery::BuildStats),
        ("SkillList", PobQuery::SkillList),
        ("Config", PobQuery::Config),
        ("Item(Weapon 1)", PobQuery::Item("Weapon 1".into())),
        ("Item(Body Armour)", PobQuery::Item("Body Armour".into())),
        ("EquippedItems", PobQuery::EquippedItems),
        ("PassiveTree", PobQuery::PassiveTree),
        (
            "PassiveStats([fire damage], r=3)",
            PobQuery::PassiveStats {
                stats: vec!["fire damage".into()],
                radius: 3,
            },
        ),
        ("UnallocatedAscendancy", PobQuery::UnallocatedAscendancy),
        (
            "GearModAnalysis(Body Armour)",
            PobQuery::GearModAnalysis("Body Armour".into()),
        ),
        (
            "SearchGems(fire, support)",
            PobQuery::SearchGems {
                query: None,
                gem_type: Some("support".into()),
                tags: vec!["fire".into()],
            },
        ),
    ];

    eprintln!("\n=== Individual query times (each reloads XML) ===");
    eprintln!("  {:.<45} {:>10}", "Query", "Time");
    eprintln!("  {}", "-".repeat(57));

    for (label, query) in queries {
        let t = Instant::now();
        let _result = parser
            .query(&xml, query.clone())
            .await
            .expect("query failed");
        let elapsed = t.elapsed();
        eprintln!("  {:.<45} {:>10.2?}", label, elapsed);
    }
}

#[tokio::test]
async fn profile_pob_reload_cost() {
    let parser = PobParser::new(Path::new(POB_PATH))
        .await
        .expect("PobParser::new failed");
    let xml = fixture_xml();

    // Simulate a typical agent tool-call sequence:
    //   get_build_stats → get_skill_list → get_item(Weapon 1)
    // Each call reloads XML from scratch.
    let sequence: &[(&str, PobQuery)] = &[
        ("get_build_stats", PobQuery::BuildStats),
        ("get_skill_list", PobQuery::SkillList),
        ("get_item(Weapon 1)", PobQuery::Item("Weapon 1".into())),
    ];

    eprintln!("\n=== Typical 3-tool sequence (sequential XML reloads) ===");

    let total = Instant::now();
    for (label, query) in sequence {
        let t = Instant::now();
        let _result = parser
            .query(&xml, query.clone())
            .await
            .expect("query failed");
        eprintln!("  {label}: {:?}", t.elapsed());
    }
    eprintln!("  ──────────────────────────────");
    eprintln!("  Total PoB time: {:?}", total.elapsed());
}

// ---------------------------------------------------------------------------
// Full agent round-trip (requires OPENAI_API_KEY)
// ---------------------------------------------------------------------------

#[tokio::test]
async fn profile_agent_roundtrip() {
    let _ = dotenvy::dotenv();

    let api_key = match std::env::var("OPENAI_API_KEY") {
        Ok(k) if !k.is_empty() => k,
        _ => {
            eprintln!("\n=== Skipping agent round-trip (no OPENAI_API_KEY) ===");
            return;
        }
    };

    let model = std::env::var("OPENAI_MODEL").unwrap_or_else(|_| "gpt-4.1-nano".into());
    eprintln!("\n=== Agent round-trip (model: {model}) ===");

    // Phase 1: Init
    let t = Instant::now();
    let parser = Arc::new(
        PobParser::new(Path::new(POB_PATH))
            .await
            .expect("PobParser::new failed"),
    );
    let init_time = t.elapsed();
    eprintln!("  PobParser init:  {init_time:?}");

    let t = Instant::now();
    let llm = ChatGptClient::new(&api_key, &model).expect("ChatGptClient::new failed");
    let llm_init_time = t.elapsed();
    eprintln!("  LLM client init: {llm_init_time:?}");

    let agent = ToolAgent::new(llm, parser, None);
    let xml = fixture_xml();
    let question = "What is my total DPS and what main skill am I using?";

    // Phase 2: Agent response
    eprintln!("  Question: \"{question}\"");
    eprintln!("  ---");

    let total = Instant::now();
    let stream = agent.respond(&xml, question, vec![]);

    use futures_lite::StreamExt;
    tokio::pin!(stream);

    let mut tool_calls = Vec::new();
    let mut first_token_time = None;
    let mut token_count = 0usize;
    let mut full_response = String::new();

    while let Some(event) = stream.next().await {
        match event.expect("stream error") {
            poe2_agent::AgentEvent::ToolCall { name } => {
                let t_call = total.elapsed();
                tool_calls.push((name.clone(), t_call));
                eprintln!("  [{t_call:?}] Tool call: {name}");
            }
            poe2_agent::AgentEvent::ToolResult { name, size_bytes } => {
                eprintln!(
                    "  [{:?}] Tool result: {name} ({size_bytes} bytes)",
                    total.elapsed()
                );
            }
            poe2_agent::AgentEvent::Token(text) => {
                if first_token_time.is_none() {
                    first_token_time = Some(total.elapsed());
                }
                token_count += 1;
                full_response.push_str(&text);
            }
            poe2_agent::AgentEvent::Usage(usage) => {
                eprintln!(
                    "  Token usage: {} input + {} output = {} total",
                    usage.input_tokens, usage.output_tokens, usage.total_tokens
                );
            }
            _ => {}
        }
    }

    let total_time = total.elapsed();

    eprintln!("  ---");
    eprintln!("  Tool calls:        {}", tool_calls.len());
    for (name, at) in &tool_calls {
        eprintln!("    {name} at {at:?}");
    }
    if let Some(ttft) = first_token_time {
        eprintln!("  Time to first token: {ttft:?}");
    }
    eprintln!("  Tokens received:   {token_count}");
    eprintln!("  Total time:        {total_time:?}");
    eprintln!("  ---");
    eprintln!(
        "  Response preview:  {}...",
        &full_response[..full_response.len().min(200)]
    );
}