indexkit 0.6.1__py3-none-any.whl

indexkit/__init__.py

@@ -0,0 +1,3 @@
+"""Local semantic and hybrid RAG: ollama, turbovec, and FTS5/BM25."""
+
+__version__ = "0.6.1"

indexkit/__main__.py

@@ -0,0 +1,4 @@
+from .cli import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())

indexkit/cli.py

@@ -0,0 +1,254 @@
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sqlite3
+import sys
+from pathlib import Path
+
+from . import __version__
+from .embed import OllamaEmbedder
+from .engine import Engine, slug
+from .storage import list_indexes, remove_index, validate_index_name
+
+
+def _first_env(*names: str) -> str | None:
+    for name in names:
+        value = os.environ.get(name)
+        if value:
+            return value
+    return None
+
+
+def _legacy_data_dir() -> Path:
+    """Where a plugin install kept index data before standalone packaging.
+
+    Kept resolvable so upgrading does not silently orphan existing indexes.
+    This is deliberate compatibility, not dead code — see the governing records.
+    """
+    # @adr 0006
+    # @adr 0007
+    return Path.home() / ".claude/plugins/data/indexkit"
+
+
+def _default_data_dir() -> Path:
+    """Host-neutral default for a plain `pip install indexkit`.
+
+    A user who installed from PyPI has no plugin host, so writing under
+    `~/.claude/` would create a directory for an unrelated tool. Follow the XDG
+    base-directory spec instead, honoring `XDG_DATA_HOME` when set.
+
+    An existing plugin-managed directory still wins when the new location has
+    not been created yet, so upgrading in place keeps working (ADR-0006).
+    """
+    xdg = os.environ.get("XDG_DATA_HOME")
+    default = (Path(xdg).expanduser() if xdg else Path.home() / ".local/share") / "indexkit"
+    if not default.exists():
+        legacy = _legacy_data_dir()
+        if legacy.is_dir():
+            return legacy
+    return default
+
+
+def _data_dir() -> Path:
+    """Resolve where indexes live.
+
+    Explicit configuration wins, then the host-provided plugin data directory,
+    then a host-neutral default. `CLAUDE_PLUGIN_DATA` is set by the plugin host,
+    so honoring it keeps plugin installs pointed at host-managed storage while a
+    standalone install never touches it.
+    """
+    configured = _first_env("CONTEXT_KIT_DATA", "PRODUCTIVITY_SKILLS_DATA", "CLAUDE_PLUGIN_DATA")
+    if configured:
+        return Path(configured).expanduser()
+    return _default_data_dir().expanduser()
+
+
+def _make_embedder(args):
+    model = (
+        getattr(args, "model", None)
+        or _first_env(
+            "CONTEXT_KIT_EMBED_MODEL",
+            "PRODUCTIVITY_SKILLS_EMBED_MODEL",
+            "CLAUDE_PLUGIN_OPTION_EMBED_MODEL",
+        )
+        or "nomic-embed-text"
+    )
+    host = (
+        _first_env(
+            "CONTEXT_KIT_OLLAMA_HOST",
+            "PRODUCTIVITY_SKILLS_OLLAMA_HOST",
+            "CLAUDE_PLUGIN_OPTION_OLLAMA_HOST",
+        )
+        or "http://localhost:11434"
+    )
+    return OllamaEmbedder(
+        model=model,
+        host=host,
+    )
+
+
+def _name_for(args, corpus=None) -> str:
+    if getattr(args, "name", None):
+        return validate_index_name(args.name)
+    return validate_index_name(slug(corpus) if corpus else "default")
+
+
+def _index_name(value: str) -> str:
+    try:
+        return validate_index_name(value)
+    except ValueError as error:
+        raise argparse.ArgumentTypeError(str(error)) from error
+
+
+def _read_allowlist(value) -> list[str] | None:
+    if not value:
+        return None
+    if value == "-":
+        return [ln.strip() for ln in sys.stdin if ln.strip()]
+    lines = Path(value).read_text(encoding="utf-8").splitlines()
+    return [ln.strip() for ln in lines if ln.strip()]
+
+
+def _missing_index(name: str) -> str:
+    return f"error: no index named '{name}'. Run 'rag index <path> --name {name}' first."
+
+
+def main(argv=None) -> int:
+    argv = list(sys.argv[1:] if argv is None else argv)
+    p = argparse.ArgumentParser(
+        prog="indexkit",
+        description="Hybrid semantic and lexical retrieval over a local index.",
+    )
+    # Exposed so an integrating adapter can pin/report a provider version the
+    # same way it does for other retrieval backends.
+    p.add_argument("--version", action="version", version=f"indexkit {__version__}")
+    sub = p.add_subparsers(dest="cmd", required=True)
+
+    pi = sub.add_parser("index")
+    pi.add_argument("path")
+    pi.add_argument("--name", type=_index_name)
+    pi.add_argument("--model")
+    pi.add_argument("--include", action="append")
+    pi.add_argument("--exclude", action="append")
+
+    pq = sub.add_parser("query")
+    pq.add_argument("text")
+    pq.add_argument("--name", type=_index_name)
+    pq.add_argument("--model")
+    pq.add_argument("--k", type=int, default=10)
+    pq.add_argument("--allowlist")
+    pq.add_argument(
+        "--hybrid",
+        action="store_true",
+        help="Fuse semantic and SQLite FTS5/BM25 candidates with reciprocal-rank fusion.",
+    )
+    pq.add_argument("--json", action="store_true")
+
+    ps = sub.add_parser("status")
+    ps.add_argument("--name", type=_index_name)
+    ps.add_argument("--model")
+    sub.add_parser("list")
+
+    pr = sub.add_parser("remove")
+    pr.add_argument("--name", required=True, type=_index_name)
+    pr.add_argument(
+        "--yes",
+        action="store_true",
+        help="Confirm non-interactive, permanent removal of the named index.",
+    )
+
+    args = p.parse_args(argv)
+    data = _data_dir()
+
+    if args.cmd == "list":
+        for name in list_indexes(data):
+            print(name)
+        return 0
+
+    if args.cmd == "remove":
+        if not args.yes:
+            print(
+                f"error: refusing to remove index '{args.name}' without --yes",
+                file=sys.stderr,
+            )
+            return 2
+        try:
+            removed = remove_index(data, args.name)
+        except (OSError, RuntimeError) as error:
+            print(f"error: {error}", file=sys.stderr)
+            return 1
+        print(f"removed={args.name} artifacts={removed}")
+        return 0
+
+    if args.cmd == "index":
+        eng = None
+        try:
+            eng = Engine(_name_for(args, args.path), data, _make_embedder(args))
+            res = eng.index(args.path, args.include, args.exclude)
+        except Exception as e:
+            print(f"error: {e}", file=sys.stderr)
+            return 1
+        finally:
+            if eng is not None:
+                eng.close()
+        print(
+            f"indexed={res['indexed']} skipped={res['skipped']} "
+            f"files={res['files']} chunks={res['chunks']}"
+        )
+        return 0
+
+    if args.cmd == "status":
+        name = _name_for(args)
+        eng = None
+        try:
+            eng = Engine(name, data, _make_embedder(args), create=False)
+            print(json.dumps(eng.store.stats()))
+        except FileNotFoundError:
+            print(_missing_index(name), file=sys.stderr)
+            return 1
+        except (OSError, RuntimeError, sqlite3.Error) as error:
+            print(f"error: {error}", file=sys.stderr)
+            return 1
+        finally:
+            if eng is not None:
+                eng.close()
+        return 0
+
+    if args.cmd == "query":
+        name = _name_for(args)
+        eng = None
+        try:
+            eng = Engine(name, data, _make_embedder(args), create=False)
+            hits = eng.query(
+                args.text,
+                k=args.k,
+                allowlist_paths=_read_allowlist(args.allowlist),
+                hybrid=args.hybrid,
+            )
+        except FileNotFoundError:
+            print(_missing_index(name), file=sys.stderr)
+            return 1
+        except Exception as e:
+            print(f"error: {e}", file=sys.stderr)
+            return 1
+        finally:
+            if eng is not None:
+                eng.close()
+        if args.json:
+            print(json.dumps(hits))
+        else:
+            for h in hits:
+                loc = f"{h['path']}" + (f" > {h['heading']}" if h["heading"] else "")
+                snippet = h["snippet"].replace("\n", " ")
+                score = f"[{h['score']:.3f}]"
+                if h["retrieval_mode"] == "hybrid":
+                    score = f"[{h['score']:.3f} hybrid]"
+                print(f"{score} {loc}\n    {snippet}")
+        return 0
+    return 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

