diff --git a/plugins/dirsql-plugin-embeddings/changelog.d/2026-09-05-one-function-per-file.md b/plugins/dirsql-plugin-embeddings/changelog.d/2026-09-05-one-function-per-file.md new file mode 100644 index 000000000..6145c2af2 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/changelog.d/2026-09-05-one-function-per-file.md @@ -0,0 +1,7 @@ +**Changed** + +Split `search/sql.py`, `search/run.py`, and `embedding/model.py` into one +module per non-trivial function, so the tree is clean at +`one-function-per-file`'s strictest threshold. Every symbol keeps its behavior; +`search.run._search` is now `search.search_rows.search_rows`, and the SQL +builders moved to modules named for them. diff --git a/plugins/dirsql-plugin-embeddings/e2e-attestations/claude-1095-ofpf-embeddings.json b/plugins/dirsql-plugin-embeddings/e2e-attestations/claude-1095-ofpf-embeddings.json new file mode 100644 index 000000000..abf6f4291 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/e2e-attestations/claude-1095-ofpf-embeddings.json @@ -0,0 +1,7 @@ +{ + "command": "uv run --with-editable ../../packages/python python -m pytest tests/e2e/ -x -q", + "ran_at": 1788575701, + "exit_code": 0, + "commit": "8f07e97547e3c145e6e6635e89839b6318d25b7c", + "branch": "claude/1095-ofpf-embeddings" +} diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/model.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/model.py index 84d0ce3a0..134334c6c 100644 --- a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/model.py +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/model.py @@ -5,11 +5,3 @@ def load_model(model_id): from model2vec import StaticModel return StaticModel.from_pretrained(model_id) - - -def model_identifier(model_id, model): - config = getattr(model, "config", None) or {} - version = config.get("model2vec_version") - if version: - return f"{model_id}@{version}" - return model_id diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/model_identifier.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/model_identifier.py new file mode 100644 index 000000000..ba1e75cc0 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/model_identifier.py @@ -0,0 +1,6 @@ +def model_identifier(model_id, model): + config = getattr(model, "config", None) or {} + version = config.get("model2vec_version") + if version: + return f"{model_id}@{version}" + return model_id diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/model_identifier_test.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/model_identifier_test.py new file mode 100644 index 000000000..829526060 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/model_identifier_test.py @@ -0,0 +1,27 @@ +from types import SimpleNamespace + +from . import model_identifier as module + + +def describe_model_identifier(): + def it_appends_the_model2vec_version_when_the_config_carries_one(): + loaded = SimpleNamespace(config={"model2vec_version": "0.3.0"}) + assert ( + module.model_identifier("some/model-id", loaded) + == "some/model-id@0.3.0" + ) + + def it_is_the_model_id_alone_when_the_config_has_no_version(): + loaded = SimpleNamespace(config={"normalize": False}) + assert module.model_identifier("some/model-id", loaded) == "some/model-id" + + def it_is_the_model_id_alone_when_the_config_is_none(): + loaded = SimpleNamespace(config=None) + assert module.model_identifier("some/model-id", loaded) == "some/model-id" + + def it_is_the_model_id_alone_when_there_is_no_config_attribute(): + assert module.model_identifier("some/model-id", object()) == "some/model-id" + + def it_treats_an_empty_version_as_absent(): + loaded = SimpleNamespace(config={"model2vec_version": ""}) + assert module.model_identifier("some/model-id", loaded) == "some/model-id" diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/model_test.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/model_test.py index bdab38c53..9bd116bd0 100644 --- a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/model_test.py +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/model_test.py @@ -1,5 +1,4 @@ import sys -from types import SimpleNamespace from unittest.mock import MagicMock, patch from . import model @@ -19,27 +18,3 @@ def it_loads_via_model2vec_from_pretrained(): "some/model-id" ) assert loaded is model2vec.StaticModel.from_pretrained.return_value - - -def describe_model_identifier(): - def it_appends_the_model2vec_version_when_the_config_carries_one(): - loaded = SimpleNamespace(config={"model2vec_version": "0.3.0"}) - assert ( - model.model_identifier("some/model-id", loaded) - == "some/model-id@0.3.0" - ) - - def it_is_the_model_id_alone_when_the_config_has_no_version(): - loaded = SimpleNamespace(config={"normalize": False}) - assert model.model_identifier("some/model-id", loaded) == "some/model-id" - - def it_is_the_model_id_alone_when_the_config_is_none(): - loaded = SimpleNamespace(config=None) - assert model.model_identifier("some/model-id", loaded) == "some/model-id" - - def it_is_the_model_id_alone_when_there_is_no_config_attribute(): - assert model.model_identifier("some/model-id", object()) == "some/model-id" - - def it_treats_an_empty_version_as_absent(): - loaded = SimpleNamespace(config={"model2vec_version": ""}) - assert model.model_identifier("some/model-id", loaded) == "some/model-id" diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/worker.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/worker.py index b0bfd6b58..5ce562c06 100644 --- a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/worker.py +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/worker.py @@ -3,6 +3,7 @@ from . import model from .cache import make_cache +from .model_identifier import model_identifier from .values import ProtocolError, decode_value MALFORMED_SHAPE = ( @@ -40,7 +41,7 @@ def _compute(self, digest, identifier): def embed(self, text, model_id): """Return the vector for ``text``, and whether the cache served it.""" loaded = self._model(model_id) - identifier = model.model_identifier(model_id, loaded) + identifier = model_identifier(model_id, loaded) digest = sha256(text.encode("utf-8")).hexdigest() self._pending = (text, loaded) self._computed = False diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/worker_test.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/worker_test.py index 87ec20a2e..66a78c549 100644 --- a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/worker_test.py +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/worker_test.py @@ -1,4 +1,5 @@ import json +from contextlib import contextmanager from hashlib import sha256 from unittest.mock import MagicMock, call, patch @@ -7,6 +8,15 @@ HELLO_DIGEST = sha256(b"hello").hexdigest() +@contextmanager +def loading(identifier="m1"): + with patch.object(worker, "model") as model: + with patch.object( + worker, "model_identifier", return_value=identifier + ) as identify: + yield model, identify + + def make_worker(): with patch.object(worker, "make_cache") as make_cache: built = worker.Worker() @@ -81,20 +91,17 @@ def describe_embed(): def it_keys_the_cache_on_the_sha256_of_the_text_and_the_identifier(): built, _ = make_worker() built._compute_cached = MagicMock(return_value=[1.0]) - with patch.object(worker, "model") as model: - model.model_identifier.return_value = "m1@9.9" + with loading("m1@9.9") as (model, identify): result = built.embed("hello", "m1") model.load_model.assert_called_once_with("m1") - model.model_identifier.assert_called_once_with( - "m1", model.load_model.return_value - ) + identify.assert_called_once_with("m1", model.load_model.return_value) built._compute_cached.assert_called_once_with(HELLO_DIGEST, "m1@9.9") assert result == ([1.0], True) def it_reports_a_hit_when_compute_never_ran(): built, _ = make_worker() built._compute_cached = MagicMock(return_value=[1.0]) - with patch.object(worker, "model"): + with loading(): _, cached = built.embed("hello", "m1") assert cached is True @@ -106,7 +113,7 @@ def compute(digest, identifier): return [1.0] built._compute_cached = MagicMock(side_effect=compute) - with patch.object(worker, "model"): + with loading(): _, cached = built.embed("hello", "m1") assert cached is False @@ -116,22 +123,21 @@ def it_clears_a_previous_calls_miss_before_consulting_the_cache(): built, _ = make_worker() built._computed = True built._compute_cached = MagicMock(return_value=[1.0]) - with patch.object(worker, "model"): + with loading(): _, cached = built.embed("hello", "m1") assert cached is True def it_stages_the_text_and_model_for_compute(): built, _ = make_worker() built._compute_cached = MagicMock(return_value=[1.0]) - with patch.object(worker, "model") as model: + with loading() as (model, _): built.embed("hello", "m1") assert built._pending == ("hello", model.load_model.return_value) def it_hashes_the_utf8_bytes_of_the_text(): built, _ = make_worker() built._compute_cached = MagicMock(return_value=[1.0]) - with patch.object(worker, "model") as model: - model.model_identifier.return_value = "m1" + with loading("m1"): built.embed("héllo", "m1") expected = sha256("héllo".encode("utf-8")).hexdigest() assert built._compute_cached.call_args == call(expected, "m1") diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/build_search_sql.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/build_search_sql.py new file mode 100644 index 000000000..fb997fbed --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/build_search_sql.py @@ -0,0 +1,18 @@ +from .embed_call import embed_call +from .normalize_glob import normalize_glob +from .quote import quote + + +def build_search_sql(glob, query, limit, model=None): + outer = embed_call(quote(query), model) + inner = embed_call("content", model) + # A matched file that is unreadable, not valid UTF-8, or deleted mid-scan + # has NULL content, so its embedding is NULL too. Dropping those rows is + # what keeps the ranking honest: a NULL distance sorts FIRST ascending in + # SQLite, so unrankable files would otherwise take the top-k slots. + return ( + f"SELECT path, vec_distance_cosine(emb, {outer}) AS distance" + f" FROM (SELECT path, {inner} AS emb FROM {quote(normalize_glob(glob))})" + f" WHERE emb IS NOT NULL" + f" ORDER BY distance LIMIT {int(limit):d}" + ) diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/build_search_sql_test.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/build_search_sql_test.py new file mode 100644 index 000000000..7742370d7 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/build_search_sql_test.py @@ -0,0 +1,50 @@ +from unittest.mock import patch + +from . import build_search_sql as module + + +def build(glob, query, limit, *args): + with patch.object(module, "quote", side_effect=lambda text: f"<{text}>"): + with patch.object( + module, "normalize_glob", side_effect=lambda g: f"N({g})" + ): + with patch.object( + module, "embed_call", side_effect=lambda a, m: f"E({a},{m})" + ) as embed: + return module.build_search_sql(glob, query, limit, *args), embed + + +def describe_build_search_sql(): + def it_assembles_the_ranking_query_from_its_collaborators(): + sql, _ = build("docs/**/*.md", "local models", 10, "my/model") + assert sql == ( + "SELECT path, vec_distance_cosine(emb, E(,my/model))" + " AS distance" + " FROM (SELECT path, E(content,my/model) AS emb" + " FROM )" + " WHERE emb IS NOT NULL" + " ORDER BY distance LIMIT 10" + ) + + def it_embeds_the_query_and_the_content_column_with_the_same_model(): + _, embed = build("g", "q", 10, "my/model") + assert [args for args, _ in embed.call_args_list] == [ + ("", "my/model"), + ("content", "my/model"), + ] + + def it_defaults_the_model_to_none(): + _, embed = build("g", "q", 10) + assert [args[1] for args, _ in embed.call_args_list] == [None, None] + + def it_drops_rows_whose_embedding_is_null(): + # A file that is unreadable or not valid UTF-8 has NULL content, so + # embed() returns NULL and vec_distance_cosine() does too. SQLite + # sorts NULLs first ascending, so without this guard those files take + # the top-k slots and then break the distance formatting. + sql, _ = build("g", "q", 10) + assert " WHERE emb IS NOT NULL ORDER BY distance" in sql + + def it_coerces_the_limit_to_a_decimal_integer(): + sql, _ = build("g", "q", "12") + assert sql.endswith("LIMIT 12") diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/count_corpus_sql.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/count_corpus_sql.py new file mode 100644 index 000000000..493335de8 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/count_corpus_sql.py @@ -0,0 +1,6 @@ +from .normalize_glob import normalize_glob +from .quote import quote + + +def count_corpus_sql(glob): + return f"SELECT COUNT(*) AS n FROM {quote(normalize_glob(glob))}" diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/count_corpus_sql_test.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/count_corpus_sql_test.py new file mode 100644 index 000000000..e56f4bb7d --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/count_corpus_sql_test.py @@ -0,0 +1,18 @@ +from unittest.mock import patch + +from . import count_corpus_sql as module + + +def count(glob): + with patch.object(module, "quote", side_effect=lambda text: f"<{text}>"): + with patch.object( + module, "normalize_glob", side_effect=lambda g: f"N({g})" + ) as normalize: + return module.count_corpus_sql(glob), normalize + + +def describe_count_corpus_sql(): + def it_counts_the_files_the_normalized_quoted_glob_matches(): + sql, normalize = count("docs/**/*.md") + assert sql == "SELECT COUNT(*) AS n FROM " + normalize.assert_called_once_with("docs/**/*.md") diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/embed_call.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/embed_call.py new file mode 100644 index 000000000..cf3672a1c --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/embed_call.py @@ -0,0 +1,7 @@ +from .quote import quote + + +def embed_call(argument, model): + if model is None: + return f"embed({argument})" + return f"embed({argument}, {quote(model)})" diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/embed_call_test.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/embed_call_test.py new file mode 100644 index 000000000..9d15bb064 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/embed_call_test.py @@ -0,0 +1,16 @@ +from unittest.mock import patch + +from . import embed_call as module + + +def call(argument, model): + with patch.object(module, "quote", side_effect=lambda text: f"<{text}>"): + return module.embed_call(argument, model) + + +def describe_embed_call(): + def it_omits_the_model_argument_when_model_is_none(): + assert call("content", None) == "embed(content)" + + def it_quotes_the_model_id_as_the_second_argument(): + assert call("content", "my/model") == "embed(content, )" diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/no_rows_message.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/no_rows_message.py new file mode 100644 index 000000000..e5bf2af35 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/no_rows_message.py @@ -0,0 +1,7 @@ +def no_rows_message(glob, matched, root): + if not matched: + return f"no files matched {glob!r} (searched from {root})" + return ( + f"{glob!r} matched {matched} file(s), but none had text content to" + f" embed -- unreadable or not valid UTF-8 (searched from {root})" + ) diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/no_rows_message_test.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/no_rows_message_test.py new file mode 100644 index 000000000..1131c994f --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/no_rows_message_test.py @@ -0,0 +1,19 @@ +from . import no_rows_message as module + + +def describe_no_rows_message(): + def it_names_the_glob_and_the_root_when_no_file_matched(): + message = module.no_rows_message("docs/**/*.md", 0, "/work") + assert message == ( + "no files matched 'docs/**/*.md' (searched from /work)" + ) + + def it_says_files_matched_but_none_were_embeddable(): + message = module.no_rows_message("docs/**/*.md", 3, "/work") + assert message == ( + "'docs/**/*.md' matched 3 file(s), but none had text content to" + " embed -- unreadable or not valid UTF-8 (searched from /work)" + ) + + def it_reports_a_none_count_as_nothing_matched(): + assert "no files matched" in module.no_rows_message("g", None, "/work") diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/normalize_glob.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/normalize_glob.py new file mode 100644 index 000000000..a9ce4cd81 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/normalize_glob.py @@ -0,0 +1,11 @@ +PATH_PREFIXES = ("./", "../", "/", "~/") + + +def normalize_glob(glob): + # The core only rescues path-shaped missing tables (./, ../, /, ~/); a + # bare relative glob like '**/*.md' would error with a "did you mean + # './...'" hint. Here GLOB is unambiguously a corpus glob, so spare the + # user the round trip. + if glob.startswith(PATH_PREFIXES): + return glob + return f"./{glob}" diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/normalize_glob_test.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/normalize_glob_test.py new file mode 100644 index 000000000..f115988c2 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/normalize_glob_test.py @@ -0,0 +1,23 @@ +from . import normalize_glob as module + + +def describe_path_prefixes(): + def it_is_exactly_the_four_prefixes_the_core_rescues(): + assert module.PATH_PREFIXES == ("./", "../", "/", "~/") + + +def describe_normalize_glob(): + def it_prefixes_a_bare_relative_glob_with_dot_slash(): + assert module.normalize_glob("**/*.md") == "./**/*.md" + + def it_keeps_a_dot_slash_glob(): + assert module.normalize_glob("./notes/*.txt") == "./notes/*.txt" + + def it_keeps_a_parent_relative_glob(): + assert module.normalize_glob("../notes/*.txt") == "../notes/*.txt" + + def it_keeps_an_absolute_glob(): + assert module.normalize_glob("/tmp/notes/*.txt") == "/tmp/notes/*.txt" + + def it_keeps_a_home_glob(): + assert module.normalize_glob("~/notes/*.txt") == "~/notes/*.txt" diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/quote.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/quote.py new file mode 100644 index 000000000..61a43d462 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/quote.py @@ -0,0 +1,3 @@ +def quote(text): + escaped = text.replace("'", "''") + return f"'{escaped}'" diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/quote_test.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/quote_test.py new file mode 100644 index 000000000..44f1bd6ca --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/quote_test.py @@ -0,0 +1,12 @@ +from . import quote as module + + +def describe_quote(): + def it_wraps_text_in_single_quotes(): + assert module.quote("hello") == "'hello'" + + def it_doubles_embedded_single_quotes(): + assert module.quote("O'Brien") == "'O''Brien'" + + def it_keeps_an_injection_shaped_value_inside_the_literal(): + assert module.quote("'; DROP TABLE x; --") == "'''; DROP TABLE x; --'" diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/run.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/run.py index 9cf71578e..203ddc6e4 100644 --- a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/run.py +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/run.py @@ -1,45 +1,17 @@ import asyncio import os -from importlib import resources - -import dirsql +from .no_rows_message import no_rows_message from .output import format_rows -from .sql import build_search_sql, count_corpus_sql +from .search_rows import search_rows class NothingToRank(Exception): """The search produced no rows, and why -- so the CLI can say which.""" -def config_fragment(): - return str(resources.files("dirsql_plugin_embeddings").joinpath("dirsql.toml")) - - -def no_rows_message(glob, matched, root): - if not matched: - return f"no files matched {glob!r} (searched from {root})" - return ( - f"{glob!r} matched {matched} file(s), but none had text content to" - f" embed -- unreadable or not valid UTF-8 (searched from {root})" - ) - - -async def _search(glob, query, limit, model): - db = dirsql.DirSQL(config=config_fragment()) - rows = await db.query(build_search_sql(glob, query, limit, model)) - if rows: - return rows, None - # Only now, on the error path, is the second scan worth its cost: it is - # what tells "no files matched" apart from "matched, none embeddable". - # One row, always: `SELECT COUNT(*)` cannot return anything else, and - # unpacking says so rather than trusting an index. - (counted,) = await db.query(count_corpus_sql(glob)) - return rows, counted["n"] - - def run_search(glob, query, limit, model=None): - rows, matched = asyncio.run(_search(glob, query, limit, model)) + rows, matched = asyncio.run(search_rows(glob, query, limit, model)) if not rows: raise NothingToRank(no_rows_message(glob, matched, os.getcwd())) return format_rows(rows) diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/run_test.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/run_test.py index 764fd6402..0f7292aac 100644 --- a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/run_test.py +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/run_test.py @@ -1,88 +1,35 @@ -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import pytest from . import run -def describe_config_fragment(): - def it_resolves_the_packaged_dirsql_toml(): - with patch.object(run, "resources") as resources: - resources.files.return_value.joinpath.return_value = "frag.toml" - assert run.config_fragment() == "frag.toml" - resources.files.assert_called_once_with("dirsql_plugin_embeddings") - resources.files.return_value.joinpath.assert_called_once_with( - "dirsql.toml" - ) - - def describe_run_search(): - def it_builds_the_sql_queries_the_sdk_and_formats_the_rows(): + def it_formats_the_rows_the_search_returned(): rows = [{"path": "a", "distance": 0.5}] - fake_dirsql = MagicMock() - fake_dirsql.DirSQL.return_value.query = AsyncMock(return_value=rows) - with patch.object(run, "build_search_sql", return_value="SQL") as build: - with patch.object(run, "dirsql", fake_dirsql): - with patch.object(run, "config_fragment", return_value="frag"): - with patch.object( - run, "format_rows", return_value=["line"] - ) as formatter: - assert run.run_search("g", "q", 5, "m") == ["line"] - build.assert_called_once_with("g", "q", 5, "m") - fake_dirsql.DirSQL.assert_called_once_with(config="frag") - fake_dirsql.DirSQL.return_value.query.assert_awaited_once_with("SQL") + found = AsyncMock(return_value=(rows, None)) + with patch.object(run, "search_rows", found): + with patch.object( + run, "format_rows", return_value=["line"] + ) as formatter: + assert run.run_search("g", "q", 5, "m") == ["line"] + found.assert_awaited_once_with("g", "q", 5, "m") formatter.assert_called_once_with(rows) def it_defaults_the_model_to_none(): - fake_dirsql = MagicMock() - fake_dirsql.DirSQL.return_value.query = AsyncMock( - return_value=[{"path": "a", "distance": 0.5}] - ) - with patch.object(run, "build_search_sql", return_value="SQL") as build: - with patch.object(run, "dirsql", fake_dirsql): - with patch.object(run, "config_fragment", return_value="frag"): - assert run.run_search("g", "q", 5) == ["a\t0.500000"] - build.assert_called_once_with("g", "q", 5, None) - - -def sdk_returning(*results): - fake_dirsql = MagicMock() - fake_dirsql.DirSQL.return_value.query = AsyncMock(side_effect=results) - return fake_dirsql - - -def search_with(fake_dirsql): - with patch.object(run, "dirsql", fake_dirsql): - with patch.object(run, "config_fragment", return_value="frag"): - return run.run_search("docs/**/*.md", "q", 5) - - -def describe_empty_results(): - def it_raises_naming_the_glob_when_no_file_matched(): - with pytest.raises(run.NothingToRank) as raised: - search_with(sdk_returning([], [{"n": 0}])) - message = str(raised.value) - assert "no files matched" in message - assert "docs/**/*.md" in message - assert "searched from" in message - - def it_says_so_when_files_matched_but_none_could_be_embedded(): - with pytest.raises(run.NothingToRank) as raised: - search_with(sdk_returning([], [{"n": 3}])) - message = str(raised.value) - assert "3" in message - assert "docs/**/*.md" in message - assert "no files matched" not in message, ( - "files did match; the message must not claim otherwise" - ) - - def it_counts_the_corpus_only_when_the_ranking_is_empty(): - fake_dirsql = sdk_returning([{"path": "a", "distance": 0.5}]) - assert search_with(fake_dirsql) == ["a\t0.500000"] - assert fake_dirsql.DirSQL.return_value.query.await_count == 1 - - def it_reuses_one_sdk_instance_for_both_queries(): - fake_dirsql = sdk_returning([], [{"n": 0}]) - with pytest.raises(run.NothingToRank): - search_with(fake_dirsql) - assert fake_dirsql.DirSQL.call_count == 1 + found = AsyncMock(return_value=([{"path": "a"}], None)) + with patch.object(run, "search_rows", found): + with patch.object(run, "format_rows", return_value=[]): + run.run_search("g", "q", 5) + found.assert_awaited_once_with("g", "q", 5, None) + + def it_raises_with_the_no_rows_message_when_nothing_ranked(): + with patch.object(run, "search_rows", AsyncMock(return_value=([], 3))): + with patch.object( + run, "no_rows_message", return_value="why" + ) as message: + with patch("os.getcwd", return_value="/work"): + with pytest.raises(run.NothingToRank, match="why"): + run.run_search("g", "q", 5) + message.assert_called_once_with("g", 3, "/work") diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/search_rows.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/search_rows.py new file mode 100644 index 000000000..89bcd4f05 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/search_rows.py @@ -0,0 +1,23 @@ +from importlib import resources + +import dirsql + +from .build_search_sql import build_search_sql +from .count_corpus_sql import count_corpus_sql + + +def config_fragment(): + return str(resources.files("dirsql_plugin_embeddings").joinpath("dirsql.toml")) + + +async def search_rows(glob, query, limit, model): + db = dirsql.DirSQL(config=config_fragment()) + rows = await db.query(build_search_sql(glob, query, limit, model)) + if rows: + return rows, None + # Only now, on the error path, is the second scan worth its cost: it is + # what tells "no files matched" apart from "matched, none embeddable". + # One row, always: `SELECT COUNT(*)` cannot return anything else, and + # unpacking says so rather than trusting an index. + (counted,) = await db.query(count_corpus_sql(glob)) + return rows, counted["n"] diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/search_rows_test.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/search_rows_test.py new file mode 100644 index 000000000..45b8043a9 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/search_rows_test.py @@ -0,0 +1,60 @@ +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +from . import search_rows as module + + +def sdk_returning(*results): + fake_dirsql = MagicMock() + fake_dirsql.DirSQL.return_value.query = AsyncMock(side_effect=results) + return fake_dirsql + + +def search(fake_dirsql, glob="docs/**/*.md", model=None): + with patch.object(module, "dirsql", fake_dirsql): + with patch.object(module, "config_fragment", return_value="frag"): + with patch.object( + module, "build_search_sql", return_value="SQL" + ) as build: + with patch.object( + module, "count_corpus_sql", return_value="COUNT" + ) as count: + rows = asyncio.run(module.search_rows(glob, "q", 5, model)) + return rows, build, count + + +def describe_config_fragment(): + def it_resolves_the_packaged_dirsql_toml(): + with patch.object(module, "resources") as resources: + resources.files.return_value.joinpath.return_value = "frag.toml" + assert module.config_fragment() == "frag.toml" + resources.files.assert_called_once_with("dirsql_plugin_embeddings") + resources.files.return_value.joinpath.assert_called_once_with( + "dirsql.toml" + ) + + +def describe_search_rows(): + def it_queries_the_sdk_with_the_built_sql_and_the_packaged_fragment(): + fake_dirsql = sdk_returning([{"path": "a"}]) + rows, build, _ = search(fake_dirsql, model="m") + assert rows == ([{"path": "a"}], None) + build.assert_called_once_with("docs/**/*.md", "q", 5, "m") + fake_dirsql.DirSQL.assert_called_once_with(config="frag") + fake_dirsql.DirSQL.return_value.query.assert_awaited_once_with("SQL") + + def it_counts_the_corpus_only_when_the_ranking_is_empty(): + _, _, count = search(sdk_returning([{"path": "a"}])) + count.assert_not_called() + + def it_returns_the_corpus_count_when_nothing_ranked(): + fake_dirsql = sdk_returning([], [{"n": 3}]) + rows, _, count = search(fake_dirsql) + assert rows == ([], 3) + count.assert_called_once_with("docs/**/*.md") + assert fake_dirsql.DirSQL.return_value.query.await_count == 2 + + def it_reuses_one_sdk_instance_for_both_queries(): + fake_dirsql = sdk_returning([], [{"n": 0}]) + search(fake_dirsql) + assert fake_dirsql.DirSQL.call_count == 1 diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/sql.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/sql.py deleted file mode 100644 index a5090919c..000000000 --- a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/sql.py +++ /dev/null @@ -1,41 +0,0 @@ -PATH_PREFIXES = ("./", "../", "/", "~/") - - -def quote(text): - escaped = text.replace("'", "''") - return f"'{escaped}'" - - -def normalize_glob(glob): - # The core only rescues path-shaped missing tables (./, ../, /, ~/); a - # bare relative glob like '**/*.md' would error with a "did you mean - # './...'" hint. Here GLOB is unambiguously a corpus glob, so spare the - # user the round trip. - if glob.startswith(PATH_PREFIXES): - return glob - return f"./{glob}" - - -def embed_call(argument, model): - if model is None: - return f"embed({argument})" - return f"embed({argument}, {quote(model)})" - - -def build_search_sql(glob, query, limit, model=None): - outer = embed_call(quote(query), model) - inner = embed_call("content", model) - # A matched file that is unreadable, not valid UTF-8, or deleted mid-scan - # has NULL content, so its embedding is NULL too. Dropping those rows is - # what keeps the ranking honest: a NULL distance sorts FIRST ascending in - # SQLite, so unrankable files would otherwise take the top-k slots. - return ( - f"SELECT path, vec_distance_cosine(emb, {outer}) AS distance" - f" FROM (SELECT path, {inner} AS emb FROM {quote(normalize_glob(glob))})" - f" WHERE emb IS NOT NULL" - f" ORDER BY distance LIMIT {int(limit):d}" - ) - - -def count_corpus_sql(glob): - return f"SELECT COUNT(*) AS n FROM {quote(normalize_glob(glob))}" diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/sql_test.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/sql_test.py deleted file mode 100644 index 491e0cae4..000000000 --- a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/search/sql_test.py +++ /dev/null @@ -1,99 +0,0 @@ -from . import sql - - -def describe_quote(): - def it_wraps_text_in_single_quotes(): - assert sql.quote("hello") == "'hello'" - - def it_doubles_embedded_single_quotes(): - assert sql.quote("O'Brien") == "'O''Brien'" - - def it_keeps_an_injection_shaped_value_inside_the_literal(): - assert sql.quote("'; DROP TABLE x; --") == "'''; DROP TABLE x; --'" - - -def describe_normalize_glob(): - def it_prefixes_a_bare_relative_glob_with_dot_slash(): - assert sql.normalize_glob("**/*.md") == "./**/*.md" - - def it_keeps_a_dot_slash_glob(): - assert sql.normalize_glob("./notes/*.txt") == "./notes/*.txt" - - def it_keeps_a_parent_relative_glob(): - assert sql.normalize_glob("../notes/*.txt") == "../notes/*.txt" - - def it_keeps_an_absolute_glob(): - assert sql.normalize_glob("/tmp/notes/*.txt") == "/tmp/notes/*.txt" - - def it_keeps_a_home_glob(): - assert sql.normalize_glob("~/notes/*.txt") == "~/notes/*.txt" - - -def describe_embed_call(): - def it_omits_the_model_argument_when_model_is_none(): - assert sql.embed_call("content", None) == "embed(content)" - - def it_quotes_the_model_id_as_the_second_argument(): - assert sql.embed_call("content", "my/model") == ( - "embed(content, 'my/model')" - ) - - def it_escapes_quotes_in_the_model_id(): - assert sql.embed_call("content", "it's") == "embed(content, 'it''s')" - - -def describe_build_search_sql(): - def it_builds_the_canonical_sql_without_a_model(): - assert sql.build_search_sql("./docs/**/*.md", "local models", 10) == ( - "SELECT path," - " vec_distance_cosine(emb, embed('local models')) AS distance" - " FROM (SELECT path, embed(content) AS emb FROM './docs/**/*.md')" - " WHERE emb IS NOT NULL" - " ORDER BY distance LIMIT 10" - ) - - def it_templates_the_model_into_both_embed_calls(): - assert sql.build_search_sql("./docs/**/*.md", "q", 3, "my/model") == ( - "SELECT path," - " vec_distance_cosine(emb, embed('q', 'my/model')) AS distance" - " FROM (SELECT path, embed(content, 'my/model') AS emb" - " FROM './docs/**/*.md')" - " WHERE emb IS NOT NULL" - " ORDER BY distance LIMIT 3" - ) - - def it_drops_rows_whose_embedding_is_null(): - # A file that is unreadable or not valid UTF-8 has NULL content, so - # embed() returns NULL and vec_distance_cosine() does too. SQLite - # sorts NULLs first ascending, so without this guard those files take - # the top-k slots and then break the distance formatting. - built = sql.build_search_sql("./docs/**/*.md", "q", 10) - assert " WHERE emb IS NOT NULL ORDER BY distance" in built - - def it_normalizes_a_bare_glob_into_the_from_clause(): - built = sql.build_search_sql("docs/**/*.md", "q", 10) - assert "FROM './docs/**/*.md'" in built - - def it_escapes_quotes_in_query_and_glob(): - built = sql.build_search_sql("it's/**", "it's a 'test'", 10) - assert "embed('it''s a ''test''')" in built - assert "FROM './it''s/**'" in built - - def it_defaults_the_model_to_none(): - assert "embed(content)" in sql.build_search_sql("g", "q", 10) - - def it_coerces_the_limit_to_a_decimal_integer(): - assert sql.build_search_sql("g", "q", "12").endswith("LIMIT 12") - - -def describe_count_corpus_sql(): - def it_counts_the_files_the_glob_matches(): - assert sql.count_corpus_sql("./docs/**/*.md") == ( - "SELECT COUNT(*) AS n FROM './docs/**/*.md'" - ) - - def it_normalizes_a_bare_glob(): - assert "FROM './docs/**/*.md'" in sql.count_corpus_sql("docs/**/*.md") - - def it_escapes_quotes_in_the_glob(): - assert "FROM './it''s/**'" in sql.count_corpus_sql("it's/**")