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.
- dirsql_plugin_embeddings-0.1.12/PKG-INFO +61 -0
- dirsql_plugin_embeddings-0.1.12/README.md +48 -0
- dirsql_plugin_embeddings-0.1.12/changelog.d/2026-08-10-embed-worker.md +14 -0
- dirsql_plugin_embeddings-0.1.12/changelog.d/2026-08-11-one-liner-cli.md +14 -0
- dirsql_plugin_embeddings-0.1.12/e2e-attestations/claude-804-embed-worker.json +7 -0
- dirsql_plugin_embeddings-0.1.12/e2e-attestations/claude-805-one-liner.json +7 -0
- {dirsql_plugin_embeddings-0.1.11 → dirsql_plugin_embeddings-0.1.12}/pyproject.toml +31 -6
- dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/cli/__init__.py +0 -0
- dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/cli/main.py +37 -0
- dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/cli/search.py +26 -0
- dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/cli/worker.py +13 -0
- dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/dirsql.toml +16 -0
- dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/embedding/__init__.py +0 -0
- dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/embedding/cache.py +19 -0
- dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/embedding/model.py +15 -0
- dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/embedding/progress.py +13 -0
- dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/embedding/values.py +28 -0
- dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/embedding/worker.py +71 -0
- dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/search/__init__.py +0 -0
- dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/search/output.py +2 -0
- dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/search/run.py +20 -0
- dirsql_plugin_embeddings-0.1.12/src/dirsql_plugin_embeddings/search/sql.py +32 -0
- dirsql_plugin_embeddings-0.1.12/tests/conftest.py +115 -0
- dirsql_plugin_embeddings-0.1.12/tests/e2e/__init__.py +0 -0
- dirsql_plugin_embeddings-0.1.12/tests/integration/__init__.py +0 -0
- dirsql_plugin_embeddings-0.1.11/PKG-INFO +0 -28
- dirsql_plugin_embeddings-0.1.11/README.md +0 -20
- dirsql_plugin_embeddings-0.1.11/src/dirsql_plugin_embeddings/dirsql.toml +0 -6
- {dirsql_plugin_embeddings-0.1.11 → dirsql_plugin_embeddings-0.1.12}/.gitignore +0 -0
- {dirsql_plugin_embeddings-0.1.11 → dirsql_plugin_embeddings-0.1.12}/changelog.d/2026-08-10-delete-old-surface.md +0 -0
- {dirsql_plugin_embeddings-0.1.11 → dirsql_plugin_embeddings-0.1.12}/migrations.d/2026-08-10-delete-old-surface.md +0 -0
- {dirsql_plugin_embeddings-0.1.11 → dirsql_plugin_embeddings-0.1.12}/src/dirsql_plugin_embeddings/__init__.py +0 -0
- {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.
|
|
@@ -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.
|
|
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
|
|
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
|
-
|
|
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 = [
|
|
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"]
|
|
File without changes
|
|
@@ -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"
|
|
File without changes
|
|
@@ -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()
|
|
File without changes
|
|
@@ -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
|
|
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"
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{dirsql_plugin_embeddings-0.1.11 → dirsql_plugin_embeddings-0.1.12}/testing-conventions.toml
RENAMED
|
File without changes
|