indexkit/embed.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+import httpx
+
+
+class EmbedError(RuntimeError):
+    pass
+
+
+class OllamaEmbedder:
+    def __init__(self, model: str = "nomic-embed-text", host: str = "http://localhost:11434"):
+        self.model = model
+        self.host = host.rstrip("/")
+        self._client = httpx.Client()
+        self._dim: int | None = None
+
+    def _embed_one(self, text: str) -> list[float]:
+        try:
+            resp = self._client.post(
+                f"{self.host}/api/embeddings",
+                json={"model": self.model, "prompt": text},
+                timeout=60.0,
+            )
+            resp.raise_for_status()
+        except httpx.ConnectError as e:
+            raise EmbedError(
+                f"Could not reach ollama at {self.host}: {e}. "
+                f"Start it with 'ollama serve' and ensure the model is pulled: "
+                f"'ollama pull {self.model}'."
+            ) from e
+        except httpx.HTTPStatusError as e:
+            raise EmbedError(
+                f"ollama returned an error for model '{self.model}'. "
+                f"Pull it with 'ollama pull {self.model}'. ({e})"
+            ) from e
+        vec = resp.json().get("embedding")
+        if not vec:
+            raise EmbedError(f"ollama returned no embedding for model '{self.model}'.")
+        return vec
+
+    def embed(self, texts: list[str]) -> list[list[float]]:
+        return [self._embed_one(t) for t in texts]
+
+    def dim(self) -> int:
+        if self._dim is None:
+            self._dim = len(self._embed_one("dimension probe"))
+        return self._dim
+
+    def close(self) -> None:
+        self._client.close()
+
+    def __enter__(self) -> "OllamaEmbedder":
+        return self
+
+    def __exit__(self, *exc) -> None:
+        self.close()

