agentic: add answer-mode retrieval core (contract, evaluator, baselines)

Author: AdnanElAssadi56

mteb/agentic/__init__.py

@@ -0,0 +1,46 @@
+"""Answer-mode retrieval benchmark (core). See README.md."""
+
+from __future__ import annotations
+
+from mteb.agentic.corpus import InMemoryCorpus
+from mteb.agentic.data import AnswerTaskData, from_mteb_retrieval
+from mteb.agentic.evaluator import AnswerEvaluationResult, AnswerEvaluator
+from mteb.agentic.interface import (
+    AnswerResult,
+    AnswerSystem,
+    ChatModel,
+    ChatResponse,
+    CorpusHandle,
+    Message,
+    Usage,
+)
+from mteb.agentic.metrics import (
+    AggregateScores,
+    ExactMatchJudge,
+    Judge,
+    LLMJudge,
+    aggregate,
+)
+from mteb.agentic.systems import ClosedBookSystem, OracleContextSystem
+
+__all__ = [
+    "AggregateScores",
+    "AnswerEvaluationResult",
+    "AnswerEvaluator",
+    "AnswerResult",
+    "AnswerSystem",
+    "AnswerTaskData",
+    "ChatModel",
+    "ChatResponse",
+    "ClosedBookSystem",
+    "CorpusHandle",
+    "ExactMatchJudge",
+    "InMemoryCorpus",
+    "Judge",
+    "LLMJudge",
+    "Message",
+    "OracleContextSystem",
+    "Usage",
+    "aggregate",
+    "from_mteb_retrieval",
+]

mteb/agentic/corpus.py

@@ -0,0 +1,14 @@
+"""Reference CorpusHandle implementations."""
+
+from __future__ import annotations
+
+
+class InMemoryCorpus:
+    """Read-only corpus backed by a dict of documents."""
+
+    def __init__(self, documents: dict[str, dict[str, str]]) -> None:
+        self._docs = documents
+
+    def get(self, doc_id: str) -> dict[str, str]:
+        """Return one document with its id."""
+        return {"id": doc_id, **self._docs[doc_id]}

mteb/agentic/data.py

@@ -0,0 +1,52 @@
+"""Adapters that turn retrieval datasets into answer-mode inputs."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+
+@dataclass
+class AnswerTaskData:
+    """Everything an AnswerEvaluator needs for one split."""
+
+    documents: dict[str, dict[str, str]]  # doc_id -> {title, text}
+    questions: dict[str, str]  # qid -> question
+    references: dict[str, str]  # qid -> reference answer
+    gold_by_qid: dict[str, list[str]]  # qid -> gold doc ids
+    gold_by_question: dict[str, list[str]]  # question text -> gold doc ids
+
+
+def from_mteb_retrieval(
+    corpus: Mapping[str, Mapping[str, str]],
+    queries: Mapping[str, str],
+    relevant_docs: Mapping[str, Mapping[str, int]],
+    answers: Mapping[str, str],
+) -> AnswerTaskData:
+    """Build answer-mode data from MTEB-style retrieval fields plus answers."""
+    documents = {
+        doc_id: {
+            "title": doc.get("title", ""),
+            "text": doc.get("text", doc.get("body", "")),
+        }
+        for doc_id, doc in corpus.items()
+    }
+    questions = dict(queries)
+    references = dict(answers)
+    gold_by_qid = {
+        qid: [doc_id for doc_id, score in rels.items() if score > 0]
+        for qid, rels in relevant_docs.items()
+    }
+    gold_by_question = {
+        questions[qid]: ids for qid, ids in gold_by_qid.items() if qid in questions
+    }
+    return AnswerTaskData(
+        documents=documents,
+        questions=questions,
+        references=references,
+        gold_by_qid=gold_by_qid,
+        gold_by_question=gold_by_question,
+    )

mteb/agentic/evaluator.py

@@ -0,0 +1,70 @@
+"""Runs one AnswerSystem over a question set and scores it on three axes."""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from mteb.agentic.metrics import aggregate
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from mteb.agentic.interface import AnswerResult, AnswerSystem, CorpusHandle
+    from mteb.agentic.metrics import AggregateScores, Judge
+
+
+@dataclass
+class AnswerEvaluationResult:
+    """Aggregate scores plus a per-question record."""
+
+    scores: AggregateScores
+    per_question: list[dict[str, Any]]
+
+
+class AnswerEvaluator:
+    """Evaluate an answer-mode system over a fixed corpus and question set."""
+
+    def __init__(
+        self,
+        questions: Mapping[str, str],
+        references: Mapping[str, str],
+        corpus: CorpusHandle,
+        judge: Judge,
+    ) -> None:
+        # questions and references are keyed by query id.
+        self.questions = questions
+        self.references = references
+        self.corpus = corpus
+        self.judge = judge
+
+    def __call__(self, system: AnswerSystem) -> AnswerEvaluationResult:
+        """Run the system over every question and return aggregate scores."""
+        results: list[AnswerResult] = []
+        correctness: list[float] = []
+        per_question: list[dict[str, Any]] = []
+        for qid, question in self.questions.items():
+            start = time.perf_counter()
+            result = system.answer(question, self.corpus)
+            elapsed = time.perf_counter() - start
+            if result.usage.latency_s is None:
+                result.usage.latency_s = elapsed
+            score = self.judge.score(question, result.answer, self.references[qid])
+            results.append(result)
+            correctness.append(score)
+            per_question.append(
+                {
+                    "query_id": qid,
+                    "answer": result.answer,
+                    "correct": score,
+                    "cited_doc_ids": result.cited_doc_ids,
+                    "latency_s": result.usage.latency_s,
+                    "cost_usd": result.usage.cost_usd,
+                    "num_llm_calls": result.usage.num_llm_calls,
+                }
+            )
+        return AnswerEvaluationResult(
+            scores=aggregate(results, correctness),
+            per_question=per_question,
+        )

mteb/agentic/interface.py

@@ -0,0 +1,73 @@
+"""Core contract for answer-mode retrieval systems."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+# Provider-agnostic chat message.
+Message = dict[str, str]
+
+
+@dataclass
+class ChatResponse:
+    """Result of one chat completion call."""
+
+    text: str
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    cost_usd: float | None = None
+
+
+@runtime_checkable
+class ChatModel(Protocol):
+    """Provider-agnostic chat interface."""
+
+    name: str
+
+    def generate(self, messages: Sequence[Message], **kwargs: Any) -> ChatResponse:
+        """Generate a single completion for a chat transcript."""
+        ...
+
+
+@runtime_checkable
+class CorpusHandle(Protocol):
+    """Read access to a fixed corpus."""
+
+    def get(self, doc_id: str) -> dict[str, str]:
+        """Fetch one document as a mapping with at least id and text."""
+        ...
+
+
+@dataclass
+class Usage:
+    """Cost and latency accounting for one answer."""
+
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    num_llm_calls: int = 0
+    cost_usd: float | None = None
+    latency_s: float | None = None
+
+
+@dataclass
+class AnswerResult:
+    """What a system returns for one question."""
+
+    answer: str
+    cited_doc_ids: list[str] = field(default_factory=list)
+    usage: Usage = field(default_factory=Usage)
+
+
+@runtime_checkable
+class AnswerSystem(Protocol):
+    """End-to-end system that produces an answer, not a ranking."""
+
+    name: str
+
+    def answer(self, question: str, corpus: CorpusHandle) -> AnswerResult:
+        """Answer a single question using the corpus."""
+        ...

