kiwimatecoder 0.1.0__py3-none-any.whl

kiwimatecoder/index/search.py

@@ -0,0 +1,247 @@
+"""BM25 ranking over the local code index, with optional vector blending.
+
+The lexical ranker is a pure function over ``{path: {term: count}}`` maps so it
+can be unit-tested on a tiny corpus. Optional embeddings only reweight files
+that already carry vectors; if the embedding call fails for any reason the
+lexical scores are returned unchanged.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from kiwimatecoder import config
+from kiwimatecoder.index import embeddings
+from kiwimatecoder.index.builder import Embedder, embeddings_configured, tokenize
+from kiwimatecoder.index.store import load_store
+
+BM25_K1 = 1.5
+BM25_B = 0.75
+SNIPPET_LINES = 6
+VECTOR_WEIGHT = 0.5
+DEFAULT_LIMIT = 10
+MAX_LIMIT = 50
+
+
+@dataclass(frozen=True)
+class IndexHit:
+    """One ranked file with the best matching snippet."""
+
+    path: str
+    score: float
+    snippet: str
+    start_line: int
+
+
+def bm25_scores(
+    query_terms: Sequence[str],
+    docs: Mapping[str, Mapping[str, int]],
+    df: Mapping[str, int],
+    *,
+    k1: float = BM25_K1,
+    b: float = BM25_B,
+) -> dict[str, float]:
+    """Score ``docs`` against ``query_terms`` with BM25.
+
+    Only documents sharing at least one query term are returned. ``df`` is the
+    document frequency of each term across the whole index.
+    """
+    total = len(docs)
+    if total == 0 or not query_terms:
+        return {}
+    lengths = {path: sum(counts.values()) for path, counts in docs.items()}
+    avgdl = sum(lengths.values()) / total
+    if avgdl <= 0:
+        return {}
+    scores: dict[str, float] = {}
+    for path, counts in docs.items():
+        score = 0.0
+        length = lengths[path]
+        for term in query_terms:
+            frequency = counts.get(term, 0)
+            if frequency <= 0:
+                continue
+            doc_freq = max(0, min(int(df.get(term, 0)), total))
+            idf = math.log(1.0 + (total - doc_freq + 0.5) / (doc_freq + 0.5))
+            denominator = frequency + k1 * (1.0 - b + b * length / avgdl)
+            score += idf * (frequency * (k1 + 1.0)) / denominator
+        if score > 0.0:
+            scores[path] = score
+    return scores
+
+
+def best_snippet(
+    lines: Sequence[str],
+    query_terms: Sequence[str],
+    *,
+    window: int = SNIPPET_LINES,
+) -> tuple[str, int]:
+    """Return the ``window``-line snippet with the most query terms.
+
+    The returned line number is 1-based. When nothing matches, the first
+    non-blank line is used.
+    """
+
+    def first_content_line() -> tuple[str, int]:
+        for index, line in enumerate(lines):
+            if line.strip():
+                return line.strip(), index + 1
+        return (lines[0].strip() if lines else "", 1)
+
+    if not lines:
+        return "", 1
+    terms = set(query_terms)
+    if not terms:
+        return first_content_line()
+    line_scores = [len(set(tokenize(line)) & terms) for line in lines]
+    best_start = 0
+    best_total = -1
+    for start in range(len(lines)):
+        total = sum(line_scores[start : start + max(1, window)])
+        if total > best_total:
+            best_total = total
+            best_start = start
+    if best_total <= 0:
+        return first_content_line()
+    snippet = "\n".join(lines[best_start : best_start + max(1, window)]).strip()
+    return snippet, best_start + 1
+
+
+def _read_lines(path: Path) -> list[str]:
+    try:
+        return path.read_text(encoding="utf-8", errors="replace").splitlines()
+    except OSError:
+        return []
+
+
+def _vector_similarity(entry: dict[str, Any], query_vector: list[float]) -> float | None:
+    vectors = entry.get("vectors")
+    chunks = entry.get("chunks")
+    if (
+        not isinstance(vectors, list)
+        or not isinstance(chunks, list)
+        or len(vectors) != len(chunks)
+        or not vectors
+    ):
+        return None
+    best = 0.0
+    for vector in vectors:
+        if isinstance(vector, list):
+            best = max(best, max(0.0, embeddings.cosine(query_vector, vector)))
+    return best
+
+
+def entry_has_vectors(entry: dict[str, Any]) -> bool:
+    """Whether an entry carries one vector per stored chunk."""
+    vectors = entry.get("vectors")
+    chunks = entry.get("chunks")
+    return (
+        isinstance(vectors, list)
+        and isinstance(chunks, list)
+        and len(vectors) == len(chunks)
+        and len(vectors) > 0
+    )
+
+
+def search_index(
+    query: str,
+    *,
+    limit: int = DEFAULT_LIMIT,
+    workspace_root: Path,
+    embedder: Embedder | None = None,
+) -> list[IndexHit]:
+    """Rank indexed files against ``query`` and return up to ``limit`` hits.
+
+    Lexical BM25 drives the ranking. When embeddings are configured (or an
+    ``embedder`` is supplied) and the query embeds successfully, each file's
+    normalized lexical score is blended 50/50 with its best chunk similarity;
+    files without vectors keep their lexical score. Any embedding failure is
+    ignored.
+    """
+    settings = config.get_index()
+    if not settings["enabled"]:
+        return []
+    try:
+        limit = max(1, min(int(limit), MAX_LIMIT))
+    except (TypeError, ValueError):
+        limit = DEFAULT_LIMIT
+    query_terms = tokenize(query)
+    if not query_terms:
+        return []
+
+    root = Path(workspace_root).resolve()
+    store = load_store(root)
+    files: dict[str, dict[str, Any]] = store["files"]
+    if not files:
+        return []
+    docs = {
+        rel: entry.get("tokens", {})
+        for rel, entry in files.items()
+        if isinstance(entry.get("tokens"), dict)
+    }
+    scores = bm25_scores(query_terms, docs, store["df"])
+    if not scores:
+        return []
+
+    embed = embedder
+    if embed is None and embeddings_configured(settings):
+        vector_targets = [rel for rel in scores if entry_has_vectors(files[rel])]
+        if vector_targets:
+            embed = _query_embedder(settings)
+    if embed is not None:
+        try:
+            vectors = embed([query])
+        except Exception:
+            vectors = []
+        if vectors:
+            scores = _blend(scores, files, vectors[0])
+
+    ranked = sorted(scores.items(), key=lambda item: (-item[1], item[0]))[:limit]
+    hits: list[IndexHit] = []
+    for rel, score in ranked:
+        lines = _read_lines(root / rel)
+        snippet, start = best_snippet(lines, query_terms)
+        hits.append(
+            IndexHit(
+                path=rel, score=round(score, 6), snippet=snippet, start_line=start
+            )
+        )
+    return hits
+
+
+def _query_embedder(settings: dict[str, Any]) -> Embedder:
+    embedding_settings = settings["embeddings"]
+    provider = str(embedding_settings["provider"])
+    model = str(embedding_settings["model"])
+    batch_size = int(embedding_settings["batch_size"])
+
+    def embedder(texts: list[str]) -> list[list[float]]:
+        return embeddings.embed_texts(
+            texts, provider, model, batch_size=batch_size
+        )
+
+    return embedder
+
+
+def _blend(
+    scores: dict[str, float],
+    files: dict[str, dict[str, Any]],
+    query_vector: list[float],
+) -> dict[str, float]:
+    """Blend lexical scores with vector similarity for vector-bearing files."""
+    largest = max(scores.values())
+    if largest <= 0:
+        return scores
+    blended: dict[str, float] = {}
+    for rel, score in scores.items():
+        lexical = score / largest
+        similarity = _vector_similarity(files[rel], query_vector)
+        if similarity is None:
+            blended[rel] = lexical
+        else:
+            blended[rel] = (1.0 - VECTOR_WEIGHT) * lexical + VECTOR_WEIGHT * similarity
+    return blended

kiwimatecoder/index/store.py

@@ -0,0 +1,251 @@
+"""Versioned JSON persistence for the codebase index.
+
+Each workspace gets its own store at
+``~/.kiwimatecoder/index/<sha1(root)[:16]>.json``. The store is written
+atomically and every read tolerates a missing, truncated, or otherwise corrupt
+file by returning an empty store, since the index is a cache that can always be
+rebuilt.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from kiwimatecoder import config
+from kiwimatecoder.tools.paths import atomic_write_text
+
+INDEX_VERSION = 1
+INDEX_DIR_NAME = "index"
+
+
+def store_path(workspace_root: Path) -> Path:
+    """Return the store path for ``workspace_root`` (directory may not exist)."""
+    root = Path(workspace_root).resolve()
+    digest = hashlib.sha1(str(root).encode("utf-8")).hexdigest()[:16]
+    return config.ensure_config_dir() / INDEX_DIR_NAME / f"{digest}.json"
+
+
+def empty_store(workspace_root: Path) -> dict[str, Any]:
+    """Return a fresh, empty store for ``workspace_root``."""
+    root = Path(workspace_root).resolve()
+    return {
+        "version": INDEX_VERSION,
+        "root": str(root),
+        "files": {},
+        "df": {},
+        "built_at": None,
+    }
+
+
+def compute_df(files: dict[str, dict[str, Any]]) -> dict[str, int]:
+    """Compute the document-frequency map for a ``rel -> entry`` mapping."""
+    df: dict[str, int] = {}
+    for entry in files.values():
+        tokens = entry.get("tokens")
+        if not isinstance(tokens, dict):
+            continue
+        for term in tokens:
+            df[str(term)] = df.get(str(term), 0) + 1
+    return df
+
+
+def _clean_tokens(raw: object) -> dict[str, int]:
+    if not isinstance(raw, dict):
+        return {}
+    tokens: dict[str, int] = {}
+    for term, count in raw.items():
+        try:
+            value = int(count)
+        except (TypeError, ValueError):
+            continue
+        if value > 0:
+            tokens[str(term)] = value
+    return tokens
+
+
+def _clean_chunks(raw: object) -> list[dict[str, Any]] | None:
+    if raw is None:
+        return None
+    if not isinstance(raw, list):
+        return None
+    chunks: list[dict[str, Any]] = []
+    for item in raw:
+        if not isinstance(item, dict) or not isinstance(item.get("text"), str):
+            return None
+        try:
+            start = int(item.get("start") or 1)
+        except (TypeError, ValueError):
+            start = 1
+        chunks.append({"start": max(1, start), "text": item["text"]})
+    return chunks
+
+
+def _clean_vectors(raw: object, expected: int) -> list[list[float]] | None:
+    if raw is None or expected <= 0:
+        return None
+    if not isinstance(raw, list) or len(raw) != expected:
+        return None
+    vectors: list[list[float]] = []
+    for vector in raw:
+        if not isinstance(vector, list) or not vector:
+            return None
+        try:
+            vectors.append([float(value) for value in vector])
+        except (TypeError, ValueError):
+            return None
+    return vectors
+
+
+def _clean_entry(raw: object) -> dict[str, Any] | None:
+    if not isinstance(raw, dict):
+        return None
+    tokens = raw.get("tokens")
+    if not isinstance(tokens, dict):
+        return None
+    try:
+        mtime = float(raw.get("mtime") or 0.0)
+        size = int(raw.get("size") or 0)
+    except (TypeError, ValueError):
+        return None
+    entry: dict[str, Any] = {
+        "mtime": mtime,
+        "size": size,
+        "tokens": _clean_tokens(tokens),
+    }
+    chunks = _clean_chunks(raw.get("chunks"))
+    if chunks is not None:
+        entry["chunks"] = chunks
+        vectors = _clean_vectors(raw.get("vectors"), len(chunks))
+        if vectors is not None:
+            entry["vectors"] = vectors
+    return entry
+
+
+def load_store(workspace_root: Path) -> dict[str, Any]:
+    """Load the store for ``workspace_root``, tolerating any corruption.
+
+    A missing file, invalid JSON, version mismatch, or wrong workspace root all
+    yield an :func:`empty_store` result so callers can simply rebuild.
+    """
+    root = Path(workspace_root).resolve()
+    path = store_path(root)
+    try:
+        stored = json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, json.JSONDecodeError, ValueError):
+        return empty_store(root)
+    if not isinstance(stored, dict) or stored.get("version") != INDEX_VERSION:
+        return empty_store(root)
+    if str(stored.get("root") or "") != str(root):
+        return empty_store(root)
+    raw_files = stored.get("files")
+    if not isinstance(raw_files, dict):
+        return empty_store(root)
+
+    files: dict[str, dict[str, Any]] = {}
+    for rel, entry in raw_files.items():
+        if not isinstance(rel, str):
+            continue
+        cleaned = _clean_entry(entry)
+        if cleaned is not None:
+            files[rel] = cleaned
+
+    df: dict[str, int] = {}
+    raw_df = stored.get("df")
+    if isinstance(raw_df, dict):
+        for term, count in raw_df.items():
+            try:
+                value = int(count)
+            except (TypeError, ValueError):
+                continue
+            if value > 0:
+                df[str(term)] = value
+    if not df and files:
+        df = compute_df(files)
+
+    built_at = stored.get("built_at")
+    if not isinstance(built_at, (int, float)):
+        built_at = None
+    return {
+        "version": INDEX_VERSION,
+        "root": str(root),
+        "files": files,
+        "df": df,
+        "built_at": built_at,
+    }
+
+
+def save_store(workspace_root: Path, store: dict[str, Any]) -> None:
+    """Write ``store`` atomically, stamping version, root, and build time."""
+    root = Path(workspace_root).resolve()
+    path = store_path(root)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    store["version"] = INDEX_VERSION
+    store["root"] = str(root)
+    store["built_at"] = time.time()
+    atomic_write_text(path, json.dumps(store, separators=(",", ":")))
+
+
+def clear_store(workspace_root: Path) -> bool:
+    """Delete the store for ``workspace_root``; returns whether it existed."""
+    path = store_path(workspace_root)
+    try:
+        path.unlink()
+    except FileNotFoundError:
+        return False
+    except OSError:
+        return False
+    return True
+
+
+@dataclass(frozen=True)
+class IndexStatus:
+    """Human-facing snapshot of the index for one workspace."""
+
+    enabled: bool
+    path: str
+    files: int
+    terms: int
+    size_bytes: int
+    stale: int
+    embeddings: bool
+    built_at: float | None
+
+
+def index_status(
+    workspace_root: Path, *, cfg: dict[str, Any] | None = None
+) -> IndexStatus:
+    """Summarize the store, counting indexed files whose stat changed."""
+    settings = config.get_index(cfg)
+    root = Path(workspace_root).resolve()
+    store = load_store(root)
+    files = store["files"]
+    stale = 0
+    for rel, entry in files.items():
+        try:
+            stat = (root / rel).stat()
+        except OSError:
+            stale += 1
+            continue
+        if stat.st_mtime != entry.get("mtime") or stat.st_size != entry.get("size"):
+            stale += 1
+    path = store_path(root)
+    try:
+        size_bytes = path.stat().st_size
+    except OSError:
+        size_bytes = 0
+    embeddings = settings["embeddings"]
+    return IndexStatus(
+        enabled=bool(settings["enabled"]),
+        path=str(path),
+        files=len(files),
+        terms=len(store["df"]),
+        size_bytes=size_bytes,
+        stale=stale,
+        embeddings=bool(embeddings["provider"] and embeddings["model"]),
+        built_at=store.get("built_at"),
+    )

