src/agent_harness/retrieval.py

"""Deterministic exact, BM25/fuzzy, and dense retrieval for the E00 pilot."""

from __future__ import annotations

from array import array
from collections import Counter
from dataclasses import dataclass
from difflib import SequenceMatcher
from hashlib import sha256
import math
from pathlib import Path
import re
import sqlite3
import time
from typing import Iterable, Sequence

from .components import Candidate
from .lm_studio_embeddings import LMStudioEmbeddingClient
from .repository import SourceChunk
from .specs import EmbeddingSpec


TOKEN_PATTERN = re.compile(r"[A-Za-z][A-Za-z0-9_./-]*")
CAMEL_BOUNDARY = re.compile(r"(?<=[a-z0-9])(?=[A-Z])")
STOP_WORDS = {
    "a", "after", "against", "and", "are", "as", "at", "be", "been", "but",
    "by", "can", "correct", "does", "ensure", "for", "from", "has", "have",
    "in", "into", "is", "it", "its", "of", "on", "or", "that", "the", "their",
    "this", "to", "under", "when", "while", "with", "without",
}


def tokenize(value: str) -> tuple[str, ...]:
    tokens: list[str] = []
    for match in TOKEN_PATTERN.findall(value):
        for slash_part in re.split(r"[./-]+", match):
            for snake_part in slash_part.split("_"):
                for camel_part in CAMEL_BOUNDARY.split(snake_part):
                    token = camel_part.lower()
                    if len(token) >= 2:
                        tokens.append(token)
    return tuple(tokens)


def query_terms(query: str) -> tuple[str, ...]:
    return tuple(dict.fromkeys(token for token in tokenize(query) if token not in STOP_WORDS))


def _candidate(chunk: SourceChunk, source: str, score: float) -> Candidate:
    return Candidate(
        path=chunk.path,
        line_start=chunk.line_start,
        line_end=chunk.line_end,
        text=chunk.text,
        source=source,
        score=score,
        metadata={"chunk_id": chunk.chunk_id},
    )


class ExactRetriever:
    """Ranks raw chunks by literal occurrences of deterministic query terms."""

    def __init__(self, chunks: Sequence[SourceChunk]):
        self.chunks = tuple(chunks)

    def retrieve(self, query: str, limit: int) -> Sequence[Candidate]:
        terms = query_terms(query)
        ranked: list[tuple[float, SourceChunk]] = []
        for chunk in self.chunks:
            haystack = f"{chunk.path}\n{chunk.text}".lower()
            matched_terms = 0
            occurrences = 0
            path_matches = 0
            lower_path = chunk.path.lower()
            for term in terms:
                count = haystack.count(term)
                if count:
                    matched_terms += 1
                    occurrences += min(count, 20)
                if term in lower_path:
                    path_matches += 1
            if matched_terms:
                score = matched_terms * 10.0 + math.log1p(occurrences) + path_matches * 5.0
                ranked.append((score, chunk))
        ranked.sort(key=lambda item: (-item[0], item[1].path, item[1].line_start))
        return tuple(_candidate(chunk, "exact", score) for score, chunk in ranked[:limit])


class BM25FuzzyRetriever:
    """BM25 over code chunks with a bounded fuzzy path-name bonus."""

    def __init__(self, chunks: Sequence[SourceChunk], k1: float = 1.2, b: float = 0.75):
        self.chunks = tuple(chunks)
        self.k1 = k1
        self.b = b
        self.term_frequencies = tuple(Counter(tokenize(f"{item.path}\n{item.text}")) for item in chunks)
        self.lengths = tuple(sum(counter.values()) for counter in self.term_frequencies)
        self.average_length = sum(self.lengths) / max(len(self.lengths), 1)
        document_frequency: Counter[str] = Counter()
        for counter in self.term_frequencies:
            document_frequency.update(counter.keys())
        self.document_frequency = document_frequency

    def retrieve(self, query: str, limit: int) -> Sequence[Candidate]:
        terms = query_terms(query)
        document_count = len(self.chunks)
        ranked: list[tuple[float, SourceChunk]] = []
        for chunk, frequencies, length in zip(self.chunks, self.term_frequencies, self.lengths):
            score = 0.0
            for term in terms:
                frequency = frequencies.get(term, 0)
                if not frequency:
                    continue
                df = self.document_frequency[term]
                inverse_document_frequency = math.log(1.0 + (document_count - df + 0.5) / (df + 0.5))
                denominator = frequency + self.k1 * (
                    1.0 - self.b + self.b * length / max(self.average_length, 1.0)
                )
                score += inverse_document_frequency * frequency * (self.k1 + 1.0) / denominator
            path_tokens = tokenize(chunk.path)
            fuzzy = max(
                (
                    SequenceMatcher(None, query_token, path_token).ratio()
                    for query_token in terms
                    for path_token in path_tokens
                ),
                default=0.0,
            )
            if fuzzy >= 0.72:
                score += (fuzzy - 0.72) * 4.0
            if score > 0.0:
                ranked.append((score, chunk))
        ranked.sort(key=lambda item: (-item[0], item[1].path, item[1].line_start))
        return tuple(_candidate(chunk, "bm25_fuzzy", score) for score, chunk in ranked[:limit])


@dataclass(frozen=True, slots=True)
class DenseIndexStats:
    total_chunks: int
    cached_chunks: int
    embedded_chunks: int
    build_seconds: float


