memory-arena 0.1.8__py3-none-any.whl

memory_arena/__init__.py

@@ -0,0 +1,19 @@
+"""Memory Arena — Knowledge Base Benchmark. Find which retrieval architecture fits your data."""
+
+__version__ = "0.1.8"
+
+from memory_arena.models.benchmark import BenchmarkResult, Question
+from memory_arena.models.document import Document, Section
+from memory_arena.models.graph import Entity, Relationship
+from memory_arena.strategies.base import Strategy
+
+__all__ = [
+    "Document",
+    "Section",
+    "Entity",
+    "Relationship",
+    "Question",
+    "BenchmarkResult",
+    "Strategy",
+    "__version__",
+]

memory_arena/arena/engine.py

@@ -0,0 +1,252 @@
+"""Arena engine - blind A/B strategy matchups with ELO rating."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import random
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from uuid import uuid4
+
+from memory_arena.settings import settings
+
+log = logging.getLogger(__name__)
+
+INITIAL_ELO = 1200.0
+K_FACTOR = 32.0
+
+
+@dataclass
+class Match:
+    """A single A/B matchup between two strategies."""
+
+    id: str
+    question: str
+    strategy_a: str
+    strategy_b: str
+    answer_a: str
+    answer_b: str
+    latency_a_ms: float = 0.0
+    latency_b_ms: float = 0.0
+    cost_a: float = 0.0
+    cost_b: float = 0.0
+    winner: str | None = None  # "a", "b", "tie", or None (pending)
+    timestamp: float = 0.0
+    sources_a: list[str] = field(default_factory=list)
+    sources_b: list[str] = field(default_factory=list)
+
+
+@dataclass
+class ArenaState:
+    """Persistent arena state with ELO ratings and match history."""
+
+    elo: dict[str, float] = field(default_factory=dict)
+    matches: list[Match] = field(default_factory=list)
+    total_votes: int = 0
+
+    def save(self, path: Path) -> None:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        data = {
+            "elo": self.elo,
+            "total_votes": self.total_votes,
+            "matches": [
+                {
+                    "id": m.id,
+                    "question": m.question,
+                    "strategy_a": m.strategy_a,
+                    "strategy_b": m.strategy_b,
+                    "answer_a": m.answer_a[:500],  # truncate for storage
+                    "answer_b": m.answer_b[:500],
+                    "latency_a_ms": m.latency_a_ms,
+                    "latency_b_ms": m.latency_b_ms,
+                    "cost_a": m.cost_a,
+                    "cost_b": m.cost_b,
+                    "winner": m.winner,
+                    "timestamp": m.timestamp,
+                }
+                for m in self.matches[-200:]  # keep last 200
+            ],
+        }
+        path.write_text(json.dumps(data, indent=2))
+
+    @classmethod
+    def load(cls, path: Path) -> ArenaState:
+        if not path.exists():
+            return cls()
+        try:
+            data = json.loads(path.read_text())
+            state = cls(
+                elo=data.get("elo", {}),
+                total_votes=data.get("total_votes", 0),
+            )
+            for m in data.get("matches", []):
+                state.matches.append(
+                    Match(
+                        id=m["id"],
+                        question=m["question"],
+                        strategy_a=m["strategy_a"],
+                        strategy_b=m["strategy_b"],
+                        answer_a=m.get("answer_a", ""),
+                        answer_b=m.get("answer_b", ""),
+                        latency_a_ms=m.get("latency_a_ms", 0),
+                        latency_b_ms=m.get("latency_b_ms", 0),
+                        cost_a=m.get("cost_a", 0),
+                        cost_b=m.get("cost_b", 0),
+                        winner=m.get("winner"),
+                        timestamp=m.get("timestamp", 0),
+                    )
+                )
+            return state
+        except (json.JSONDecodeError, KeyError):
+            log.warning("Corrupt arena state, starting fresh")
+            return cls()
+
+
+class ArenaEngine:
+    """Manages blind A/B matches between retrieval strategies."""
+
+    def __init__(self, strategies: dict) -> None:
+        self.strategies = strategies
+        self._state_path = Path(settings.results_path) / "arena_state.json"
+        self.state = ArenaState.load(self._state_path)
+        # Initialize ELO for new strategies
+        for name in strategies:
+            if name not in self.state.elo:
+                self.state.elo[name] = INITIAL_ELO
+
+    async def create_match(self, question: str, corpus: str = "") -> Match:
+        """Pick two random strategies, query both, return a blind match."""
+        names = list(self.strategies.keys())
+        if len(names) < 2:
+            raise ValueError("Need at least 2 strategies for arena mode")
+        a_name, b_name = random.sample(names, 2)
+
+        result_a, result_b = await asyncio.gather(
+            self.strategies[a_name].query(question),
+            self.strategies[b_name].query(question),
+        )
+
+        match = Match(
+            id=uuid4().hex[:8],
+            question=question,
+            strategy_a=a_name,
+            strategy_b=b_name,
+            answer_a=result_a.answer,
+            answer_b=result_b.answer,
+            latency_a_ms=result_a.latency_ms,
+            latency_b_ms=result_b.latency_ms,
+            cost_a=result_a.cost_usd,
+            cost_b=result_b.cost_usd,
+            sources_a=result_a.sources,
+            sources_b=result_b.sources,
+            timestamp=time.time(),
+        )
+        self.state.matches.append(match)
+        return match
+
+    def vote(self, match_id: str, winner: str) -> dict:
+        """Record a vote and update ELO. winner: 'a', 'b', or 'tie'."""
+        if winner not in ("a", "b", "tie"):
+            return {"error": f"Invalid winner: {winner}. Must be 'a', 'b', or 'tie'"}
+
+        match = next((m for m in self.state.matches if m.id == match_id), None)
+        if not match:
+            return {"error": "Match not found"}
+        if match.winner is not None:
+            return {"error": "Match already voted on"}
+
+        match.winner = winner
+        self.state.total_votes += 1
+        self._update_elo(match)
+        self.state.save(self._state_path)
+        self._append_vote_jsonl(match)
+
+        return {
+            "strategy_a": match.strategy_a,
+            "strategy_b": match.strategy_b,
+            "winner": winner,
+            "elo": dict(self.state.elo),
+            "total_votes": self.state.total_votes,
+        }
+
+    def _update_elo(self, match: Match) -> None:
+        """Standard ELO rating update."""
+        ea = self.state.elo.get(match.strategy_a, INITIAL_ELO)
+        eb = self.state.elo.get(match.strategy_b, INITIAL_ELO)
+        expected_a = 1.0 / (1.0 + 10.0 ** ((eb - ea) / 400.0))
+
+        if match.winner == "a":
+            score_a = 1.0
+        elif match.winner == "b":
+            score_a = 0.0
+        else:  # tie
+            score_a = 0.5
+
+        self.state.elo[match.strategy_a] = ea + K_FACTOR * (score_a - expected_a)
+        self.state.elo[match.strategy_b] = eb + K_FACTOR * ((1 - score_a) - (1 - expected_a))
+
+    def leaderboard(self) -> list[dict]:
+        """Return strategies sorted by ELO rating."""
+        board = []
+        for name, elo in self.state.elo.items():
+            wins = sum(
+                1
+                for m in self.state.matches
+                if m.winner
+                and (
+                    (m.strategy_a == name and m.winner == "a")
+                    or (m.strategy_b == name and m.winner == "b")
+                )
+            )
+            losses = sum(
+                1
+                for m in self.state.matches
+                if m.winner
+                and (
+                    (m.strategy_a == name and m.winner == "b")
+                    or (m.strategy_b == name and m.winner == "a")
+                )
+            )
+            ties = sum(
+                1
+                for m in self.state.matches
+                if m.winner == "tie" and (m.strategy_a == name or m.strategy_b == name)
+            )
+            board.append(
+                {
+                    "strategy": name,
+                    "elo": round(elo, 1),
+                    "wins": wins,
+                    "losses": losses,
+                    "ties": ties,
+                    "matches": wins + losses + ties,
+                }
+            )
+        return sorted(board, key=lambda x: x["elo"], reverse=True)
+
+    def _append_vote_jsonl(self, match: Match) -> None:
+        """Append-only JSONL log of all votes (survives state resets)."""
+        jsonl_path = self._state_path.parent / "arena_votes.jsonl"
+        jsonl_path.parent.mkdir(parents=True, exist_ok=True)
+        record = {
+            "match_id": match.id,
+            "question": match.question[:200],
+            "strategy_a": match.strategy_a,
+            "strategy_b": match.strategy_b,
+            "winner": match.winner,
+            "latency_a_ms": round(match.latency_a_ms, 1),
+            "latency_b_ms": round(match.latency_b_ms, 1),
+            "cost_a": match.cost_a,
+            "cost_b": match.cost_b,
+            "timestamp": match.timestamp,
+            "elo_snapshot": {k: round(v, 1) for k, v in self.state.elo.items()},
+        }
+        with open(jsonl_path, "a") as f:
+            f.write(json.dumps(record) + "\n")
+
+    def get_pending_match(self, match_id: str) -> Match | None:
+        """Get a match by ID."""
+        return next((m for m in self.state.matches if m.id == match_id), None)

memory_arena/benchmark/__init__.py

@@ -0,0 +1,15 @@
+"""Memory Arena benchmark engine."""
+
+from memory_arena.benchmark.evaluator import (
+    MemoryScore,
+    evaluate_memory_answer,
+)
+from memory_arena.benchmark.questions import load_memory_questions
+from memory_arena.benchmark.recall_metrics import compute_memory_recall_metrics
+
+__all__ = [
+    "MemoryScore",
+    "compute_memory_recall_metrics",
+    "evaluate_memory_answer",
+    "load_memory_questions",
+]

memory_arena/benchmark/evaluator.py

@@ -0,0 +1,315 @@
+"""7-axis evaluator for memory-arena.
+
+Axes (4 lifted from kb-arena + 3 memory-specific):
+  1. structural — must_mention / must_not_claim / max_tokens
+  2. sources    — was at least one supporting_session_id cited?
+  3. judge      — Opus rates accuracy 0..100 against the reference
+  4. memo       — identical (answer, reference) pairs cached in-process
+  5. temporal_correctness — answer's claimed timestamp is inside ground_truth.valid_as_of window
+  6. update_precision     — answer reflects the latest fact version, not an earlier one
+  7. abstention_f1        — abstention questions get F1 over an abstain/no-abstain classifier
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from memory_arena.llm.client import LLMClient
+from memory_arena.sessions.schema import Constraints, GroundTruth, QuestionRecord
+
+_JUDGE_SYSTEM = (
+    "You are a strict judge comparing a candidate answer to a reference answer for a question "
+    "about a chat history. Output a JSON object with this exact shape:\n"
+    '{"accuracy": <0-100>, "completeness": <0-100>, "rationale": "<one sentence>"}\n'
+    "Be strict: an answer that misses key facts should score below 60. An answer that "
+    "contradicts the reference should score below 30. Reply with only the JSON object."
+)
+
+
+_ABSTAIN_SYSTEM = (
+    "Classify whether the following answer abstains from answering (says it does not know, "
+    "cannot find the information, or refuses to answer). Reply with only YES or NO."
+)
+
+
+_TEMPORAL_SYSTEM = (
+    "Extract any explicit time marker mentioned in the answer (date, week-of, month, year, "
+    "or relative phrase like 'last week'). Reply with only the time marker, or NONE."
+)
+
+
+class MemoryScore(BaseModel):
+    """7-axis evaluation score."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    accuracy: float = 0.0
+    completeness: float = 0.0
+    structural_pass: bool = True
+    structural_fails: list[str] = Field(default_factory=list)
+    sources_pass: bool = False
+    judge_score: float = 0.0
+    judge_rationale: str = ""
+    temporal_correct: bool = False
+    # update_precision_correct can be None when the question does not
+    # carry fact_versions to check — runner aggregates only non-None.
+    update_precision_correct: bool | None = False
+    abstained: bool = False
+    abstention_correct: bool = False
+    cost_usd: float = 0.0
+    tokens_used: int = 0
+
+
+# Module-level memo cache. Keyed by (question_id, answer-hash, reference-hash)
+# so long answers that diverge after the first 500 characters do not collide.
+_judge_cache: dict[tuple[str, str, str], dict] = {}
+
+
+def _evaluate_structural(
+    answer: str,
+    constraints: Constraints,
+) -> tuple[bool, list[str]]:
+    fails: list[str] = []
+    answer_lower = answer.lower()
+    for required in constraints.must_mention:
+        if required.lower() not in answer_lower:
+            fails.append(f"missing must_mention: {required}")
+    for forbidden in constraints.must_not_claim:
+        if forbidden.lower() in answer_lower:
+            fails.append(f"contains must_not_claim: {forbidden}")
+    if constraints.max_tokens > 0:
+        approx_tokens = max(1, len(answer) // 4)
+        if approx_tokens > constraints.max_tokens * 1.5:
+            fails.append(f"exceeds max_tokens by 50%: {approx_tokens} > {constraints.max_tokens}")
+    return (len(fails) == 0), fails
+
+
+def _evaluate_sources(
+    answer: str,
+    supporting_session_ids: list[str],
+    ground_truth: GroundTruth,
+) -> bool:
+    expected = set(ground_truth.supporting_session_ids)
+    if not expected:
+        return True  # no labeled sources -> can't fail
+    return bool(set(supporting_session_ids) & expected)
+
+
+async def _evaluate_judge(
+    answer: str,
+    reference: str,
+    llm: LLMClient,
+    question_id: str = "",
+) -> tuple[float, float, str, float, int]:
+    import hashlib
+
+    a_hash = hashlib.sha256(answer.strip().encode("utf-8")).hexdigest()
+    r_hash = hashlib.sha256(reference.strip().encode("utf-8")).hexdigest()
+    key = (question_id, a_hash, r_hash)
+    if key in _judge_cache:
+        cached = _judge_cache[key]
+        return (
+            cached["accuracy"],
+            cached["completeness"],
+            cached["rationale"],
+            0.0,
+            0,
+        )
+    if not reference:
+        return 0.0, 0.0, "no reference provided", 0.0, 0
+    resp = await llm.judge(answer=answer, reference=reference, system_prompt=_JUDGE_SYSTEM)
+    text = resp.text.strip()
+    accuracy = 0.0
+    completeness = 0.0
+    rationale = ""
+    try:
+        import json
+
+        match = re.search(r"\{.*\}", text, re.DOTALL)
+        if match:
+            data = json.loads(match.group(0))
+            accuracy = float(data.get("accuracy", 0))
+            completeness = float(data.get("completeness", 0))
+            rationale = data.get("rationale", "")
+    except Exception:
+        rationale = f"judge parse failure: {text[:80]}"
+    _judge_cache[key] = {
+        "accuracy": accuracy,
+        "completeness": completeness,
+        "rationale": rationale,
+    }
+    return accuracy, completeness, rationale, resp.cost_usd, resp.total_tokens
+
+
+async def _classify_abstention(answer: str, llm: LLMClient) -> tuple[bool, float, int]:
+    if _quick_abstain_match(answer):
+        return True, 0.0, 0
+    resp = await llm.classify(
+        query=answer[:1000], system_prompt=_ABSTAIN_SYSTEM, allowed_values=["YES", "NO"]
+    )
+    return resp.upper().startswith("Y"), 0.0, 0
+
+
+def _quick_abstain_match(answer: str) -> bool:
+    a = answer.lower()
+    needles = (
+        "i do not have that information",
+        "i don't have that information",
+        "i don't know",
+        "i do not know",
+        "i'm not sure",
+        "i am not sure",
+        "no information",
+        "cannot answer",
+        "unable to answer",
+    )
+    return any(n in a for n in needles)
+
+
+async def _check_temporal(
+    answer: str,
+    valid_as_of: str | None,
+    llm: LLMClient,
+) -> bool:
+    if not valid_as_of:
+        return True
+    resp = await llm.classify(query=answer[:1000], system_prompt=_TEMPORAL_SYSTEM)
+    if resp.upper() == "NONE" or not resp:
+        return False
+    return _temporal_overlap(resp, valid_as_of)
+
+
+def _temporal_overlap(claimed: str, valid_as_of: str) -> bool:
+    """Naive substring overlap on year/month/day numbers. Tighten in v0.2."""
+    claimed_nums = set(re.findall(r"\d+", claimed))
+    expected_nums = set(re.findall(r"\d+", valid_as_of))
+    return bool(claimed_nums & expected_nums) if expected_nums else False
+
+
+async def _check_update_precision(
+    answer: str,
+    ground_truth: GroundTruth,
+) -> bool | None:
+    """Answer must reflect the latest fact version, not an earlier one.
+
+    Returns None when fact_versions is empty/missing — there's nothing to
+    check, and reporting True would inflate the per-category metric to
+    1.0 across every strategy regardless of whether the question
+    actually exercised an update. The runner aggregates only non-None
+    values.
+    """
+    if not ground_truth.fact_versions:
+        return None
+    a = answer.lower()
+    latest = ground_truth.fact_versions[-1].value.lower()
+    if latest and latest in a:
+        return True
+    earlier = [v.value.lower() for v in ground_truth.fact_versions[:-1]]
+    if any(v and v in a for v in earlier):
+        return False
+    return False
+
+
+async def evaluate_memory_answer(
+    answer: str,
+    ground_truth: GroundTruth,
+    constraints: Constraints,
+    question: QuestionRecord,
+    llm: LLMClient,
+    supporting_session_ids: list[str] | None = None,
+) -> MemoryScore:
+    """Run all 7 axes and return a MemoryScore."""
+    score = MemoryScore()
+    if supporting_session_ids is None:
+        supporting_session_ids = []
+
+    # 1. structural
+    structural_pass, fails = _evaluate_structural(answer, constraints)
+    score.structural_pass = structural_pass
+    score.structural_fails = fails
+
+    # 2. sources
+    score.sources_pass = _evaluate_sources(answer, supporting_session_ids, ground_truth)
+
+    # 3. judge
+    accuracy, completeness, rationale, cost, toks = await _evaluate_judge(
+        answer, ground_truth.answer, llm, question_id=question.id
+    )
+    score.judge_score = accuracy
+    score.judge_rationale = rationale
+    score.completeness = completeness / 100.0 if completeness > 1 else completeness
+    score.cost_usd += cost
+    score.tokens_used += toks
+
+    # 5. temporal
+    if question.category == "temporal":
+        score.temporal_correct = await _check_temporal(answer, ground_truth.valid_as_of, llm)
+    else:
+        score.temporal_correct = True
+
+    # 6. update precision. _check_update_precision returns None when the
+    # question has no fact_versions to verify; preserve that so the
+    # runner can aggregate only the questions that actually exercised an
+    # update instead of polluting the metric with structural-trues.
+    if question.category == "knowledge_update":
+        score.update_precision_correct = await _check_update_precision(answer, ground_truth)
+    else:
+        score.update_precision_correct = True
+
+    # 7. abstention
+    abstained, _, _ = await _classify_abstention(answer, llm)
+    score.abstained = abstained
+    score.abstention_correct = abstained == constraints.abstention_expected
+
+    # Composite accuracy: judge score normalized to 0-1, dampened by structural and sources
+    base = accuracy / 100.0 if accuracy > 1 else accuracy
+    if not structural_pass:
+        base *= 0.5
+    if not score.sources_pass and ground_truth.supporting_session_ids:
+        base *= 0.8
+    if question.category == "abstention":
+        base = 1.0 if score.abstention_correct else 0.0
+    if question.category == "temporal" and not score.temporal_correct:
+        base *= 0.5
+    if question.category == "knowledge_update" and score.update_precision_correct is False:
+        base *= 0.5
+    score.accuracy = max(0.0, min(1.0, base))
+
+    return score
+
+
+def clear_eval_cache() -> None:
+    _judge_cache.clear()
+
+
+@dataclass
+class EvaluatorBundle:
+    """Convenience wrapper for callers that batch many evaluations."""
+
+    llm: LLMClient = field(default_factory=LLMClient)
+
+    async def evaluate(
+        self,
+        answer: str,
+        question: QuestionRecord,
+        supporting_session_ids: list[str] | None = None,
+    ) -> MemoryScore:
+        return await evaluate_memory_answer(
+            answer=answer,
+            ground_truth=question.ground_truth,
+            constraints=question.constraints,
+            question=question,
+            llm=self.llm,
+            supporting_session_ids=supporting_session_ids,
+        )
+
+
+__all__ = [
+    "EvaluatorBundle",
+    "MemoryScore",
+    "clear_eval_cache",
+    "evaluate_memory_answer",
+]