diff --git a/docs/docs/cli/llm-stamp-themes.md b/docs/docs/cli/llm-stamp-themes.md index 1ebaf77..a8ca76c 100644 --- a/docs/docs/cli/llm-stamp-themes.md +++ b/docs/docs/cli/llm-stamp-themes.md @@ -21,7 +21,7 @@ Usage: stack llm stamp-themes [OPTIONS] │ --key TEXT Only this item key. │ │ --top INTEGER Themes kept per chunk. [default: 2] │ │ --min-score FLOAT Cosine similarity a theme must reach. │ -│ [default: 0.55] │ +│ [default: 0.6] │ │ --force Re-stamp chunks already on the current │ │ vocabulary version. │ │ --help Show this message and exit. │ diff --git a/src/cli/llm.py b/src/cli/llm.py index 6bb5f61..fccc8ea 100644 --- a/src/cli/llm.py +++ b/src/cli/llm.py @@ -489,7 +489,7 @@ def stamp_themes( key: str = typer.Option("", "--key", help="Only this item key."), top: int = typer.Option(2, "--top", help="Themes kept per chunk."), min_score: float = typer.Option( - 0.55, "--min-score", help="Cosine similarity a theme must reach." + 0.60, "--min-score", help="Cosine similarity a theme must reach." ), force: bool = typer.Option( False, diff --git a/src/llm/search.py b/src/llm/search.py index 6fab6d9..cd8cc6a 100644 --- a/src/llm/search.py +++ b/src/llm/search.py @@ -54,13 +54,16 @@ def _collections_for(collection: str) -> tuple[str, ...]: return (collection,) -def _store_filter(filters: dict[str, str]) -> dict[str, str]: +def _store_filter(filters: dict[str, str]) -> dict[str, Any]: """The langchain PGVector ``filter=`` metadata dict for *filters* — docket/item_key/kind and the manual coordinates (doctype, manual, chapter, iom_section) are pushed down to the store; ``year`` is applied in Python (:func:`_year_ok`) since it's a prefix of the - ``date`` field, not its own metadata key.""" - return { + ``date`` field, not its own metadata key. A ``theme`` becomes an + ``$or`` of equalities over the ``theme_1..theme_3`` rank keys + ``llm.themes`` writes, so the overfetch is not spent on off-theme + chunks (the comma-joined ``themes`` is re-checked in Python).""" + out: dict[str, Any] = { key: filters[key] for key in ( "docket", @@ -73,6 +76,10 @@ def _store_filter(filters: dict[str, str]) -> dict[str, str]: ) if filters.get(key) } + theme = str(filters.get("theme") or "") + if theme: + out["$or"] = [{f"theme_{i}": theme} for i in range(1, 4)] + return out def _year_ok(md: dict[str, Any], year: str) -> bool: diff --git a/src/llm/themes.py b/src/llm/themes.py index 881e466..adb461f 100644 --- a/src/llm/themes.py +++ b/src/llm/themes.py @@ -93,7 +93,7 @@ def stamp( doctype: str = "manual", item_key: str = "", top: int = 2, - min_score: float = 0.55, + min_score: float = 0.60, batch: int = 1000, force: bool = False, ) -> dict[str, int]: @@ -136,6 +136,12 @@ def stamp( "themes": ",".join(s for s, _ in chosen), "theme_scores": ",".join(f"{s}:{sc}" for s, sc in chosen), "themes_vocab": version, + # one scalar key per rank so the store can filter with an + # equality ($or over theme_1..theme_N) instead of a LIKE + **{ + f"theme_{i}": (chosen[i - 1][0] if i <= len(chosen) else "") + for i in range(1, top + 1) + }, } patches.append({"id": cid, "patch": json.dumps(patch)}) if patches: diff --git a/tests/llm/test_search.py b/tests/llm/test_search.py index ab208c2..78a7a0b 100644 --- a/tests/llm/test_search.py +++ b/tests/llm/test_search.py @@ -391,3 +391,12 @@ class TestThemeFilter: assert not _theme_ok({"themes": "telehealth,supervision"}, "drugs") assert not _theme_ok({}, "drugs") assert _theme_ok({}, "") + + def test_theme_is_pushed_down_as_an_or_of_rank_keys(self): + out = _store_filter({"theme": "telehealth", "doctype": "manual"}) + assert out["doctype"] == "manual" + assert out["$or"] == [ + {"theme_1": "telehealth"}, + {"theme_2": "telehealth"}, + {"theme_3": "telehealth"}, + ] diff --git a/tests/llm/test_themes.py b/tests/llm/test_themes.py index 2503a75..c12538c 100644 --- a/tests/llm/test_themes.py +++ b/tests/llm/test_themes.py @@ -64,7 +64,14 @@ class TestStamp: 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 by["a"]["theme_1"] == "telehealth" and by["a"]["theme_2"] == "mixed" + assert by["c"] == { + "themes": "", + "theme_scores": "", + "themes_vocab": "2", + "theme_1": "", + "theme_2": "", + } assert "b" not in by def test_force_restamps_everything_and_pages_by_id(self):