indexkit/engine.py

@@ -0,0 +1,272 @@
+from __future__ import annotations
+
+from contextlib import ExitStack
+import hashlib
+import re
+from pathlib import Path
+
+from .embed import OllamaEmbedder
+from .index import VecIndex
+from .loaders.markdown import iter_corpus, load_markdown
+from .storage import IndexLock, ensure_index_dir, existing_index_dir
+from .store import MetaStore
+
+RRF_K = 60
+SEMANTIC_RRF_WEIGHT = 1.0
+LEXICAL_RRF_WEIGHT = 1.0
+HYBRID_CANDIDATE_MULTIPLIER = 3
+
+
+def slug(path) -> str:
+    s = re.sub(r"[^a-zA-Z0-9]+", "-", str(Path(path).resolve())).strip("-").lower()
+    return s[-80:].lstrip("-") or "default"
+
+
+def reciprocal_rank_fusion(
+    semantic_hits: list[tuple[int, float]],
+    lexical_hits: list[tuple[int, float]],
+) -> list[dict]:
+    """Fuse ranked sources with ``weight / (RRF_K + rank)``.
+
+    The returned order breaks equal RRF scores by best source rank, then chunk id,
+    making fusion reproducible even when a source assigns equal scores.
+    """
+    candidates: dict[int, dict] = {}
+    for rank, (chunk_id, score) in enumerate(semantic_hits, start=1):
+        hit = candidates.setdefault(chunk_id, {"id": chunk_id})
+        hit["semantic_rank"] = rank
+        hit["semantic_score"] = score
+    for rank, (chunk_id, score) in enumerate(lexical_hits, start=1):
+        hit = candidates.setdefault(chunk_id, {"id": chunk_id})
+        hit["lexical_rank"] = rank
+        hit["lexical_score"] = score
+    for hit in candidates.values():
+        hit["fused_score"] = (
+            SEMANTIC_RRF_WEIGHT / (RRF_K + hit["semantic_rank"]) if "semantic_rank" in hit else 0.0
+        ) + (LEXICAL_RRF_WEIGHT / (RRF_K + hit["lexical_rank"]) if "lexical_rank" in hit else 0.0)
+    return sorted(
+        candidates.values(),
+        key=lambda hit: (
+            -hit["fused_score"],
+            min(hit.get("semantic_rank", float("inf")), hit.get("lexical_rank", float("inf"))),
+            hit["id"],
+        ),
+    )
+
+
+class Engine:
+    def __init__(self, name: str, data_dir, embedder=None, *, create: bool = True):
+        self.name = name
+        embedder_instance = embedder or OllamaEmbedder()
+        with ExitStack() as cleanup:
+            cleanup.callback(embedder_instance.close)
+            lock = IndexLock(data_dir, name)
+            lock.acquire()
+            cleanup.callback(lock.release)
+            directory = (
+                ensure_index_dir(data_dir, name) if create else existing_index_dir(data_dir, name)
+            )
+            store = MetaStore(directory / "meta.sqlite")
+            cleanup.callback(store.close)
+            store.init_schema()
+            self.dir = directory
+            self.embedder = embedder_instance
+            self.store = store
+            self.index_path = self.dir / "index.tvim"
+            self._lock = lock
+            cleanup.pop_all()
+
+    def close(self) -> None:
+        try:
+            try:
+                store = getattr(self, "store", None)
+                if store is not None:
+                    store.close()
+            finally:
+                embedder = getattr(self, "embedder", None)
+                if embedder is not None:
+                    embedder.close()
+        finally:
+            lock = getattr(self, "_lock", None)
+            if lock is not None:
+                lock.release()
+
+    def _dim(self) -> int:
+        d = self.embedder.dim()
+        stored = self.store.get_meta("dim")
+        if stored is None:
+            self.store.set_meta("dim", str(d))
+            self.store.set_meta("model", self.embedder.model)
+        elif int(stored) != d:
+            raise ValueError(
+                f"Embedding dim {d} (model '{self.embedder.model}') != index dim {stored} "
+                f"(model '{self.store.get_meta('model')}'). Reindex with --name NEW or a matching model."
+            )
+        return d
+
+    def _load_index(self, dim: int) -> VecIndex:
+        if self.index_path.exists():
+            return VecIndex.load(dim=dim, path=self.index_path)
+        return VecIndex(dim=dim, path=self.index_path)
+
+    def index(self, root, include=None, exclude=None) -> dict:
+        dim = self._dim()
+        idx = self._load_index(dim)
+        root = Path(root)
+        indexed = skipped = remove_failures = 0
+        seen: set[str] = set()
+        for fp in iter_corpus(root, include, exclude):
+            rel = fp.relative_to(root).as_posix()
+            seen.add(rel)
+            raw = fp.read_text(encoding="utf-8", errors="replace")
+            h = hashlib.sha256(raw.encode("utf-8")).hexdigest()
+            if self.store.file_hash(rel) == h:
+                skipped += 1
+                continue
+            for cid in self.store.chunk_ids_for_paths([rel]):
+                try:
+                    idx.remove(cid)
+                except Exception:
+                    remove_failures += 1
+            chunks = load_markdown(raw, rel)
+            if not chunks:
+                self.store.upsert_file(rel, h, [])
+                continue
+            vecs = self.embedder.embed([c.text for c in chunks])
+            ids = self.store.upsert_file(rel, h, chunks)
+            idx.add(ids, vecs)
+            indexed += 1
+        for gone in set(self.store.all_paths()) - seen:
+            for cid in self.store.chunk_ids_for_paths([gone]):
+                try:
+                    idx.remove(cid)
+                except Exception:
+                    remove_failures += 1
+            self.store.delete_file(gone)
+        self.store.set_meta("root", str(Path(root).resolve()))
+        idx.save()
+        st = self.store.stats()
+        return {
+            "indexed": indexed,
+            "skipped": skipped,
+            "chunks": st["chunks"],
+            "files": st["files"],
+            "remove_failures": remove_failures,
+        }
+
+    def _normalize_allowlist(self, paths: list[str]) -> list[str]:
+        """Map incoming file paths (absolute, $VAULT-prefixed, relative, or
+        basename) to the corpus-relative keys used in the store."""
+        stored = self.store.all_paths()
+        have = set(stored)
+        root_meta = self.store.get_meta("root")
+        root = Path(root_meta) if root_meta else None
+        out: list[str] = []
+        for raw in paths:
+            if raw in have:  # already a stored relative key
+                out.append(raw)
+                continue
+            p = Path(raw)
+            matched = False
+            if root is not None:
+                try:
+                    rel = p.resolve().relative_to(root).as_posix()
+                except ValueError:
+                    rel = None
+                if rel and rel in have:
+                    out.append(rel)
+                    matched = True
+            if not matched:
+                # basename fallback (handles $VAULT-relative prefixes)
+                base = p.name
+                for h in stored:
+                    if h == base or Path(h).name == base or h.endswith("/" + base):
+                        out.append(h)
+                        matched = True
+        # de-dup preserving order
+        seen: set[str] = set()
+        return [x for x in out if not (x in seen or seen.add(x))]
+
+    @staticmethod
+    def _result(
+        ch: dict,
+        *,
+        score: float,
+        retrieval_mode: str,
+        semantic_rank: int | None = None,
+        semantic_score: float | None = None,
+        lexical_rank: int | None = None,
+        lexical_score: float | None = None,
+        fused_rank: int | None = None,
+        fused_score: float | None = None,
+    ) -> dict:
+        return {
+            "path": ch["path"],
+            "heading": ch["heading"],
+            "score": score,
+            "snippet": ch["text"][:240],
+            "start": ch["start"],
+            "end": ch["end"],
+            "retrieval_mode": retrieval_mode,
+            "semantic_rank": semantic_rank,
+            "semantic_score": semantic_score,
+            "lexical_rank": lexical_rank,
+            "lexical_score": lexical_score,
+            "fused_rank": fused_rank,
+            "fused_score": fused_score,
+        }
+
+    def query(
+        self, text: str, k: int = 10, allowlist_paths=None, hybrid: bool = False
+    ) -> list[dict]:
+        if k <= 0:
+            return []
+        if hybrid and not self.store.fts5_available:
+            raise RuntimeError("SQLite FTS5 is unavailable; --hybrid cannot be used.")
+        dim = int(self.store.get_meta("dim") or self.embedder.dim())
+        idx = self._load_index(dim)
+        qv = self.embedder.embed([text])[0]
+        allow = None
+        if allowlist_paths is not None:
+            allow = self.store.chunk_ids_for_paths(self._normalize_allowlist(allowlist_paths))
+        candidate_depth = k * HYBRID_CANDIDATE_MULTIPLIER if hybrid else k
+        semantic_hits = idx.search(qv, k=candidate_depth, allowlist=allow)
+        if hybrid:
+            lexical_hits = self.store.lexical_search(text, k=candidate_depth, allowlist=allow)
+            fused_hits = reciprocal_rank_fusion(semantic_hits, lexical_hits)[:k]
+            out = []
+            for fused_rank, hit in enumerate(fused_hits, start=1):
+                try:
+                    ch = self.store.get_chunk(hit["id"])
+                except KeyError:
+                    continue
+                out.append(
+                    self._result(
+                        ch,
+                        score=hit["fused_score"],
+                        retrieval_mode="hybrid",
+                        semantic_rank=hit.get("semantic_rank"),
+                        semantic_score=hit.get("semantic_score"),
+                        lexical_rank=hit.get("lexical_rank"),
+                        lexical_score=hit.get("lexical_score"),
+                        fused_rank=fused_rank,
+                        fused_score=hit["fused_score"],
+                    )
+                )
+            return out
+        out = []
+        for semantic_rank, (cid, score) in enumerate(semantic_hits, start=1):
+            try:
+                ch = self.store.get_chunk(cid)
+            except KeyError:
+                continue
+            out.append(
+                self._result(
+                    ch,
+                    score=score,
+                    retrieval_mode="semantic",
+                    semantic_rank=semantic_rank,
+                    semantic_score=score,
+                )
+            )
+        return out

