boring-agent-memory 0.2.2__py3-none-any.whl

boring_agent_memory/embeddings.py

@@ -0,0 +1,253 @@
+from __future__ import annotations
+
+import math
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass, replace
+from pathlib import Path
+from typing import Any, Protocol, cast
+
+from .privacy import redact_secrets
+from .query import QueryResult
+from .schema import connect, init_db, require_index
+
+
+class EmbeddingAdapter(Protocol):
+    model_id: str
+
+    def embed_documents(self, texts: Sequence[str]) -> list[list[float]]: ...
+
+    def embed_query(self, text: str) -> list[float]: ...
+
+
+class FastEmbedAdapter:
+    """Optional FastEmbed adapter with a strict no-download default."""
+
+    def __init__(
+        self,
+        model_name: str,
+        model_path: Path | str | None = None,
+        allow_download: bool = False,
+        cache_dir: Path | str | None = None,
+    ) -> None:
+        if model_path is None and not allow_download:
+            raise ValueError(
+                "a local embedding model path is required unless allow_download=True"
+            )
+        resolved_model_path = Path(model_path).expanduser().resolve() if model_path else None
+        if resolved_model_path is not None and not resolved_model_path.exists():
+            raise FileNotFoundError(f"embedding model path does not exist: {resolved_model_path}")
+        try:
+            from fastembed import TextEmbedding  # type: ignore[import-not-found]
+        except ImportError as exc:
+            raise RuntimeError(
+                "optional embeddings require: pip install 'boring-agent-memory[embeddings]'"
+            ) from exc
+
+        resolved_cache_dir = (
+            Path(cache_dir).expanduser().resolve().as_posix()
+            if cache_dir is not None
+            else None
+        )
+        self._model = TextEmbedding(
+            model_name=model_name,
+            cache_dir=resolved_cache_dir,
+            local_files_only=not allow_download,
+            specific_model_path=(
+                resolved_model_path.as_posix()
+                if resolved_model_path is not None
+                else None
+            ),
+        )
+        self.model_id = model_name
+        self.allow_download = allow_download
+        self.model_path = resolved_model_path
+
+    def embed_documents(self, texts: Sequence[str]) -> list[list[float]]:
+        redacted_texts = [redact_secrets(text)[0] for text in texts]
+        return [_vector(values) for values in self._model.embed(redacted_texts)]
+
+    def embed_query(self, text: str) -> list[float]:
+        text = redact_secrets(text)[0]
+        query_embed = getattr(self._model, "query_embed", None)
+        values = next(iter(query_embed(text) if query_embed else self._model.embed([text])))
+        return _vector(values)
+
+
+@dataclass(frozen=True)
+class EmbeddingDocument:
+    result: QueryResult
+    text: str
+
+
+@dataclass(frozen=True)
+class DenseIndex:
+    adapter: EmbeddingAdapter
+    documents: tuple[EmbeddingDocument, ...]
+    vectors: tuple[tuple[float, ...], ...]
+
+    @classmethod
+    def build(
+        cls,
+        adapter: EmbeddingAdapter,
+        documents: Sequence[EmbeddingDocument],
+    ) -> DenseIndex:
+        vectors = adapter.embed_documents(
+            [redact_secrets(document.text)[0] for document in documents]
+        )
+        if len(vectors) != len(documents):
+            raise ValueError("embedding adapter returned the wrong number of vectors")
+        dimensions = {len(vector) for vector in vectors}
+        if len(dimensions) > 1 or dimensions == {0}:
+            raise ValueError("embedding adapter returned inconsistent vector dimensions")
+        return cls(
+            adapter=adapter,
+            documents=tuple(documents),
+            vectors=tuple(tuple(vector) for vector in vectors),
+        )
+
+    def query(self, text: str, limit: int = 5) -> list[QueryResult]:
+        query_vector = self.adapter.embed_query(redact_secrets(text)[0])
+        if self.vectors and len(query_vector) != len(self.vectors[0]):
+            raise ValueError("query vector dimension does not match the dense index")
+        scored = [
+            (cosine_similarity(query_vector, vector), document.result)
+            for document, vector in zip(self.documents, self.vectors)
+        ]
+        scored.sort(key=lambda item: (-item[0], item[1].chunk_id))
+        seen_documents: set[str] = set()
+        results: list[QueryResult] = []
+        for score, result in scored:
+            if result.id in seen_documents:
+                continue
+            seen_documents.add(result.id)
+            results.append(replace(result, score=-score, strategy="dense"))
+            if len(results) >= limit:
+                break
+        return results
+
+
+def load_embedding_documents(
+    db_path: Path | str,
+    source_type: str | None = None,
+    workspace: Path | str | None = None,
+) -> list[EmbeddingDocument]:
+    conn = connect(require_index(db_path))
+    init_db(conn)
+    try:
+        filters: list[str] = []
+        params: list[object] = []
+        if source_type:
+            filters.append("d.source_type = ?")
+            params.append(source_type)
+        if workspace:
+            filters.append("d.workspace = ?")
+            params.append(Path(workspace).expanduser().resolve().as_posix())
+        where = "WHERE " + " AND ".join(filters) if filters else ""
+        rows = conn.execute(
+            f"""
+            SELECT
+              d.id, c.id AS chunk_id, d.source_type, d.source_path,
+              d.workspace, d.title, c.heading, c.start_line, c.end_line,
+              c.content
+            FROM chunks c
+            JOIN documents d ON d.id = c.document_id
+            {where}
+            ORDER BY d.source_path, c.ordinal
+            """,
+            params,
+        ).fetchall()
+    finally:
+        conn.close()
+
+    documents: list[EmbeddingDocument] = []
+    for row in rows:
+        location = f"L{row['start_line']}-L{row['end_line']}"
+        citation = f"{row['source_path']}:{location}"
+        if row["heading"]:
+            citation = f"{row['source_path']}#{row['heading']}:{location}"
+        result = QueryResult(
+            id=row["id"],
+            chunk_id=row["chunk_id"],
+            source_type=row["source_type"],
+            source_path=row["source_path"],
+            workspace=row["workspace"],
+            title=row["title"],
+            heading=row["heading"],
+            start_line=int(row["start_line"]),
+            end_line=int(row["end_line"]),
+            citation=citation,
+            score=0.0,
+            snippet=row["content"],
+            strategy="dense",
+        )
+        embedding_fields = (
+            redact_secrets(value)[0]
+            for value in (row["title"], row["heading"], row["content"])
+            if value
+        )
+        text = "\n".join(embedding_fields)
+        documents.append(EmbeddingDocument(result=result, text=text))
+    return documents
+
+
+def reciprocal_rank_fusion(
+    lexical: Sequence[QueryResult],
+    dense: Sequence[QueryResult],
+    limit: int,
+    lexical_weight: float = 0.6,
+    dense_weight: float = 0.4,
+    rrf_k: int = 60,
+) -> list[QueryResult]:
+    scores: dict[str, float] = {}
+    results: dict[str, QueryResult] = {}
+    tie_breaks: dict[str, tuple[int, int, str]] = {}
+    missing_rank = 10**9
+    lexical_ranks = {result.chunk_id: rank for rank, result in enumerate(lexical, start=1)}
+    dense_ranks = {result.chunk_id: rank for rank, result in enumerate(dense, start=1)}
+    for result in [*lexical, *dense]:
+        results.setdefault(result.chunk_id, result)
+    for chunk_id, result in results.items():
+        lexical_rank = lexical_ranks.get(chunk_id)
+        dense_rank = dense_ranks.get(chunk_id)
+        score = 0.0
+        if lexical_rank is not None:
+            score += lexical_weight / (rrf_k + lexical_rank)
+        if dense_rank is not None:
+            score += dense_weight / (rrf_k + dense_rank)
+        scores[chunk_id] = score
+        tie_breaks[chunk_id] = (
+            lexical_rank or missing_rank,
+            dense_rank or missing_rank,
+            chunk_id,
+        )
+    ordered = sorted(results, key=lambda key: (-scores[key], tie_breaks[key]))
+    seen_documents: set[str] = set()
+    fused: list[QueryResult] = []
+    for chunk_id in ordered:
+        result = results[chunk_id]
+        if result.id in seen_documents:
+            continue
+        seen_documents.add(result.id)
+        fused.append(replace(result, score=-scores[chunk_id], strategy="hybrid_rrf"))
+        if len(fused) >= limit:
+            break
+    return fused
+
+
+def cosine_similarity(left: Sequence[float], right: Sequence[float]) -> float:
+    if len(left) != len(right):
+        raise ValueError("vector dimensions do not match")
+    numerator = sum(a * b for a, b in zip(left, right))
+    left_norm = math.sqrt(sum(value * value for value in left))
+    right_norm = math.sqrt(sum(value * value for value in right))
+    if left_norm == 0 or right_norm == 0:
+        return 0.0
+    return numerator / (left_norm * right_norm)
+
+
+def _vector(values: object) -> list[float]:
+    if hasattr(values, "tolist"):
+        values = values.tolist()
+    if not isinstance(values, Iterable):
+        raise TypeError("embedding adapter returned a non-iterable vector")
+    return [float(value) for value in cast(Iterable[Any], values)]