class SQLiteEmbeddingCache:
    def __init__(self, path: Path, spec: EmbeddingSpec):
        path.parent.mkdir(parents=True, exist_ok=True)
        self.spec = spec
        self.connection = sqlite3.connect(path)
        self.connection.execute(
            """
            CREATE TABLE IF NOT EXISTS embeddings (
                cache_key TEXT PRIMARY KEY,
                dimension INTEGER NOT NULL,
                vector BLOB NOT NULL
            )
            """
        )
        self.connection.commit()

    def get(self, cache_key: str) -> array | None:
        row = self.connection.execute(
            "SELECT dimension, vector FROM embeddings WHERE cache_key = ?",
            (cache_key,),
        ).fetchone()
        if row is None:
            return None
        dimension, payload = row
        if dimension != self.spec.vector_dimension:
            raise ValueError(f"Cached embedding dimension mismatch for {cache_key}")
        vector = array("f")
        vector.frombytes(payload)
        if len(vector) != self.spec.vector_dimension:
            raise ValueError(f"Cached embedding payload is malformed for {cache_key}")
        return vector

    def put_many(self, values: Sequence[tuple[str, array]]) -> None:
        self.connection.executemany(
            "INSERT OR REPLACE INTO embeddings(cache_key, dimension, vector) VALUES (?, ?, ?)",
            ((key, self.spec.vector_dimension, vector.tobytes()) for key, vector in values),
        )
        self.connection.commit()

    def close(self) -> None:
        self.connection.close()

    def __enter__(self) -> "SQLiteEmbeddingCache":
        return self

    def __exit__(self, exc_type: object, exc: object, traceback: object) -> None:
        self.close()


class DenseRetriever:
    def __init__(
        self,
        chunks: Sequence[SourceChunk],
        vectors: Sequence[array],
        spec: EmbeddingSpec,
        client: LMStudioEmbeddingClient,
        cache: SQLiteEmbeddingCache | None = None,
    ):
        if len(chunks) != len(vectors):
            raise ValueError("dense chunks and vectors must have equal length")
        self.chunks = tuple(chunks)
        self.vectors = tuple(vectors)
        self.spec = spec
        self.client = client
        self.cache = cache

    @classmethod
    def build(
        cls,
        chunks: Sequence[SourceChunk],
        spec: EmbeddingSpec,
        client: LMStudioEmbeddingClient,
        cache: SQLiteEmbeddingCache,
    ) -> tuple["DenseRetriever", DenseIndexStats]:
        started = time.monotonic()
        vectors: list[array | None] = []
        cache_keys: list[str] = []
        missing_indices: list[int] = []
        for index, chunk in enumerate(chunks):
            document = spec.document_prefix_template.format(path=chunk.path) + chunk.text
            cache_key = sha256(
                f"{spec.config_hash}\0{document}".encode("utf-8")
            ).hexdigest()
            cache_keys.append(cache_key)
            vector = cache.get(cache_key)
            vectors.append(vector)
            if vector is None:
                missing_indices.append(index)

        for offset in range(0, len(missing_indices), spec.batch_size):
            batch_indices = missing_indices[offset : offset + spec.batch_size]
            documents = [
                spec.document_prefix_template.format(path=chunks[index].path) + chunks[index].text
                for index in batch_indices
            ]
            embedded = client.embed(documents)
            cached_batch: list[tuple[str, array]] = []
            for index, values in zip(batch_indices, embedded):
                vector = array("f", values)
                vectors[index] = vector
                cached_batch.append((cache_keys[index], vector))
            cache.put_many(cached_batch)

        resolved_vectors = tuple(vector for vector in vectors if vector is not None)
        if len(resolved_vectors) != len(chunks):
            raise RuntimeError("dense index construction left missing vectors")
        stats = DenseIndexStats(
            total_chunks=len(chunks),
            cached_chunks=len(chunks) - len(missing_indices),
            embedded_chunks=len(missing_indices),
            build_seconds=time.monotonic() - started,
        )
        return cls(chunks, resolved_vectors, spec, client, cache), stats

    def query_vector(self, query: str) -> array:
        instructed_query = f"Instruct: {self.spec.query_instruction}\nQuery: {query}"
        cache_key = sha256(
            f"{self.spec.config_hash}\0query\0{instructed_query}".encode("utf-8")
        ).hexdigest()
        if self.cache is not None:
            cached = self.cache.get(cache_key)
            if cached is not None:
                return cached
        vector = array("f", self.client.embed([instructed_query])[0])
        if self.cache is not None:
            self.cache.put_many([(cache_key, vector)])
        return vector

    def retrieve(self, query: str, limit: int) -> Sequence[Candidate]:
        query_vector = self.query_vector(query)
        ranked = [
            (sum(left * right for left, right in zip(query_vector, vector)), chunk)
            for chunk, vector in zip(self.chunks, self.vectors)
        ]
        ranked.sort(key=lambda item: (-item[0], item[1].path, item[1].line_start))
        return tuple(_candidate(chunk, "dense", score) for score, chunk in ranked[:limit])


def unique_file_ranking(candidates: Iterable[Candidate]) -> tuple[Candidate, ...]:
    seen: set[str] = set()
    result: list[Candidate] = []
    for candidate in candidates:
        if candidate.path not in seen:
            seen.add(candidate.path)
            result.append(candidate)
    return tuple(result)