dirsql-plugin-embeddings 0.1.0__py3-none-any.whl

dirsql_plugin_embeddings/dirsql.toml

@@ -0,0 +1,18 @@
+# Fragment shipped by the dirsql-plugin-embeddings plugin.
+# Discovered via the package's [project.entry-points.dirsql] declaration and
+# injected by the launcher as an ordinary -c flag ("installed = active", #529).
+#
+# The table name here MUST match pre_query.TABLE_NAME (the SQL the pre-query
+# hook prints queries this table).
+[dirsql]
+pre-query = "dirsql-embeddings-pre-query {args}"
+hook-timeout = 300
+
+[[dirsql.extension]]
+path = "sqlite_vec"
+entrypoint = "sqlite3_vec_init"
+
+[[table]]
+ddl = "CREATE TABLE documents (path TEXT, text TEXT, embedding TEXT)"
+glob = "**/*.md"
+on-file = "dirsql-embeddings-on-file {path}"

dirsql_plugin_embeddings/embedder.py

@@ -0,0 +1,71 @@
+"""Embed text via an OpenAI-compatible ``/v1/embeddings`` endpoint.
+
+Configuration comes from three environment variables (base URL, model, API
+key). The HTTP call is behind an injected ``post`` seam so the unit tests drive
+it without a real network; ``_urllib_post`` is the production seam.
+
+Annotations are evaluated at runtime (no ``from __future__ import annotations``)
+so a mutated ``X | None`` union in a signature fails at import instead of being
+an inert string -- otherwise every annotation-union mutant would survive.
+"""
+
+import json
+import os
+import urllib.error
+import urllib.request
+from collections.abc import Mapping
+
+ENV_BASE_URL = "DIRSQL_EMBEDDINGS_BASE_URL"
+ENV_MODEL = "DIRSQL_EMBEDDINGS_MODEL"
+ENV_API_KEY = "DIRSQL_EMBEDDINGS_API_KEY"
+
+
+class EmbeddingError(RuntimeError):
+    """A configuration, transport, or response error while embedding."""
+
+
+def _urllib_post(url: str, data: bytes, headers: Mapping[str, str]):
+    request = urllib.request.Request(
+        url, data=data, headers=dict(headers), method="POST"
+    )
+    try:
+        with urllib.request.urlopen(request) as response:
+            return response.status, response.read()
+    except urllib.error.HTTPError as error:
+        return error.code, error.read()
+
+
+def _require(env: Mapping[str, str], name: str) -> str:
+    value = env.get(name, "")
+    if not value:
+        raise EmbeddingError(f"missing required environment variable {name}")
+    return value
+
+
+def embed(
+    text: str, env: Mapping[str, str] | None = None, post=_urllib_post
+) -> list[float]:
+    if env is None:
+        env = os.environ
+    base_url = _require(env, ENV_BASE_URL).rstrip("/")
+    model = _require(env, ENV_MODEL)
+    api_key = _require(env, ENV_API_KEY)
+
+    data = json.dumps({"model": model, "input": [text]}).encode("utf-8")
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {api_key}",
+    }
+    status, body = post(f"{base_url}/v1/embeddings", data=data, headers=headers)
+    if status != 200:
+        raise EmbeddingError(f"embeddings endpoint returned status {status}: {body!r}")
+
+    try:
+        entry = json.loads(body)["data"][0]
+    except (json.JSONDecodeError, KeyError, IndexError, TypeError) as exc:
+        raise EmbeddingError(f"malformed embeddings response: {body!r}") from exc
+
+    vector = entry.get("embedding")
+    if not isinstance(vector, list) or not vector:
+        raise EmbeddingError(f"embeddings response carried no vector: {body!r}")
+    return [float(component) for component in vector]

dirsql_plugin_embeddings/on_file.py

