kc-cli 0.4.0__py3-none-any.whl

kc/search/fts.py

@@ -0,0 +1,284 @@
+"""Range retrieval over SQLite FTS5 plus semantic vectors."""
+
+from __future__ import annotations
+
+import sqlite3
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+from urllib.parse import quote
+
+import orjson
+
+from kc.errors import KcError
+from kc.models.source import SourceRecord
+from kc.models.source_range import SourceRangeRecord
+from kc.paths import ensure_data_dir_exists
+from kc.search.semantic import (
+    assert_semantic_index_ready,
+    build_semantic_index,
+    semantic_index_status,
+    semantic_rankings,
+)
+from kc.store.jsonl import read_jsonl
+from kc.store.sqlite import index_status, init_db, rebuild_index
+
+
+@dataclass(frozen=True)
+class Bm25Rank:
+    range_id: str
+    rank: int
+    score: float
+
+
+@dataclass(frozen=True)
+class CombinedRank:
+    range_id: str
+    bm25_rank: int | None = None
+    bm25_score: float | None = None
+    semantic_rank: int | None = None
+    semantic_score: float | None = None
+    rrf_score: float | None = None
+
+
+def ensure_index(db_path: Path, sources_path: Path, ranges_path: Path) -> None:
+    ensure_data_dir_exists()
+    sources = read_jsonl(sources_path, SourceRecord)
+    ranges = read_jsonl(ranges_path, SourceRangeRecord)
+    status = index_status(db_path, sources, ranges)
+    if not status["sqlite_exists"] or status["stale"]:
+        rebuild_index(db_path, sources, ranges)
+    semantic_status = semantic_index_status(db_path, ranges)
+    if (
+        not semantic_status["index_metadata"]
+        or not semantic_status["metadata_match"]
+        or semantic_status["missing_vectors"]
+        or semantic_status["stale_vectors"]
+    ):
+        try:
+            build_semantic_index(db_path, ranges)
+        except KcError as exc:
+            if exc.code != "KC_RETRIEVAL_MODEL_UNAVAILABLE":
+                raise
+
+
+def _build_fts_query(query: str) -> str:
+    terms = [t.replace('"', "").strip() for t in query.split() if t.strip()]
+    return " OR ".join(f'"{term}"' for term in terms)
+
+
+def citation_token(
+    source_id: str,
+    locator: dict[str, Any],
+    *,
+    range_id: str | None = None,
+    legacy: bool = False,
+) -> str:
+    prefix = f"[kc:{source_id}"
+    if range_id and not legacy:
+        prefix += f":{range_id}"
+    if locator.get("kind") == "line_range":
+        return f"{prefix}:L{locator.get('start_line')}-L{locator.get('end_line')}]"
+    if locator.get("kind") == "json_pointer":
+        pointer = quote(str(locator.get("pointer", "/")), safe="/~")
+        return f"{prefix}:JP:{pointer}]"
+    if locator.get("kind") == "csv_row_range":
+        return f"{prefix}:CSV:R{locator.get('start_row')}-R{locator.get('end_row')}]"
+    return f"{prefix}]"
+
+
+def rrf_score(ranks: list[int], *, k: int = 60) -> float:
+    return sum(1.0 / (k + rank) for rank in ranks)
+
+
+def _connect(db_path: Path) -> sqlite3.Connection:
+    conn = sqlite3.connect(str(db_path))
+    conn.row_factory = sqlite3.Row
+    return conn
+
+
+def _bm25_rankings(
+    conn: sqlite3.Connection,
+    query: str,
+    *,
+    domain: str | None = None,
+    limit: int = 100,
+) -> list[Bm25Rank]:
+    fts_query = _build_fts_query(query)
+    if not fts_query:
+        return []
+    sql = """
+        SELECT
+          f.range_id,
+          bm25(source_ranges_fts) AS bm25_score
+        FROM source_ranges_fts f
+        JOIN source_ranges r ON r.range_id = f.range_id
+        JOIN sources s ON s.source_id = f.source_id
+        WHERE source_ranges_fts MATCH ?
+    """
+    params: list[Any] = [fts_query]
+    if domain:
+        sql += " AND f.domain LIKE ?"
+        params.append(f"%{domain}%")
+    sql += " ORDER BY bm25_score, f.range_id LIMIT ?"
+    params.append(limit)
+    try:
+        rows = conn.execute(sql, params).fetchall()
+    except sqlite3.OperationalError as exc:
+        raise KcError(
+            code="KC_INDEX_BUILD_FAILED",
+            message=f"Search failed: {exc}",
+            details={"query": query},
+        ) from exc
+    return [
+        Bm25Rank(range_id=row["range_id"], rank=rank, score=float(row["bm25_score"]))
+        for rank, row in enumerate(rows, start=1)
+    ]
+
+
+def _rows_for_ranges(conn: sqlite3.Connection, range_ids: list[str]) -> dict[str, sqlite3.Row]:
+    if not range_ids:
+        return {}
+    placeholders = ", ".join("?" for _ in range_ids)
+    rows = conn.execute(
+        f"""
+        SELECT
+          r.range_id,
+          r.source_id,
+          r.locator_json,
+          r.source_fingerprint,
+          r.record_json AS range_record_json,
+          r.excerpt,
+          s.display_name,
+          s.status AS source_status,
+          s.authority_json,
+          s.fingerprint AS current_source_fingerprint
+        FROM source_ranges r
+        JOIN sources s ON s.source_id = r.source_id
+        WHERE r.range_id IN ({placeholders})
+        """,
+        range_ids,
+    ).fetchall()
+    return {row["range_id"]: row for row in rows}
+
+
+def _format_result(
+    row: sqlite3.Row,
+    rank: CombinedRank,
+    *,
+    hybrid_rank: int,
+) -> dict[str, Any]:
+    locator = orjson.loads(row["locator_json"])
+    authority = orjson.loads(row["authority_json"])
+    return {
+        "range_id": row["range_id"],
+        "source_id": row["source_id"],
+        "display_name": row["display_name"],
+        "locator": locator,
+        "excerpt": row["excerpt"],
+        "scores": {
+            "bm25_rank": rank.bm25_rank,
+            "bm25_score": rank.bm25_score,
+            "semantic_rank": rank.semantic_rank,
+            "semantic_score": rank.semantic_score,
+            "hybrid_rank": hybrid_rank,
+            "rrf_score": rank.rrf_score,
+        },
+        "citation_token": citation_token(row["source_id"], locator, range_id=row["range_id"]),
+        "legacy_citation_token": citation_token(row["source_id"], locator, legacy=True),
+        "source_authority": authority,
+        "source_status": row["source_status"],
+        "source_fingerprint": row["source_fingerprint"],
+        "current_source_fingerprint": row["current_source_fingerprint"],
+    }
+
+
+def _combine_hybrid(
+    bm25: list[Bm25Rank],
+    semantic: list[Any],
+    *,
+    rrf_k: int,
+    limit: int,
+) -> list[CombinedRank]:
+    by_range: dict[str, CombinedRank] = {}
+    for item in bm25:
+        by_range[item.range_id] = CombinedRank(
+            range_id=item.range_id,
+            bm25_rank=item.rank,
+            bm25_score=item.score,
+        )
+    for item in semantic:
+        existing = by_range.get(item.range_id)
+        by_range[item.range_id] = CombinedRank(
+            range_id=item.range_id,
+            bm25_rank=existing.bm25_rank if existing else None,
+            bm25_score=existing.bm25_score if existing else None,
+            semantic_rank=item.rank,
+            semantic_score=item.score,
+        )
+
+    merged = []
+    for item in by_range.values():
+        ranks = [rank for rank in [item.bm25_rank, item.semantic_rank] if rank is not None]
+        merged.append(
+            CombinedRank(
+                range_id=item.range_id,
+                bm25_rank=item.bm25_rank,
+                bm25_score=item.bm25_score,
+                semantic_rank=item.semantic_rank,
+                semantic_score=item.semantic_score,
+                rrf_score=rrf_score(ranks, k=rrf_k),
+            )
+        )
+    sentinel = 1_000_000
+    merged.sort(
+        key=lambda item: (
+            -(item.rrf_score or 0.0),
+            item.bm25_rank if item.bm25_rank is not None else sentinel,
+            item.semantic_rank if item.semantic_rank is not None else sentinel,
+            item.range_id,
+        )
+    )
+    return merged[:limit]
+
+
+def search_ranges(
+    db_path: Path,
+    query: str,
+    *,
+    domain: str | None = None,
+    limit: int = 10,
+    rrf_k: int = 60,
+    ranges: list[SourceRangeRecord] | None = None,
+    metadata: dict[str, Any] | None = None,
+) -> list[dict[str, Any]]:
+    init_db(db_path)
+    conn = _connect(db_path)
+    try:
+        candidate_limit = max(limit * 5, 100)
+        source_ranges = ranges if ranges is not None else []
+        bm25 = _bm25_rankings(conn, query, domain=domain, limit=candidate_limit)
+        mode = "hybrid"
+        semantic_unavailable: str | None = None
+        try:
+            model = assert_semantic_index_ready(db_path, source_ranges)
+            semantic = semantic_rankings(conn, query, model, domain=domain, limit=candidate_limit)
+        except KcError as exc:
+            if exc.code != "KC_RETRIEVAL_MODEL_UNAVAILABLE":
+                raise
+            semantic = []
+            mode = "fts_fallback"
+            semantic_unavailable = exc.message
+        combined = _combine_hybrid(bm25, semantic, rrf_k=rrf_k, limit=limit)
+        if metadata is not None:
+            metadata["mode"] = mode
+            metadata["semantic_unavailable_reason"] = semantic_unavailable
+        rows = _rows_for_ranges(conn, [item.range_id for item in combined])
+        results = []
+        for hybrid_rank, item in enumerate(combined, start=1):
+            row = rows.get(item.range_id)
+            if row is not None:
+                results.append(_format_result(row, item, hybrid_rank=hybrid_rank))
+        return results
+    finally:
+        conn.close()

