diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 00000000..5b102d99 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,13 @@ +repos: + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.11.0 + hooks: + - id: ruff + args: [--fix] + - id: ruff-format + + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.15.0 + hooks: + - id: mypy + additional_dependencies: [types-pyyaml] diff --git a/pyproject.toml b/pyproject.toml index 5f39c75b..f6463c64 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,25 +19,22 @@ classifiers = [ dependencies = [ "pydantic>=2.13.4,<3", "click>=8.4.0,<9", - "click>=8.1,<9", "pyyaml>=6,<7", "mcp>=1.0,<2", ] [project.optional-dependencies] +embeddings = [ + "sentence-transformers>=3,<4", + "numpy>=1.26,<2", +] dev = [ "pytest>=9.0.3,<10", "pytest-cov>=5,<6", "mypy>=2.1.0", "ruff>=0.15.13", - "mypy>=1.10", "types-pyyaml", ] -embeddings = [ - "sentence-transformers>=2.7,<4", - "numpy>=1.26,<3", - "sqlite-vec>=0.1,<1", -] embeddings-fast = [ "fastembed>=0.3,<1", "onnxruntime>=1.18,<2", diff --git a/src/vouch/bundle.py b/src/vouch/bundle.py index 7a1959a0..1faefd95 100644 --- a/src/vouch/bundle.py +++ b/src/vouch/bundle.py @@ -209,8 +209,6 @@ class ImportCheckResult: def _validate_content(path: str, data: bytes, issues: list[str]) -> None: - if not any(path.endswith(ext) for ext in (".yaml", ".yml", ".md")): - return subdir = path.split("/")[0] # Source artifacts have two file kinds: # sources//meta.yaml -- the Source pydantic model (validate) @@ -222,6 +220,8 @@ def _validate_content(path: str, data: bytes, issues: list[str]) -> None: validator = VALIDATORS.get(subdir) if validator is None: return + if not any(path.lower().endswith(ext) for ext in (".yaml", ".yml", ".md")): + return try: validator(data) except Exception as e: diff --git a/src/vouch/capabilities.py b/src/vouch/capabilities.py index c54df48f..4d4fc087 100644 --- a/src/vouch/capabilities.py +++ b/src/vouch/capabilities.py @@ -57,10 +57,18 @@ def capabilities() -> Capabilities: + retrieval = ["fts5", "substring"] + try: + from .embeddings import get_embedder + get_embedder() + retrieval.append("embedding") + retrieval.append("hybrid") + except Exception: + pass return Capabilities( version=__version__, methods=METHODS, - retrieval=["fts5", "substring"], + retrieval=retrieval, review_gated=True, transports=["mcp", "jsonl"], ) diff --git a/src/vouch/cli.py b/src/vouch/cli.py index c00877d3..86e09b91 100644 --- a/src/vouch/cli.py +++ b/src/vouch/cli.py @@ -465,20 +465,36 @@ def crystallize(session_id: str, no_page: bool) -> None: @cli.command() @click.argument("query") @click.option("--limit", default=10, show_default=True, type=int) -def search(query: str, limit: int) -> None: +@click.option("--embedding", "use_embedding", is_flag=True, + help="Use embedding-based semantic search") +def search(query: str, limit: int, use_embedding: bool) -> None: """Search claims, pages, and entities (embedding → fts5 → substring).""" from . import index_db store = _load_store() - try: - hits = index_db.search(store.kb_dir, query, limit=limit) - if not hits: + + if use_embedding: + try: + from .embeddings import get_embedder + embedder = get_embedder() + except Exception: + click.echo("Error: sentence-transformers not installed. " + "pip install vouch[embeddings]", err=True) + raise SystemExit(1) from None + vec = embedder.encode(query).tolist() + hits = index_db.search_embedding(store.kb_dir, query_vec=vec, limit=limit) + backend = "embedding" + else: + try: + hits = index_db.search(store.kb_dir, query, limit=limit) + if not hits: + hits = store.search_substring(query, limit=limit) + backend = "substring" + else: + backend = "fts5" + except Exception: hits = store.search_substring(query, limit=limit) backend = "substring" - else: - backend = "fts5" - except Exception: - hits = store.search_substring(query, limit=limit) - backend = "substring" + for kind, hid, snippet, score in hits: click.echo(f"[{kind}] {hid} score={score:.3f} ({backend})") click.echo(f" {snippet[:200]}") diff --git a/src/vouch/health.py b/src/vouch/health.py index 0bc684cb..d83b99b9 100644 --- a/src/vouch/health.py +++ b/src/vouch/health.py @@ -164,9 +164,35 @@ def rebuild_index(store: KBStore) -> dict: conn, id=e.id, name=e.name, description=e.description, type=e.type.value, aliases=e.aliases, ) + _rebuild_embeddings(store) return index_db.stats(store.kb_dir) +def _rebuild_embeddings(store: KBStore) -> None: + try: + from .embeddings import get_embedder + embedder = get_embedder() + except Exception: + return + with index_db.open_db(store.kb_dir) as conn: + texts: list[tuple[str, str, str]] = [] + for c in store.list_claims(): + texts.append(("claim", c.id, c.text)) + for p in store.list_pages(): + texts.append(("page", p.id, f"{p.title} {p.body}")) + for e in store.list_entities(): + texts.append(("entity", e.id, f"{e.name} {e.description or ''}")) + if not texts: + return + batch_size = 64 + for i in range(0, len(texts), batch_size): + batch = texts[i:i + batch_size] + vecs = embedder.encode_batch([t[2] for t in batch]) + for (kind, eid, _), row in zip(batch, vecs, strict=True): + index_db.index_embedding(conn, kind=kind, id=eid, + vec=row.tolist()) + + # --- helpers used by `vouch discover` (CLI) ------------------------------- def hash_path(p: Path) -> str: diff --git a/src/vouch/index_db.py b/src/vouch/index_db.py index 71c7734a..bb350c6a 100644 --- a/src/vouch/index_db.py +++ b/src/vouch/index_db.py @@ -12,6 +12,7 @@ from __future__ import annotations import datetime as _dt +import json import sqlite3 from collections.abc import Iterable from contextlib import contextmanager, suppress @@ -33,6 +34,14 @@ id UNINDEXED, name, description, type UNINDEXED, aliases ); +CREATE TABLE IF NOT EXISTS embeddings ( + kind TEXT NOT NULL, + id TEXT NOT NULL, + vec BLOB NOT NULL, + model TEXT NOT NULL DEFAULT 'all-MiniLM-L6-v2', + PRIMARY KEY (kind, id) +); + CREATE TABLE IF NOT EXISTS index_meta ( key TEXT PRIMARY KEY, value TEXT NOT NULL ); @@ -108,6 +117,52 @@ def reset(kb_dir: Path) -> None: ) +def index_embedding(conn: sqlite3.Connection, *, kind: str, id: str, + vec: list[float]) -> None: + conn.execute( + "INSERT OR REPLACE INTO embeddings (kind, id, vec) VALUES (?, ?, ?)", + (kind, id, json.dumps(vec)), + ) + + +def search_embeddings(kb_dir: Path, query_vec: list[float], *, + limit: int = 10 + ) -> list[tuple[str, str, str, float]]: + """Return (kind, id, snippet, cosine_score) via brute-force NumPy scan.""" + import numpy as np # type: ignore[import-not-found] + out: list[tuple[str, str, str, float]] = [] + if not query_vec: + return out + q = np.array(query_vec, dtype=np.float32) + q_norm = np.linalg.norm(q) + if q_norm == 0.0: + return out + q = q / q_norm + with open_db(kb_dir) as conn: + rows = conn.execute( + "SELECT kind, id, vec FROM embeddings" + ).fetchall() + for kind, eid, vec_json in rows: + v = np.array(json.loads(vec_json), dtype=np.float32) + if v.ndim != 1 or v.shape[0] != q.shape[0]: + continue + score = float(np.dot(q, v)) + snippet = _snippet_for(kb_dir, kind, eid) + out.append((kind, eid, snippet, score)) + out.sort(key=lambda x: x[3], reverse=True) + return out[:limit] + + +def _snippet_for(kb_dir: Path, kind: str, eid: str) -> str: + path = kb_dir / kind / f"{eid}.yaml" + if not path.exists(): + path = kb_dir / kind / f"{eid}.md" + if not path.exists(): + return eid + text = path.read_text() + return text[:200].replace("\n", " ") + + def index_claim(conn: sqlite3.Connection, *, id: str, text: str, type: str, status: str, tags: Iterable[str]) -> None: conn.execute("DELETE FROM claims_fts WHERE id = ?", (id,)) diff --git a/src/vouch/jsonl_server.py b/src/vouch/jsonl_server.py index f6f6a293..aaf7e472 100644 --- a/src/vouch/jsonl_server.py +++ b/src/vouch/jsonl_server.py @@ -81,9 +81,6 @@ def _h_search(p: dict) -> list[dict]: hits: list[tuple[str, str, str, float]] = [] used = backend_arg - # Reject unknown backends with a clear error rather than silently - # returning []. Falling through hides client typos and diverges from - # the MCP transport, which raises ValueError on the same input. valid_backends = {"auto", "embedding", "fts5", "substring", "hybrid"} if backend_arg not in valid_backends: raise ValueError( @@ -128,6 +125,7 @@ def _h_search(p: dict) -> list[dict]: ] + def _h_context(p: dict) -> dict: return build_context_pack( _store(),