Spaces:
Sleeping
Sleeping
| """Tests for the self-verification tools: check_result + compare_candidates. | |
| Use a scripted fake LLM so assertions are deterministic.""" | |
| from __future__ import annotations | |
| from pathlib import Path | |
| import duckdb | |
| import pandas as pd | |
| from lexsi_ds.agent.context import AgentContext, ColumnInfo, DatasetHandle, TableInfo | |
| from lexsi_ds.agent.tools.check_result import CheckResultArgs | |
| from lexsi_ds.agent.tools.check_result import _run as check_run | |
| from lexsi_ds.agent.tools.compare_candidates import CompareCandidatesArgs | |
| from lexsi_ds.agent.tools.compare_candidates import _run as cmp_run | |
| from lexsi_ds.llm.client import LLMResult | |
| class FakeLLM: | |
| name = "fake" | |
| def __init__(self, *texts): | |
| self.texts = list(texts) | |
| self.i = 0 | |
| def complete(self, system, user): | |
| t = self.texts[min(self.i, len(self.texts) - 1)] | |
| self.i += 1 | |
| return LLMResult(text=t) | |
| def _ctx(tmp_path, llm, tables=None): | |
| ds = DatasetHandle(id="t", kind="attached", | |
| duckdb_path=tmp_path / "x.duckdb", tables=tables or []) | |
| return AgentContext(dataset=ds, run_id="t", llm=llm) | |
| # ---- check_result ---- | |
| def test_check_result_flags_wrong(tmp_path): | |
| ctx = _ctx(tmp_path, FakeLLM( | |
| '{"verdict":"wrong","issues":["grouped by track not song"],"fix":"GROUP BY song"}')) | |
| ctx.cache["last_sql_result"] = pd.DataFrame({"track": ["a"], "rev": [1]}) | |
| ctx.cache["last_sql"] = "SELECT track, SUM(rev) FROM sales GROUP BY track" | |
| res = check_run(CheckResultArgs(question="which song earned the most?"), ctx) | |
| assert res.ok | |
| assert res.payload["verdict"] == "wrong" | |
| assert "track" in res.payload["issues"][0] | |
| assert res.payload["fix"] | |
| def test_check_result_ok_verdict(tmp_path): | |
| ctx = _ctx(tmp_path, FakeLLM('{"verdict":"ok","issues":[],"fix":""}')) | |
| ctx.cache["last_sql_result"] = pd.DataFrame({"n": [42]}) | |
| ctx.cache["last_sql"] = "SELECT count(*) AS n FROM t" | |
| res = check_run(CheckResultArgs(question="how many?"), ctx) | |
| assert res.ok and res.payload["verdict"] == "ok" | |
| def test_check_result_no_result_errors(tmp_path): | |
| res = check_run(CheckResultArgs(question="x"), _ctx(tmp_path, FakeLLM("{}"))) | |
| assert not res.ok and res.error == "no_result" | |
| def test_check_result_non_json_is_unsure(tmp_path): | |
| ctx = _ctx(tmp_path, FakeLLM("Looks fine to me.")) | |
| ctx.cache["last_sql_result"] = pd.DataFrame({"n": [1]}) | |
| res = check_run(CheckResultArgs(question="x"), ctx) | |
| assert res.ok and res.payload["verdict"] == "unsure" | |
| # ---- compare_candidates ---- | |
| def _duck(tmp_path): | |
| p = tmp_path / "x.duckdb" | |
| con = duckdb.connect(str(p)) | |
| con.execute("CREATE TABLE t (x INTEGER)") | |
| con.executemany("INSERT INTO t VALUES (?)", [(1,), (2,), (3,)]) | |
| con.close() | |
| return [TableInfo(name="t", columns=[ColumnInfo("x", "INTEGER")], n_rows=3)] | |
| def test_compare_candidates_agree(tmp_path): | |
| tables = _duck(tmp_path) | |
| ctx = _ctx(tmp_path, FakeLLM("SELECT count(*) FROM t", "SELECT count(*) FROM t"), tables) | |
| res = cmp_run(CompareCandidatesArgs(question="how many rows?", n=2), ctx) | |
| assert res.ok and res.payload["agree"] is True | |
| assert len(res.payload["candidates"]) == 2 | |
| def test_compare_candidates_disagree(tmp_path): | |
| tables = _duck(tmp_path) | |
| ctx = _ctx(tmp_path, FakeLLM("SELECT count(*) FROM t", "SELECT sum(x) FROM t"), tables) | |
| res = cmp_run(CompareCandidatesArgs(question="how big is t?", n=2), ctx) | |
| assert res.ok and res.payload["agree"] is False | |
| def test_compare_candidates_surfaces_sql_error(tmp_path): | |
| tables = _duck(tmp_path) | |
| ctx = _ctx(tmp_path, FakeLLM("SELECT count(*) FROM t", "SELECT * FROM nonexistent"), tables) | |
| res = cmp_run(CompareCandidatesArgs(question="?", n=2), ctx) | |
| assert res.ok | |
| errs = [c["error"] for c in res.payload["candidates"]] | |
| assert any(e for e in errs) # the bad candidate's error is captured, not raised | |