Spaces:
Sleeping
Sleeping
File size: 4,303 Bytes
f619a51 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """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
|