tests/test_b_tools.py

"""Tests for the tool surface

What we test, what we deliberately skip:

  * `inspect_data`, `run_sql` — pure local tools. Tested against PKDD with
    no LLM dependency. Skip if PKDD DuckDB isn't bootstrapped.
  * `run_sql` cache mechanics — verifies that the `sql_result:<label>`
    handoff to `train_tabular_model` / `predict` works.
  * `text_to_sql`, `summarize_result` — need a real LLM. Gated behind the
    `lexsi_env` fixture; skip cleanly if `SDK_ACCESS_TOKEN` + project env
    vars aren't set.
  * Full agent-loop integration on the analytic samples is the last test;
    also LLM-gated.

"""

from __future__ import annotations

import pandas as pd
import pytest

from lexsi_ds.agent.loop import AgentLoop
from lexsi_ds.agent.samples import SAMPLES, by_category, by_id
from lexsi_ds.agent.tools.inspect_data import InspectDataArgs
from lexsi_ds.agent.tools.inspect_data import TOOL as inspect_data_tool
from lexsi_ds.agent.tools.query_kg import QueryKgArgs
from lexsi_ds.agent.tools.query_kg import TOOL as query_kg_tool
from lexsi_ds.agent.tools.run_sql import RunSqlArgs
from lexsi_ds.agent.tools.run_sql import TOOL as run_sql_tool
from lexsi_ds.agent.tools.sample_values import SampleValuesArgs
from lexsi_ds.agent.tools.sample_values import TOOL as sample_values_tool
from lexsi_ds.agent.tools.summarize_result import SummarizeResultArgs
from lexsi_ds.agent.tools.summarize_result import TOOL as summarize_result_tool
from lexsi_ds.agent.tools.text_to_sql import TextToSqlArgs
from lexsi_ds.agent.tools.text_to_sql import TOOL as text_to_sql_tool


# ----- inspect_data -----

class TestInspectData:
    def test_returns_pkdd_schema(self, pkdd_ctx):
        result = inspect_data_tool.run(InspectDataArgs(), pkdd_ctx)
        assert result.ok, result.summary
        # All 8 PKDD tables should appear in the rendered schema
        for tbl in ("fin_account", "fin_loan", "fin_client", "fin_district",
                    "fin_disp", "fin_card", "fin_trans", "fin_order"):
            assert tbl in result.summary
        assert result.payload["has_kg"] is True
        # KG content moved to `query_kg`; inspect_data should now nudge toward it.
        assert "query_kg" in result.summary

    def test_table_filter(self, pkdd_ctx):
        result = inspect_data_tool.run(
            InspectDataArgs(table_filter=["fin_loan", "fin_card"]),
            pkdd_ctx,
        )
        assert result.ok
        assert "fin_loan" in result.summary
        assert "fin_card" in result.summary
        # Filtered-out tables should NOT appear in the per-table headings
        assert "### `fin_trans`" not in result.summary

    def test_bad_filter(self, pkdd_ctx):
        result = inspect_data_tool.run(
            InspectDataArgs(table_filter=["does_not_exist"]),
            pkdd_ctx,
        )
        assert result.ok is False
        assert result.error == "no_matching_tables"


# ----- sample_values -----

class TestSampleValues:
    def test_district_name_includes_prague_in_czech(self, pkdd_ctx):
        """The m02 fix: searching the column should surface 'Hl.m. Praha'
        so the planner can recover after a 0-row WHERE = 'Prague' attempt."""
        result = sample_values_tool.run(
            SampleValuesArgs(table="fin_district", column="district_name", k=80),
            pkdd_ctx,
        )
        assert result.ok, result.summary
        values = result.payload["values"]
        assert "Hl.m. Praha" in values, values

    def test_card_type_values(self, pkdd_ctx):
        result = sample_values_tool.run(
            SampleValuesArgs(table="fin_card", column="card_type"),
            pkdd_ctx,
        )
        assert result.ok
        # Card types in PKDD are single letters: J / C / G
        assert set(result.payload["values"]).issubset({"J", "C", "G"})

    def test_unknown_table_fails_softly(self, pkdd_ctx):
        result = sample_values_tool.run(
            SampleValuesArgs(table="nope", column="x"),
            pkdd_ctx,
        )
        assert result.ok is False
        assert result.error == "unknown_table"

    def test_bad_identifier_rejected(self, pkdd_ctx):
        # Injection attempt should be caught by the identifier validator,
        # not by DuckDB downstream.
        result = sample_values_tool.run(
            SampleValuesArgs(table="fin_district; DROP TABLE foo;",
                             column="district_name"),
            pkdd_ctx,
        )
        assert result.ok is False
        # Could be `bad_identifier` (preferred) or `unknown_table` (cheaper check
        # ran first). Either way it failed safely.
        assert result.error in ("bad_identifier", "unknown_table")


