diff --git a/.gitignore b/.gitignore index 1da3f5ad..2aacff62 100644 --- a/.gitignore +++ b/.gitignore @@ -11,7 +11,7 @@ wheels/ # tests .coverage* -tests/data/ +evaluations/data/ .pytest_cache/ .ruff_cache/ diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 9705e677..07f0c6df 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -2,8 +2,9 @@ We use the [repliqa](https://huggingface.co/datasets/ServiceNow/repliqa) dataset for the evaluation of `haiku.rag`. -You can perform your own evaluations using as example the script found at -`tests/generate_benchmark_db.py`. The evaluation flow is orchestrated with +You can perform your own evaluations with the Typer CLI in +`evaluations/benchmark.py`, for example `python -m evaluations.benchmark repliqa`. +The evaluation flow is orchestrated with [`pydantic-evals`](https://github.com/pydantic/pydantic-ai/tree/main/libs/pydantic-evals), which we leverage for dataset management, scoring, and report generation. @@ -16,13 +17,11 @@ The recall obtained is ~0.79 for matching in the top result, raising to ~0.91 fo | Embedding Model | Document in top 1 | Document in top 3 | Reranker | |---------------------------------------|-------------------|-------------------|------------------------| +| Ollama / `qwen3-embedding` | 0.81 | 0.95 | None | +| Ollama / `qwen3-embedding` | 0.91 | 0.98 | `mxbai-rerank-base-v2` | | Ollama / `mxbai-embed-large` | 0.79 | 0.91 | None | | Ollama / `mxbai-embed-large` | 0.90 | 0.95 | `mxbai-rerank-base-v2` | | Ollama / `nomic-embed-text-v1.5` | 0.74 | 0.90 | None | -| Ollama / `qwen3-embedding` | 0.81 | 0.95 | None | - ## Question/Answer evaluation @@ -38,5 +37,29 @@ determine whether the answer is correct. The obtained accuracy is as follows: | Ollama / `mxbai-embed-large` | Ollama / `qwen3:0.6b` | 0.28 | None | Note the significant degradation when very small models are used such as `qwen3:0.6b`. - + +## Wix dataset + +We also track retrieval performance on [WixQA](https://huggingface.co/datasets/Wix/WixQA), +a dataset of real customer support questions paired with curated answers from +Wix. The benchmark follows the evaluation protocol described in the +[WixQA paper](https://arxiv.org/abs/2505.08643) and gives us a view into how the +system handles conversational, product-specific support queries. + +For recall, we index the reference answer passages shipped with the dataset and +run retrieval against each user question. Each sample supplies one or more +relevant passage URIs; we count how many of those URIs land inside the top *k* +retrieved documents, divide by the number of relevant passages for that query, +and average across all queries. + +The results for recall using the `WixQA` dataset are as follows: + +| Embedding Model | Document in top 1 | Document in top 3 | Reranker | +|----------------------------|-------------------|-------------------|------------------------| +| `qwen3-embedding` | 0.36 | 0.57 | `mxbai-rerank-base-v2` | + +And for QA accuracy, + +| Embedding Model | QA Model | Accuracy | Reranker | +|----------------------------|-----------|----------|------------------------| +| `qwen3-embedding` | `gpt-oss` | 0.75 | `mxbai-rerank-base-v2` | diff --git a/evaluations/__init__.py b/evaluations/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/evaluations/benchmark.py b/evaluations/benchmark.py new file mode 100644 index 00000000..f309bd7e --- /dev/null +++ b/evaluations/benchmark.py @@ -0,0 +1,320 @@ +import asyncio +from collections.abc import Mapping +from typing import Any, cast + +import logfire +import typer +from pydantic_ai.models.openai import OpenAIChatModel +from pydantic_ai.providers.ollama import OllamaProvider +from pydantic_evals import Dataset as EvalDataset +from pydantic_evals.evaluators import IsInstance, LLMJudge +from pydantic_evals.reporting import ReportCaseFailure +from rich.console import Console +from rich.progress import Progress + +from evaluations.config import DatasetSpec, RetrievalSample +from evaluations.datasets import DATASETS +from evaluations.llm_judge import ANSWER_EQUIVALENCE_RUBRIC +from haiku.rag import logging # noqa: F401 +from haiku.rag.client import HaikuRAG +from haiku.rag.config import Config +from haiku.rag.logging import configure_cli_logging +from haiku.rag.qa import get_qa_agent + +QA_JUDGE_MODEL = "qwen3" + +logfire.configure(send_to_logfire="if-token-present", service_name="evals") +logfire.instrument_pydantic_ai() +configure_cli_logging() +console = Console() + + +async def populate_db(spec: DatasetSpec) -> None: + spec.db_path.parent.mkdir(parents=True, exist_ok=True) + corpus = spec.document_loader() + if spec.document_limit is not None: + corpus = corpus.select(range(min(spec.document_limit, len(corpus)))) + + with Progress() as progress: + task = progress.add_task("[green]Populating database...", total=len(corpus)) + async with HaikuRAG(spec.db_path) as rag: + for doc in corpus: + doc_mapping = cast(Mapping[str, Any], doc) + payload = spec.document_mapper(doc_mapping) + if payload is None: + progress.advance(task) + continue + + existing = await rag.get_document_by_uri(payload.uri) + if existing is not None: + assert existing.id + chunks = await rag.chunk_repository.get_by_document_id(existing.id) + if chunks: + progress.advance(task) + continue + await rag.document_repository.delete(existing.id) + + await rag.create_document( + content=payload.content, + uri=payload.uri, + title=payload.title, + metadata=payload.metadata, + ) + progress.advance(task) + rag.store.vacuum() + + +def _is_relevant_match(retrieved_uri: str | None, sample: RetrievalSample) -> bool: + return retrieved_uri is not None and retrieved_uri in sample.expected_uris + + +async def run_retrieval_benchmark(spec: DatasetSpec) -> dict[str, float] | None: + if spec.retrieval_loader is None or spec.retrieval_mapper is None: + console.print("Skipping retrieval benchmark; no retrieval config.") + return None + + corpus = spec.retrieval_loader() + + recall_totals = { + 1: 0.0, + 3: 0.0, + 5: 0.0, + } + total_queries = 0 + + with Progress() as progress: + task = progress.add_task( + "[blue]Running retrieval benchmark...", total=len(corpus) + ) + async with HaikuRAG(spec.db_path) as rag: + for doc in corpus: + doc_mapping = cast(Mapping[str, Any], doc) + sample = spec.retrieval_mapper(doc_mapping) + if sample is None or sample.skip: + progress.advance(task) + continue + + matches = await rag.search(query=sample.question, limit=5) + if not matches: + progress.advance(task) + continue + + total_queries += 1 + + retrieved_uris: list[str] = [] + for chunk, _ in matches: + if chunk.document_id is None: + continue + retrieved_doc = await rag.get_document_by_id(chunk.document_id) + if retrieved_doc and retrieved_doc.uri: + retrieved_uris.append(retrieved_doc.uri) + + # Compute per-query recall@K by counting how many relevant + # documents are retrieved within the first K results and + # averaging these fractions across all queries. + for cutoff in (1, 3, 5): + top_k = set(retrieved_uris[:cutoff]) + relevant = set(sample.expected_uris) + if relevant: + matched = len(top_k & relevant) + recall_totals[cutoff] += matched / len(relevant) + + progress.advance(task) + + if total_queries == 0: + console.print("No retrieval cases to evaluate.") + return None + + recall_at_1 = recall_totals[1] / total_queries + recall_at_3 = recall_totals[3] / total_queries + recall_at_5 = recall_totals[5] / total_queries + + console.print("\n=== Retrieval Benchmark Results ===", style="bold cyan") + console.print(f"Total queries: {total_queries}") + console.print(f"Recall@1: {recall_at_1:.4f}") + console.print(f"Recall@3: {recall_at_3:.4f}") + console.print(f"Recall@5: {recall_at_5:.4f}") + + return { + "recall@1": recall_at_1, + "recall@3": recall_at_3, + "recall@5": recall_at_5, + } + + +async def run_qa_benchmark( + spec: DatasetSpec, qa_limit: int | None = None +) -> ReportCaseFailure[str, str, dict[str, str]] | None: + corpus = spec.qa_loader() + if qa_limit is not None: + corpus = corpus.select(range(min(qa_limit, len(corpus)))) + + cases = [ + spec.qa_case_builder(index, cast(Mapping[str, Any], doc)) + for index, doc in enumerate(corpus, start=1) + ] + + judge_model = OpenAIChatModel( + model_name=QA_JUDGE_MODEL, + provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"), + ) + + evaluation_dataset = EvalDataset[str, str, dict[str, str]]( + cases=cases, + evaluators=[ + IsInstance(type_name="str"), + LLMJudge( + rubric=ANSWER_EQUIVALENCE_RUBRIC, + include_input=True, + include_expected_output=True, + model=judge_model, + assertion={ + "evaluation_name": "answer_equivalent", + "include_reason": True, + }, + ), + ], + ) + + total_processed = 0 + passing_cases = 0 + failures: list[ReportCaseFailure[str, str, dict[str, str]]] = [] + + with Progress(console=console) as progress: + qa_task = progress.add_task( + "[yellow]Evaluating QA cases...", + total=len(evaluation_dataset.cases), + ) + + async with HaikuRAG(spec.db_path) as rag: + qa = get_qa_agent(rag) + + async def answer_question(question: str) -> str: + return await qa.answer(question) + + for case in evaluation_dataset.cases: + progress.console.print(f"\n[bold]Evaluating case:[/bold] {case.name}") + + single_case_dataset = EvalDataset[str, str, dict[str, str]]( + cases=[case], + evaluators=evaluation_dataset.evaluators, + ) + + report = await single_case_dataset.evaluate( + answer_question, + name="qa_answer", + max_concurrency=1, + progress=False, + ) + + total_processed += 1 + + if report.cases: + result_case = report.cases[0] + + equivalence = result_case.assertions.get("answer_equivalent") + progress.console.print(f"Question: {result_case.inputs}") + progress.console.print(f"Expected: {result_case.expected_output}") + progress.console.print(f"Generated: {result_case.output}") + if equivalence is not None: + progress.console.print( + f"Equivalent: {equivalence.value}" + + (f" — {equivalence.reason}" if equivalence.reason else "") + ) + if equivalence.value: + passing_cases += 1 + + progress.console.print("") + + if report.failures: + failures.extend(report.failures) + failure = report.failures[0] + progress.console.print( + "[red]Failure encountered during case evaluation:[/red]" + ) + progress.console.print(f"Question: {failure.inputs}") + progress.console.print(f"Error: {failure.error_message}") + progress.console.print("") + + progress.console.print( + f"[green]Accuracy: {(passing_cases / total_processed):.4f} " + f"{passing_cases}/{total_processed}[/green]" + ) + progress.advance(qa_task) + + total_cases = total_processed + accuracy = passing_cases / total_cases if total_cases > 0 else 0 + + console.print("\n=== QA Benchmark Results ===", style="bold cyan") + console.print(f"Total questions: {total_cases}") + console.print(f"Correct answers: {passing_cases}") + console.print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)") + + if failures: + console.print("[red]\nSummary of failures:[/red]") + for failure in failures: + console.print(f"Case: {failure.name}") + console.print(f"Question: {failure.inputs}") + console.print(f"Error: {failure.error_message}") + console.print("") + + return failures[0] if failures else None + + +async def evaluate_dataset( + spec: DatasetSpec, + skip_db: bool, + skip_retrieval: bool, + skip_qa: bool, + qa_limit: int | None, +) -> None: + if not skip_db: + console.print(f"Using dataset: {spec.key}", style="bold magenta") + await populate_db(spec) + + if not skip_retrieval: + console.print("Running retrieval benchmarks...", style="bold blue") + await run_retrieval_benchmark(spec) + + if not skip_qa: + console.print("\nRunning QA benchmarks...", style="bold yellow") + await run_qa_benchmark(spec, qa_limit=qa_limit) + + +app = typer.Typer(help="Run retrieval and QA benchmarks for configured datasets.") + + +@app.command() +def run( + dataset: str = typer.Argument(..., help="Dataset key to evaluate."), + skip_db: bool = typer.Option( + False, "--skip-db", help="Skip updateing the evaluation db." + ), + skip_retrieval: bool = typer.Option( + False, "--skip-retrieval", help="Skip retrieval benchmark." + ), + skip_qa: bool = typer.Option(False, "--skip-qa", help="Skip QA benchmark."), + qa_limit: int | None = typer.Option( + None, "--qa-limit", help="Limit number of QA cases." + ), +) -> None: + spec = DATASETS.get(dataset.lower()) + if spec is None: + valid_datasets = ", ".join(sorted(DATASETS)) + raise typer.BadParameter( + f"Unknown dataset '{dataset}'. Choose from: {valid_datasets}" + ) + + asyncio.run( + evaluate_dataset( + spec=spec, + skip_db=skip_db, + skip_retrieval=skip_retrieval, + skip_qa=skip_qa, + qa_limit=qa_limit, + ) + ) + + +if __name__ == "__main__": + app() diff --git a/evaluations/config.py b/evaluations/config.py new file mode 100644 index 00000000..8e230dda --- /dev/null +++ b/evaluations/config.py @@ -0,0 +1,46 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from datasets import Dataset +from pydantic_evals import Case + + +@dataclass +class DocumentPayload: + uri: str + content: str + title: str | None = None + metadata: dict[str, Any] | None = None + + +@dataclass +class RetrievalSample: + question: str + expected_uris: tuple[str, ...] + skip: bool = False + + +DocumentLoader = Callable[[], Dataset] +DocumentMapper = Callable[[Mapping[str, Any]], DocumentPayload | None] +RetrievalLoader = Callable[[], Dataset] +RetrievalMapper = Callable[[Mapping[str, Any]], RetrievalSample | None] +CaseBuilder = Callable[[int, Mapping[str, Any]], Case[str, str, dict[str, str]]] + + +@dataclass +class DatasetSpec: + key: str + db_filename: str + document_loader: DocumentLoader + document_mapper: DocumentMapper + qa_loader: DocumentLoader + qa_case_builder: CaseBuilder + retrieval_loader: RetrievalLoader | None = None + retrieval_mapper: RetrievalMapper | None = None + document_limit: int | None = None + + @property + def db_path(self) -> Path: + return Path(__file__).parent / "data" / self.db_filename diff --git a/evaluations/datasets/__init__.py b/evaluations/datasets/__init__.py new file mode 100644 index 00000000..94a00c34 --- /dev/null +++ b/evaluations/datasets/__init__.py @@ -0,0 +1,8 @@ +from evaluations.config import DatasetSpec + +from .repliqa import REPLIQ_SPEC +from .wix import WIX_SPEC + +DATASETS: dict[str, DatasetSpec] = {spec.key: spec for spec in (REPLIQ_SPEC, WIX_SPEC)} + +__all__ = ["DATASETS"] diff --git a/evaluations/datasets/repliqa.py b/evaluations/datasets/repliqa.py new file mode 100644 index 00000000..2fc21d0f --- /dev/null +++ b/evaluations/datasets/repliqa.py @@ -0,0 +1,58 @@ +from collections.abc import Mapping +from typing import Any, cast + +from datasets import Dataset, DatasetDict, load_dataset +from pydantic_evals import Case + +from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample + + +def load_repliqa_corpus() -> Dataset: + dataset_dict = cast(DatasetDict, load_dataset("ServiceNow/repliqa")) + dataset = cast(Dataset, dataset_dict["repliqa_3"]) + return dataset.filter(lambda doc: doc["document_topic"] == "News Stories") + + +def map_repliqa_document(doc: Mapping[str, Any]) -> DocumentPayload: + return DocumentPayload( + uri=str(doc["document_id"]), + content=doc["document_extracted"], + ) + + +def map_repliqa_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None: + expected_answer = doc["answer"] + if expected_answer == "The answer is not found in the document.": + return None + return RetrievalSample( + question=doc["question"], + expected_uris=(str(doc["document_id"]),), + ) + + +def build_repliqa_case( + index: int, doc: Mapping[str, Any] +) -> Case[str, str, dict[str, str]]: + document_id = doc["document_id"] + case_name = f"{index}_{document_id}" if document_id is not None else f"case_{index}" + return Case( + name=case_name, + inputs=doc["question"], + expected_output=doc["answer"], + metadata={ + "document_id": str(document_id), + "case_index": str(index), + }, + ) + + +REPLIQ_SPEC = DatasetSpec( + key="repliqa", + db_filename="repliqa.lancedb", + document_loader=load_repliqa_corpus, + document_mapper=map_repliqa_document, + qa_loader=load_repliqa_corpus, + qa_case_builder=build_repliqa_case, + retrieval_loader=load_repliqa_corpus, + retrieval_mapper=map_repliqa_retrieval, +) diff --git a/evaluations/datasets/wix.py b/evaluations/datasets/wix.py new file mode 100644 index 00000000..112ba171 --- /dev/null +++ b/evaluations/datasets/wix.py @@ -0,0 +1,81 @@ +import json +from collections.abc import Iterable, Mapping +from typing import Any, cast + +from datasets import Dataset, DatasetDict, load_dataset +from pydantic_evals import Case + +from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample + + +def load_wix_corpus() -> Dataset: + dataset_dict = cast(DatasetDict, load_dataset("Wix/WixQA", "wix_kb_corpus")) + return cast(Dataset, dataset_dict["train"]) + + +def map_wix_document(doc: Mapping[str, Any]) -> DocumentPayload: + article_id = doc.get("id") + url = doc.get("url") + uri = str(article_id) if article_id is not None else str(url) + + metadata: dict[str, str] = {} + if article_id is not None: + metadata["article_id"] = str(article_id) + if url: + metadata["url"] = str(url) + + return DocumentPayload( + uri=uri, + content=doc["contents"], + title=doc.get("title"), + metadata=metadata or None, + ) + + +def load_wix_qa() -> Dataset: + dataset_dict = cast(DatasetDict, load_dataset("Wix/WixQA", "wixqa_expertwritten")) + return cast(Dataset, dataset_dict["train"]) + + +def map_wix_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None: + article_ids: Iterable[int | str] | None = doc.get("article_ids") + if not article_ids: + return None + + expected_uris = tuple(str(article_id) for article_id in article_ids) + return RetrievalSample( + question=doc["question"], + expected_uris=expected_uris, + ) + + +def build_wix_case( + index: int, doc: Mapping[str, Any] +) -> Case[str, str, dict[str, str]]: + article_ids = tuple(str(article_id) for article_id in doc.get("article_ids") or []) + joined_ids = "-".join(article_ids) + case_name = f"{index}_{joined_ids}" if joined_ids else f"case_{index}" + + metadata = { + "case_index": str(index), + "document_ids": json.dumps(article_ids), + } + + return Case( + name=case_name, + inputs=doc["question"], + expected_output=doc["answer"], + metadata=metadata, + ) + + +WIX_SPEC = DatasetSpec( + key="wix", + db_filename="wix.lancedb", + document_loader=load_wix_corpus, + document_mapper=map_wix_document, + qa_loader=load_wix_qa, + qa_case_builder=build_wix_case, + retrieval_loader=load_wix_qa, + retrieval_mapper=map_wix_retrieval, +) diff --git a/tests/llm_judge.py b/evaluations/llm_judge.py similarity index 97% rename from tests/llm_judge.py rename to evaluations/llm_judge.py index 435e517b..7dcfccac 100644 --- a/tests/llm_judge.py +++ b/evaluations/llm_judge.py @@ -37,7 +37,7 @@ class LLMJudgeResponseSchema(BaseModel): class LLMJudge: """LLM-as-judge for evaluating answer equivalence using Pydantic AI.""" - def __init__(self, model: str = "qwen3"): + def __init__(self, model: str = "gpt-oss"): # Create Ollama model ollama_model = OpenAIChatModel( model_name=model, @@ -49,6 +49,7 @@ class LLMJudge: model=ollama_model, output_type=LLMJudgeResponseSchema, system_prompt=ANSWER_EQUIVALENCE_RUBRIC, + retries=3, ) async def judge_answers( diff --git a/pyproject.toml b/pyproject.toml index 67ddddc5..09727f05 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,9 @@ haiku-rag = "haiku.rag.cli:cli" requires = ["hatchling"] build-backend = "hatchling.build" +[tool.hatch.build] +exclude = ["/docs", "/evaluations", "/tests", "/.github"] + [tool.hatch.build.targets.wheel] packages = ["src/haiku"] diff --git a/src/haiku/rag/embeddings/ollama.py b/src/haiku/rag/embeddings/ollama.py index 2dbd8ea4..a7303ea7 100644 --- a/src/haiku/rag/embeddings/ollama.py +++ b/src/haiku/rag/embeddings/ollama.py @@ -7,6 +7,8 @@ from haiku.rag.embeddings.base import EmbedderBase class Embedder(EmbedderBase): async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]: client = AsyncOpenAI(base_url=f"{Config.OLLAMA_BASE_URL}/v1", api_key="dummy") + if not text: + return [] response = await client.embeddings.create( model=self._model, input=text, diff --git a/src/haiku/rag/embeddings/openai.py b/src/haiku/rag/embeddings/openai.py index 14d9129a..5b0ea2ff 100644 --- a/src/haiku/rag/embeddings/openai.py +++ b/src/haiku/rag/embeddings/openai.py @@ -6,6 +6,8 @@ from haiku.rag.embeddings.base import EmbedderBase class Embedder(EmbedderBase): async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]: client = AsyncOpenAI() + if not text: + return [] response = await client.embeddings.create( model=self._model, input=text, diff --git a/src/haiku/rag/embeddings/vllm.py b/src/haiku/rag/embeddings/vllm.py index cae33398..2d2f77bd 100644 --- a/src/haiku/rag/embeddings/vllm.py +++ b/src/haiku/rag/embeddings/vllm.py @@ -9,6 +9,8 @@ class Embedder(EmbedderBase): client = AsyncOpenAI( base_url=f"{Config.VLLM_EMBEDDINGS_BASE_URL}/v1", api_key="dummy" ) + if not text: + return [] response = await client.embeddings.create( model=self._model, input=text, diff --git a/src/haiku/rag/embeddings/voyageai.py b/src/haiku/rag/embeddings/voyageai.py index 4c0e7e89..60c8c6b3 100644 --- a/src/haiku/rag/embeddings/voyageai.py +++ b/src/haiku/rag/embeddings/voyageai.py @@ -6,6 +6,8 @@ try: class Embedder(EmbedderBase): async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]: client = Client() + if not text: + return [] if isinstance(text, str): res = client.embed([text], model=self._model, output_dtype="float") return res.embeddings[0] # type: ignore[return-value] diff --git a/src/haiku/rag/reranking/__init__.py b/src/haiku/rag/reranking/__init__.py index f63453c6..f4753d50 100644 --- a/src/haiku/rag/reranking/__init__.py +++ b/src/haiku/rag/reranking/__init__.py @@ -1,3 +1,5 @@ +import os + from haiku.rag.config import Config from haiku.rag.reranking.base import RerankerBase @@ -17,6 +19,7 @@ def get_reranker() -> RerankerBase | None: try: from haiku.rag.reranking.mxbai import MxBAIReranker + os.environ["TOKENIZERS_PARALLELISM"] = "true" _reranker = MxBAIReranker() return _reranker except ImportError: diff --git a/tests/generate_benchmark_db.py b/tests/generate_benchmark_db.py deleted file mode 100644 index f1209ac5..00000000 --- a/tests/generate_benchmark_db.py +++ /dev/null @@ -1,258 +0,0 @@ -import asyncio -from pathlib import Path - -import logfire -from datasets import Dataset, load_dataset -from llm_judge import ANSWER_EQUIVALENCE_RUBRIC -from pydantic_ai.models.openai import OpenAIChatModel -from pydantic_ai.providers.ollama import OllamaProvider -from pydantic_evals import Case -from pydantic_evals import Dataset as EvalDataset -from pydantic_evals.evaluators import IsInstance, LLMJudge -from pydantic_evals.reporting import ReportCaseFailure -from rich.console import Console -from rich.progress import Progress - -from haiku.rag import logging # noqa -from haiku.rag.client import HaikuRAG -from haiku.rag.config import Config -from haiku.rag.logging import configure_cli_logging -from haiku.rag.qa import get_qa_agent - -logfire.configure(send_to_logfire="if-token-present", service_name="evals") -logfire.instrument_pydantic_ai() -configure_cli_logging() -console = Console() - -QA_JUDGE_MODEL: str = "qwen3" -db_path = Path(__file__).parent / "data" / "benchmark.lancedb" - - -async def populate_db(): - ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"] # type: ignore - corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories") - - with Progress() as progress: - task = progress.add_task("[green]Populating database...", total=len(corpus)) - - async with HaikuRAG(db_path) as rag: - for doc in corpus: - uri = doc["document_id"] # type: ignore - existing_doc = await rag.get_document_by_uri(uri) - if existing_doc is not None: - progress.advance(task) - continue - - await rag.create_document( - content=doc["document_extracted"], # type: ignore - uri=uri, - ) - progress.advance(task) - rag.store.vacuum() - - -async def run_match_benchmark(): - ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"] # type: ignore - corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories") - - correct_at_1 = 0 - correct_at_2 = 0 - correct_at_3 = 0 - total_queries = 0 - - with Progress() as progress: - task = progress.add_task( - "[blue]Running retrieval benchmark...", total=len(corpus) - ) - - async with HaikuRAG(db_path) as rag: - for doc in corpus: - doc_id = doc["document_id"] # type: ignore - expected_answer = doc["answer"] # type: ignore - if expected_answer == "The answer is not found in the document.": - progress.advance(task) - continue - matches = await rag.search( - query=doc["question"], # type: ignore - limit=3, - ) - - total_queries += 1 - - # Check position of correct document in results - for position, (chunk, _) in enumerate(matches): - assert chunk.document_id is not None, ( - "Chunk document_id should not be None" - ) - retrieved = await rag.get_document_by_id(chunk.document_id) - if retrieved and retrieved.uri == doc_id: - if position == 0: # First position - correct_at_1 += 1 - correct_at_2 += 1 - correct_at_3 += 1 - elif position == 1: # Second position - correct_at_2 += 1 - correct_at_3 += 1 - elif position == 2: # Third position - correct_at_3 += 1 - break - - progress.advance(task) - - # Calculate recall metrics - recall_at_1 = correct_at_1 / total_queries - recall_at_2 = correct_at_2 / total_queries - recall_at_3 = correct_at_3 / total_queries - - console.print("\n=== Retrieval Benchmark Results ===", style="bold cyan") - console.print(f"Total queries: {total_queries}") - console.print(f"Recall@1: {recall_at_1:.4f}") - console.print(f"Recall@2: {recall_at_2:.4f}") - console.print(f"Recall@3: {recall_at_3:.4f}") - - return {"recall@1": recall_at_1, "recall@2": recall_at_2, "recall@3": recall_at_3} - - -async def run_qa_benchmark(k: int | None = None): - """Run QA benchmarking on the corpus.""" - ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"] # type: ignore - corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories") - - if k is not None: - corpus = corpus.select(range(min(k, len(corpus)))) - - cases: list[Case[str, str, dict[str, str]]] = [] - for index, doc in enumerate(corpus, start=1): - question = doc["question"] # type: ignore[index] - expected_answer = doc["answer"] # type: ignore[index] - doc_id = doc["document_id"] # type: ignore[index] - case_name = f"{index}_{doc_id}" if doc_id is not None else f"case_{index}" - - cases.append( - Case( - name=case_name, - inputs=question, - expected_output=expected_answer, - metadata={ - "document_id": str(doc_id), - "case_index": str(index), - }, - ) - ) - - judge_model = OpenAIChatModel( - model_name=QA_JUDGE_MODEL, - provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"), - ) - - evaluation_dataset = EvalDataset[str, str, dict[str, str]]( - cases=cases, - evaluators=[ - IsInstance(type_name="str"), - LLMJudge( - rubric=ANSWER_EQUIVALENCE_RUBRIC, - include_input=True, - include_expected_output=True, - model=judge_model, - assertion={ - "evaluation_name": "answer_equivalent", - "include_reason": True, - }, - ), - ], - ) - - total_processed = 0 - passing_cases = 0 - failures: list[ReportCaseFailure[str, str, dict[str, str]]] = [] - - with Progress(console=console) as progress: - qa_task = progress.add_task( - "[yellow]Evaluating QA cases...", - total=len(evaluation_dataset.cases), - ) - - async with HaikuRAG(db_path) as rag: - qa = get_qa_agent(rag) - - async def answer_question(question: str) -> str: - return await qa.answer(question) - - for case in evaluation_dataset.cases: - progress.console.print(f"\n[bold]Evaluating case:[/bold] {case.name}") - - single_case_dataset = EvalDataset[str, str, dict[str, str]]( - cases=[case], - evaluators=evaluation_dataset.evaluators, - ) - - report = await single_case_dataset.evaluate( - answer_question, - name="qa_answer", - max_concurrency=1, - progress=False, - ) - - total_processed += 1 - - if report.cases: - result_case = report.cases[0] - - equivalence = result_case.assertions.get("answer_equivalent") - progress.console.print(f"Question: {result_case.inputs}") - progress.console.print(f"Expected: {result_case.expected_output}") - progress.console.print(f"Generated: {result_case.output}") - if equivalence is not None: - progress.console.print( - f"Equivalent: {equivalence.value}" - + (f" — {equivalence.reason}" if equivalence.reason else "") - ) - if equivalence.value: - passing_cases += 1 - - progress.console.print("") - - if report.failures: - failures.extend(report.failures) - failure = report.failures[0] - progress.console.print( - "[red]Failure encountered during case evaluation:[/red]" - ) - progress.console.print(f"Question: {failure.inputs}") - progress.console.print(f"Error: {failure.error_message}") - progress.console.print("") - - progress.console.print( - f"[green]Accuracy: {(passing_cases / total_processed):.4f} " - f"{passing_cases}/{total_processed}[/green]" - ) - progress.advance(qa_task) - total_cases = total_processed - accuracy = passing_cases / total_cases if total_cases > 0 else 0 - - console.print("\n=== QA Benchmark Results ===", style="bold cyan") - console.print(f"Total questions: {total_cases}") - console.print(f"Correct answers: {passing_cases}") - console.print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)") - - if failures: - console.print("[red]\nSummary of failures:[/red]") - for failure in failures: - console.print(f"Case: {failure.name}") - console.print(f"Question: {failure.inputs}") - console.print(f"Error: {failure.error_message}") - console.print("") - - -async def main(): - await populate_db() - - console.print("Running retrieval benchmarks...", style="bold blue") - await run_match_benchmark() - - console.print("\nRunning QA benchmarks...", style="bold yellow") - await run_qa_benchmark() - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/tests/test_qa.py b/tests/test_qa.py index 2fa2d417..1f2f4022 100644 --- a/tests/test_qa.py +++ b/tests/test_qa.py @@ -1,12 +1,11 @@ import pytest from datasets import Dataset +from evaluations.llm_judge import LLMJudge from haiku.rag.client import HaikuRAG from haiku.rag.config import Config from haiku.rag.qa.agent import QuestionAnswerAgent -from .llm_judge import LLMJudge - OPENAI_AVAILABLE = bool(Config.OPENAI_API_KEY) ANTHROPIC_AVAILABLE = bool(Config.ANTHROPIC_API_KEY) VLLM_QA_AVAILABLE = bool(Config.VLLM_QA_BASE_URL)