mteb/agentic/metrics.py

@@ -0,0 +1,92 @@
+"""Correctness judges and three-axis aggregation for answer-mode eval."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from mteb.agentic.interface import AnswerResult, ChatModel
+
+
+@runtime_checkable
+class Judge(Protocol):
+    """Scores answer correctness in the range 0 to 1."""
+
+    def score(self, question: str, predicted: str, reference: str) -> float:
+        """Grade a predicted answer against a reference answer."""
+        ...
+
+
+def _normalize(text: str) -> str:
+    # Lowercase, drop articles and punctuation, collapse whitespace.
+    text = text.lower()
+    text = re.sub(r"\b(a|an|the)\b", " ", text)
+    text = re.sub(r"[^a-z0-9 ]", " ", text)
+    return " ".join(text.split())
+
+
+class ExactMatchJudge:
+    """Normalized exact match."""
+
+    def score(self, question: str, predicted: str, reference: str) -> float:  # noqa: PLR6301
+        """Return 1.0 on a normalized exact match, else 0.0."""
+        return 1.0 if _normalize(predicted) == _normalize(reference) else 0.0
+
+
+_JUDGE_PROMPT = (
+    "You grade a predicted answer against a reference answer. "
+    "Reply with only YES if the prediction is correct, otherwise NO.\n\n"
+    "Question: {question}\nReference: {reference}\nPrediction: {predicted}"
+)
+
+
+class LLMJudge:
+    """Grades open ended answers with a ChatModel."""
+
+    def __init__(self, model: ChatModel) -> None:
+        self.model = model
+
+    def score(self, question: str, predicted: str, reference: str) -> float:
+        """Return 1.0 if the judge model answers YES, else 0.0."""
+        prompt = _JUDGE_PROMPT.format(
+            question=question, reference=reference, predicted=predicted
+        )
+        out = self.model.generate([{"role": "user", "content": prompt}])
+        return 1.0 if out.text.strip().lower().startswith("yes") else 0.0
+
+
+@dataclass
+class AggregateScores:
+    """Three-axis summary over a question set: quality, cost, latency."""
+
+    accuracy: float
+    mean_cost_usd: float | None
+    total_cost_usd: float | None
+    mean_latency_s: float | None
+    mean_llm_calls: float
+    n: int
+
+
+def aggregate(
+    results: Sequence[AnswerResult], correctness: Sequence[float]
+) -> AggregateScores:
+    """Reduce per-question results into the three reported axes."""
+    n = len(results)
+    if n == 0:
+        return AggregateScores(0.0, None, None, None, 0.0, 0)
+    costs = [r.usage.cost_usd for r in results if r.usage.cost_usd is not None]
+    latencies = [r.usage.latency_s for r in results if r.usage.latency_s is not None]
+    total_cost = sum(costs) if costs else None
+    mean_cost = sum(costs) / len(costs) if costs else None
+    return AggregateScores(
+        accuracy=sum(correctness) / n,
+        mean_cost_usd=mean_cost,
+        total_cost_usd=total_cost,
+        mean_latency_s=(sum(latencies) / len(latencies)) if latencies else None,
+        mean_llm_calls=sum(r.usage.num_llm_calls for r in results) / n,
+        n=n,
+    )

mteb/agentic/systems/__init__.py

@@ -0,0 +1,10 @@
+"""Answer-mode reference systems: floor and ceiling baselines."""
+
+from __future__ import annotations
+
+from mteb.agentic.systems.baselines import ClosedBookSystem, OracleContextSystem
+
+__all__ = [
+    "ClosedBookSystem",
+    "OracleContextSystem",
+]