indexkit/index.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+from turbovec import IdMapIndex
+
+
+class VecIndex:
+    """Thin wrapper around turbovec's ``IdMapIndex``.
+
+    Isolates all turbovec-specific quirks (2-D batch queries, numpy
+    dtypes, return shapes) so the rest of ``indexkit`` works with plain
+    Python ``int``/``float`` ids and scores.
+    """
+
+    def __init__(self, dim: int, path) -> None:
+        self.dim = dim
+        self.path = Path(path)
+        self._idx = IdMapIndex(dim=dim, bit_width=4)
+
+    @classmethod
+    def load(cls, dim: int, path) -> "VecIndex":
+        obj = cls.__new__(cls)
+        obj.dim = dim
+        obj.path = Path(path)
+        obj._idx = IdMapIndex.load(str(Path(path)))
+        return obj
+
+    def add(self, ids: list[int], vectors: list[list[float]]) -> None:
+        v = np.asarray(vectors, dtype=np.float32)
+        i = np.asarray(ids, dtype=np.uint64)
+        self._idx.add_with_ids(v, i)
+
+    def remove(self, chunk_id: int) -> None:
+        self._idx.remove(int(chunk_id))
+
+    def search(
+        self,
+        vector: list[float],
+        k: int,
+        allowlist: list[int] | None = None,
+    ) -> list[tuple[int, float]]:
+        # turbovec expects a 2-D (nq, dim) query and returns
+        # (scores, ids) as (nq, effective_k) arrays. We submit a single
+        # row and read row 0 back.
+        q = np.asarray([vector], dtype=np.float32)
+        if allowlist is not None:
+            if not allowlist:
+                return []
+            allow = np.asarray(allowlist, dtype=np.uint64)
+            scores, ids = self._idx.search(q, k=k, allowlist=allow)
+        else:
+            scores, ids = self._idx.search(q, k=k)
+        return [(int(i), float(s)) for i, s in zip(ids[0], scores[0])]
+
+    def save(self) -> None:
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        self._idx.write(str(self.path))

indexkit/loaders/__init__.py

@@ -0,0 +1,3 @@
+from .markdown import Chunk, load_markdown, iter_corpus
+
+__all__ = ["Chunk", "load_markdown", "iter_corpus"]