# ----- query_kg -----

class TestQueryKg:
    def test_outline(self, pkdd_ctx):
        result = query_kg_tool.run(QueryKgArgs(), pkdd_ctx)
        assert result.ok, result.summary
        assert result.payload["mode"] == "outline"
        # PKDD has at least concepts + disambigs + metrics + join_paths
        counts = result.payload["counts"]
        for section in ("concept", "disambig", "metric", "join_path"):
            assert counts.get(section, 0) > 0, f"missing section: {section}"
        # Outline should suggest next moves
        assert "search" in result.summary.lower()

    def test_search_finds_district_disambig_from_prague_question(self, pkdd_ctx):
        """The actual m02 failure case: 'customers in Prague' should surface
        the district disambig + customer concept WITHOUT the planner already
        knowing the disambig name."""
        result = query_kg_tool.run(
            QueryKgArgs(search="how many customers live in prague"),
            pkdd_ctx,
        )
        assert result.ok, result.summary
        assert result.payload["n_results"] > 0
        body = result.summary.lower()
        # Either disambig name should be in the top results
        assert any(
            t in body
            for t in ("district_disambig", "customer_vs_account", "home_district")
        ), result.summary

    def test_search_with_section_filter(self, pkdd_ctx):
        result = query_kg_tool.run(
            QueryKgArgs(search="loan default", section="disambigs"),
            pkdd_ctx,
        )
        assert result.ok
        # All returned node ids should be in the disambig section
        for nid in result.payload["node_ids"]:
            assert nid.startswith("disambig:"), nid

    def test_node_fetch_by_label(self, pkdd_ctx):
        result = query_kg_tool.run(
            QueryKgArgs(node="district_disambig"),
            pkdd_ctx,
        )
        assert result.ok, result.summary
        assert result.payload["mode"] == "node"
        # Rule keywords should be in the rendered output
        body = result.summary.lower()
        assert "customer" in body or "branch" in body

    def test_unknown_node_fails_softly(self, pkdd_ctx):
        result = query_kg_tool.run(
            QueryKgArgs(node="totally_made_up_node"),
            pkdd_ctx,
        )
        assert result.ok is False
        assert result.error == "unknown_node"

    def test_section_listing(self, pkdd_ctx):
        result = query_kg_tool.run(
            QueryKgArgs(section="disambigs"),
            pkdd_ctx,
        )
        assert result.ok
        assert result.payload["n_results"] > 0
        # Every node_id in payload should be a disambig
        assert all(n.startswith("disambig:") for n in result.payload["node_ids"])


# ----- run_sql -----

class TestRunSql:
    @pytest.mark.parametrize("sample", by_category("analytic"), ids=lambda s: s.id)
    def test_each_analytic_gold_sql_executes(self, pkdd_ctx, sample):
        """Every analytic gold SQL should execute against the PKDD DuckDB."""
        result = run_sql_tool.run(RunSqlArgs(sql=sample.gold_sql), pkdd_ctx)
        assert result.ok, f"{sample.id} failed: {result.summary}"
        df = result.payload["df"]
        assert isinstance(df, pd.DataFrame)
        assert len(df) > 0, f"{sample.id}: gold SQL returned 0 rows"

    @pytest.mark.parametrize("sample", by_category("ambiguity"), ids=lambda s: s.id)
    def test_each_ambiguity_gold_sql_executes(self, pkdd_ctx, sample):
        result = run_sql_tool.run(RunSqlArgs(sql=sample.gold_sql), pkdd_ctx)
        assert result.ok, f"{sample.id} failed: {result.summary}"

    def test_label_stashes_into_cache(self, pkdd_ctx):
        sample = by_id("a02")  # avg loan principal
        result = run_sql_tool.run(
            RunSqlArgs(sql=sample.gold_sql, label="avg_loan"),
            pkdd_ctx,
        )
        assert result.ok
        # Both the labeled slot and `last_sql_result` should be populated
        assert isinstance(pkdd_ctx.cache["sql_result:avg_loan"], pd.DataFrame)
        assert pkdd_ctx.cache["last_sql_result"] is result.payload["df"]
        assert pkdd_ctx.cache["last_sql"] == sample.gold_sql.strip().rstrip(";")

    def test_empty_sql_fails_softly(self, pkdd_ctx):
        result = run_sql_tool.run(RunSqlArgs(sql=""), pkdd_ctx)
        assert result.ok is False
        assert result.error == "empty_sql"

    def test_bad_sql_fails_softly(self, pkdd_ctx):
        result = run_sql_tool.run(RunSqlArgs(sql="SELECT * FROM no_such_table"), pkdd_ctx)
        assert result.ok is False
        # Catalog error should be surfaced; loop will let the LLM try again
        assert "no_such_table" in result.error.lower() or "catalog" in result.summary.lower()

    def test_known_row_count_a01(self, pkdd_ctx):
        """A01: Prague accounts. Gold result is a single COUNT row."""
        sample = by_id("a01")
        result = run_sql_tool.run(RunSqlArgs(sql=sample.gold_sql), pkdd_ctx)
        df = result.payload["df"]
        assert len(df) == 1
        n_accounts = int(df.iloc[0, 0])
        # PKDD: ~554 accounts at the Hl.m. Praha branch. Don't pin the exact
        # number, just sanity-check the order of magnitude.
        assert 100 < n_accounts < 5000