kc/search/semantic.py

@@ -0,0 +1,346 @@
+"""Local semantic retrieval support for source ranges."""
+
+from __future__ import annotations
+
+import hashlib
+import sqlite3
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from functools import lru_cache
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+import orjson
+
+from kc.errors import KcError
+from kc.models.source_range import SourceRangeRecord
+from kc.store.sqlite import init_db
+
+MODEL_PROVIDER = "model2vec"
+BUNDLED_MODEL_NAME = "potion-base-8M"
+EXPECTED_DIMENSION = 256
+EXPECTED_CHECKSUM = "sha256:1630533e875a4725d11a703f48cb9598e59f3cfe3fa0d5a19a54045d52922848"
+SEMANTIC_METADATA_KEY = "semantic_model"
+TEXT_MODEL_SUFFIXES = {".json", ".md"}
+
+
+@dataclass(frozen=True)
+class SemanticRank:
+    range_id: str
+    rank: int
+    score: float
+
+
+def bundled_model_dir() -> Path:
+    return Path(__file__).resolve().parent.parent / "embedding_models" / BUNDLED_MODEL_NAME
+
+
+def model_directory_checksum(model_dir: Path) -> str:
+    root = Path(model_dir)
+    digest = hashlib.sha256()
+    paths = (path for path in root.rglob("*") if path.is_file())
+    for path in sorted(paths, key=lambda item: item.relative_to(root).as_posix()):
+        rel = path.relative_to(root).as_posix().encode("utf-8")
+        digest.update(rel)
+        digest.update(b"\0")
+        digest.update(_checksum_file_bytes(path))
+        digest.update(b"\0")
+    return f"sha256:{digest.hexdigest()}"
+
+
+def _checksum_file_bytes(path: Path) -> bytes:
+    data = path.read_bytes()
+    if path.suffix.lower() in TEXT_MODEL_SUFFIXES:
+        return data.replace(b"\r\n", b"\n").replace(b"\r", b"\n")
+    return data
+
+
+def _model_unavailable(message: str, **details: Any) -> KcError:
+    return KcError(
+        code="KC_RETRIEVAL_MODEL_UNAVAILABLE",
+        message=message,
+        details=details,
+        suggested_action="run kc index build after fixing the local model",
+    )
+
+
+def semantic_model_metadata(model: Any | None = None) -> dict[str, Any]:
+    dimension = int(getattr(model, "dim", EXPECTED_DIMENSION) or EXPECTED_DIMENSION)
+    return {
+        "provider": MODEL_PROVIDER,
+        "model": BUNDLED_MODEL_NAME,
+        "dimension": dimension,
+        "checksum": model_directory_checksum(bundled_model_dir()),
+        "purpose": "ranking_only",
+    }
+
+
+@lru_cache(maxsize=1)
+def load_semantic_model() -> Any:
+    model_dir = bundled_model_dir()
+    if not model_dir.exists():
+        raise _model_unavailable("Bundled semantic model directory is missing.", model_dir=str(model_dir))
+    checksum = model_directory_checksum(model_dir)
+    if checksum != EXPECTED_CHECKSUM:
+        raise _model_unavailable(
+            "Bundled semantic model checksum does not match the configured checksum.",
+            model_dir=str(model_dir),
+            expected_checksum=EXPECTED_CHECKSUM,
+            actual_checksum=checksum,
+        )
+    try:
+        from model2vec import StaticModel
+    except ImportError as exc:
+        raise _model_unavailable("model2vec is not installed.", dependency="model2vec") from exc
+    try:
+        model = StaticModel.from_pretrained(str(model_dir))
+    except Exception as exc:
+        raise _model_unavailable("Failed to load bundled semantic model.", model_dir=str(model_dir)) from exc
+    dimension = int(getattr(model, "dim", 0) or 0)
+    if dimension and dimension != EXPECTED_DIMENSION:
+        raise _model_unavailable(
+            "Bundled semantic model dimension does not match the configured dimension.",
+            expected_dimension=EXPECTED_DIMENSION,
+            actual_dimension=dimension,
+        )
+    return model
+
+
+def is_model_available() -> tuple[bool, dict[str, Any] | None, str | None]:
+    try:
+        model = load_semantic_model()
+        return True, semantic_model_metadata(model), None
+    except KcError as exc:
+        return False, None, exc.message
+
+
+def _encode(model: Any, texts: str | list[str]) -> npt.NDArray[np.float32]:
+    try:
+        encoded = model.encode(texts, show_progress_bar=False, use_multiprocessing=False)
+    except TypeError:
+        encoded = model.encode(texts)
+    return np.asarray(encoded, dtype=np.float32)
+
+
+def embed_text(model: Any, text: str) -> npt.NDArray[np.float32]:
+    return _encode(model, text).reshape(-1).astype(np.float32)
+
+
+def embed_texts(model: Any, texts: list[str]) -> npt.NDArray[np.float32]:
+    if not texts:
+        return np.empty((0, EXPECTED_DIMENSION), dtype=np.float32)
+    encoded = _encode(model, texts)
+    if encoded.ndim == 1:
+        encoded = encoded.reshape(1, -1)
+    return encoded.astype(np.float32)
+
+
+def embedding_to_blob(embedding: npt.NDArray[np.float32]) -> bytes:
+    return np.asarray(embedding, dtype=np.float32).reshape(-1).tobytes()
+
+
+def blob_to_embedding(blob: bytes) -> npt.NDArray[np.float32]:
+    return np.frombuffer(blob, dtype=np.float32)
+
+
+def cosine_similarity(a: npt.NDArray[np.float32], b: npt.NDArray[np.float32]) -> float:
+    left = np.asarray(a, dtype=np.float32).reshape(-1)
+    right = np.asarray(b, dtype=np.float32).reshape(-1)
+    norm = np.linalg.norm(left) * np.linalg.norm(right)
+    if norm == 0:
+        return 0.0
+    return float(np.dot(left, right) / norm)
+
+
+def _now() -> str:
+    return datetime.now(UTC).isoformat()
+
+
+def _connect(db_path: Path) -> sqlite3.Connection:
+    conn = sqlite3.connect(str(db_path))
+    conn.row_factory = sqlite3.Row
+    return conn
+
+
+def _table_exists(conn: sqlite3.Connection, table: str) -> bool:
+    row = conn.execute(
+        "SELECT name FROM sqlite_master WHERE type IN ('table', 'view') AND name = ?",
+        (table,),
+    ).fetchone()
+    return row is not None
+
+
+def _metadata_from_db(conn: sqlite3.Connection) -> dict[str, Any] | None:
+    row = conn.execute("SELECT value_json FROM index_metadata WHERE key = ?", (SEMANTIC_METADATA_KEY,)).fetchone()
+    if row is None:
+        return None
+    return orjson.loads(row["value_json"])
+
+
+def build_semantic_index(db_path: Path, ranges: list[SourceRangeRecord]) -> dict[str, Any]:
+    init_db(db_path)
+    model = load_semantic_model()
+    metadata = semantic_model_metadata(model)
+    vectors = embed_texts(model, [r.excerpt for r in ranges])
+    if len(vectors) != len(ranges):
+        raise _model_unavailable(
+            "Semantic model returned an unexpected number of embeddings.",
+            expected=len(ranges),
+            actual=len(vectors),
+        )
+    dimension = int(vectors.shape[1]) if vectors.ndim == 2 and len(ranges) > 0 else int(metadata["dimension"])
+    if dimension != int(metadata["dimension"]):
+        raise _model_unavailable(
+            "Semantic model returned embeddings with an unexpected dimension.",
+            expected_dimension=metadata["dimension"],
+            actual_dimension=dimension,
+        )
+
+    conn = _connect(db_path)
+    try:
+        timestamp = _now()
+        conn.execute("DELETE FROM source_range_embeddings")
+        conn.executemany(
+            """
+            INSERT INTO source_range_embeddings(
+              range_id, source_id, source_fingerprint, text_hash, model_name,
+              model_checksum, dimension, embedding, updated_at
+            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+            """,
+            [
+                (
+                    source_range.range_id,
+                    source_range.source_id,
+                    source_range.source_fingerprint,
+                    source_range.text_hash,
+                    metadata["model"],
+                    metadata["checksum"],
+                    metadata["dimension"],
+                    embedding_to_blob(vectors[index]),
+                    timestamp,
+                )
+                for index, source_range in enumerate(ranges)
+            ],
+        )
+        stored_metadata = {**metadata, "built_at": timestamp, "ranges": len(ranges)}
+        conn.execute(
+            "INSERT OR REPLACE INTO index_metadata(key, value_json, updated_at) VALUES (?, ?, ?)",
+            (SEMANTIC_METADATA_KEY, orjson.dumps(stored_metadata).decode(), timestamp),
+        )
+        conn.commit()
+        return {"enabled": True, "model": stored_metadata, "embeddings": len(ranges)}
+    finally:
+        conn.close()
+
+
+def semantic_index_status(db_path: Path, ranges: list[SourceRangeRecord] | None = None) -> dict[str, Any]:
+    available, model_metadata, unavailable_reason = is_model_available()
+    status: dict[str, Any] = {
+        "model_available": available,
+        "model": model_metadata,
+        "unavailable_reason": unavailable_reason,
+        "index_metadata": None,
+        "metadata_match": False,
+        "vector_count": 0,
+        "missing_vectors": None,
+        "stale_vectors": None,
+    }
+    if not db_path.exists():
+        return status
+    conn = _connect(db_path)
+    try:
+        if not _table_exists(conn, "source_range_embeddings"):
+            return status
+        status["vector_count"] = int(conn.execute("SELECT COUNT(*) FROM source_range_embeddings").fetchone()[0])
+        index_metadata = _metadata_from_db(conn)
+        status["index_metadata"] = index_metadata
+        if model_metadata and index_metadata:
+            status["metadata_match"] = all(
+                index_metadata.get(key) == model_metadata.get(key)
+                for key in ("provider", "model", "dimension", "checksum", "purpose")
+            )
+        if ranges is not None:
+            rows = {
+                row["range_id"]: row
+                for row in conn.execute(
+                    """
+                    SELECT range_id, source_fingerprint, text_hash, model_checksum, dimension
+                    FROM source_range_embeddings
+                    """
+                ).fetchall()
+            }
+            missing = 0
+            stale = 0
+            checksum = model_metadata.get("checksum") if model_metadata else None
+            dimension_value = model_metadata.get("dimension") if model_metadata else None
+            dimension = int(dimension_value) if dimension_value is not None else None
+            for source_range in ranges:
+                row = rows.get(source_range.range_id)
+                if row is None:
+                    missing += 1
+                    continue
+                if (
+                    row["source_fingerprint"] != source_range.source_fingerprint
+                    or row["text_hash"] != source_range.text_hash
+                    or row["model_checksum"] != checksum
+                    or dimension is None
+                    or int(row["dimension"]) != dimension
+                ):
+                    stale += 1
+            status["missing_vectors"] = missing
+            status["stale_vectors"] = stale
+        return status
+    finally:
+        conn.close()
+
+
+def assert_semantic_index_ready(
+    db_path: Path,
+    ranges: list[SourceRangeRecord],
+) -> Any:
+    model = load_semantic_model()
+    status = semantic_index_status(db_path, ranges)
+    if not status["index_metadata"] or not status["metadata_match"]:
+        raise _model_unavailable(
+            "Semantic index metadata is missing or stale. Run kc index build.",
+            status=status,
+        )
+    if status["missing_vectors"] or status["stale_vectors"]:
+        raise _model_unavailable(
+            "Semantic index vectors are missing or stale. Run kc index build.",
+            missing_vectors=status["missing_vectors"],
+            stale_vectors=status["stale_vectors"],
+        )
+    return model
+
+
+def semantic_rankings(
+    conn: sqlite3.Connection,
+    query: str,
+    model: Any,
+    *,
+    domain: str | None = None,
+    limit: int = 100,
+) -> list[SemanticRank]:
+    query_embedding = embed_text(model, query)
+    sql = """
+        SELECT e.range_id, e.embedding
+        FROM source_range_embeddings e
+        JOIN sources s ON s.source_id = e.source_id
+    """
+    params: list[Any] = []
+    if domain:
+        sql += " WHERE s.domain_json LIKE ?"
+        params.append(f"%{domain}%")
+    rows = conn.execute(sql, params).fetchall()
+    scored: list[tuple[str, float]] = []
+    for row in rows:
+        score = cosine_similarity(query_embedding, blob_to_embedding(row["embedding"]))
+        scored.append((row["range_id"], score))
+    scored.sort(key=lambda item: (-item[1], item[0]))
+    return [
+        SemanticRank(range_id=range_id, rank=rank, score=score)
+        for rank, (range_id, score) in enumerate(scored[:limit], start=1)
+    ]

