99 lines
2.7 KiB
Python
99 lines
2.7 KiB
Python
"""
|
|
Unit tests for the RAG agent and auto-check graph.
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import tempfile
|
|
import shutil
|
|
import pytest
|
|
|
|
from src.index import RAGAgent, auto_check_graph
|
|
|
|
# Helper to create a temporary agent with fake embeddings/LLM
|
|
def create_temp_agent(tmp_path):
|
|
# Ensure no OpenAI key
|
|
os.environ.pop("OPENAI_API_KEY", None)
|
|
agent = RAGAgent(
|
|
vector_store_path=tmp_path / "vector_store.faiss",
|
|
documents_dir=tmp_path / "docs",
|
|
)
|
|
return agent
|
|
|
|
def test_add_and_query():
|
|
tmp_dir = tempfile.mkdtemp()
|
|
try:
|
|
agent = create_temp_agent(tmp_dir)
|
|
docs = [
|
|
"The capital of France is Paris.",
|
|
"William Shakespeare wrote Hamlet.",
|
|
]
|
|
agent.add_documents(docs)
|
|
# Query for first doc
|
|
answer = agent.query("What is the capital of France?")
|
|
assert "Paris" in answer
|
|
# Query for second doc
|
|
answer2 = agent.query("Who wrote Hamlet?")
|
|
assert "Shakespeare" in answer2
|
|
finally:
|
|
shutil.rmtree(tmp_dir)
|
|
|
|
def test_auto_check_pass():
|
|
tmp_dir = tempfile.mkdtemp()
|
|
try:
|
|
agent = create_temp_agent(tmp_dir)
|
|
docs = [
|
|
"The capital of France is Paris.",
|
|
"William Shakespeare wrote Hamlet.",
|
|
]
|
|
agent.add_documents(docs)
|
|
ground_truth = {
|
|
"What is the capital of France?": "Paris",
|
|
"Who wrote Hamlet?": "William Shakespeare",
|
|
}
|
|
result = auto_check_graph(
|
|
"What is the capital of France?", agent, ground_truth
|
|
)
|
|
assert result["verdict_row"] == "PASS"
|
|
assert "Paris" in result["answer"]
|
|
finally:
|
|
shutil.rmtree(tmp_dir)
|
|
|
|
def test_auto_check_fail():
|
|
tmp_dir = tempfile.mkdtemp()
|
|
try:
|
|
agent = create_temp_agent(tmp_dir)
|
|
docs = [
|
|
"The capital of France is Paris.",
|
|
]
|
|
agent.add_documents(docs)
|
|
ground_truth = {
|
|
"What is the capital of France?": "Berlin",
|
|
}
|
|
result = auto_check_graph(
|
|
"What is the capital of France?", agent, ground_truth
|
|
)
|
|
assert result["verdict_row"] == "FAIL"
|
|
finally:
|
|
shutil.rmtree(tmp_dir)
|
|
|
|
def test_auto_check_unknown():
|
|
tmp_dir = tempfile.mkdtemp()
|
|
try:
|
|
agent = create_temp_agent(tmp_dir)
|
|
docs = [
|
|
"The capital of France is Paris.",
|
|
]
|
|
agent.add_documents(docs)
|
|
ground_truth = {
|
|
"What is the capital of Germany?": "Berlin",
|
|
}
|
|
result = auto_check_graph(
|
|
"What is the capital of Germany?", agent, ground_truth
|
|
)
|
|
assert result["verdict_row"] == "UNKNOWN"
|
|
finally:
|
|
shutil.rmtree(tmp_dir)
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__]) |