kiwimatecoder/instructions.py

@@ -0,0 +1,83 @@
+"""Project instruction files (``AGENTS.md`` and friends).
+
+Instruction files are project-owned guidance that the agent should follow:
+build commands, style rules, testing conventions. They are loaded from the
+workspace root and injected into the system prompt with a size budget so a
+pathological file cannot crowd out the conversation.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+# Checked in order; the first matching file for each name wins.
+INSTRUCTION_FILES: tuple[str, ...] = (
+    "AGENTS.md",
+    "CLAUDE.md",
+    ".kiwimatecoder/AGENTS.md",
+    ".kiwimatecoder/instructions.md",
+)
+
+MAX_INSTRUCTION_FILE_BYTES = 32 * 1024
+MAX_INSTRUCTION_TOTAL_BYTES = 48 * 1024
+
+
+def discover_instruction_files(workspace_root: Path) -> list[Path]:
+    """Return existing instruction files under ``workspace_root``, in order."""
+    found: list[Path] = []
+    for relative in INSTRUCTION_FILES:
+        candidate = workspace_root / relative
+        if candidate.is_file():
+            found.append(candidate)
+    return found
+
+
+def load_instructions(
+    workspace_root: Path,
+    *,
+    file_budget: int = MAX_INSTRUCTION_FILE_BYTES,
+    total_budget: int = MAX_INSTRUCTION_TOTAL_BYTES,
+) -> list[tuple[str, str]]:
+    """Load instruction files as ``(relative_path, text)`` pairs.
+
+    Returns an empty list when there are none. Missing/unreadable/binary files
+    are skipped. Total and per-file byte budgets are enforced.
+    """
+    loaded: list[tuple[str, str]] = []
+    remaining = max(0, total_budget)
+    for path in discover_instruction_files(workspace_root):
+        if remaining <= 0:
+            break
+        limit = min(file_budget, remaining)
+        try:
+            raw = path.read_bytes()
+        except OSError:
+            continue
+        if b"\x00" in raw[:1024]:
+            continue
+        clipped = raw[:limit]
+        text = clipped.decode("utf-8", "replace").strip()
+        if not text:
+            continue
+        truncated = len(raw) > len(clipped)
+        if truncated:
+            text += f"\n\n[truncated: showing {len(clipped)} of {len(raw)} bytes]"
+        loaded.append((path.relative_to(workspace_root).as_posix(), text))
+        remaining -= len(clipped)
+    return loaded
+
+
+def instructions_section(workspace_root: Path) -> str:
+    """Render the project-instructions block for the system prompt."""
+    loaded = load_instructions(workspace_root)
+    if not loaded:
+        return ""
+    blocks = [
+        "Project instructions follow. These come from the repository and should "
+        "be followed for work in this workspace:"
+    ]
+    for relative, text in loaded:
+        blocks.append(
+            f'<project_instructions path="{relative}">\n{text}\n</project_instructions>'
+        )
+    return "\n\n" + "\n\n".join(blocks)