@@ -0,0 +1,33 @@
+"""``on-file`` console script: embed one matched file into a dirsql row.
+
+Reads the file at ``argv[1]``, embeds its text, and prints a one-line JSON row
+array (``path``, ``text``, ``embedding``) -- the embedding stored as JSON text,
+which ``sqlite-vec`` accepts directly.
+
+Annotations are evaluated at runtime (no ``from __future__ import annotations``)
+so a mutated ``X | None`` union in a signature fails at import rather than
+surviving as an inert string.
+"""
+
+import json
+import sys
+
+from .embedder import embed
+
+
+def _read_text(path: str) -> str:
+    with open(path, encoding="utf-8") as handle:
+        return handle.read()
+
+
+def build_rows(path: str, text: str, vector: list[float]) -> list[dict]:
+    return [{"path": path, "text": text, "embedding": json.dumps(vector)}]
+
+
+def main(argv: list[str] | None = None) -> int:
+    if argv is None:
+        argv = sys.argv
+    path = argv[1]
+    text = _read_text(path)
+    print(json.dumps(build_rows(path, text, embed(text))))
+    return 0

dirsql_plugin_embeddings/pre_query.py

@@ -0,0 +1,42 @@
+"""``pre-query`` console script: turn a ``{"q": ...}`` body into search SQL.
+
+Accepts both a verbatim server body (``{"q": ...}``) and the CLI ``query``
+subcommand's ``{"sql": <arg>}`` wrapper, embeds the question, and prints the
+nearest-neighbor SQL over the ``documents`` table (ordered by
+``vec_distance_cosine``). The hook owns SQL safety: the only interpolated value
+is a numeric vector this script produced.
+
+Annotations are evaluated at runtime (no ``from __future__ import annotations``)
+so a mutated ``X | None`` union in a signature fails at import rather than
+surviving as an inert string.
+"""
+
+import json
+import sys
+
+from .embedder import embed
+
+TABLE_NAME = "documents"
+RESULT_LIMIT = 3
+
+
+def question(raw_body: str) -> str:
+    body = json.loads(raw_body)
+    if "q" in body:
+        return body["q"]
+    return json.loads(body["sql"])["q"]
+
+
+def build_sql(vector: list[float]) -> str:
+    needle = json.dumps(vector)
+    return (
+        f"SELECT path, ROUND(vec_distance_cosine(embedding, '{needle}'), 3) "
+        f"AS distance FROM {TABLE_NAME} ORDER BY distance LIMIT {RESULT_LIMIT}"
+    )
+
+
+def main(argv: list[str] | None = None) -> int:
+    if argv is None:
+        argv = sys.argv
+    print(build_sql(embed(question(argv[1]))))
+    return 0

dirsql_plugin_embeddings-0.1.0.dist-info/METADATA

@@ -0,0 +1,5 @@
+Metadata-Version: 2.4
+Name: dirsql-plugin-embeddings
+Version: 0.1.0
+Summary: First-party dirsql plugin: semantic search via an OpenAI-compatible embeddings endpoint.
+Requires-Python: >=3.11

dirsql_plugin_embeddings-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+dirsql_plugin_embeddings/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+dirsql_plugin_embeddings/embedder.py,sha256=4nNii0sGhskBwH7ppSMFgVKw2Z4yQT3KsR-CNYFNui8,2540
+dirsql_plugin_embeddings/on_file.py,sha256=OojSpkUSLeU4j2CbAdy97k_WcfKQ9pS3EzxEe7UbD-4,994
+dirsql_plugin_embeddings/pre_query.py,sha256=x4fOxHKmGlxBp3ojcPOjsIcM2kiQOC5veAFQgUjXcqQ,1265
+dirsql_plugin_embeddings/dirsql.toml,sha256=G3vdR9fpLt53IoPCkmd0KX9RqUCn3jPTXsphBk5W4FA,624
+dirsql_plugin_embeddings-0.1.0.dist-info/METADATA,sha256=nnaduB4QSPOi8CpJeCv65wIP4dOYDQ6FnTq8NzJau5I,190
+dirsql_plugin_embeddings-0.1.0.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
+dirsql_plugin_embeddings-0.1.0.dist-info/entry_points.txt,sha256=JhSvJnpAdXe7FoQ1vtbBgYmLiVPBr1AHPAYla-x93_s,202
+dirsql_plugin_embeddings-0.1.0.dist-info/RECORD,,

dirsql_plugin_embeddings-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.31.0
+Root-Is-Purelib: true
+Tag: py3-none-any

dirsql_plugin_embeddings-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,6 @@
+[console_scripts]
+dirsql-embeddings-on-file = dirsql_plugin_embeddings.on_file:main
+dirsql-embeddings-pre-query = dirsql_plugin_embeddings.pre_query:main
+
+[dirsql]
+embeddings = dirsql_plugin_embeddings