kc/store/__init__.py

@@ -0,0 +1 @@
+"""Persistence helpers for kc."""

kc/store/jsonl.py

@@ -0,0 +1,55 @@
+"""JSONL canonical store helpers."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+from pathlib import Path
+
+import orjson
+from pydantic import BaseModel
+
+from kc.atomic_write import atomic_write_bytes
+from kc.errors import KcError
+
+
+def read_jsonl[T: BaseModel](path: Path, model: type[T]) -> list[T]:
+    if not path.exists():
+        return []
+    records: list[T] = []
+    for line_no, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
+        if not line.strip():
+            continue
+        try:
+            records.append(model.model_validate(orjson.loads(line)))
+        except Exception as exc:
+            raise KcError(
+                code="KC_CONFIG_INVALID",
+                message=f"Invalid JSONL record in {path} at line {line_no}: {exc}",
+                details={"path": str(path), "line": line_no},
+            ) from exc
+    return records
+
+
+def write_jsonl(path: Path, records: Sequence[BaseModel]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    chunks = [
+        orjson.dumps(record.model_dump(mode="json"), option=orjson.OPT_APPEND_NEWLINE)
+        for record in records
+    ]
+    atomic_write_bytes(path, b"".join(chunks))
+
+
+def append_jsonl(path: Path, records: Iterable[BaseModel]) -> None:
+    existing: list[BaseModel] = []
+    if path.exists():
+        existing_lines = path.read_bytes().splitlines()
+        existing_data = [orjson.loads(line) for line in existing_lines if line.strip()]
+        existing = [_RawModel(value=item) for item in existing_data]
+    write_jsonl(path, [*existing, *records])
+
+
+class _RawModel(BaseModel):
+    value: dict
+
+    def model_dump(self, *args, **kwargs):  # type: ignore[no-untyped-def]
+        return self.value