dirsql-plugin-embeddings 0.1.11__tar.gz → 0.1.12__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (33) hide show
  1. dirsql_plugin_embeddings-0.1.12/PKG-INFO +61 -0
  2. dirsql_plugin_embeddings-0.1.12/README.md +48 -0
  3. dirsql_plugin_embeddings-0.1.12/changelog.d/2026-08-10-embed-worker.md +14 -0
  4. dirsql_plugin_embeddings-0.1.12/changelog.d/2026-08-11-one-liner-cli.md +14 -0
  5. dirsql_plugin_embeddings-0.1.12/e2e-attestations/claude-804-embed-worker.json +7 -0
  6. dirsql_plugin_embeddings-0.1.12/e2e-attestations/claude-805-one-liner.json +7 -0
  7. {dirsql_plugin_embeddings-0.1.11 → dirsql_plugin_embeddings-0.1.12}/pyproject.toml +31 -6
  8. dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/cli/__init__.py +0 -0
  9. dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/cli/main.py +37 -0
  10. dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/cli/search.py +26 -0
  11. dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/cli/worker.py +13 -0
  12. dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/dirsql.toml +16 -0
  13. dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/embedding/__init__.py +0 -0
  14. dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/embedding/cache.py +19 -0
  15. dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/embedding/model.py +15 -0
  16. dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/embedding/progress.py +13 -0
  17. dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/embedding/values.py +28 -0
  18. dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/embedding/worker.py +71 -0
  19. dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/search/__init__.py +0 -0
  20. dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/search/output.py +2 -0
  21. dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/search/run.py +20 -0
  22. dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/search/sql.py +32 -0
  23. dirsql_plugin_embeddings-0.1.12/tests/conftest.py +115 -0
  24. dirsql_plugin_embeddings-0.1.12/tests/e2e/__init__.py +0 -0
  25. dirsql_plugin_embeddings-0.1.12/tests/integration/__init__.py +0 -0
  26. dirsql_plugin_embeddings-0.1.11/PKG-INFO +0 -28
  27. dirsql_plugin_embeddings-0.1.11/README.md +0 -20
  28. dirsql_plugin_embeddings-0.1.11/src/dirsql_plugin_embeddings/dirsql.toml +0 -6
  29. {dirsql_plugin_embeddings-0.1.11 → dirsql_plugin_embeddings-0.1.12}/.gitignore +0 -0
  30. {dirsql_plugin_embeddings-0.1.11 → dirsql_plugin_embeddings-0.1.12}/changelog.d/2026-08-10-delete-old-surface.md +0 -0
  31. {dirsql_plugin_embeddings-0.1.11 → dirsql_plugin_embeddings-0.1.12}/migrations.d/2026-08-10-delete-old-surface.md +0 -0
  32. {dirsql_plugin_embeddings-0.1.11 → dirsql_plugin_embeddings-0.1.12}/src/dirsql_plugin_embeddings/__init__.py +0 -0
  33. {dirsql_plugin_embeddings-0.1.11 → dirsql_plugin_embeddings-0.1.12}/testing-conventions.toml +0 -0
