Use cassette to record and reply llms in qa tests

This commit is contained in:
Yiorgis Gozadinos 2025-12-26 15:41:11 +02:00
parent 1f654917b8
commit 4a82e47d03
No known key found for this signature in database
8 changed files with 3158 additions and 14 deletions

View file

@ -75,6 +75,7 @@ dev = [
"pytest>=9.0.2", "pytest>=9.0.2",
"pytest-asyncio>=1.3.0", "pytest-asyncio>=1.3.0",
"pytest-cov>=7.0.0", "pytest-cov>=7.0.0",
"pytest-recording>=0.13.2",
"ruff>=0.14.8", "ruff>=0.14.8",
] ]

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View file

@ -1,6 +1,8 @@
import logging
import os import os
import tempfile import tempfile
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Any
# Prevent tests from loading user's local haiku.rag.yaml by setting env var # Prevent tests from loading user's local haiku.rag.yaml by setting env var
# to an empty config file BEFORE any haiku.rag imports. # to an empty config file BEFORE any haiku.rag imports.
@ -10,10 +12,17 @@ _test_config_path = Path(_test_config_dir) / "test-defaults.yaml"
_test_config_path.write_text("{}") # Empty YAML = use all defaults _test_config_path.write_text("{}") # Empty YAML = use all defaults
os.environ["HAIKU_RAG_CONFIG_PATH"] = str(_test_config_path) os.environ["HAIKU_RAG_CONFIG_PATH"] = str(_test_config_path)
import pydantic_ai.models # noqa: E402
import pytest # noqa: E402 import pytest # noqa: E402
import yaml # noqa: E402 import yaml # noqa: E402
from datasets import Dataset, load_dataset, load_from_disk # noqa: E402 from datasets import Dataset, load_dataset, load_from_disk # noqa: E402
if TYPE_CHECKING:
from vcr import VCR
pydantic_ai.models.ALLOW_MODEL_REQUESTS = False
logging.getLogger("vcr.cassette").setLevel(logging.WARNING)
@pytest.fixture(scope="session") @pytest.fixture(scope="session")
def qa_corpus() -> Dataset: def qa_corpus() -> Dataset:
@ -71,3 +80,33 @@ def temp_yaml_config(tmp_path, monkeypatch):
monkeypatch.setenv("HAIKU_RAG_CONFIG_PATH", str(config_file)) monkeypatch.setenv("HAIKU_RAG_CONFIG_PATH", str(config_file))
yield config_file yield config_file
@pytest.fixture
def allow_model_requests():
with pydantic_ai.models.override_allow_model_requests(True):
yield
@pytest.fixture(autouse=True)
def set_mock_api_keys(monkeypatch):
"""Set mock API keys for providers that require them during initialization."""
if not os.getenv("OPENAI_API_KEY"):
monkeypatch.setenv("OPENAI_API_KEY", "sk-mock-key-for-vcr-playback")
if not os.getenv("ANTHROPIC_API_KEY"):
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-mock-key-for-vcr-playback")
def pytest_recording_configure(config: Any, vcr: "VCR"):
from . import json_body_serializer
vcr.register_serializer("yaml", json_body_serializer)
@pytest.fixture(scope="module")
def vcr_config():
return {
"ignore_localhost": False,
"filter_headers": ["authorization", "x-api-key"],
"decode_compressed_response": True,
}

View file

