"""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 == ""