boring_agent_memory/eval.py

@@ -0,0 +1,271 @@
+from __future__ import annotations
+
+import json
+import re
+import shutil
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from .canonical import verify_canonical_source
+from .index import build_index
+from .privacy import SECRET_PATTERNS
+from .query import QueryResult, query_memory
+
+
+@dataclass(frozen=True)
+class EvalCase:
+    id: str
+    query: str
+    expected_source: str
+    expected_terms: tuple[str, ...] = ()
+    forbidden_sources: tuple[str, ...] = ()
+    forbidden_terms: tuple[str, ...] = ()
+    expected_rank_at_most: int = 3
+    limit: int = 3
+    mutate_after_build: dict[str, str] | None = None
+    expect_stale: bool = False
+
+
+def run_eval(
+    fixture_dir: Path | str,
+    golden_path: Path | str | None = None,
+    limit: int = 3,
+) -> dict[str, Any]:
+    fixture_dir = Path(fixture_dir).expanduser().resolve()
+    golden_path = Path(golden_path).expanduser().resolve() if golden_path else fixture_dir.parent / "golden.jsonl"
+    cases = load_cases(golden_path)
+
+    with tempfile.TemporaryDirectory() as tmp:
+        work_fixture = Path(tmp) / "fixtures"
+        shutil.copytree(fixture_dir, work_fixture)
+        eval_db_path = Path(tmp) / "memory.db"
+
+        build_stats = build_index(
+            db_path=eval_db_path,
+            includes=[work_fixture.as_posix()],
+            workspace=work_fixture,
+            source_type="eval_fixture",
+        )
+        apply_mutations(work_fixture, cases)
+
+        rows = [
+            evaluate_case(
+                case=case,
+                db_path=eval_db_path,
+                workspace=work_fixture,
+                limit=max(limit, case.limit, case.expected_rank_at_most),
+            )
+            for case in cases
+        ]
+
+    metrics = summarize(rows)
+    return {
+        "fixture_dir": fixture_dir.as_posix(),
+        "golden_path": golden_path.as_posix(),
+        "build": build_stats,
+        "metrics": metrics,
+        "cases": rows,
+    }
+
+
+def evaluate_gates(
+    report: dict[str, Any],
+    min_recall_at_1: float | None = None,
+    min_recall_at_3: float | None = None,
+    min_source_accuracy: float | None = None,
+    min_snippet_term_rate: float | None = None,
+    min_stale_detection_rate: float | None = None,
+    max_privacy_leaks: int | None = None,
+) -> list[str]:
+    metrics = report["metrics"]
+    failures: list[str] = []
+    threshold_checks = {
+        "recall_at_1": min_recall_at_1,
+        "recall_at_3": min_recall_at_3,
+        "source_accuracy": min_source_accuracy,
+        "snippet_term_rate": min_snippet_term_rate,
+        "stale_detection_rate": min_stale_detection_rate,
+    }
+    for metric, minimum in threshold_checks.items():
+        if minimum is not None and float(metrics[metric]) < minimum:
+            failures.append(f"{metric}={metrics[metric]} is below required {minimum}")
+
+    if max_privacy_leaks is not None and int(metrics["privacy_leak_count"]) > max_privacy_leaks:
+        failures.append(
+            f"privacy_leak_count={metrics['privacy_leak_count']} exceeds allowed {max_privacy_leaks}"
+        )
+    return failures
+
+
+def load_cases(golden_path: Path) -> list[EvalCase]:
+    cases: list[EvalCase] = []
+    for line_number, line in enumerate(golden_path.read_text(encoding="utf-8").splitlines(), start=1):
+        if not line.strip():
+            continue
+        payload = json.loads(line)
+        try:
+            cases.append(
+                EvalCase(
+                    id=payload["id"],
+                    query=payload["query"],
+                    expected_source=payload["expected_source"],
+                    expected_terms=tuple(payload.get("expected_terms", ())),
+                    forbidden_sources=tuple(payload.get("forbidden_sources", ())),
+                    forbidden_terms=tuple(payload.get("forbidden_terms", ())),
+                    expected_rank_at_most=int(payload.get("expected_rank_at_most", payload.get("limit", 3))),
+                    limit=int(payload.get("limit", 3)),
+                    mutate_after_build=payload.get("mutate_after_build"),
+                    expect_stale=bool(payload.get("expect_stale", False)),
+                )
+            )
+        except KeyError as exc:
+            raise ValueError(f"{golden_path}:{line_number} missing required field {exc}") from exc
+    if not cases:
+        raise ValueError(f"{golden_path} does not contain any eval cases")
+    return cases
+
+
+def apply_mutations(work_fixture: Path, cases: list[EvalCase]) -> None:
+    work_fixture = work_fixture.resolve()
+    for case in cases:
+        if not case.mutate_after_build:
+            continue
+        source_path = case.mutate_after_build["source_path"]
+        content = case.mutate_after_build["content"]
+        target = (work_fixture / source_path).resolve()
+        if not target.is_relative_to(work_fixture):
+            raise ValueError(f"mutation target escapes fixture directory: {source_path}")
+        target.write_text(content, encoding="utf-8")
+
+
+def evaluate_case(
+    case: EvalCase,
+    db_path: Path,
+    workspace: Path,
+    limit: int,
+) -> dict[str, Any]:
+    workspace = workspace.resolve()
+    results = query_memory(
+        db_path=db_path,
+        query=case.query,
+        limit=limit,
+        source_type="eval_fixture",
+        workspace=workspace,
+    )
+    expected_path = (workspace / case.expected_source).resolve().as_posix()
+    rank = rank_of_expected(results, expected_path)
+    top = results[0] if results else None
+    snippet_terms_present = terms_present(top, case.expected_terms) if top else False
+    privacy_leaks = count_privacy_leaks(results, case.forbidden_terms)
+    forbidden_source_ranks = ranked_forbidden_sources(results, workspace, case.forbidden_sources)
+    stale_detected = False
+    if case.expect_stale or case.forbidden_sources:
+        stale_detected = stale_sources_did_not_outrank_expected(rank, forbidden_source_ranks)
+        if case.expect_stale and top:
+            verification = verify_canonical_source(db_path, top.source_path)
+            stale_detected = stale_detected and verification.get("content_hash_match") is False
+
+    return {
+        "id": case.id,
+        "query": case.query,
+        "expected_source": case.expected_source,
+        "top_source": relative_source(top.source_path, workspace) if top else None,
+        "rank": rank,
+        "recall_at_1": rank == 1,
+        "recall_at_k": rank is not None and rank <= case.expected_rank_at_most,
+        "source_accuracy": top is not None and top.source_path == expected_path,
+        "snippet_terms_present": snippet_terms_present,
+        "privacy_leak_count": privacy_leaks,
+        "forbidden_source_ranks": forbidden_source_ranks,
+        "stale_expected": case.expect_stale or bool(case.forbidden_sources),
+        "stale_detected": stale_detected,
+    }
+
+
+def rank_of_expected(results: list[QueryResult], expected_path: str) -> int | None:
+    for index, result in enumerate(results, start=1):
+        if result.source_path == expected_path:
+            return index
+    return None
+
+
+def terms_present(result: QueryResult, terms: tuple[str, ...]) -> bool:
+    if not terms:
+        return True
+    haystack = normalize_text(f"{result.title}\n{result.snippet}")
+    return all(normalize_text(term) in haystack for term in terms)
+
+
+def count_privacy_leaks(results: list[QueryResult], forbidden_terms: tuple[str, ...] = ()) -> int:
+    count = 0
+    for result in results:
+        for pattern in SECRET_PATTERNS:
+            for match in pattern.finditer(result.snippet):
+                if "[REDACTED]" not in match.group(0):
+                    count += 1
+        lowered = result.snippet.lower()
+        for term in forbidden_terms:
+            if term.lower() in lowered:
+                count += 1
+    return count
+
+
+def ranked_forbidden_sources(
+    results: list[QueryResult],
+    workspace: Path,
+    sources: tuple[str, ...],
+) -> dict[str, int | None]:
+    forbidden = set(sources)
+    ranks: dict[str, int | None] = {source: None for source in sources}
+    for result in results:
+        relative = relative_source(result.source_path, workspace)
+        if relative in forbidden:
+            ranks[relative] = rank_of_expected(results, result.source_path)
+    return ranks
+
+
+def stale_sources_did_not_outrank_expected(
+    expected_rank: int | None,
+    forbidden_source_ranks: dict[str, int | None],
+) -> bool:
+    if expected_rank is None:
+        return False
+    return all(rank is None or rank > expected_rank for rank in forbidden_source_ranks.values())
+
+
+def relative_source(source_path: str, workspace: Path) -> str:
+    try:
+        return Path(source_path).resolve().relative_to(workspace).as_posix()
+    except ValueError:
+        return source_path
+
+
+def summarize(rows: list[dict[str, Any]]) -> dict[str, Any]:
+    total = len(rows)
+    stale_cases = [row for row in rows if row["stale_expected"]]
+    mrr_total = sum((1 / row["rank"]) for row in rows if row["rank"])
+    return {
+        "cases": total,
+        "recall_at_1": ratio(sum(1 for row in rows if row["recall_at_1"]), total),
+        "recall_at_3": ratio(sum(1 for row in rows if row["recall_at_k"]), total),
+        "mrr": ratio(mrr_total, total),
+        "source_accuracy": ratio(sum(1 for row in rows if row["source_accuracy"]), total),
+        "snippet_term_rate": ratio(sum(1 for row in rows if row["snippet_terms_present"]), total),
+        "privacy_leak_count": sum(row["privacy_leak_count"] for row in rows),
+        "stale_detection_rate": ratio(
+            sum(1 for row in stale_cases if row["stale_detected"]),
+            len(stale_cases),
+        ),
+    }
+
+
+def ratio(numerator: float, denominator: int) -> float:
+    if denominator == 0:
+        return 0.0
+    return round(numerator / denominator, 4)
+
+
+def normalize_text(value: str) -> str:
+    return re.sub(r"[^a-z0-9]+", " ", value.lower()).strip()