@ -0,0 +1,110 @@
# Adapted from pydantic-ai: https://github.com/pydantic/pydantic-ai/blob/main/tests/json_body_serializer.py
# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false
import json
import urllib.parse
from typing import TYPE_CHECKING, Any
import yaml
if TYPE_CHECKING:
from yaml import Dumper, SafeLoader
else:
try:
from yaml import CDumper as Dumper
from yaml import CSafeLoader as SafeLoader
except ImportError:
from yaml import Dumper, SafeLoader
FILTERED_HEADER_PREFIXES = ["anthropic-", "cf-", "x-"]
FILTERED_HEADERS = {
"authorization",
"date",
"request-id",
"server",
"user-agent",
"via",
"set-cookie",
"api-key",
}
ALLOWED_HEADER_PREFIXES: set[str] = set()
ALLOWED_HEADERS: set[str] = set()
ALLOWED_LOCALHOST_PATHS = ["/api/", "/v1/"]
class LiteralDumper(Dumper):
pass
def str_presenter(dumper: Dumper, data: str):
if "\n" in data:
return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|")
return dumper.represent_scalar("tag:yaml.org,2002:str", data)
LiteralDumper.add_representer(str, str_presenter)
def _is_filtered_localhost(uri: str) -> bool:
parsed = urllib.parse.urlparse(uri)
if parsed.hostname not in ("localhost", "127.0.0.1"):
return False
return not any(parsed.path.startswith(p) for p in ALLOWED_LOCALHOST_PATHS)
def deserialize(cassette_string: str):
cassette_dict = yaml.load(cassette_string, Loader=SafeLoader)
for interaction in cassette_dict["interactions"]:
for kind, data in interaction.items():
parsed_body = data.pop("parsed_body", None)
if parsed_body is not None:
dumped_body = json.dumps(parsed_body)
data["body"] = (
{"string": dumped_body} if kind == "response" else dumped_body
)
return cassette_dict
def serialize(cassette_dict: Any):
cassette_dict["interactions"] = [
i
for i in cassette_dict["interactions"]
if not _is_filtered_localhost(i["request"]["uri"])
]
for interaction in cassette_dict["interactions"]:
for _kind, data in interaction.items():
headers: dict[str, list[str]] = data.get("headers", {})
headers = {k.lower(): v for k, v in headers.items()}
headers = {k: v for k, v in headers.items() if k not in FILTERED_HEADERS}
headers = {
k: v
for k, v in headers.items()
if not any(k.startswith(prefix) for prefix in FILTERED_HEADER_PREFIXES)
or k in ALLOWED_HEADERS
or any(k.startswith(prefix) for prefix in ALLOWED_HEADER_PREFIXES)
}
data["headers"] = headers
content_type = headers.get("content-type", [])
if any(
isinstance(header, str) and header.startswith("application/json")
for header in content_type
):
body = data.get("body", None)
assert body is not None, data
if isinstance(body, dict):
body = body.get("string")
if body:
data["parsed_body"] = json.loads(body)
if "access_token" in data["parsed_body"]:
data["parsed_body"]["access_token"] = "scrubbed"
del data["body"]
if content_type == ["application/x-www-form-urlencoded"]:
query_params = urllib.parse.parse_qs(data["body"])
for key in ["client_id", "client_secret", "refresh_token"]:
if key in query_params:
query_params[key] = ["scrubbed"]
data["body"] = urllib.parse.urlencode(query_params)
return yaml.dump(cassette_dict, Dumper=LiteralDumper, allow_unicode=True, width=120)

View file

