82fd10e0ee
Convert all LanceDB operations from sync calls wrapped in async functions to the native async API (connect_async, AsyncConnection, AsyncTable, AsyncQuery). Database I/O no longer blocks the event loop. - Store and HaikuRAG use async context managers (async with). Store initialization is deferred to __aenter__; direct construction without async with is no longer supported. - Index creation uses config objects (FTS, BTree, IvfPq) instead of string-based index_type parameter. - Upgrade callbacks are async. - HaikuRAG tracks background vacuum tasks and awaits them in __aexit__ and before destructive rebuild operations to avoid races with concurrent table mutations. - temp_db_path fixture uses pytest's tmp_path for reliable async cleanup.
69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
from pathlib import Path
|
|
|
|
import pytest
|
|
from datasets import Dataset
|
|
from evaluations.evaluators import LLMJudge
|
|
|
|
from haiku.rag.agents.qa.agent import QuestionAnswerAgent
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config import Config
|
|
from haiku.rag.config.models import ModelConfig
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def vcr_cassette_dir():
|
|
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_qa")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_qa_agent_factory(temp_db_path):
|
|
"""Test get_qa_agent factory function creates a properly configured agent."""
|
|
from haiku.rag.agents.qa import get_qa_agent
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
agent = get_qa_agent(client, Config)
|
|
|
|
assert agent is not None
|
|
assert isinstance(agent, QuestionAnswerAgent)
|
|
# Verify internal client is set correctly
|
|
assert agent._client is client
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_qa_agent_with_custom_prompt(temp_db_path):
|
|
"""Test get_qa_agent factory with custom system prompt."""
|
|
from haiku.rag.agents.qa import get_qa_agent
|
|
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
custom_prompt = "You are a custom QA assistant."
|
|
agent = get_qa_agent(client, Config, system_prompt=custom_prompt)
|
|
|
|
assert agent is not None
|
|
assert isinstance(agent, QuestionAnswerAgent)
|
|
assert agent._system_prompt == custom_prompt
|
|
|
|
|
|
@pytest.mark.vcr()
|
|
async def test_qa_ollama(allow_model_requests, qa_corpus: Dataset, temp_db_path):
|
|
"""Test Ollama QA with LLM judge (VCR recorded)."""
|
|
async with HaikuRAG(temp_db_path, create=True) as client:
|
|
qa = QuestionAnswerAgent(
|
|
client,
|
|
ModelConfig(provider="ollama", name="gpt-oss", enable_thinking=True),
|
|
)
|
|
llm_judge = LLMJudge()
|
|
|
|
doc = qa_corpus[1]
|
|
await client.create_document(
|
|
content=doc["document_extracted"], uri=doc["document_id"]
|
|
)
|
|
|
|
question = doc["question"]
|
|
expected_answer = doc["answer"]
|
|
|
|
answer, _ = await qa.answer(question)
|
|
is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer)
|
|
|
|
assert is_equivalent, (
|
|
f"Generated answer not equivalent to expected answer.\nQuestion: {question}\nGenerated: {answer}\nExpected: {expected_answer}"
|
|
)
|