diff --git a/docs/plugins.md b/docs/plugins.md index 3123f5056..dbf7fd6a5 100644 --- a/docs/plugins.md +++ b/docs/plugins.md @@ -93,9 +93,10 @@ scope. ### Vector cache Computed vectors are cached at `~/.cache/dirsql/embeddings/` (respecting -`XDG_CACHE_HOME`), keyed on the SHA-256 of the value bytes plus the model -identifier. Changing either recomputes — switching models never serves stale -vectors — and re-running a query over unchanged files is cache hits all the -way. There is no eviction: the directory is **safe to wipe at any time**; the -only cost is re-embedding. The cache never lives inside a queried tree — the -worker writes nothing into the directories you query. +`XDG_CACHE_HOME`): one SQLite database per model identifier, each row keyed +on the SHA-256 of the value bytes. Changing the value or the model recomputes +— switching models never serves stale vectors — and re-running a query over +unchanged files is cache hits all the way, a whole batch answered by one +query. There is no eviction: the directory is **safe to wipe at any time**; +the only cost is re-embedding. The cache never lives inside a queried tree — +the worker writes nothing into the directories you query. diff --git a/plugins/dirsql-plugin-embeddings/README.md b/plugins/dirsql-plugin-embeddings/README.md index 68d375715..d43cfa971 100644 --- a/plugins/dirsql-plugin-embeddings/README.md +++ b/plugins/dirsql-plugin-embeddings/README.md @@ -53,9 +53,10 @@ Results print one `pathdistance` line per match, closest first. `embed()` is inert until a query calls it: no worker process is spawned and no model is loaded for queries that never use it. On the first call, dirsql spawns the plugin's worker process (`dirsql-plugin-embeddings worker`), which -serves every call of the invocation over stdin/stdout. Only the values the -query actually selects are embedded — the worker receives values, not paths, -and never opens files itself. +serves every call of the invocation over stdin/stdout. A statement's values +reach the worker in batches of up to 4096, each encoded in one model call. +Only the values the query actually selects are embedded — the worker receives +values, not paths, and never opens files itself. ## Model @@ -79,12 +80,14 @@ The one-liner's `--model` flag templates the same second argument. ## Vector cache Computed vectors are cached at `~/.cache/dirsql/embeddings/` (or -`$XDG_CACHE_HOME/dirsql/embeddings/` when `XDG_CACHE_HOME` is set), keyed by -the SHA-256 of the value bytes plus the model identifier — changing either -recomputes; switching models never serves stale vectors. There is no -eviction: **the directory is safe to wipe at any time**; the only cost is -re-embedding. The cache never lives inside a queried tree — the worker -receives values, not paths, and writes nothing anywhere else. +`$XDG_CACHE_HOME/dirsql/embeddings/` when `XDG_CACHE_HOME` is set): one +SQLite database per model identifier, each row keyed by the SHA-256 of the +value bytes. Changing the value or the model recomputes; switching models +never serves stale vectors. A batch is looked up in one query and its misses +written in one transaction. There is no eviction: **the directory is safe to +wipe at any time**; the only cost is re-embedding. The cache never lives +inside a queried tree — the worker receives values, not paths, and writes +nothing anywhere else. ## Docs diff --git a/plugins/dirsql-plugin-embeddings/changelog.d/2026-10-02-batched-embed-sqlite-cache.md b/plugins/dirsql-plugin-embeddings/changelog.d/2026-10-02-batched-embed-sqlite-cache.md new file mode 100644 index 000000000..748a660a4 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/changelog.d/2026-10-02-batched-embed-sqlite-cache.md @@ -0,0 +1,4 @@ +**Changed** + +- **`embed()` now encodes a statement's values in batches.** The packaged fragment sets `batch = 4096`, so dirsql hands the worker up to 4096 values per request and the worker answers them with one `model.encode()` call per model instead of one per value. Embedding every abstract in a corpus, startup excluded, went from 4.5x the time of a bare batch encode (32768 abstracts) to 0.83x of it (65536 abstracts). The worker still answers single `{"call": [...]}` requests. (#1204) +- **The vector cache is one SQLite database per model.** `~/.cache/dirsql/embeddings/` now holds `.db` files, each a `vectors(digest, vector)` table; a batch is looked up in one query and its misses stored in one transaction. The `cachetta` dependency is gone. Vectors written by earlier releases (one file per value) are ignored, never read; wipe the directory to reclaim their space. (#1204) diff --git a/plugins/dirsql-plugin-embeddings/e2e-attestations/claude-1204-embed-plugin-batch.json b/plugins/dirsql-plugin-embeddings/e2e-attestations/claude-1204-embed-plugin-batch.json new file mode 100644 index 000000000..0aaa96ef4 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/e2e-attestations/claude-1204-embed-plugin-batch.json @@ -0,0 +1,7 @@ +{ + "command": "uv run --with-editable ../../packages/python python -m pytest tests/e2e/ -x -q", + "ran_at": 1790960575, + "exit_code": 0, + "commit": "64524926303ffd11d9f107fd843b6c2f21648bbd", + "branch": "claude/1204-embed-plugin-batch" +} diff --git a/plugins/dirsql-plugin-embeddings/migrations.d/2026-10-02-batched-embed-sqlite-cache.md b/plugins/dirsql-plugin-embeddings/migrations.d/2026-10-02-batched-embed-sqlite-cache.md new file mode 100644 index 000000000..b8208ae79 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/migrations.d/2026-10-02-batched-embed-sqlite-cache.md @@ -0,0 +1,44 @@ +### dirsql-plugin-embeddings: vector cache moves to SQLite, worker requests are batched + +#### Summary + +The on-disk vector cache changes format: `~/.cache/dirsql/embeddings/` (or +`$XDG_CACHE_HOME/dirsql/embeddings/`) now holds one SQLite database per model +identifier instead of one `cachetta` file per cached value. Nothing reads the +old files. The first query after upgrading re-embeds everything it touches; +later runs hit the new store. The worker protocol gains the batched +`{"calls": [...]}` request alongside the single `{"call": [...]}` request, +and the packaged fragment opts in with `batch = 4096`, which needs dirsql +0.4.60 or later. + +#### Required changes + +| Surface | Before | After | +| ------- | ------ | ----- | +| Cache directory contents | One `cachetta` file per cached value | One `.db` SQLite database per model; old files are dead weight — wipe the directory to reclaim their space | +| `dirsql` version with the packaged fragment | Any release with `[[dirsql.function]]` | 0.4.60 or later (`batch` is rejected as an unknown field by older releases) | +| Worker protocol | `{"call": [value, model_id?]}` only | Also `{"calls": [[value, model_id?], ...]}` answered by `{"results": [...]}` in request order | + +#### Deprecations removed + +_None._ + +#### Behavior changes without code changes + +- The first query after upgrading reports every `embed()` call as computed + rather than cached, even over files embedded before, because the old + per-value files are never read. +- A statement's `embed()` values reach the worker in batches of up to 4096; + a failing model fails every call in its batch group with the same error + instead of one call at a time. + +#### Verification + +```bash +cd "$(mktemp -d)" +printf 'hello world\n' > note.md +uvx --with dirsql-plugin-embeddings dirsql query "SELECT path, vec_length(embed(content)) AS dims FROM '*.md'" +# expected: [{"path":"note.md","dims":512}] +ls "${XDG_CACHE_HOME:-$HOME/.cache}/dirsql/embeddings/" +# expected: exactly one <64 hex chars>.db file for the default model +``` diff --git a/plugins/dirsql-plugin-embeddings/pyproject.toml b/plugins/dirsql-plugin-embeddings/pyproject.toml index 8dbdac896..b959ffd97 100644 --- a/plugins/dirsql-plugin-embeddings/pyproject.toml +++ b/plugins/dirsql-plugin-embeddings/pyproject.toml @@ -17,7 +17,8 @@ requires-python = ">=3.10" # # model2vec (NOT sentence-transformers/torch) is the inference stack: static # embeddings whose runtime deps are numpy + tokenizers, viable in uvx -# ephemeral envs. cachetta is the vector cache; tqdm the stderr progress UI. +# ephemeral envs. The vector cache is stdlib sqlite3; tqdm the stderr +# progress UI. # # dirsql is the SDK the one-liner search command runs its generated SQL # through. Functionally it needs >= 0.4.17 (the first release with the core @@ -30,7 +31,6 @@ dependencies = [ "dirsql", "sqlite-vec>=0.1", "model2vec>=0.6", - "cachetta>=0.7", "tqdm>=4", "click>=8", ] diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/dirsql.toml b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/dirsql.toml index 14c032043..660ae6d6e 100644 --- a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/dirsql.toml +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/dirsql.toml @@ -7,10 +7,12 @@ entrypoint = "sqlite3_vec_init" # The embed() SQL scalar function, served by the plugin's worker process over # newline-delimited JSON on stdin/stdout. Inert until a query calls it. The -# generous timeout absorbs the first call's model download. +# generous timeout absorbs the first call's model download. A statement's +# values reach the worker in batches, each encoded in one model call. [[dirsql.function]] name = "embed" args = [1, 2] command = "dirsql-plugin-embeddings worker" deterministic = true timeout = "600s" +batch = 4096 diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/cache.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/cache.py index 5489fab91..8a8ec72c1 100644 --- a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/cache.py +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/cache.py @@ -1,19 +1,8 @@ import os -from datetime import timedelta from pathlib import Path -from cachetta import Cachetta - -# No eviction: entries stay until the user wipes the directory (documented as -# safe -- the only cost is re-embedding). -NO_EVICTION = timedelta(days=365000) - def cache_dir(): xdg = os.environ.get("XDG_CACHE_HOME", "") base = Path(xdg) if xdg else Path.home() / ".cache" return base / "dirsql" / "embeddings" - - -def make_cache(): - return Cachetta(path=cache_dir(), hashed=True, duration=NO_EVICTION) diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/cache_test.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/cache_test.py index 6b10c7ea6..9632f9363 100644 --- a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/cache_test.py +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/cache_test.py @@ -1,4 +1,3 @@ -from datetime import timedelta from pathlib import Path from unittest.mock import patch @@ -19,20 +18,3 @@ def it_treats_an_empty_xdg_cache_home_as_unset(): with patch.dict(cache.os.environ, {"XDG_CACHE_HOME": ""}): with patch.object(cache.Path, "home", return_value=Path("/home/u")): assert cache.cache_dir() == Path("/home/u/.cache/dirsql/embeddings") - - -def describe_make_cache(): - def it_builds_a_hashed_cachetta_at_the_cache_dir_with_no_eviction(): - with patch.object(cache, "Cachetta") as cachetta: - with patch.object( - cache, "cache_dir", return_value=Path("/somewhere") - ) as directory: - built = cache.make_cache() - directory.assert_called_once_with() - cachetta.assert_called_once_with( - path=Path("/somewhere"), hashed=True, duration=cache.NO_EVICTION - ) - assert built is cachetta.return_value - - def it_pins_no_eviction_to_a_thousand_years(): - assert cache.NO_EVICTION == timedelta(days=365000) diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/store.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/store.py new file mode 100644 index 000000000..3e9e0d010 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/store.py @@ -0,0 +1,53 @@ +import sqlite3 +from array import array +from hashlib import sha256 +from json import dumps + +SCHEMA = ( + "CREATE TABLE IF NOT EXISTS vectors" + " (digest TEXT PRIMARY KEY, vector BLOB NOT NULL)" +) +LOOKUP = ( + "SELECT digest, vector FROM vectors" + " WHERE digest IN (SELECT value FROM json_each(?))" +) +INSERT = "INSERT OR REPLACE INTO vectors (digest, vector) VALUES (?, ?)" + + +def pack(vector): + return array("d", vector).tobytes() + + +def unpack(blob): + doubles = array("d") + doubles.frombytes(blob) + return doubles.tolist() + + +class VectorStore: + """One SQLite database per model under ``directory``, keyed by text digest.""" + + def __init__(self, directory): + self._directory = directory + self._connections = {} + + def _connection(self, identifier): + if identifier not in self._connections: + self._directory.mkdir(parents=True, exist_ok=True) + name = sha256(identifier.encode("utf-8")).hexdigest() + ".db" + connection = sqlite3.connect(self._directory / name) + connection.execute(SCHEMA) + self._connections[identifier] = connection + return self._connections[identifier] + + def lookup(self, identifier, digests): + """The stored vectors among ``digests``, keyed by digest.""" + rows = self._connection(identifier).execute(LOOKUP, (dumps(list(digests)),)) + return {digest: unpack(blob) for digest, blob in rows} + + def insert(self, identifier, vectors): + """Store ``vectors`` (digest to vector) in one transaction.""" + connection = self._connection(identifier) + rows = [(digest, pack(vector)) for digest, vector in vectors.items()] + with connection: + connection.executemany(INSERT, rows) diff --git a/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/store_test.py b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/store_test.py new file mode 100644 index 000000000..f4b1acbd1 --- /dev/null +++ b/plugins/dirsql-plugin-embeddings/src/dirsql_plugin_embeddings/embedding/store_test.py @@ -0,0 +1,123 @@ +from hashlib import sha256 +from unittest.mock import MagicMock, call, patch + +from . import store + + +def make_store(): + directory = MagicMock() + return store.VectorStore(directory), directory + + +def describe_databases(): + def it_keeps_one_database_per_model_named_by_the_identifier_digest(): + built, directory = make_store() + with patch.object(store.sqlite3, "connect") as connect: + built.lookup("m1", ["d"]) + built.lookup("m2", ["d"]) + assert directory.__truediv__.call_args_list == [ + call(sha256(b"m1").hexdigest() + ".db"), + call(sha256(b"m2").hexdigest() + ".db"), + ] + assert connect.call_args_list == [call(directory / "x")] * 2 + + def it_creates_the_cache_directory_before_opening(): + built, directory = make_store() + with patch.object(store.sqlite3, "connect"): + built.lookup("m1", []) + directory.mkdir.assert_called_once_with(parents=True, exist_ok=True) + + def it_creates_the_table_on_open(): + built, _ = make_store() + with patch.object(store.sqlite3, "connect") as connect: + built.lookup("m1", []) + assert connect.return_value.execute.call_args_list[0] == call(store.SCHEMA) + + def it_opens_each_model_once(): + built, _ = make_store() + with patch.object(store.sqlite3, "connect") as connect: + built.lookup("m1", ["a"]) + built.insert("m1", {"a": [1.0]}) + built.lookup("m1", ["a"]) + assert connect.call_count == 1 + + +def describe_lookup(): + def it_finds_nothing_on_a_cold_lookup(): + built, _ = make_store() + with patch.object(store.sqlite3, "connect") as connect: + connect.return_value.execute.return_value = [] + assert built.lookup("m1", ["a", "b"]) == {} + + def it_unpacks_every_stored_row_keyed_by_digest(): + built, _ = make_store() + with patch.object(store.sqlite3, "connect") as connect: + connect.return_value.execute.return_value = [ + ("a", store.pack([1.0, -0.5])), + ("c", store.pack([3.0])), + ] + found = built.lookup("m1", ["a", "b", "c"]) + assert found == {"a": [1.0, -0.5], "c": [3.0]} + + def it_asks_for_the_whole_batch_in_one_query(): + built, _ = make_store() + with patch.object(store.sqlite3, "connect") as connect: + connect.return_value.execute.return_value = [] + built.lookup("m1", ["a", "b", "c"]) + selects = [ + args + for args, _ in connect.return_value.execute.call_args_list + if args[0] is store.LOOKUP + ] + assert selects == [(store.LOOKUP, ('["a", "b", "c"]',))] + + def it_accepts_any_iterable_of_digests(): + built, _ = make_store() + with patch.object(store.sqlite3, "connect") as connect: + connect.return_value.execute.return_value = [] + built.lookup("m1", iter(["a", "b"])) + connect.return_value.execute.assert_called_with( + store.LOOKUP, ('["a", "b"]',) + ) + + +def describe_insert(): + def it_writes_the_packed_batch_in_one_transaction(): + built, _ = make_store() + with patch.object(store.sqlite3, "connect") as connect: + connection = connect.return_value + built.insert("m1", {"a": [1.0], "b": [2.0, 3.0]}) + connection.__enter__.assert_called_once() + connection.__exit__.assert_called_once() + connection.executemany.assert_called_once_with( + store.INSERT, + [("a", store.pack([1.0])), ("b", store.pack([2.0, 3.0]))], + ) + + def it_writes_inside_the_transaction_not_before_it(): + built, _ = make_store() + events = [] + with patch.object(store.sqlite3, "connect") as connect: + connection = connect.return_value + connection.__enter__.side_effect = lambda: events.append("enter") + connection.executemany.side_effect = lambda *_: events.append("write") + built.insert("m1", {"a": [1.0]}) + assert events == ["enter", "write"] + + def it_accepts_an_empty_batch(): + built, _ = make_store() + with patch.object(store.sqlite3, "connect") as connect: + built.insert("m1", {}) + connect.return_value.executemany.assert_called_once_with(store.INSERT, []) + + +def describe_pack(): + def it_encodes_doubles_and_unpack_restores_them_exactly(): + vector = [0.1, -2.5e-8, 1e300, 0.0, 3.0] + assert store.unpack(store.pack(vector)) == vector + + def it_encodes_eight_bytes_per_component(): + assert len(store.pack([1.0, 2.0, 3.0])) == 24 + + def it_unpacks_an_empty_blob_to_an_empty_vector(): + assert store.unpack(b"") == [] 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 5ce562c06..52501a88a 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 @@ -2,67 +2,54 @@ from hashlib import sha256 from . import model -from .cache import make_cache +from .cache import cache_dir from .model_identifier import model_identifier +from .store import VectorStore from .values import ProtocolError, decode_value MALFORMED_SHAPE = ( - 'malformed request: expected {"call": [value, model_id?]} on one line' + 'malformed request: expected {"call": [value, model_id?]}' + ' or {"calls": [[value, model_id?], ...]} on one line' ) -MALFORMED_ARITY = 'malformed request: "call" must carry 1 or 2 arguments' +MALFORMED_CALLS = 'malformed request: "calls" must be a list of argument lists' +MALFORMED_ARITY = "malformed request: a call must carry 1 or 2 arguments" MALFORMED_MODEL_ID = "the model id must be TEXT" class Worker: def __init__(self): self._models = {} - self._cache = make_cache() - self._compute_cached = self._cache.wrap(self._compute) - self._pending = None - self._computed = False + self._store = VectorStore(cache_dir()) def _model(self, model_id): if model_id not in self._models: self._models[model_id] = model.load_model(model_id) return self._models[model_id] - def _compute(self, digest, identifier): - # cachetta's wrapper returns a hit and a miss identically, so the only - # place a miss is observable is here: _compute runs on a miss and never - # on a hit. embed() clears the flag and reads it back afterwards. - self._computed = True - text, loaded = self._pending - # No progress bar: the protocol embeds one value per round trip, so a - # per-call bar can only ever say 1/1. The model *download* bar, which - # measures something, is TTY-gated in `progress.configure`. - (vector,) = loaded.encode([text], show_progress_bar=False) - return [float(component) for component in vector] - - def embed(self, text, model_id): - """Return the vector for ``text``, and whether the cache served it.""" + def embed(self, texts, model_id): + """The vector for each of ``texts``, with whether the cache served it.""" loaded = self._model(model_id) identifier = model_identifier(model_id, loaded) - digest = sha256(text.encode("utf-8")).hexdigest() - self._pending = (text, loaded) - self._computed = False - vector = self._compute_cached(digest, identifier) - return vector, not self._computed - - def _embed_response(self, text, model_id): - try: - vector, cached = self.embed(text, model_id) - except Exception as error: - return {"err": f"embed({model_id!r}) failed: {error}"} - return {"ok": vector, "meta": {"cached": cached}} + digests = [sha256(text.encode("utf-8")).hexdigest() for text in texts] + found = self._store.lookup(identifier, digests) + missing = {} + for digest, text in zip(digests, texts): + if digest not in found: + missing[digest] = text + computed = {} + if missing: + # No progress bar: dirsql's own progress line counts the calls. The + # model *download* bar is TTY-gated in `progress.configure`. + vectors = loaded.encode(list(missing.values()), show_progress_bar=False) + computed = dict(zip(missing, vectors.tolist())) + self._store.insert(identifier, computed) + return [ + (found[digest], True) if digest in found else (computed[digest], False) + for digest in digests + ] - def handle(self, line): - try: - request = json.loads(line) - except json.JSONDecodeError as error: - return {"err": f"malformed request: invalid JSON: {error}"} - if not isinstance(request, dict) or "call" not in request: - return {"err": MALFORMED_SHAPE} - call = request["call"] + def _decode(self, call): + """A call's ``(text, model_id)``, or the response that answers it.""" if not isinstance(call, list) or len(call) not in (1, 2): return {"err": MALFORMED_ARITY} value, *rest = call @@ -75,7 +62,56 @@ def handle(self, line): (model_id,) = rest or [model.DEFAULT_MODEL_ID] if not isinstance(model_id, str): return {"err": MALFORMED_MODEL_ID} - return self._embed_response(text, model_id) + return text, model_id + + def _embed_responses(self, calls): + responses = [None] * len(calls) + by_model = {} + for index, (_, model_id) in enumerate(calls): + by_model.setdefault(model_id, []).append(index) + for model_id, indices in by_model.items(): + texts = [calls[index][0] for index in indices] + try: + group = [ + {"ok": vector, "meta": {"cached": cached}} + for vector, cached in self.embed(texts, model_id) + ] + except Exception as error: + group = [{"err": f"embed({model_id!r}) failed: {error}"}] * len(texts) + for index, response in zip(indices, group): + responses[index] = response + return responses + + def _responses(self, calls): + responses = [None] * len(calls) + pending = {} + for index, call in enumerate(calls): + decoded = self._decode(call) + if isinstance(decoded, dict): + responses[index] = decoded + else: + pending[index] = decoded + embedded = self._embed_responses(list(pending.values())) + for index, response in zip(pending, embedded): + responses[index] = response + return responses + + def handle(self, line): + try: + request = json.loads(line) + except json.JSONDecodeError as error: + return {"err": f"malformed request: invalid JSON: {error}"} + if not isinstance(request, dict): + return {"err": MALFORMED_SHAPE} + if "calls" in request: + calls = request["calls"] + if not isinstance(calls, list): + return {"err": MALFORMED_CALLS} + return {"results": self._responses(calls)} + if "call" not in request: + return {"err": MALFORMED_SHAPE} + (response,) = self._responses([request["call"]]) + return response def serve(self, stdin, stdout): for line in stdin: 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 66a78c549..7d5ec0118 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,11 +1,11 @@ -import json from contextlib import contextmanager from hashlib import sha256 from unittest.mock import MagicMock, call, patch from . import worker -HELLO_DIGEST = sha256(b"hello").hexdigest() +HELLO = sha256(b"hello").hexdigest() +WORLD = sha256(b"world").hexdigest() @contextmanager @@ -17,28 +17,34 @@ def loading(identifier="m1"): yield model, identify -def make_worker(): - with patch.object(worker, "make_cache") as make_cache: - built = worker.Worker() - return built, make_cache.return_value +def make_worker(found=None): + with patch.object(worker, "VectorStore") as store_class: + with patch.object(worker, "cache_dir") as cache_dir: + built = worker.Worker() + store = store_class.return_value + store.lookup.return_value = {} if found is None else found + return built, store, store_class, cache_dir -def describe_init(): - def it_wraps_compute_with_the_cache(): - built, cache = make_worker() - cache.wrap.assert_called_once_with(built._compute) - assert built._compute_cached is cache.wrap.return_value +def encoding(loaded, vectors): + loaded.encode.return_value.tolist.return_value = vectors + - def it_starts_with_no_loaded_models_and_nothing_pending(): - built, _ = make_worker() +def describe_init(): + def it_opens_the_vector_store_at_the_cache_dir(): + built, store, store_class, cache_dir = make_worker() + cache_dir.assert_called_once_with() + store_class.assert_called_once_with(cache_dir.return_value) + assert built._store is store + + def it_starts_with_no_loaded_models(): + built, *_ = make_worker() assert built._models == {} - assert built._pending is None - assert built._computed is False def describe_model_loading(): def it_loads_a_model_once_and_memoizes_it(): - built, _ = make_worker() + built, *_ = make_worker() with patch.object(worker, "model") as model: first = built._model("m1") second = built._model("m1") @@ -46,160 +52,133 @@ def it_loads_a_model_once_and_memoizes_it(): assert first is second is model.load_model.return_value def it_loads_each_distinct_model_id(): - built, _ = make_worker() + built, *_ = make_worker() with patch.object(worker, "model") as model: built._model("m1") built._model("m2") assert model.load_model.call_args_list == [call("m1"), call("m2")] -def describe_compute(): - def it_records_that_it_ran_so_a_miss_is_distinguishable_from_a_hit(): - # _compute runs only on a cache miss, so it is the one place a miss is - # observable -- cachetta's wrapper returns both cases identically. - built, _ = make_worker() - loaded = MagicMock() - loaded.encode.return_value = [[1.0]] - built._pending = ("hello", loaded) - built._computed = False - built._compute(HELLO_DIGEST, "m1") - assert built._computed is True - - def it_encodes_the_pending_text_and_returns_plain_floats(): - built, _ = make_worker() - loaded = MagicMock() - loaded.encode.return_value = [[1, 2.5]] - built._pending = ("hello", loaded) - vector = built._compute(HELLO_DIGEST, "m1") - loaded.encode.assert_called_once_with(["hello"], show_progress_bar=False) - assert vector == [1.0, 2.5] - assert all(isinstance(component, float) for component in vector) - - def it_never_shows_a_per_call_progress_bar_even_on_a_tty(): - # The protocol embeds exactly one value per round trip, so a per-call - # bar can only ever say 1/1 -- one meaningless bar per embedded value. - built, _ = make_worker() - loaded = MagicMock() - loaded.encode.return_value = [[0.0]] - built._pending = ("hello", loaded) - with patch("sys.stderr.isatty", return_value=True): - built._compute(HELLO_DIGEST, "m1") - loaded.encode.assert_called_once_with(["hello"], show_progress_bar=False) - - 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]) + def it_looks_the_whole_batch_up_under_the_model_identifier(): + built, store, *_ = make_worker() with loading("m1@9.9") as (model, identify): - result = built.embed("hello", "m1") + encoding(model.load_model.return_value, [[1.0], [2.0]]) + built.embed(["hello", "world"], "m1") model.load_model.assert_called_once_with("m1") 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 loading(): - _, cached = built.embed("hello", "m1") - assert cached is True - - def it_reports_a_miss_when_compute_ran(): - built, _ = make_worker() - - def compute(digest, identifier): - built._computed = True - return [1.0] - - built._compute_cached = MagicMock(side_effect=compute) - with loading(): - _, cached = built.embed("hello", "m1") - assert cached is False - - def it_clears_a_previous_calls_miss_before_consulting_the_cache(): - # Without the reset, one miss would make every later hit read as a - # miss for the rest of the process. - built, _ = make_worker() - built._computed = True - built._compute_cached = MagicMock(return_value=[1.0]) - 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]) + store.lookup.assert_called_once_with("m1@9.9", [HELLO, WORLD]) + + def it_hashes_the_utf8_bytes_of_each_text(): + built, store, *_ = make_worker() 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 loading("m1"): - built.embed("héllo", "m1") + encoding(model.load_model.return_value, [[1.0]]) + built.embed(["héllo"], "m1") expected = sha256("héllo".encode("utf-8")).hexdigest() - assert built._compute_cached.call_args == call(expected, "m1") + store.lookup.assert_called_once_with("m1", [expected]) + + def it_encodes_only_the_misses_in_one_call_and_stores_them(): + built, store, *_ = make_worker(found={HELLO: [1.0, 0.0]}) + with loading() as (model, _): + loaded = model.load_model.return_value + encoding(loaded, [[0.0, 1.0]]) + result = built.embed(["hello", "world"], "m1") + loaded.encode.assert_called_once_with(["world"], show_progress_bar=False) + store.insert.assert_called_once_with("m1", {WORLD: [0.0, 1.0]}) + assert result == [([1.0, 0.0], True), ([0.0, 1.0], False)] + + def it_reports_a_cold_batch_as_computed_in_request_order(): + built, store, *_ = make_worker() + with loading() as (model, _): + encoding(model.load_model.return_value, [[1.0], [2.0], [3.0]]) + result = built.embed(["a", "b", "c"], "m1") + assert result == [([1.0], False), ([2.0], False), ([3.0], False)] + + def it_never_encodes_when_every_text_is_cached(): + built, store, *_ = make_worker(found={HELLO: [1.0], WORLD: [2.0]}) + with loading() as (model, _): + loaded = model.load_model.return_value + result = built.embed(["world", "hello"], "m1") + loaded.encode.assert_not_called() + store.insert.assert_not_called() + assert result == [([2.0], True), ([1.0], True)] + + def it_encodes_a_repeated_miss_once(): + built, store, *_ = make_worker() + with loading() as (model, _): + loaded = model.load_model.return_value + encoding(loaded, [[1.0]]) + result = built.embed(["hello", "hello"], "m1") + loaded.encode.assert_called_once_with(["hello"], show_progress_bar=False) + assert result == [([1.0], False), ([1.0], False)] + + def it_returns_the_vectors_as_plain_lists(): + built, *_ = make_worker() + with loading() as (model, _): + loaded = model.load_model.return_value + encoding(loaded, [[1.0, 2.5]]) + ((vector, _),) = built.embed(["hello"], "m1") + loaded.encode.return_value.tolist.assert_called_once_with() + assert vector == [1.0, 2.5] def describe_handle(): def it_answers_ok_with_the_embedding(): - built, _ = make_worker() + built, *_ = make_worker() with patch.object( - built, "embed", return_value=([1.0, 0.0], False) + built, "embed", return_value=[([1.0, 0.0], False)] ) as embed: response = built.handle('{"call": ["hello", "m1"]}') - embed.assert_called_once_with("hello", "m1") + embed.assert_called_once_with(["hello"], "m1") assert response == {"ok": [1.0, 0.0], "meta": {"cached": False}} def it_reports_a_cache_hit_in_the_response_metadata(): - built, _ = make_worker() - with patch.object(built, "embed", return_value=([1.0], True)): + built, *_ = make_worker() + with patch.object(built, "embed", return_value=[([1.0], True)]): response = built.handle('{"call": ["hello", "m1"]}') assert response == {"ok": [1.0], "meta": {"cached": True}} def it_uses_the_default_model_for_a_single_argument_call(): - built, _ = make_worker() - with patch.object(built, "embed", return_value=([1.0], False)) as embed: + built, *_ = make_worker() + with patch.object(built, "embed", return_value=[([1.0], False)]) as embed: with patch.object(worker, "model") as model: model.DEFAULT_MODEL_ID = "default/model" response = built.handle('{"call": ["hello"]}') - embed.assert_called_once_with("hello", "default/model") + embed.assert_called_once_with(["hello"], "default/model") assert response == {"ok": [1.0], "meta": {"cached": False}} def it_answers_err_on_invalid_json(): - built, _ = make_worker() + built, *_ = make_worker() response = built.handle("{not json") assert set(response) == {"err"} assert response["err"].startswith("malformed request: invalid JSON:") def it_answers_err_on_a_non_object_request(): - built, _ = make_worker() + built, *_ = make_worker() assert built.handle("[1, 2]") == {"err": worker.MALFORMED_SHAPE} def it_answers_err_when_call_is_missing(): - built, _ = make_worker() + built, *_ = make_worker() assert built.handle('{"other": []}') == {"err": worker.MALFORMED_SHAPE} def it_answers_err_when_call_is_not_a_list(): - built, _ = make_worker() + built, *_ = make_worker() assert built.handle('{"call": "hello"}') == { "err": worker.MALFORMED_ARITY } def it_answers_err_on_an_empty_call(): - built, _ = make_worker() + built, *_ = make_worker() assert built.handle('{"call": []}') == {"err": worker.MALFORMED_ARITY} def it_answers_err_on_three_arguments(): - built, _ = make_worker() + built, *_ = make_worker() assert built.handle('{"call": ["a", "b", "c"]}') == { "err": worker.MALFORMED_ARITY } def it_answers_err_with_the_protocol_message_for_a_bad_value(): - built, _ = make_worker() + built, *_ = make_worker() with patch.object( worker, "decode_value", @@ -210,21 +189,21 @@ def it_answers_err_with_the_protocol_message_for_a_bad_value(): assert response == {"err": "bad value"} def it_answers_ok_null_for_a_null_value(): - built, _ = make_worker() + built, *_ = make_worker() with patch.object(built, "embed") as embed: response = built.handle('{"call": [null]}') embed.assert_not_called() assert response == {"ok": None} def it_answers_err_for_a_non_text_model_id(): - built, _ = make_worker() + built, *_ = make_worker() with patch.object(built, "embed") as embed: response = built.handle('{"call": ["hello", 42]}') embed.assert_not_called() assert response == {"err": worker.MALFORMED_MODEL_ID} def it_answers_err_naming_the_model_when_embedding_fails(): - built, _ = make_worker() + built, *_ = make_worker() with patch.object( built, "embed", side_effect=RuntimeError("model exploded") ): @@ -234,9 +213,107 @@ def it_answers_err_naming_the_model_when_embedding_fails(): } +def describe_handle_batched(): + def it_answers_results_in_request_order(): + built, *_ = make_worker() + with patch.object( + built, "embed", return_value=[([1.0], False), ([2.0], True)] + ) as embed: + response = built.handle( + '{"calls": [["hello", "m1"], ["world", "m1"]]}' + ) + embed.assert_called_once_with(["hello", "world"], "m1") + assert response == { + "results": [ + {"ok": [1.0], "meta": {"cached": False}}, + {"ok": [2.0], "meta": {"cached": True}}, + ] + } + + def it_embeds_once_per_model_and_restores_the_request_order(): + built, *_ = make_worker() + by_model = { + "m1": [([1.0], False), ([3.0], False)], + "m2": [([2.0], False)], + } + with patch.object( + built, "embed", side_effect=lambda texts, model_id: by_model[model_id] + ) as embed: + response = built.handle( + '{"calls": [["a", "m1"], ["b", "m2"], ["c", "m1"]]}' + ) + assert embed.call_args_list == [ + call(["a", "c"], "m1"), + call(["b"], "m2"), + ] + assert [result["ok"] for result in response["results"]] == [ + [1.0], + [2.0], + [3.0], + ] + + def it_answers_each_malformed_call_in_place_and_embeds_the_rest(): + built, *_ = make_worker() + with patch.object( + built, "embed", return_value=[([1.0], False)] + ) as embed: + response = built.handle( + '{"calls": [[null, "m1"], ["a", "b", "c"], ["hello", 42],' + ' ["hello", "m1"]]}' + ) + embed.assert_called_once_with(["hello"], "m1") + assert response == { + "results": [ + {"ok": None}, + {"err": worker.MALFORMED_ARITY}, + {"err": worker.MALFORMED_MODEL_ID}, + {"ok": [1.0], "meta": {"cached": False}}, + ] + } + + def it_fails_every_call_of_a_model_that_fails_and_keeps_serving_the_others(): + built, *_ = make_worker() + + def embed(texts, model_id): + if model_id == "bad": + raise RuntimeError("model exploded") + return [([1.0], False)] * len(texts) + + with patch.object(built, "embed", side_effect=embed): + response = built.handle( + '{"calls": [["a", "bad"], ["b", "good"], ["c", "bad"]]}' + ) + failure = {"err": "embed('bad') failed: model exploded"} + assert response == { + "results": [failure, {"ok": [1.0], "meta": {"cached": False}}, failure] + } + + def it_uses_the_default_model_for_single_argument_calls(): + built, *_ = make_worker() + with patch.object(built, "embed", return_value=[([1.0], False)]) as embed: + with patch.object(worker, "model") as model: + model.DEFAULT_MODEL_ID = "default/model" + built.handle('{"calls": [["hello"]]}') + embed.assert_called_once_with(["hello"], "default/model") + + def it_answers_no_results_for_an_empty_batch(): + built, *_ = make_worker() + with patch.object(built, "embed") as embed: + response = built.handle('{"calls": []}') + embed.assert_not_called() + assert response == {"results": []} + + def it_answers_err_when_calls_is_not_a_list(): + built, *_ = make_worker() + with patch.object(built, "embed") as embed: + response = built.handle('{"calls": "hello"}') + embed.assert_not_called() + assert response == {"err": worker.MALFORMED_CALLS} + + def describe_serve(): def it_writes_one_compact_json_line_per_request_and_flushes(): - built, _ = make_worker() + built, *_ = make_worker() stdout = MagicMock() with patch.object( built, "handle", side_effect=[{"ok": [1.0]}, {"err": "x"}] @@ -253,7 +330,7 @@ def it_writes_one_compact_json_line_per_request_and_flushes(): assert stdout.flush.call_count == 2 def it_flushes_after_each_line_not_only_at_the_end(): - built, _ = make_worker() + built, *_ = make_worker() stdout = MagicMock() events = [] stdout.write.side_effect = lambda line: events.append(("write", line)) @@ -268,7 +345,7 @@ def it_flushes_after_each_line_not_only_at_the_end(): ] def it_skips_blank_lines_and_keeps_serving_later_requests(): - built, _ = make_worker() + built, *_ = make_worker() stdout = MagicMock() with patch.object(built, "handle", return_value={"ok": None}) as handle: built.serve(["\n", " \n", '{"call": ["a"]}\n'], stdout) diff --git a/plugins/dirsql-plugin-embeddings/tests/e2e/embed_function_test.py b/plugins/dirsql-plugin-embeddings/tests/e2e/embed_function_test.py index 7cc53e961..da372a75a 100644 --- a/plugins/dirsql-plugin-embeddings/tests/e2e/embed_function_test.py +++ b/plugins/dirsql-plugin-embeddings/tests/e2e/embed_function_test.py @@ -24,9 +24,11 @@ import json import os import shutil +import sqlite3 import stat import subprocess import sys +from contextlib import closing from importlib import resources import pytest @@ -63,7 +65,12 @@ def _cache_entries(cache_home): embeddings = cache_home / "dirsql" / "embeddings" if not embeddings.is_dir(): return 0 - return sum(1 for path in embeddings.iterdir() if path.is_file()) + total = 0 + for path in embeddings.glob("*.db"): + with closing(sqlite3.connect(path)) as connection: + (count,) = connection.execute("SELECT count(*) FROM vectors").fetchone() + total += count + return total def describe_glob_scoping(): diff --git a/plugins/dirsql-plugin-embeddings/tests/e2e/speed_of_light_test.py b/plugins/dirsql-plugin-embeddings/tests/e2e/speed_of_light_test.py index 5fce8cd4f..c2e4c2636 100644 --- a/plugins/dirsql-plugin-embeddings/tests/e2e/speed_of_light_test.py +++ b/plugins/dirsql-plugin-embeddings/tests/e2e/speed_of_light_test.py @@ -83,10 +83,6 @@ def root(tmp_path): finally: shutil.rmtree(tree, ignore_errors=True) - @pytest.mark.xfail( - strict=True, - reason="embed() makes one worker round trip per value instead of a batch", - ) def it_matches_the_batch_encode_within_the_bar(root, tiny_model, tmp_path): startup = startup_seconds([*LAUNCHER, "--help"], cwd=tmp_path) diff --git a/plugins/dirsql-plugin-embeddings/tests/e2e/worker_cli_test.py b/plugins/dirsql-plugin-embeddings/tests/e2e/worker_cli_test.py index bff134f77..60c3a56a2 100644 --- a/plugins/dirsql-plugin-embeddings/tests/e2e/worker_cli_test.py +++ b/plugins/dirsql-plugin-embeddings/tests/e2e/worker_cli_test.py @@ -1,8 +1,8 @@ """E2E: the installed ``dirsql-plugin-embeddings worker`` CLI, nothing mocked. Spawns the real console script (as installed on PATH) and speaks the wire -protocol over real pipes: real argparse dispatch, real model2vec inference, -real cachetta writes under a real ``XDG_CACHE_HOME``. +protocol over real pipes: real click dispatch, real model2vec inference, +real SQLite cache writes under a real ``XDG_CACHE_HOME``. The model is the real on-disk model2vec model from conftest, passed through the worker's ordinary model-override argument. The default @@ -16,6 +16,7 @@ uv run python -m pytest tests/e2e -q """ +import json import shutil @@ -30,6 +31,15 @@ def it_serves_an_embed_request_over_real_pipes(spawn_worker, tiny_model): worker = spawn_worker(argv=worker_argv()) assert worker.request("hello", tiny_model)["ok"] == [1.0, 0.0] + def it_serves_a_batched_request_over_real_pipes(spawn_worker, tiny_model): + worker = spawn_worker(argv=worker_argv()) + line = json.dumps({"calls": [["hello", tiny_model], ["world", tiny_model]]}) + response = worker.send_line(line) + assert [result["ok"] for result in response["results"]] == [ + [1.0, 0.0], + [0.0, 1.0], + ] + def it_survives_a_malformed_request_and_keeps_serving( spawn_worker, tiny_model ): diff --git a/plugins/dirsql-plugin-embeddings/tests/integration/cache_behavior_test.py b/plugins/dirsql-plugin-embeddings/tests/integration/cache_behavior_test.py index 011e9c55d..68cc8734e 100644 --- a/plugins/dirsql-plugin-embeddings/tests/integration/cache_behavior_test.py +++ b/plugins/dirsql-plugin-embeddings/tests/integration/cache_behavior_test.py @@ -1,16 +1,20 @@ """Integration: the on-disk vector cache, observed through real workers. -Each (value bytes, model identifier) pair owns exactly one cache entry under -``$XDG_CACHE_HOME/dirsql/embeddings/``, so entry counts across real worker -runs pin hit/miss behavior: repeats (same process or a fresh one) add nothing, -content changes and model changes each add one. The default-model requests are -driven through a wrapper process that points the default at the on-disk test -model (the sandboxed suite cannot download the real default); the worker -itself runs unmodified. +The cache is one SQLite database per model under +``$XDG_CACHE_HOME/dirsql/embeddings/``, holding one row per (value bytes, +model identifier) pair, so row counts across real worker runs pin hit/miss +behavior: repeats (same process or a fresh one) add nothing, content changes +and model changes each add one. The default-model requests are driven through +a wrapper process that points the default at the on-disk test model (the +sandboxed suite cannot download the real default); the worker itself runs +unmodified. """ import base64 +import json +import sqlite3 import sys +from contextlib import closing DEFAULT_PATCH_ARGV_PREFIX = [ sys.executable, @@ -24,11 +28,20 @@ ] -def entries(cache_home): +def databases(cache_home): embeddings = cache_home / "dirsql" / "embeddings" if not embeddings.is_dir(): - return 0 - return sum(1 for path in embeddings.iterdir() if path.is_file()) + return [] + return sorted(embeddings.glob("*.db")) + + +def entries(cache_home): + total = 0 + for path in databases(cache_home): + with closing(sqlite3.connect(path)) as connection: + (count,) = connection.execute("SELECT count(*) FROM vectors").fetchone() + total += count + return total def describe_cache_location(): @@ -39,6 +52,15 @@ def it_writes_under_xdg_cache_home_dirsql_embeddings( worker.request("hello", tiny_model) assert entries(cache_home) == 1 + def it_keeps_one_database_per_model( + spawn_worker, cache_home, tiny_model, other_model + ): + worker = spawn_worker() + worker.request("hello", tiny_model) + assert len(databases(cache_home)) == 1 + worker.request("hello", other_model) + assert len(databases(cache_home)) == 2 + def it_never_writes_into_the_working_directory( spawn_worker, tiny_model, tmp_path ): @@ -111,11 +133,52 @@ def it_shares_one_entry_between_default_and_explicit_same_id( assert entries(cache_home) == 1 +def describe_batched_requests(): + def _batch(worker, *texts, model): + line = json.dumps({"calls": [[text, model] for text in texts]}) + return worker.send_line(line)["results"] + + def it_adds_one_entry_per_distinct_value(spawn_worker, cache_home, tiny_model): + worker = spawn_worker() + _batch(worker, "hello", "world", "hello", model=tiny_model) + assert entries(cache_home) == 2 + + def it_serves_a_repeated_batch_from_the_cache( + spawn_worker, cache_home, tiny_model + ): + worker = spawn_worker() + cold = _batch(worker, "hello", "world", model=tiny_model) + warm = _batch(worker, "hello", "world", model=tiny_model) + assert [result["meta"] for result in cold] == [{"cached": False}] * 2 + assert [result["meta"] for result in warm] == [{"cached": True}] * 2 + assert [result["ok"] for result in warm] == [[1.0, 0.0], [0.0, 1.0]] + assert entries(cache_home) == 2 + + def it_computes_only_the_misses_of_a_mixed_batch( + spawn_worker, cache_home, tiny_model + ): + worker = spawn_worker() + worker.request("hello", tiny_model) + hit, miss = _batch(worker, "hello", "world", model=tiny_model) + assert hit["meta"] == {"cached": True} + assert miss["meta"] == {"cached": False} + assert entries(cache_home) == 2 + + def it_survives_the_process(spawn_worker, cache_home, tiny_model): + first = spawn_worker() + _batch(first, "hello", "world", model=tiny_model) + first.close() + warm = _batch(spawn_worker(), "world", "hello", model=tiny_model) + assert [result["ok"] for result in warm] == [[0.0, 1.0], [1.0, 0.0]] + assert [result["meta"] for result in warm] == [{"cached": True}] * 2 + assert entries(cache_home) == 2 + + def describe_cache_reporting(): - """The worker tells dirsql which answers cost nothing (#1034). + """The worker tells dirsql which answers cost nothing. - ``cachetta``'s wrapper returns a hit and a miss identically, so the worker - has to recover the signal itself; these pin that it recovers the right one. + A hit is a row the store already held; a miss is a vector the model + computed this call. These pin that the flag says which. """ def it_flags_a_recomputed_value_as_not_cached(spawn_worker, tiny_model): diff --git a/plugins/dirsql-plugin-embeddings/tests/integration/config_fragment_test.py b/plugins/dirsql-plugin-embeddings/tests/integration/config_fragment_test.py index 3a17d71be..b97421ae9 100644 --- a/plugins/dirsql-plugin-embeddings/tests/integration/config_fragment_test.py +++ b/plugins/dirsql-plugin-embeddings/tests/integration/config_fragment_test.py @@ -2,8 +2,8 @@ Reads the packaged ``dirsql.toml`` exactly as the launcher does and pins the whole surface it declares: the sqlite-vec ``[[dirsql.extension]]`` entry -(unchanged) plus the ``[[dirsql.function]]`` entry for ``embed``, exactly as -specified in #801. +(unchanged) plus the ``[[dirsql.function]]`` entry for ``embed``, batched so +a statement's values reach the worker in bulk. """ import sys @@ -45,6 +45,7 @@ def it_declares_the_embed_function_exactly_per_spec(): "command": "dirsql-plugin-embeddings worker", "deterministic": True, "timeout": "600s", + "batch": 4096, } diff --git a/plugins/dirsql-plugin-embeddings/tests/integration/worker_protocol_test.py b/plugins/dirsql-plugin-embeddings/tests/integration/worker_protocol_test.py index d62af2562..9868e8393 100644 --- a/plugins/dirsql-plugin-embeddings/tests/integration/worker_protocol_test.py +++ b/plugins/dirsql-plugin-embeddings/tests/integration/worker_protocol_test.py @@ -1,10 +1,12 @@ """Integration: the worker's stdin/stdout protocol over real pipes. Spawns the real ``worker`` subcommand as a subprocess and speaks the -newline-delimited JSON protocol from #801: requests are -``{"call": [value, model_id?]}``, responses ``{"ok": [floats...]}`` (plus an -advisory ``"meta"``, covered in cache_behavior_test) or ``{"err": "message"}``. The model is a real model2vec model on disk (see -conftest), passed through the ordinary model-override argument. +newline-delimited JSON protocol: requests are ``{"call": [value, model_id?]}`` +or, batched, ``{"calls": [[value, model_id?], ...]}``; responses +``{"ok": [floats...]}`` (plus an advisory ``"meta"``, covered in +cache_behavior_test) or ``{"err": "message"}``, one per call, wrapped in +``{"results": [...]}`` for a batch. The model is a real model2vec model on +disk (see conftest), passed through the ordinary model-override argument. """ import base64 @@ -108,3 +110,80 @@ def it_answers_err_naming_the_model_and_stays_alive( assert "err" in response assert "no-such-model" in response["err"] assert worker.request("hello", tiny_model)["ok"] == [1.0, 0.0] + + +def describe_batched_requests(): + def _batch(worker, *calls): + return worker.send_line(json.dumps({"calls": [list(c) for c in calls]})) + + def it_answers_results_in_request_order(spawn_worker, tiny_model): + worker = spawn_worker() + response = _batch( + worker, + ("hello", tiny_model), + ("world", tiny_model), + ("hello world", tiny_model), + ) + assert [result["ok"] for result in response["results"]] == [ + [1.0, 0.0], + [0.0, 1.0], + [0.5, 0.5], + ] + + def it_serves_two_models_in_one_request( + spawn_worker, tiny_model, other_model + ): + worker = spawn_worker() + response = _batch( + worker, + ("hello", tiny_model), + ("hello", other_model), + ("world", tiny_model), + ) + assert [result["ok"] for result in response["results"]] == [ + [1.0, 0.0], + [0.0, 2.0], + [0.0, 1.0], + ] + + def it_answers_each_malformed_call_in_place(spawn_worker, tiny_model): + worker = spawn_worker() + response = _batch( + worker, + ("hello", tiny_model), + (7, tiny_model), + (None, tiny_model), + ("hello", 42), + ) + ok, bad_value, null, bad_model = response["results"] + assert ok["ok"] == [1.0, 0.0] + assert "err" in bad_value + assert null == {"ok": None} + assert "err" in bad_model + + def it_fails_only_the_calls_of_an_unknown_model( + spawn_worker, tiny_model, tmp_path + ): + worker = spawn_worker() + missing = str(tmp_path / "no-such-model") + response = _batch(worker, ("hello", missing), ("hello", tiny_model)) + failed, ok = response["results"] + assert "no-such-model" in failed["err"] + assert ok["ok"] == [1.0, 0.0] + + def it_answers_an_empty_batch_with_no_results(spawn_worker): + worker = spawn_worker() + assert _batch(worker) == {"results": []} + + def it_rejects_calls_that_is_not_a_list_and_stays_alive( + spawn_worker, tiny_model + ): + worker = spawn_worker() + response = worker.send_line(json.dumps({"calls": "hello"})) + assert "err" in response + assert worker.request("hello", tiny_model)["ok"] == [1.0, 0.0] + + def it_keeps_serving_single_calls_after_a_batch(spawn_worker, tiny_model): + worker = spawn_worker() + _batch(worker, ("hello", tiny_model), ("world", tiny_model)) + assert worker.request("hello world", tiny_model)["ok"] == [0.5, 0.5]