@ -1,4 +1,4 @@
import os from pathlib import Path
import pytest import pytest
from datasets import Dataset from datasets import Dataset
@ -8,13 +8,15 @@ from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import ModelConfig from haiku.rag.config.models import ModelConfig
from haiku.rag.qa.agent import QuestionAnswerAgent from haiku.rag.qa.agent import QuestionAnswerAgent
OPENAI_AVAILABLE = bool(os.getenv("OPENAI_API_KEY"))
ANTHROPIC_AVAILABLE = bool(os.getenv("ANTHROPIC_API_KEY")) @pytest.fixture(scope="module")
def vcr_cassette_dir():
return str(Path(__file__).parent / "cassettes" / "test_qa")
@pytest.mark.asyncio @pytest.mark.vcr()
async def test_qa_ollama(qa_corpus: Dataset, temp_db_path): async def test_qa_ollama(allow_model_requests, qa_corpus: Dataset, temp_db_path):
"""Test Ollama QA with LLM judge.""" """Test Ollama QA with LLM judge (VCR recorded)."""
client = HaikuRAG(temp_db_path, create=True) client = HaikuRAG(temp_db_path, create=True)
qa = QuestionAnswerAgent( qa = QuestionAnswerAgent(
client, ModelConfig(provider="ollama", name="gpt-oss", enable_thinking=True) client, ModelConfig(provider="ollama", name="gpt-oss", enable_thinking=True)
@ -37,10 +39,9 @@ async def test_qa_ollama(qa_corpus: Dataset, temp_db_path):
) )
@pytest.mark.asyncio @pytest.mark.vcr()
@pytest.mark.skipif(not OPENAI_AVAILABLE, reason="OpenAI not available") async def test_qa_openai(allow_model_requests, qa_corpus: Dataset, temp_db_path):
async def test_qa_openai(qa_corpus: Dataset, temp_db_path): """Test OpenAI QA with LLM judge (VCR recorded)."""
"""Test OpenAI QA with LLM judge."""
client = HaikuRAG(temp_db_path, create=True) client = HaikuRAG(temp_db_path, create=True)
qa = QuestionAnswerAgent(client, ModelConfig(provider="openai", name="gpt-4o-mini")) qa = QuestionAnswerAgent(client, ModelConfig(provider="openai", name="gpt-4o-mini"))
llm_judge = LLMJudge() llm_judge = LLMJudge()
@ -61,10 +62,9 @@ async def test_qa_openai(qa_corpus: Dataset, temp_db_path):
) )
@pytest.mark.asyncio @pytest.mark.vcr()
@pytest.mark.skipif(not ANTHROPIC_AVAILABLE, reason="Anthropic not available") async def test_qa_anthropic(allow_model_requests, qa_corpus: Dataset, temp_db_path):
async def test_qa_anthropic(qa_corpus: Dataset, temp_db_path): """Test Anthropic QA with LLM judge (VCR recorded)."""
"""Test Anthropic QA with LLM judge."""
client = HaikuRAG(temp_db_path, create=True) client = HaikuRAG(temp_db_path, create=True)
qa = QuestionAnswerAgent( qa = QuestionAnswerAgent(
client, ModelConfig(provider="anthropic", name="claude-3-5-haiku-20241022") client, ModelConfig(provider="anthropic", name="claude-3-5-haiku-20241022")

28
uv.lock
View file

@ -1270,6 +1270,7 @@ dev = [
{ name = "pytest" }, { name = "pytest" },
{ name = "pytest-asyncio" }, { name = "pytest-asyncio" },
{ name = "pytest-cov" }, { name = "pytest-cov" },
{ name = "pytest-recording" },
{ name = "ruff" }, { name = "ruff" },
] ]
@ -1291,6 +1292,7 @@ dev = [
{ name = "pytest", specifier = ">=9.0.2" }, { name = "pytest", specifier = ">=9.0.2" },
{ name = "pytest-asyncio", specifier = ">=1.3.0" }, { name = "pytest-asyncio", specifier = ">=1.3.0" },
{ name = "pytest-cov", specifier = ">=7.0.0" }, { name = "pytest-cov", specifier = ">=7.0.0" },
{ name = "pytest-recording", specifier = ">=0.13.2" },
{ name = "ruff", specifier = ">=0.14.8" }, { name = "ruff", specifier = ">=0.14.8" },
] ]
@ -3765,6 +3767,19 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl", hash = "sha256:3b8e9558b16cc1479da72058bdecf8073661c7f57f7d3c5f22a1c23507f2d861", size = 22424, upload-time = "2025-09-09T10:57:00.695Z" }, { url = "https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl", hash = "sha256:3b8e9558b16cc1479da72058bdecf8073661c7f57f7d3c5f22a1c23507f2d861", size = 22424, upload-time = "2025-09-09T10:57:00.695Z" },
] ]
[[package]]
name = "pytest-recording"
version = "0.13.4"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pytest" },
{ name = "vcrpy" },
]
sdist = { url = "https://files.pythonhosted.org/packages/32/9c/f4027c5f1693847b06d11caf4b4f6bb09f22c1581ada4663877ec166b8c6/pytest_recording-0.13.4.tar.gz", hash = "sha256:568d64b2a85992eec4ae0a419c855d5fd96782c5fb016784d86f18053792768c", size = 26576, upload-time = "2025-05-08T10:41:11.231Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/42/c2/ce34735972cc42d912173e79f200fe66530225190c06655c5632a9d88f1e/pytest_recording-0.13.4-py3-none-any.whl", hash = "sha256:ad49a434b51b1c4f78e85b1e6b74fdcc2a0a581ca16e52c798c6ace971f7f439", size = 13723, upload-time = "2025-05-08T10:41:09.684Z" },
]
[[package]] [[package]]
name = "python-dateutil" name = "python-dateutil"
version = "2.9.0.post0" version = "2.9.0.post0"
@ -4919,6 +4934,19 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/ee/d9/d88e73ca598f4f6ff671fb5fde8a32925c2e08a637303a1d12883c7305fa/uvicorn-0.38.0-py3-none-any.whl", hash = "sha256:48c0afd214ceb59340075b4a052ea1ee91c16fbc2a9b1469cca0e54566977b02", size = 68109, upload-time = "2025-10-18T13:46:42.958Z" }, { url = "https://files.pythonhosted.org/packages/ee/d9/d88e73ca598f4f6ff671fb5fde8a32925c2e08a637303a1d12883c7305fa/uvicorn-0.38.0-py3-none-any.whl", hash = "sha256:48c0afd214ceb59340075b4a052ea1ee91c16fbc2a9b1469cca0e54566977b02", size = 68109, upload-time = "2025-10-18T13:46:42.958Z" },
] ]
[[package]]
name = "vcrpy"
version = "8.1.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pyyaml" },
{ name = "wrapt" },
]
sdist = { url = "https://files.pythonhosted.org/packages/23/74/4200cb68d59e86849992eb6512d969cd561051f034e017428e040b113974/vcrpy-8.1.0.tar.gz", hash = "sha256:e585ca3cd9bb751e402728a00394847561250588eebc047b4d3c8948d5487733", size = 85930, upload-time = "2025-12-08T16:46:13.049Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/09/77/892bcd82445ac949816205b51ab80deb86a492a315f2e290ed4eab35021c/vcrpy-8.1.0-py3-none-any.whl", hash = "sha256:fc4fb6e954c6d082ba6d329c6f3d1228f5b1b1d2836f9022c301b587cfad7378", size = 42748, upload-time = "2025-12-08T16:46:12.08Z" },
]
[[package]] [[package]]
name = "virtualenv" name = "virtualenv"
version = "20.35.4" version = "20.35.4"