256 lines
9.2 KiB
Python
256 lines
9.2 KiB
Python
"""Phase 1 tests: RAG service logic.
|
|
|
|
Covers:
|
|
- ChromaDB document ingestion with metadata
|
|
- Retrieval with query keywords
|
|
- Response generation with strict RAG prompt
|
|
- Metadata handling per chunk
|
|
|
|
All tests use real ChromaDB via tmp_path. Only the LLM client (external API)
|
|
is mocked where needed.
|
|
"""
|
|
import pytest
|
|
import chromadb
|
|
from unittest.mock import AsyncMock
|
|
|
|
|
|
class _MockLLM:
|
|
"""Minimal mock for the external LLM API."""
|
|
|
|
def __init__(self, response: str = "mock answer"):
|
|
self._response = response
|
|
self.last_prompt: str | None = None
|
|
|
|
async def complete(self, prompt: str, temperature: float = 0.7, step_name: str = "LLM") -> str: # type: ignore[override]
|
|
self.last_prompt = prompt
|
|
return self._response
|
|
|
|
|
|
def _setup_chroma(tmp_path, monkeypatch, collection_name: str = "documents"):
|
|
"""Create an isolated real ChromaDB client + collection for a test.
|
|
|
|
Returns (client, collection, service_kwargs) where service_kwargs can be
|
|
unpacked into RAGService().
|
|
"""
|
|
from app.core.config import get_settings
|
|
|
|
monkeypatch.setenv("CHROMA_DB_PATH", str(tmp_path / "test_chroma"))
|
|
get_settings.cache_clear()
|
|
|
|
client = chromadb.PersistentClient(path=str(tmp_path / "test_chroma"))
|
|
collection = client.get_or_create_collection(name=collection_name)
|
|
return client, collection
|
|
|
|
|
|
class TestRAGService:
|
|
"""RAG retrieval and prompt logic tests."""
|
|
|
|
def test_ingest_document_adds_chunks(self, tmp_path, monkeypatch):
|
|
"""Should add chunks with metadata to real ChromaDB collection."""
|
|
from app.services.rag import RAGService
|
|
|
|
client, collection = _setup_chroma(tmp_path, monkeypatch)
|
|
|
|
service = RAGService(chroma_client=client)
|
|
|
|
chunks = ["chunk one", "chunk two"]
|
|
metadata = [
|
|
{"filename": "test.txt", "upload_date": "2024-01-01", "content_summary": "summary 1", "chunk_index": 0},
|
|
{"filename": "test.txt", "upload_date": "2024-01-01", "content_summary": "summary 2", "chunk_index": 1},
|
|
]
|
|
|
|
doc_id = service.ingest_document("test.txt", chunks, metadata)
|
|
|
|
assert doc_id != ""
|
|
assert collection.count() == 2
|
|
|
|
stored = collection.get(include=["documents", "metadatas"])
|
|
assert len(stored["documents"]) == 2
|
|
assert stored["documents"] == chunks
|
|
for i, meta in enumerate(stored["metadatas"]):
|
|
assert meta["filename"] == "test.txt"
|
|
assert meta["content_summary"] == f"summary {i + 1}"
|
|
|
|
def test_ingest_document_empty_chunks(self, tmp_path, monkeypatch):
|
|
"""Should not add anything when chunks list is empty."""
|
|
from app.services.rag import RAGService
|
|
|
|
client, collection = _setup_chroma(tmp_path, monkeypatch)
|
|
|
|
service = RAGService(chroma_client=client)
|
|
result = service.ingest_document("test.txt", [], [])
|
|
|
|
assert result == ""
|
|
assert collection.count() == 0
|
|
|
|
def test_retrieve_returns_chunks(self, tmp_path, monkeypatch):
|
|
"""Should retrieve chunks from real ChromaDB by query."""
|
|
from app.services.rag import RAGService
|
|
|
|
client, collection = _setup_chroma(tmp_path, monkeypatch)
|
|
|
|
collection.add(
|
|
documents=["chunk one about apples", "chunk two about bananas"],
|
|
metadatas=[
|
|
{"filename": "test.txt"},
|
|
{"filename": "test.txt"},
|
|
],
|
|
ids=["id1", "id2"],
|
|
)
|
|
|
|
service = RAGService(chroma_client=client)
|
|
results = service.retrieve(["apples"], n_results=5)
|
|
|
|
assert len(results) >= 1
|
|
assert "apples" in results[0][0]
|
|
assert results[0][1]["filename"] == "test.txt"
|
|
assert isinstance(results[0][2], float)
|
|
|
|
def test_retrieve_no_results(self, tmp_path, monkeypatch):
|
|
"""Should return empty list when querying an empty collection."""
|
|
from app.services.rag import RAGService
|
|
|
|
client, _ = _setup_chroma(tmp_path, monkeypatch)
|
|
|
|
service = RAGService(chroma_client=client)
|
|
results = service.retrieve(["nonexistent query terms xyz"])
|
|
|
|
assert results == []
|
|
|
|
async def test_generate_response_calls_llm(self, tmp_path, monkeypatch, mock_prompt_service):
|
|
"""Should call LLM with strict RAG prompt."""
|
|
from app.services.rag import RAGService
|
|
|
|
client, _ = _setup_chroma(tmp_path, monkeypatch)
|
|
|
|
mock_llm = _MockLLM(response="- Bullet point answer")
|
|
|
|
service = RAGService(
|
|
chroma_client=client,
|
|
llm_client=mock_llm,
|
|
prompt_service=mock_prompt_service,
|
|
)
|
|
|
|
chunks = ["relevant chunk"]
|
|
metadata = [{"filename": "test.txt", "content_summary": "summary"}]
|
|
|
|
answer, gen_prompt = await service.generate_response("What is this?", chunks, metadata)
|
|
|
|
assert mock_llm.last_prompt is not None
|
|
assert "What is this?" in mock_llm.last_prompt
|
|
assert "relevant chunk" in mock_llm.last_prompt
|
|
assert "test.txt" in mock_llm.last_prompt
|
|
assert "only these document chunks" in mock_llm.last_prompt.lower()
|
|
assert answer == "- Bullet point answer"
|
|
assert gen_prompt == mock_llm.last_prompt
|
|
|
|
async def test_generate_response_no_chunks(self, tmp_path, monkeypatch):
|
|
"""Should return fallback message when no chunks provided."""
|
|
from app.services.rag import RAGService
|
|
|
|
client, _ = _setup_chroma(tmp_path, monkeypatch)
|
|
mock_llm = _MockLLM()
|
|
|
|
service = RAGService(chroma_client=client, llm_client=mock_llm)
|
|
|
|
answer, gen_prompt = await service.generate_response("What is this?", [], [])
|
|
|
|
assert "no relevant" in answer.lower() or "could not find" in answer.lower()
|
|
assert gen_prompt == ""
|
|
|
|
def test_retrieve_per_subquestion_returns_per_query(self, tmp_path, monkeypatch):
|
|
"""Each sub-question retrieves its own chunks independently."""
|
|
from app.services.rag import RAGService
|
|
|
|
client, collection = _setup_chroma(tmp_path, monkeypatch)
|
|
|
|
collection.add(
|
|
documents=["Alpha content about apples", "Alpha extra about apples"],
|
|
metadatas=[{"filename": "a.pdf"}, {"filename": "a2.pdf"}],
|
|
ids=["a1", "a2"],
|
|
)
|
|
collection.add(
|
|
documents=["Beta content about bananas"],
|
|
metadatas=[{"filename": "b.pdf"}],
|
|
ids=["b1"],
|
|
)
|
|
|
|
service = RAGService(chroma_client=client)
|
|
results = service.retrieve_per_subquestion(["apples", "bananas"], n_results=5)
|
|
|
|
assert len(results) == 2
|
|
assert results[0][0] == "apples"
|
|
assert len(results[0][1]) >= 1
|
|
assert results[1][0] == "bananas"
|
|
assert len(results[1][1]) >= 1
|
|
|
|
def test_retrieve_per_subquestion_empty_list(self, tmp_path, monkeypatch):
|
|
"""Empty sub_questions list returns empty list without querying."""
|
|
from app.services.rag import RAGService
|
|
|
|
client, _ = _setup_chroma(tmp_path, monkeypatch)
|
|
|
|
service = RAGService(chroma_client=client)
|
|
results = service.retrieve_per_subquestion([], n_results=5)
|
|
|
|
assert results == []
|
|
|
|
async def test_generate_response_per_subquestion_calls_llm(self, tmp_path, monkeypatch, mock_prompt_service):
|
|
"""LLM should receive sub-question-organized context."""
|
|
from app.services.rag import RAGService
|
|
|
|
client, _ = _setup_chroma(tmp_path, monkeypatch)
|
|
|
|
mock_llm = _MockLLM(response="## Sub-question 1: Q?\n- Answer")
|
|
|
|
service = RAGService(
|
|
chroma_client=client,
|
|
llm_client=mock_llm,
|
|
prompt_service=mock_prompt_service,
|
|
)
|
|
|
|
answer, gen_prompt, grouped_sources = await service.generate_response_per_subquestion(
|
|
["What is X?"],
|
|
[["chunk data"]],
|
|
[[{"filename": "f.txt", "content_summary": "sum"}]],
|
|
)
|
|
|
|
assert mock_llm.last_prompt is not None
|
|
assert "chunk data" in mock_llm.last_prompt
|
|
assert answer == "## Sub-question 1: Q?\n- Answer"
|
|
assert len(grouped_sources) == 1
|
|
assert grouped_sources[0][0]["filename"] == "f.txt"
|
|
|
|
async def test_generate_response_per_subquestion_no_subquestions(self, tmp_path, monkeypatch):
|
|
"""Should return fallback when sub_questions is empty."""
|
|
from app.services.rag import RAGService
|
|
|
|
client, _ = _setup_chroma(tmp_path, monkeypatch)
|
|
mock_llm = _MockLLM()
|
|
|
|
service = RAGService(chroma_client=client, llm_client=mock_llm)
|
|
|
|
answer, gen_prompt, grouped_sources = await service.generate_response_per_subquestion(
|
|
[], [], [],
|
|
)
|
|
|
|
assert "could not find" in answer.lower()
|
|
assert gen_prompt == ""
|
|
assert grouped_sources == []
|
|
|
|
async def test_generate_response_per_subquestion_no_chunks(self, tmp_path, monkeypatch):
|
|
"""Should return fallback when all chunk lists are empty."""
|
|
from app.services.rag import RAGService
|
|
|
|
client, _ = _setup_chroma(tmp_path, monkeypatch)
|
|
mock_llm = _MockLLM()
|
|
|
|
service = RAGService(chroma_client=client, llm_client=mock_llm)
|
|
|
|
answer, gen_prompt, grouped_sources = await service.generate_response_per_subquestion(
|
|
["Q?"], [[]], [[]],
|
|
)
|
|
|
|
assert "could not find" in answer.lower()
|
|
assert gen_prompt == ""
|