Build the skills rag_db fixture once per session
This commit is contained in:
parent
e4b1f16f97
commit
ac64aa9c17
1 changed files with 44 additions and 32 deletions
|
|
@ -1,5 +1,5 @@
|
||||||
import random
|
import random
|
||||||
from unittest.mock import MagicMock
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from pydantic_ai import RunContext
|
from pydantic_ai import RunContext
|
||||||
|
|
@ -12,6 +12,19 @@ from haiku.rag.skills._deps import AnalysisRunDeps, RAGRunDeps
|
||||||
VECTOR_DIM = 2560
|
VECTOR_DIM = 2560
|
||||||
|
|
||||||
|
|
||||||
|
def _seeded_vector(text: str) -> list[float]:
|
||||||
|
random.seed(hash(text) % (2**32))
|
||||||
|
return [random.random() for _ in range(VECTOR_DIM)]
|
||||||
|
|
||||||
|
|
||||||
|
async def _fake_embed_query(self, text: str) -> list[float]:
|
||||||
|
return _seeded_vector(text)
|
||||||
|
|
||||||
|
|
||||||
|
async def _fake_embed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||||
|
return [_seeded_vector(t) for t in texts]
|
||||||
|
|
||||||
|
|
||||||
def _make_ctx(state=None, rag=None, sandbox=None):
|
def _make_ctx(state=None, rag=None, sandbox=None):
|
||||||
"""Create a mock RunContext with RAGRunDeps (or AnalysisRunDeps when state is AnalysisState)."""
|
"""Create a mock RunContext with RAGRunDeps (or AnalysisRunDeps when state is AnalysisState)."""
|
||||||
from haiku.rag.skills.analysis import AnalysisState
|
from haiku.rag.skills.analysis import AnalysisState
|
||||||
|
|
@ -35,20 +48,8 @@ def _get_tool(skill, name):
|
||||||
@pytest.fixture(autouse=True)
|
@pytest.fixture(autouse=True)
|
||||||
def mock_embedder(monkeypatch):
|
def mock_embedder(monkeypatch):
|
||||||
"""Monkeypatch the embedder to return deterministic vectors."""
|
"""Monkeypatch the embedder to return deterministic vectors."""
|
||||||
|
monkeypatch.setattr(EmbedderWrapper, "embed_query", _fake_embed_query)
|
||||||
async def fake_embed_query(self, text):
|
monkeypatch.setattr(EmbedderWrapper, "embed_documents", _fake_embed_documents)
|
||||||
random.seed(hash(text) % (2**32))
|
|
||||||
return [random.random() for _ in range(VECTOR_DIM)]
|
|
||||||
|
|
||||||
async def fake_embed_documents(self, texts):
|
|
||||||
result = []
|
|
||||||
for t in texts:
|
|
||||||
random.seed(hash(t) % (2**32))
|
|
||||||
result.append([random.random() for _ in range(VECTOR_DIM)])
|
|
||||||
return result
|
|
||||||
|
|
||||||
monkeypatch.setattr(EmbedderWrapper, "embed_query", fake_embed_query)
|
|
||||||
monkeypatch.setattr(EmbedderWrapper, "embed_documents", fake_embed_documents)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
|
|
@ -56,23 +57,34 @@ def test_app_config():
|
||||||
return AppConfig(environment="skills-test")
|
return AppConfig(environment="skills-test")
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture(scope="session")
|
||||||
async def rag_db(temp_db_path):
|
async def rag_db(tmp_path_factory):
|
||||||
"""Create a test database with sample documents."""
|
"""Sample database with two documents, built once and shared read-only.
|
||||||
async with HaikuRAG(temp_db_path, create=True) as rag:
|
|
||||||
await rag.create_document(
|
Consumers (``rag_client``, ``sandbox_factory``) only read, so the docling
|
||||||
"Artificial intelligence is transforming industries worldwide. "
|
conversion + ingest is paid once per session instead of per test. Document
|
||||||
"Deep learning models are used in healthcare, finance, and transportation.",
|
vectors use the same seeded fakes as ``mock_embedder`` so search stays
|
||||||
title="AI Overview",
|
consistent with query-time embeddings.
|
||||||
uri="test://ai-overview",
|
"""
|
||||||
)
|
db_path = tmp_path_factory.mktemp("skills_rag_db") / "rag.lancedb"
|
||||||
await rag.create_document(
|
with (
|
||||||
"Machine learning is a subset of artificial intelligence. "
|
patch.object(EmbedderWrapper, "embed_query", _fake_embed_query),
|
||||||
"It includes supervised learning, unsupervised learning, and reinforcement learning.",
|
patch.object(EmbedderWrapper, "embed_documents", _fake_embed_documents),
|
||||||
title="ML Basics",
|
):
|
||||||
uri="test://ml-basics",
|
async with HaikuRAG(db_path, create=True) as rag:
|
||||||
)
|
await rag.create_document(
|
||||||
return temp_db_path
|
"Artificial intelligence is transforming industries worldwide. "
|
||||||
|
"Deep learning models are used in healthcare, finance, and transportation.",
|
||||||
|
title="AI Overview",
|
||||||
|
uri="test://ai-overview",
|
||||||
|
)
|
||||||
|
await rag.create_document(
|
||||||
|
"Machine learning is a subset of artificial intelligence. "
|
||||||
|
"It includes supervised learning, unsupervised learning, and reinforcement learning.",
|
||||||
|
title="ML Basics",
|
||||||
|
uri="test://ml-basics",
|
||||||
|
)
|
||||||
|
return db_path
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue