scripts/retrieval.py

"""Shared retrieval for the CLI and the Space, so the two cannot drift apart.

Beyond loading the index, this does one thing worth explaining: it down-weights
release notes.

Zephyr's release notes are a quarter of the corpus — 1,688 chunks of the 6,949 —
and they are written in the same vocabulary as the documentation they describe.
A question like "what does west build actually do?" retrieved
`develop/west/release-notes - v0.7.0` above `develop/west/build-flash-debug`,
because a changelog entry about west build looks a lot like documentation about
west build to an embedding model.

For a coding assistant that is the wrong answer nearly every time: a developer
asking how something works wants the reference page, not what changed in v0.7.0.
So release notes are penalised, unless the question is actually about a version
or a change — in which case they are exactly right and the penalty is lifted.
"""

from __future__ import annotations

import json
import re
from functools import lru_cache
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
DEFAULT_INDEX = ROOT / "data" / "index"

# Historical rather than instructional. Not excluded — "what changed in 3.5?" is
# a fair question — just made to earn its place.
HISTORICAL = re.compile(r"^releases/|/release-notes", re.IGNORECASE)
HISTORICAL_PENALTY = 0.08

# When the question is about a version or a change, the penalty is lifted.
VERSION_QUERY = re.compile(
    r"\b(v?\d+\.\d+|release|changelog|changed|new in|migration|deprecat\w*|"
    r"breaking|since|upgrade|version)\b",
    re.IGNORECASE,
)


@lru_cache(maxsize=4)
def load(index_dir: str):
    import faiss
    from sentence_transformers import SentenceTransformer

    path = Path(index_dir)
    if not (path / "docs.faiss").exists():
        raise FileNotFoundError(
            f"no index at {path} - run scripts/fetch_docs.py then scripts/build_index.py"
        )
    meta = json.loads((path / "meta.json").read_text(encoding="utf-8"))
    chunks = [
        json.loads(line)
        for line in (path / "chunks.jsonl").read_text(encoding="utf-8").splitlines()
    ]
    embedder = SentenceTransformer(meta["embedding_model"])
    return faiss.read_index(str(path / "docs.faiss")), chunks, embedder, meta


def search(question: str, k: int = 6, index_dir: Path | str = DEFAULT_INDEX) -> list[dict]:
    index, chunks, embedder, _ = load(str(index_dir))

    vector = embedder.encode(
        [question], convert_to_numpy=True, normalize_embeddings=True
    ).astype("float32")

    # Over-fetch so re-ranking has something to promote from. Without this the
    # penalty could only demote within an already-decided top k.
    scores, ids = index.search(vector, min(k * 4, len(chunks)))

    wants_history = bool(VERSION_QUERY.search(question))
    hits: list[dict] = []
    for score, position in zip(scores[0], ids[0]):
        if position < 0:
            continue
        hit = dict(chunks[int(position)])
        raw = float(score)
        adjusted = raw
        if not wants_history and HISTORICAL.search(hit.get("source", "")):
            adjusted -= HISTORICAL_PENALTY
        hit["score"] = round(adjusted, 4)
        hit["raw_score"] = round(raw, 4)
        hits.append(hit)

    hits.sort(key=lambda item: item["score"], reverse=True)
    return hits[:k]


def meta(index_dir: Path | str = DEFAULT_INDEX) -> dict:
    return load(str(index_dir))[3]