Spaces:
Sleeping
Sleeping
| """transform_column: derive a structured column from free text via the LLM, | |
| materialize it, then query with SQL. Scripted fake LLM for determinism.""" | |
| from __future__ import annotations | |
| import duckdb | |
| import pytest | |
| from lexsi_ds.agent.context import AgentContext, ColumnInfo, DatasetHandle, TableInfo | |
| from lexsi_ds.agent.tools.transform_column import TransformColumnArgs, _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, rows, llm): | |
| p = tmp_path / "d.duckdb" | |
| con = duckdb.connect(str(p)) | |
| con.execute("CREATE TABLE proj (repo VARCHAR, prose VARCHAR)") | |
| con.executemany("INSERT INTO proj VALUES (?, ?)", rows) | |
| con.close() | |
| ds = DatasetHandle(id="t", kind="attached", duckdb_path=p, tables=[ | |
| TableInfo(name="proj", columns=[ColumnInfo("repo", "VARCHAR"), | |
| ColumnInfo("prose", "VARCHAR")])]) | |
| return AgentContext(dataset=ds, run_id="t", llm=llm) | |
| def test_extracts_number_from_prose_as_integer(tmp_path): | |
| rows = [("a/b", "has 38715 stars and 10 forks"), ("c/d", "currently 0 stars")] | |
| ctx = _ctx(tmp_path, rows, FakeLLM('["38715", "0"]')) | |
| res = _run(TransformColumnArgs(table="proj", source_column="prose", | |
| new_column="stars", instruction="extract the integer star count", | |
| output_type="integer"), ctx) | |
| assert res.ok, res.summary | |
| con = duckdb.connect(str(ctx.dataset.duckdb_path), read_only=True) | |
| rows_out = con.execute('SELECT repo, stars FROM "proj__stars" ORDER BY stars DESC').fetchall() | |
| typ = con.execute("SELECT data_type FROM information_schema.columns " | |
| "WHERE table_name='proj__stars' AND column_name='stars'").fetchone()[0] | |
| con.close() | |
| assert rows_out[0] == ("a/b", 38715) # numeric, sortable | |
| assert "INT" in typ.upper() or "BIGINT" in typ.upper() | |
| def test_where_scopes_rows(tmp_path): | |
| rows = [("a/b", "100 stars"), ("c/d", "200 stars")] | |
| ctx = _ctx(tmp_path, rows, FakeLLM('["100"]')) | |
| res = _run(TransformColumnArgs(table="proj", source_column="prose", new_column="stars", | |
| instruction="extract stars", where="repo = 'a/b'", | |
| output_type="integer"), ctx) | |
| assert res.ok and res.payload["n_rows"] == 1 | |
| def test_classification_labels(tmp_path): | |
| rows = [("x", "stocks tumble on wall street"), ("y", "team wins the cup")] | |
| ctx = _ctx(tmp_path, rows, FakeLLM('["Business", "Sports"]')) | |
| res = _run(TransformColumnArgs(table="proj", source_column="prose", new_column="category", | |
| instruction="classify into World/Sports/Business/SciTech"), ctx) | |
| assert res.ok | |
| con = duckdb.connect(str(ctx.dataset.duckdb_path), read_only=True) | |
| cats = dict(con.execute('SELECT repo, category FROM "proj__category"').fetchall()) | |
| con.close() | |
| assert cats == {"x": "Business", "y": "Sports"} | |
| def test_no_llm_errors(tmp_path): | |
| ctx = _ctx(tmp_path, [("a", "x")], FakeLLM("[]")) | |
| ctx.llm = None | |
| res = _run(TransformColumnArgs(table="proj", source_column="prose", | |
| new_column="c", instruction="x"), ctx) | |
| assert not res.ok and res.error == "no_llm" | |
| def test_bad_new_column_rejected(tmp_path): | |
| ctx = _ctx(tmp_path, [("a", "x")], FakeLLM("[]")) | |
| res = _run(TransformColumnArgs(table="proj", source_column="prose", | |
| new_column="bad name", instruction="x"), ctx) | |
| assert not res.ok and res.error == "bad_identifier" | |
| def test_short_array_padded_with_none(tmp_path): | |
| # LLM returns fewer outputs than inputs → missing rows become null, no crash | |
| rows = [("a", "10 stars"), ("b", "20 stars")] | |
| ctx = _ctx(tmp_path, rows, FakeLLM('["10"]')) | |
| res = _run(TransformColumnArgs(table="proj", source_column="prose", new_column="stars", | |
| instruction="extract stars", output_type="integer"), ctx) | |
| assert res.ok and res.payload["n_nonnull"] == 1 | |