# ----- text_to_sql (LLM-gated) -----

class TestTextToSqlWithLexsi:
    @pytest.mark.parametrize(
        "sample",
        [by_id("a02"), by_id("a03"), by_id("a04")],
        ids=lambda s: s.id,
    )
    def test_emits_runnable_sql(self, pkdd_ctx_lexsi, sample):
        result = text_to_sql_tool.run(
            TextToSqlArgs(question=sample.question, mode="analytic"),
            pkdd_ctx_lexsi,
        )
        assert result.ok, result.summary
        sql = result.payload["sql"]
        assert sql.strip().lower().startswith("select")
        # The emitted SQL should at least execute (correctness checked
        # separately in `TestAgentLoopOnSamples`).
        exec_result = run_sql_tool.run(RunSqlArgs(sql=sql), pkdd_ctx_lexsi)
        assert exec_result.ok, (
            f"{sample.id}: text_to_sql emitted SQL that didn't execute: {exec_result.summary}"
        )

    def test_predictive_mode_emits_two_sqls(self, pkdd_ctx_lexsi):
        sample = by_id("p01")
        result = text_to_sql_tool.run(
            TextToSqlArgs(question=sample.question, mode="predictive"),
            pkdd_ctx_lexsi,
        )
        assert result.ok, result.summary
        assert "context_sql" in result.payload
        assert "predict_sql" in result.payload
        assert result.payload["task"].get("task_type") in {"classification", "regression"}


# ----- summarize_result (LLM-gated) -----

class TestSummarizeResultWithLexsi:
    def test_writes_narrative_from_cached_df(self, pkdd_ctx_lexsi):
        # Stage a result in cache by running run_sql first
        sample = by_id("a03")
        run_sql_tool.run(
            RunSqlArgs(sql=sample.gold_sql, label="top_districts"),
            pkdd_ctx_lexsi,
        )
        result = summarize_result_tool.run(
            SummarizeResultArgs(
                question=sample.question,
                sql_result_label="top_districts",
            ),
            pkdd_ctx_lexsi,
        )
        assert result.ok, result.summary
        narrative = result.payload["narrative"]
        assert len(narrative) > 40
        # Narrative should mention at least one district name from the result.
        # Don't pin the wording; check substring overlap with the gold df.
        gold_df = pkdd_ctx_lexsi.cache["sql_result:top_districts"]
        district_names = gold_df.iloc[:, 0].astype(str).tolist()
        assert any(name in narrative for name in district_names), narrative


# ----- full loop (LLM-gated) -----

class TestAgentLoopOnSamples:
    @pytest.mark.parametrize(
        "sample",
        [by_id("a02"), by_id("a03")],
        ids=lambda s: s.id,
    )
    def test_analytic_end_to_end(self, pkdd_ctx_lexsi, sample):
        loop = AgentLoop(ctx=pkdd_ctx_lexsi, max_steps=6)
        run = loop.run(sample.question)
        assert run.ok, f"{sample.id} run failed: {run.error}\nsteps: {[s.tool for s in run.steps]}"
        assert run.final_answer, "no final answer emitted"
        # The agent should have called run_sql at some point.
        tools_called = [s.tool for s in run.steps if s.tool]
        assert "run_sql" in tools_called, f"{sample.id}: tools were {tools_called}"