102234e132
Some checks failed
CI / lint (push) Has been cancelled
CI / test (push) Has been cancelled
CI / notebooks-smoke (push) Has been cancelled
Deploy / notebooks (push) Has been cancelled
Deploy / docs (push) Has been cancelled
Deploy / api (push) Has been cancelled
Deploy / llm (push) Has been cancelled
Deploy / mc (push) Has been cancelled
Deploy / report (push) Has been cancelled
Infra CI / notebooks (push) Has been cancelled
Infra CI / zotero (push) Has been cancelled
Infra CI / docs (push) Has been cancelled
Infra CI / api (push) Has been cancelled
Infra CI / llm (push) Has been cancelled
Infra CI / mc (push) Has been cancelled
Deploy / zotero (push) Has been cancelled
77 lines
2.7 KiB
Python
77 lines
2.7 KiB
Python
"""llm.themes — cosine theme stamps on indexed chunks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from unittest.mock import MagicMock
|
|
|
|
from llm.themes import pick, stamp
|
|
|
|
CARDS = {"telehealth": [1.0, 0.0], "drugs": [0.0, 1.0], "mixed": [0.7071, 0.7071]}
|
|
|
|
|
|
class TestPick:
|
|
def test_top_by_similarity_with_threshold(self):
|
|
assert pick([1.0, 0.0], CARDS, top=2, min_score=0.5) == [
|
|
("telehealth", 1.0),
|
|
("mixed", 0.7071),
|
|
]
|
|
assert pick([1.0, 0.0], CARDS, top=1, min_score=0.5) == [("telehealth", 1.0)]
|
|
assert pick([1.0, 0.0], CARDS, top=3, min_score=0.9) == [("telehealth", 1.0)]
|
|
assert pick([0.0, 0.0], CARDS, top=2, min_score=0.1) == []
|
|
|
|
|
|
def _engine(pages):
|
|
"""A fake engine whose SELECT returns *pages* in turn, then nothing;
|
|
UPDATEs are recorded."""
|
|
engine = MagicMock()
|
|
conn = engine.begin.return_value.__enter__.return_value
|
|
selects = list(pages) + [[]]
|
|
updates = []
|
|
|
|
def execute(sql, params=None):
|
|
if "SELECT" in str(sql):
|
|
r = MagicMock()
|
|
r.fetchall.return_value = selects.pop(0)
|
|
return r
|
|
updates.extend(params) # an UPDATE with a list of params (executemany)
|
|
return MagicMock()
|
|
|
|
conn.execute.side_effect = execute
|
|
return engine, updates
|
|
|
|
|
|
class TestStamp:
|
|
def test_stamps_unstamped_chunks_and_skips_current_ones(self):
|
|
rows = [
|
|
("a", json.dumps([1.0, 0.0]), None),
|
|
(
|
|
"b",
|
|
json.dumps([0.0, 1.0]),
|
|
"2",
|
|
), # already stamped with this vocab version
|
|
(
|
|
"c",
|
|
json.dumps([0.0, 0.0]),
|
|
"1",
|
|
), # stale version, no theme clears the bar
|
|
]
|
|
engine, updates = _engine([rows])
|
|
stats = stamp(engine, CARDS, vocab_version=2, top=2, min_score=0.5)
|
|
assert stats == {"scanned": 3, "stamped": 2, "skipped": 1, "unthemed": 1}
|
|
by = {u["id"]: json.loads(u["patch"]) for u in updates}
|
|
assert (
|
|
by["a"]["themes"] == "telehealth,mixed" and by["a"]["themes_vocab"] == "2"
|
|
)
|
|
assert by["a"]["theme_scores"].startswith("telehealth:1.0,mixed:0.7071")
|
|
assert by["c"] == {"themes": "", "theme_scores": "", "themes_vocab": "2"}
|
|
assert "b" not in by
|
|
|
|
def test_force_restamps_everything_and_pages_by_id(self):
|
|
page1 = [("a", json.dumps([1.0, 0.0]), "2")]
|
|
page2 = [("b", json.dumps([0.0, 1.0]), "2")]
|
|
engine, updates = _engine([page1, page2])
|
|
stats = stamp(engine, CARDS, vocab_version=2, force=True, batch=1)
|
|
assert stats["stamped"] == 2 and stats["skipped"] == 0
|
|
assert [u["id"] for u in updates] == ["a", "b"]
|