#!/usr/bin/env -S uv run --script # /// script # requires-python = ">=3.11" # dependencies = ["requests", "pyyaml", "markdown-it-py", "sqlite-vec==0.1.6"] # /// """Query side of the wiki skill — three layers, cheapest first. grep live ripgrep over the files on disk. No index, never stale, and not limited to the indexed extensions, so it is the layer that covers source code. toc directory -> file -> title + tags, read straight from the `files` catalog. search FTS5 (BM25) and vec0 (KNN) over the same `chunks` rows, merged with RRF. Both halves rank the same unit, which is what makes the merge meaningful. The parameters below are module constants on purpose: `--limit` is the only knob worth exposing, and a config key that never changes is a config key that rots. """ from __future__ import annotations import argparse import json import shutil import sqlite3 import subprocess import sys from pathlib import Path import wiki_store as store from wiki_config import ( GIT_KIND, ConfigError, SourceConfig, WikiConfig, db_path, load_config, source_root, ) from wiki_embed import ( EmbeddingUnavailable, IndexIdentityMismatch, OllamaEmbedder, require_matching_index, ) CANDIDATE_LIMIT = 50 # KNN k, and the LIMIT for BM25 RRF_K = 60 # Cormack et al. 2009 DEFAULT_LIMIT = 10 PREFIX_MIN_LENGTH = 3 # a one- or two-character prefix matches too widely to carry signal EXCERPT_CHARS = 320 GREP_TIMEOUT = 30 GREP_MAX_PER_FILE = 5 def rrf_merge(ranked_lists: list[list[int]]) -> list[tuple[int, float]]: """score(id) = sum over lists of 1 / (RRF_K + rank). No weights: both halves count equally.""" scores: dict[int, float] = {} for ids in ranked_lists: for rank, chunk_id in enumerate(ids, start=1): scores[chunk_id] = scores.get(chunk_id, 0.0) + 1.0 / (RRF_K + rank) return sorted(scores.items(), key=lambda item: (-item[1], item[0])) def fts_match_expression(query: str) -> str: """Turn free text into an FTS5 OR-query of quoted prefix terms. Terms are quoted so reserved words (`and`, `not`, `near`) and punctuation cannot be read as operators. The trailing `*` covers Czech inflection, which `unicode61` does not stem — `záloh*` finds záloha/zálohování/zálohy. """ terms = [] for word in _tokenize(query): wildcard = "*" if len(word) >= PREFIX_MIN_LENGTH else "" terms.append(f'"{word}"{wildcard}') return " OR ".join(terms) def _tokenize(text: str) -> list[str]: return "".join(char if char.isalnum() else " " for char in text).split() # --------------------------------------------------------------------------- # search # --------------------------------------------------------------------------- def run_search(config: WikiConfig, query: str, limit: int) -> int: expression = fts_match_expression(query) if not expression: print("(empty query)") return 0 with store.connection(db_path()) as conn: try: require_matching_index(store.read_meta(conn), config.embedding) except IndexIdentityMismatch as exc: print(str(exc), file=sys.stderr) return 1 bm25_ids = store.bm25_ranked_ids(conn, expression, CANDIDATE_LIMIT) vector_ids: list[int] = [] degraded: str | None = None try: vector = OllamaEmbedder(config.embedding).embed_query(query) vector_ids = store.knn_ranked_ids(conn, vector, CANDIDATE_LIMIT) except EmbeddingUnavailable as exc: degraded = f"note: embeddings unavailable ({exc}) — FTS-only results" pending = store.index_stats(conn)["pending"] merged = rrf_merge([ids for ids in (bm25_ids, vector_ids) if ids])[:limit] rows = store.fetch_chunks(conn, [chunk_id for chunk_id, _ in merged]) if degraded: print(degraded) elif pending: print(f"note: {pending} chunks still awaiting vectors — semantic half is incomplete") if not merged: print("(no matches)") return 0 bm25_rank = {chunk_id: rank for rank, chunk_id in enumerate(bm25_ids, start=1)} vector_rank = {chunk_id: rank for rank, chunk_id in enumerate(vector_ids, start=1)} for position, (chunk_id, score) in enumerate(merged, start=1): row = rows.get(chunk_id) if row is None: continue origin = _origin_label(bm25_rank.get(chunk_id), vector_rank.get(chunk_id)) print(f"{position}. {row['source_id']}:{row['path']} (rrf {score:.4f}, {origin})") print(f" {row['breadcrumb']}") print(_indent(_excerpt(row["text"], row["breadcrumb"]))) print() return 0 def _origin_label(bm25: int | None, vector: int | None) -> str: parts = [] if bm25 is not None: parts.append(f"bm25 #{bm25}") if vector is not None: parts.append(f"vec #{vector}") return ", ".join(parts) def _excerpt(text: str, breadcrumb: str) -> str: body = text[len(breadcrumb) :].lstrip("\n") if text.startswith(breadcrumb) else text body = " ".join(body.split()) return body[:EXCERPT_CHARS] + ("…" if len(body) > EXCERPT_CHARS else "") def _indent(text: str) -> str: return "\n".join(f" {line}" for line in text.splitlines()) # --------------------------------------------------------------------------- # toc # --------------------------------------------------------------------------- def run_toc(source_id: str | None, tag: str | None) -> int: with store.connection(db_path()) as conn: rows = store.list_toc_files(conn, source_id=source_id, tag=tag) if not rows: print("(no indexed files match)") return 0 by_source: dict[str, list[sqlite3.Row]] = {} for row in rows: by_source.setdefault(row["source_id"], []).append(row) for source, files in by_source.items(): print(f"{source} ({len(files)} files)") print() directory = None for row in files: path = row["path"] parent, _, name = path.rpartition("/") if parent != directory: directory = parent print(f"{parent}/" if parent else "./") print(f" {name:<28} {row['title'] or '':<32} {_tag_label(row['tags'])}".rstrip()) print() return 0 def _tag_label(raw: str | None) -> str: try: tags = json.loads(raw) if raw else [] except json.JSONDecodeError: return "" return f"[{', '.join(tags)}]" if tags else "" # --------------------------------------------------------------------------- # grep # --------------------------------------------------------------------------- def grep_roots(source: SourceConfig) -> list[Path]: """Directories ripgrep should walk for one source. A git clone is searched whole. A workspace source is bounded by the literal prefix of each `paths` glob, so grep stays inside what the source claims without being narrowed to the indexed extensions. """ root = source_root(source) if source.kind == GIT_KIND: return [root] if root.is_dir() else [] roots = [] for pattern in source.paths: prefix = _literal_prefix(pattern) candidate = root / prefix if prefix else root if candidate.is_dir() and candidate not in roots: roots.append(candidate) return roots def _literal_prefix(pattern: str) -> str: parts = [] for segment in pattern.split("/"): if any(char in segment for char in "*?["): break parts.append(segment) return "/".join(parts) def run_grep(config: WikiConfig, pattern: str, source_id: str | None) -> int: if shutil.which("rg") is None: print("ripgrep (rg) not found", file=sys.stderr) return 1 sources = [s for s in config.sources if not source_id or s.source_id == source_id] if not sources: print(f"unknown source {source_id!r}", file=sys.stderr) return 1 roots = [path for source in sources for path in grep_roots(source)] if not roots: print("(nothing on disk to grep — has wiki_sync.py run?)") return 0 command = [ "rg", "--line-number", "--no-heading", "--color", "never", "--smart-case", "--max-count", str(GREP_MAX_PER_FILE), "--glob", "!.git", pattern, *[str(path) for path in roots], ] try: result = subprocess.run(command, capture_output=True, text=True, timeout=GREP_TIMEOUT, check=False) except subprocess.TimeoutExpired: print(f"grep timed out after {GREP_TIMEOUT}s", file=sys.stderr) return 1 if result.returncode not in (0, 1): print(result.stderr.strip(), file=sys.stderr) return 1 output = result.stdout.strip() print(output if output else "(no matches)") return 0 # --------------------------------------------------------------------------- # cli # --------------------------------------------------------------------------- def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description="Search the wiki index over your notes.") subparsers = parser.add_subparsers(dest="command", required=True) search = subparsers.add_parser("search", help="hybrid BM25 + vector search over chunks") search.add_argument("query", help="free text; Czech inflection is covered by prefix matching") search.add_argument("--limit", type=int, default=DEFAULT_LIMIT, help="chunks to return") toc = subparsers.add_parser("toc", help="directory -> file -> title + tags from the catalog") toc.add_argument("--source", help="limit to one source id") toc.add_argument("--tag", help="only files carrying this frontmatter tag") grep = subparsers.add_parser("grep", help="live ripgrep over the files on disk, index-free") grep.add_argument("pattern", help="ripgrep regex") grep.add_argument("--source", help="limit to one source id") return parser.parse_args(argv) def main(argv: list[str] | None = None) -> int: args = parse_args(argv) try: config = load_config() except ConfigError as exc: print(f"config error: {exc}", file=sys.stderr) return 1 if args.command == "search": return run_search(config, args.query, args.limit) if args.command == "toc": return run_toc(args.source, args.tag) return run_grep(config, args.pattern, args.source) if __name__ == "__main__": sys.exit(main())