@@ -0,0 +1,61 @@
1
+ Metadata-Version: 2.5
2
+ Name: dirsql-plugin-embeddings
3
+ Version: 0.1.12
4
+ Summary: First-party dirsql plugin: semantic search over files.
5
+ Requires-Python: >=3.10
6
+ Requires-Dist: cachetta>=0.7
7
+ Requires-Dist: click>=8
8
+ Requires-Dist: dirsql
9
+ Requires-Dist: model2vec>=0.6
10
+ Requires-Dist: sqlite-vec>=0.1
11
+ Requires-Dist: tqdm>=4
12
+ Description-Content-Type: text/markdown
13
+
14
+ # dirsql-plugin-embeddings
15
+
16
+ A first-party [`dirsql`](https://github.com/thekevinscott/dirsql) plugin for
17
+ semantic search over files.
18
+
19
+ Installing the plugin loads the
20
+ [`sqlite-vec`](https://github.com/asg017/sqlite-vec) extension (for
21
+ `vec_distance_cosine()` and friends) and declares an `embed()` SQL scalar
22
+ function that turns TEXT or BLOB values into embedding vectors:
23
+
24
+ ```sh
25
+ uvx --with dirsql-plugin-embeddings dirsql "
26
+ SELECT path
27
+ FROM (SELECT path, embed(content ->> 'abstract') AS emb
28
+ FROM 'arxiv-firehose/data/**/metadata.json')
29
+ ORDER BY vec_distance_cosine(emb, embed('local private models'))
30
+ LIMIT 10"
31
+ ```
32
+
33
+ `embed()` is inert until a query calls it: no worker process is spawned and no
34
+ model is loaded for queries that never use it. On the first call, dirsql
35
+ spawns the plugin's worker process (`dirsql-plugin-embeddings worker`), which
36
+ serves every call of the invocation over stdin/stdout.
37
+
38
+ ## Model
39
+
40
+ Embeddings come from [model2vec](https://github.com/MinishLab/model2vec)
41
+ (static embeddings — numpy + tokenizers, no torch), defaulting to
42
+ [`minishlab/potion-retrieval-32M`](https://huggingface.co/minishlab/potion-retrieval-32M).
43
+ The model downloads to the standard Hugging Face cache on the first ever run,
44
+ with progress on stderr when stderr is a TTY.
45
+
46
+ An optional second argument overrides the model per call — the id must be
47
+ model2vec-loadable (sentence-transformers/torch models are out of scope):
48
+
49
+ ```sql
50
+ SELECT embed('some text', 'minishlab/potion-base-8M')
51
+ ```
52
+
53
+ ## Vector cache
54
+
55
+ Computed vectors are cached at `~/.cache/dirsql/embeddings/` (or
56
+ `$XDG_CACHE_HOME/dirsql/embeddings/` when `XDG_CACHE_HOME` is set), keyed by
57
+ the SHA-256 of the value bytes plus the model identifier — changing either
58
+ recomputes; switching models never serves stale vectors. There is no
59
+ eviction: **the directory is safe to wipe at any time**; the only cost is
60
+ re-embedding. The cache never lives inside a queried tree — the worker
61
+ receives values, not paths, and writes nothing anywhere else.
@@ -0,0 +1,48 @@
1
+ # dirsql-plugin-embeddings
2
+
3
+ A first-party [`dirsql`](https://github.com/thekevinscott/dirsql) plugin for
4
+ semantic search over files.
5
+
6
+ Installing the plugin loads the
7
+ [`sqlite-vec`](https://github.com/asg017/sqlite-vec) extension (for
8
+ `vec_distance_cosine()` and friends) and declares an `embed()` SQL scalar
9
+ function that turns TEXT or BLOB values into embedding vectors:
10
+
11
+ ```sh
12
+ uvx --with dirsql-plugin-embeddings dirsql "
13
+ SELECT path
14
+ FROM (SELECT path, embed(content ->> 'abstract') AS emb
15
+ FROM 'arxiv-firehose/data/**/metadata.json')
16
+ ORDER BY vec_distance_cosine(emb, embed('local private models'))
17
+ LIMIT 10"
18
+ ```
19
+
20
+ `embed()` is inert until a query calls it: no worker process is spawned and no
21
+ model is loaded for queries that never use it. On the first call, dirsql
22
+ spawns the plugin's worker process (`dirsql-plugin-embeddings worker`), which
23
+ serves every call of the invocation over stdin/stdout.
24
+
25
+ ## Model
26
+
27
+ Embeddings come from [model2vec](https://github.com/MinishLab/model2vec)
28
+ (static embeddings — numpy + tokenizers, no torch), defaulting to
29
+ [`minishlab/potion-retrieval-32M`](https://huggingface.co/minishlab/potion-retrieval-32M).
30
+ The model downloads to the standard Hugging Face cache on the first ever run,
31
+ with progress on stderr when stderr is a TTY.
32
+
33
+ An optional second argument overrides the model per call — the id must be
34
+ model2vec-loadable (sentence-transformers/torch models are out of scope):
35
+
36
+ ```sql
37
+ SELECT embed('some text', 'minishlab/potion-base-8M')
38
+ ```
39
+
40
+ ## Vector cache
41
+
42
+ Computed vectors are cached at `~/.cache/dirsql/embeddings/` (or
43
+ `$XDG_CACHE_HOME/dirsql/embeddings/` when `XDG_CACHE_HOME` is set), keyed by
44
+ the SHA-256 of the value bytes plus the model identifier — changing either
45
+ recomputes; switching models never serves stale vectors. There is no
46
+ eviction: **the directory is safe to wipe at any time**; the only cost is
47
+ re-embedding. The cache never lives inside a queried tree — the worker
48
+ receives values, not paths, and writes nothing anywhere else.
@@ -0,0 +1,14 @@
1
+ **Added** the `dirsql-plugin-embeddings worker` subcommand: a persistent
2
+ stdin/stdout process serving `embed()` requests as newline-delimited JSON
3
+ (`{"call": [value, model_id?]}` → `{"ok": [floats...]}` / `{"err": "message"}`).
4
+ Values are SQL TEXT (JSON strings) or BLOB (`{"$bytes": "<base64>"}`, decoded
5
+ as utf-8); embeddings come from model2vec (default
6
+ `minishlab/potion-retrieval-32M`, loaded lazily on the first request, one load
7
+ per process; an optional second argument overrides the model id). Vectors are
8
+ cached on disk via cachetta at `~/.cache/dirsql/embeddings/` (respects
9
+ `XDG_CACHE_HOME`; safe to wipe — the only cost is re-embedding), keyed by the
10
+ SHA-256 of the value bytes plus the model identifier. Progress is reported on
11
+ stderr only when stderr is a TTY. The packaged `dirsql.toml` now declares the
12
+ `embed` SQL function via `[[dirsql.function]]` (name `embed`, arities 1–2,
13
+ deterministic, 600s per-call timeout) alongside the existing sqlite-vec
14
+ `[[dirsql.extension]]` entry.
@@ -0,0 +1,14 @@
1
+ **Added** the one-liner search CLI as the default command:
2
+ `dirsql-plugin-embeddings '<glob>' '<query>' [-k/--limit N] [--model ID]`.
3
+ The corpus glob is a required first positional (no default corpus), the query
4
+ text the second; `-k`/`--limit` (default 10) is exactly the SQL `LIMIT`, and
5
+ `--model` templates the model id as `embed()`'s second argument in the
6
+ generated SQL. The command builds the canonical search SQL (an `embed()`
7
+ subquery over the glob's rows, `vec_distance_cosine` against
8
+ `embed('<query>')`, `ORDER BY distance LIMIT k`), runs it via the dirsql
9
+ Python SDK (now a declared dependency; needs dirsql >= 0.4.17, the first
10
+ release with `[[dirsql.function]]`) with the packaged
11
+ config fragment, and prints ranked `path<TAB>distance` lines. Query text,
12
+ model id, and glob are SQL-escaped. The `worker` subcommand is unchanged;
13
+ there is no explicit `search` spelling — a literal `search` first token is a
14
+ corpus glob.
@@ -0,0 +1,7 @@
1
+ {
2
+ "command": "uv run --with-editable ../../packages/python python -m pytest tests/e2e/ -x -q",
3
+ "ran_at": 1786455708,
4
+ "exit_code": 0,
5
+ "commit": "9d8bf1c91b82cb4f6648789e0d4df9b1aa932921",
6
+ "branch": "claude/804-embed-worker"
7
+ }
@@ -0,0 +1,7 @@
1
+ {
2
+ "command": "uv run --with-editable ../../packages/python python -m pytest tests/e2e/ -x -q",
3
+ "ran_at": 1786459020,
4
+ "exit_code": 0,
5
+ "commit": "1bd863d143acbb821b3e6e594e1e1688a5e1178d",
6
+ "branch": "claude/805-one-liner"
7
+ }
@@ -3,12 +3,10 @@ requires = ["hatchling", "hatch-vcs"]
3
3
  build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
- # First-party dirsql plugin, published independently to PyPI. Mid-rebuild:
7
- # today it only loads sqlite-vec; the embed() SQL function surface is landing
8
- # in follow-up slices.
6
+ # First-party dirsql plugin, published independently to PyPI.
9
7
  name = "dirsql-plugin-embeddings"
10
8
  dynamic = ["version"]
11
- description = "First-party dirsql plugin: semantic search over files (rebuild in progress)."
9
+ description = "First-party dirsql plugin: semantic search over files."
12
10
  readme = "README.md"
13
11
  requires-python = ">=3.10"
14
12
  # sqlite-vec is a runtime dependency, not a docs suggestion: the shipped
@@ -16,7 +14,29 @@ requires-python = ">=3.10"
16
14
  # only resolve when the package is installed. Declaring it here makes the
17
15
  # quickstart (`uvx --with dirsql-plugin-embeddings dirsql`) work without a
18
16
  # separate `--with sqlite-vec`.
19
- dependencies = ["sqlite-vec>=0.1"]
17
+ #
18
+ # model2vec (NOT sentence-transformers/torch) is the inference stack: static
19
+ # embeddings whose runtime deps are numpy + tokenizers, viable in uvx
20
+ # ephemeral envs. cachetta is the vector cache; tqdm the stderr progress UI.
21
+ #
22
+ # dirsql is the SDK the one-liner search command runs its generated SQL
23
+ # through. Functionally it needs >= 0.4.17 (the first release with the core
24
+ # [[dirsql.function]] mechanism), but the requirement stays unversioned: the
25
+ # in-repo editable SDK builds under the stale on-disk version literal (the
26
+ # release pipeline rewrites it at publish time), so any floor would make the
27
+ # local e2e overlay (`uv run --with-editable ../../packages/python`)
28
+ # unresolvable. Fresh installs resolve the latest release regardless.
29
+ dependencies = [
30
+ "dirsql",
31
+ "sqlite-vec>=0.1",
32
+ "model2vec>=0.6",
33
+ "cachetta>=0.7",
34
+ "tqdm>=4",
35
+ "click>=8",
36
+ ]
37
+
38
+ [project.scripts]
39
+ dirsql-plugin-embeddings = "dirsql_plugin_embeddings.cli.main:main"
20
40
 
21
41
  [tool.hatch.version]
22
42
  source = "vcs"
@@ -47,7 +67,12 @@ git_describe_command = ["git", "describe", "--dirty", "--tags", "--long", "--mat
47
67
  embeddings = "dirsql_plugin_embeddings"
48
68
 
49
69
  [dependency-groups]
50
- dev = ["pytest>=8", "pytest-cov>=5", "pytest-describe>=2"]
70
+ dev = [
71
+ "pytest>=8",
72
+ "pytest-cov>=5",
73
+ "pytest-describe>=2",
74
+ "tomli>=2; python_version < '3.11'",
75
+ ]
51
76
 
52
77
  [tool.hatch.build]
53
78
  exclude = ["**/*_test.py"]
@@ -0,0 +1,37 @@
1
+ import click
2
+
3
+ from .search import search
4
+ from .worker import worker
5
+
6
+
7
+ class DefaultCommandGroup(click.Group):
8
+ """Route non-`worker` invocations to the hidden search command.
9
+
10
+ Bare positionals are the plugin's only search interface
11
+ (`dirsql-plugin-embeddings '<glob>' '<query>'`): a first token counts as
12
+ a subcommand only when it names a *visible* command, so a literal
13
+ 'search' first token is a corpus glob, not a spelling of the command.
14
+ """
15
+
16
+ def parse_args(self, ctx, args):
17
+ visible = [
18
+ name
19
+ for name, command in self.commands.items()
20
+ if not command.hidden
21
+ ]
22
+ if not args or (args[0] != "--help" and args[0] not in visible):
23
+ args = [search.name, *args]
24
+ return super().parse_args(ctx, args)
25
+
26
+
27
+ @click.group(cls=DefaultCommandGroup)
28
+ def main():
29
+ """dirsql embeddings plugin.
30
+
31
+ Bare arguments run the semantic search:
32
+ dirsql-plugin-embeddings '<glob>' '<query>' [-k N] [--model ID]
33
+ """
34
+
35
+
36
+ main.add_command(worker, name="worker")
37
+ main.add_command(search, name="search")
@@ -0,0 +1,26 @@
1
+ import click
2
+
3
+ from ..search.run import run_search
4
+
5
+
6
+ @click.command(hidden=True)
7
+ @click.argument("glob")
8
+ @click.argument("query")
9
+ @click.option(
10
+ "-k",
11
+ "--limit",
12
+ "limit",
13
+ type=int,
14
+ default=10,
15
+ show_default=True,
16
+ help="Maximum results; exactly the SQL LIMIT.",
17
+ )
18
+ @click.option(
19
+ "--model",
20
+ default=None,
21
+ help="model2vec model id, templated as embed()'s second argument.",
22
+ )
23
+ def search(glob, query, limit, model):
24
+ """Rank files matching GLOB by semantic similarity to QUERY."""
25
+ for line in run_search(glob, query, limit, model):
26
+ click.echo(line)
@@ -0,0 +1,13 @@
1
+ import sys
2
+
3
+ import click
4
+
5
+ from ..embedding.progress import configure
6
+ from ..embedding.worker import Worker
7
+
8
+
9
+ @click.command()
10
+ def worker():
11
+ """Serve embed() requests as newline-delimited JSON on stdin/stdout."""
12
+ configure()
13
+ Worker().serve(sys.stdin, sys.stdout)
@@ -0,0 +1,16 @@
1
+ # Fragment shipped by the dirsql-plugin-embeddings plugin.
2
+ # Discovered via the package's [project.entry-points.dirsql] declaration and
3
+ # injected by the launcher as an ordinary -c flag ("installed = active").
4
+ [[dirsql.extension]]
5
+ path = "sqlite_vec"
6
+ entrypoint = "sqlite3_vec_init"
7
+
8
+ # The embed() SQL scalar function, served by the plugin's worker process over
9
+ # newline-delimited JSON on stdin/stdout. Inert until a query calls it. The
10
+ # generous timeout absorbs the first call's model download.
11
+ [[dirsql.function]]
12
+ name = "embed"
13
+ args = [1, 2]
14
+ command = "dirsql-plugin-embeddings worker"
15
+ deterministic = true
16
+ timeout = "600s"
@@ -0,0 +1,19 @@
1
+ import os
2
+ from datetime import timedelta
3
+ from pathlib import Path
4
+
5
+ from cachetta import Cachetta
6
+
7
+ # No eviction: entries stay until the user wipes the directory (documented as
8
+ # safe -- the only cost is re-embedding).
9
+ NO_EVICTION = timedelta(days=365000)
10
+
11
+
12
+ def cache_dir():
13
+ xdg = os.environ.get("XDG_CACHE_HOME", "")
14
+ base = Path(xdg) if xdg else Path.home() / ".cache"
15
+ return base / "dirsql" / "embeddings"
16
+
17
+
18
+ def make_cache():
19
+ return Cachetta(path=cache_dir(), hashed=True, duration=NO_EVICTION)
@@ -0,0 +1,15 @@
1
+ DEFAULT_MODEL_ID = "minishlab/potion-retrieval-32M"
2
+
3
+
4
+ def load_model(model_id):
5
+ from model2vec import StaticModel
6
+
7
+ return StaticModel.from_pretrained(model_id)
8
+
9
+
10
+ def model_identifier(model_id, model):
11
+ config = getattr(model, "config", None) or {}
12
+ version = config.get("model2vec_version")
13
+ if version:
14
+ return f"{model_id}@{version}"
15
+ return model_id
@@ -0,0 +1,13 @@
1
+ import os
2
+ import sys
3
+
4
+
5
+ def stderr_is_tty():
6
+ return sys.stderr.isatty()
7
+
8
+
9
+ def configure():
10
+ # Hugging Face hub download bars honor this variable; setdefault keeps an
11
+ # explicit user choice in force.
12
+ if not stderr_is_tty():
13
+ os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
@@ -0,0 +1,28 @@
1
+ import base64
2
+ import binascii
3
+
4
+
5
+ class ProtocolError(ValueError):
6
+ """A request value that does not follow the worker protocol."""
7
+
8
+
9
+ def decode_value(value):
10
+ if value is None:
11
+ return None
12
+ if isinstance(value, str):
13
+ return value
14
+ if isinstance(value, dict) and set(value) == {"$bytes"}:
15
+ encoded = value["$bytes"]
16
+ if not isinstance(encoded, str):
17
+ raise ProtocolError('the "$bytes" value must be a base64 string')
18
+ try:
19
+ raw = base64.b64decode(encoded, validate=True)
20
+ except binascii.Error as error:
21
+ raise ProtocolError(f'invalid base64 in "$bytes": {error}') from error
22
+ try:
23
+ return raw.decode("utf-8")
24
+ except UnicodeDecodeError as error:
25
+ raise ProtocolError(f"BLOB is not valid utf-8 text: {error}") from error
26
+ raise ProtocolError(
27
+ f"embed() accepts TEXT or BLOB values, got {type(value).__name__}"
28
+ )
@@ -0,0 +1,71 @@
1
+ import json
2
+ from hashlib import sha256
3
+
4
+ from . import model
5
+ from .cache import make_cache
6
+ from .progress import stderr_is_tty
7
+ from .values import ProtocolError, decode_value
8
+
9
+ MALFORMED_SHAPE = (
10
+ 'malformed request: expected {"call": [value, model_id?]} on one line'
11
+ )
12
+ MALFORMED_ARITY = 'malformed request: "call" must carry 1 or 2 arguments'
13
+ MALFORMED_MODEL_ID = "the model id must be TEXT"
14
+
15
+
16
+ class Worker:
17
+ def __init__(self):
18
+ self._models = {}
19
+ self._cache = make_cache()
20
+ self._compute_cached = self._cache.wrap(self._compute)
21
+ self._pending = None
22
+
23
+ def _model(self, model_id):
24
+ if model_id not in self._models:
25
+ self._models[model_id] = model.load_model(model_id)
26
+ return self._models[model_id]
27
+
28
+ def _compute(self, digest, identifier):
29
+ text, loaded = self._pending
30
+ (vector,) = loaded.encode([text], show_progress_bar=stderr_is_tty())
31
+ return [float(component) for component in vector]
32
+
33
+ def embed(self, text, model_id):
34
+ loaded = self._model(model_id)
35
+ identifier = model.model_identifier(model_id, loaded)
36
+ digest = sha256(text.encode("utf-8")).hexdigest()
37
+ self._pending = (text, loaded)
38
+ return self._compute_cached(digest, identifier)
39
+
40
+ def handle(self, line):
41
+ try:
42
+ request = json.loads(line)
43
+ except json.JSONDecodeError as error:
44
+ return {"err": f"malformed request: invalid JSON: {error}"}
45
+ if not isinstance(request, dict) or "call" not in request:
46
+ return {"err": MALFORMED_SHAPE}
47
+ call = request["call"]
48
+ if not isinstance(call, list) or len(call) not in (1, 2):
49
+ return {"err": MALFORMED_ARITY}
50
+ value, *rest = call
51
+ try:
52
+ text = decode_value(value)
53
+ except ProtocolError as error:
54
+ return {"err": str(error)}
55
+ if text is None:
56
+ return {"ok": None}
57
+ (model_id,) = rest or [model.DEFAULT_MODEL_ID]
58
+ if not isinstance(model_id, str):
59
+ return {"err": MALFORMED_MODEL_ID}
60
+ try:
61
+ return {"ok": self.embed(text, model_id)}
62
+ except Exception as error:
63
+ return {"err": f"embed({model_id!r}) failed: {error}"}
64
+
65
+ def serve(self, stdin, stdout):
66
+ for line in stdin:
67
+ if not line.strip():
68
+ continue
69
+ response = self.handle(line)
70
+ stdout.write(json.dumps(response, separators=(",", ":")) + "\n")
71
+ stdout.flush()
@@ -0,0 +1,2 @@
1
+ def format_rows(rows):
2
+ return [f"{row['path']}\t{row['distance']:.6f}" for row in rows]
@@ -0,0 +1,20 @@
1
+ import asyncio
2
+ from importlib import resources
3
+
4
+ import dirsql
5
+
6
+ from .output import format_rows
7
+ from .sql import build_search_sql
8
+
9
+
10
+ def config_fragment():
11
+ return str(resources.files("dirsql_plugin_embeddings").joinpath("dirsql.toml"))
12
+
13
+
14
+ async def _query(sql):
15
+ return await dirsql.DirSQL(config=config_fragment()).query(sql)
16
+
17
+
18
+ def run_search(glob, query, limit, model=None):
19
+ rows = asyncio.run(_query(build_search_sql(glob, query, limit, model)))
20
+ return format_rows(rows)
@@ -0,0 +1,32 @@
1
+ PATH_PREFIXES = ("./", "../", "/", "~/")
2
+
3
+
4
+ def quote(text):
5
+ escaped = text.replace("'", "''")
6
+ return f"'{escaped}'"
7
+
8
+
9
+ def normalize_glob(glob):
10
+ # The core only rescues path-shaped missing tables (./, ../, /, ~/); a
11
+ # bare relative glob like '**/*.md' would error with a "did you mean
12
+ # './...'" hint. Here GLOB is unambiguously a corpus glob, so spare the
13
+ # user the round trip.
14
+ if glob.startswith(PATH_PREFIXES):
15
+ return glob
16
+ return f"./{glob}"
17
+
18
+
19
+ def embed_call(argument, model):
20
+ if model is None:
21
+ return f"embed({argument})"
22
+ return f"embed({argument}, {quote(model)})"
23
+
24
+
25
+ def build_search_sql(glob, query, limit, model=None):
26
+ outer = embed_call(quote(query), model)
27
+ inner = embed_call("content", model)
28
+ return (
29
+ f"SELECT path, vec_distance_cosine(emb, {outer}) AS distance"
30
+ f" FROM (SELECT path, {inner} AS emb FROM {quote(normalize_glob(glob))})"
31
+ f" ORDER BY distance LIMIT {int(limit):d}"
32
+ )
@@ -0,0 +1,115 @@
1
+ """Shared fixtures for the integration and e2e tiers.
2
+
3
+ ``make_model`` builds a real model2vec model on disk: a three-token WordLevel
4
+ vocabulary with hand-picked vectors, so embeddings are deterministic and no
5
+ network or Hugging Face download is involved. The worker loads it through the
6
+ ordinary model-override argument (`{"call": [text, "<path>"]}`), exercising
7
+ the exact code path a hub model id takes.
8
+
9
+ ``worker_process`` spawns the real worker subcommand as a subprocess over real
10
+ pipes, with ``XDG_CACHE_HOME`` pointed at a per-test temp dir so the vector
11
+ cache is isolated and its location is observable.
12
+ """
13
+
14
+ import json
15
+ import os
16
+ import subprocess
17
+ import sys
18
+
19
+ import numpy as np
20
+ import pytest
21
+ from model2vec import StaticModel
22
+ from tokenizers import Tokenizer
23
+ from tokenizers.models import WordLevel
24
+ from tokenizers.pre_tokenizers import Whitespace
25
+
26
+ VOCAB = {"[UNK]": 0, "hello": 1, "world": 2}
27
+
28
+ WORKER_ARGV = [
29
+ sys.executable,
30
+ "-c",
31
+ "import sys; from dirsql_plugin_embeddings.cli.main import main; sys.exit(main())",
32
+ "worker",
33
+ ]
34
+
35
+
36
+ def build_model(directory, rows):
37
+ tokenizer = Tokenizer(WordLevel(VOCAB, unk_token="[UNK]"))
38
+ tokenizer.pre_tokenizer = Whitespace()
39
+ vectors = np.array(rows, dtype=np.float32)
40
+ model = StaticModel(vectors, tokenizer, config={"normalize": False})
41
+ model.save_pretrained(directory)
42
+ return str(directory)
43
+
44
+
45
+ @pytest.fixture(scope="session")
46
+ def tiny_model(tmp_path_factory):
47
+ # [UNK] -> [0, 0], hello -> [1, 0], world -> [0, 1].
48
+ return build_model(
49
+ tmp_path_factory.mktemp("model-a"),
50
+ [[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]],
51
+ )
52
+
53
+
54
+ @pytest.fixture(scope="session")
55
+ def other_model(tmp_path_factory):
56
+ # Same vocabulary, different vectors: hello -> [0, 2].
57
+ return build_model(
58
+ tmp_path_factory.mktemp("model-b"),
59
+ [[0.0, 0.0], [0.0, 2.0], [2.0, 0.0]],
60
+ )
61
+
62
+
63
+ class WorkerProcess:
64
+ def __init__(self, argv, cache_home, cwd=None):
65
+ self.process = subprocess.Popen(
66
+ argv,
67
+ stdin=subprocess.PIPE,
68
+ stdout=subprocess.PIPE,
69
+ stderr=subprocess.PIPE,
70
+ text=True,
71
+ cwd=cwd,
72
+ env={**os.environ, "XDG_CACHE_HOME": str(cache_home)},
73
+ )
74
+
75
+ def send_line(self, line):
76
+ self.process.stdin.write(line + "\n")
77
+ self.process.stdin.flush()
78
+ response = self.process.stdout.readline()
79
+ assert response, (
80
+ f"worker produced no response line; exited with"
81
+ f" {self.process.poll()}, stderr: {self.process.stderr.read()!r}"
82
+ )
83
+ return json.loads(response)
84
+
85
+ def request(self, *call):
86
+ return self.send_line(json.dumps({"call": list(call)}))
87
+
88
+ def close(self):
89
+ self.process.stdin.close()
90
+ code = self.process.wait(timeout=30)
91
+ stderr = self.process.stderr.read()
92
+ self.process.stdout.close()
93
+ self.process.stderr.close()
94
+ return code, stderr
95
+
96
+
97
+ @pytest.fixture
98
+ def cache_home(tmp_path):
99
+ return tmp_path / "xdg-cache"
100
+
101
+
102
+ @pytest.fixture
103
+ def spawn_worker(cache_home):
104
+ workers = []
105
+
106
+ def spawn(argv=WORKER_ARGV, cwd=None):
107
+ worker = WorkerProcess(argv, cache_home, cwd=cwd)
108
+ workers.append(worker)
109
+ return worker
110
+
111
+ yield spawn
112
+ for worker in workers:
113
+ if worker.process.poll() is None:
114
+ worker.process.kill()
115
+ worker.process.wait()
File without changes
@@ -1,28 +0,0 @@
1
- Metadata-Version: 2.4
2
- Name: dirsql-plugin-embeddings
3
- Version: 0.1.11
4
- Summary: First-party dirsql plugin: semantic search over files (rebuild in progress).
5
- Requires-Python: >=3.10
6
- Requires-Dist: sqlite-vec>=0.1
7
- Description-Content-Type: text/markdown
8
-
9
- # dirsql-plugin-embeddings
10
-
11
- A first-party [`dirsql`](https://github.com/thekevinscott/dirsql) plugin for
12
- semantic search over files.
13
-
14
- **Rebuild in progress**
15
- ([#800](https://github.com/thekevinscott/dirsql/issues/800)): the previous
16
- surface — a declared `documents` table built eagerly over the working
17
- directory, per-file `on-file` embedding hooks, a `pre-query` hook, and
18
- `DIRSQL_EMBEDDINGS_*` endpoint configuration — has been removed. The plugin is
19
- being rebuilt around a plugin-provided `embed()` SQL function that is inert
20
- until a query calls it.
21
-
22
- Today, installing the plugin loads the
23
- [`sqlite-vec`](https://github.com/asg017/sqlite-vec) extension (for
24
- `vec_distance_cosine()` and friends) and nothing else.
25
-
26
- ```sh
27
- uvx --with dirsql-plugin-embeddings dirsql "SELECT vec_distance_cosine('[1, 0]', '[0, 1]') AS d"
28
- ```
@@ -1,20 +0,0 @@
1
- # dirsql-plugin-embeddings
2
-
3
- A first-party [`dirsql`](https://github.com/thekevinscott/dirsql) plugin for
4
- semantic search over files.
5
-
6
- **Rebuild in progress**
7
- ([#800](https://github.com/thekevinscott/dirsql/issues/800)): the previous
8
- surface — a declared `documents` table built eagerly over the working
9
- directory, per-file `on-file` embedding hooks, a `pre-query` hook, and
10
- `DIRSQL_EMBEDDINGS_*` endpoint configuration — has been removed. The plugin is
11
- being rebuilt around a plugin-provided `embed()` SQL function that is inert
12
- until a query calls it.
13
-
14
- Today, installing the plugin loads the
15
- [`sqlite-vec`](https://github.com/asg017/sqlite-vec) extension (for
16
- `vec_distance_cosine()` and friends) and nothing else.
17
-
18
- ```sh
19
- uvx --with dirsql-plugin-embeddings dirsql "SELECT vec_distance_cosine('[1, 0]', '[0, 1]') AS d"
20
- ```
@@ -1,6 +0,0 @@
1
- # Fragment shipped by the dirsql-plugin-embeddings plugin.
2
- # Discovered via the package's [project.entry-points.dirsql] declaration and
3
- # injected by the launcher as an ordinary -c flag ("installed = active").
4
- [[dirsql.extension]]
5
- path = "sqlite_vec"
6
- entrypoint = "sqlite3_vec_init"