from unittest.mock import AsyncMock from haiku.rag.agents.research.models import Citation, ResearchReport from haiku.rag.agents.rlm.models import RLMResult from haiku.rag.client import HaikuRAG from haiku.rag.store.models.chunk import SearchResult from haiku.rag.tools.document import DocumentInfo from haiku.rag.tools.qa import QAHistoryEntry from .conftest import _get_tool, _make_ctx class TestRAGSkillCreation: def test_create_skill_returns_valid_skill(self, temp_db_path): from haiku.rag.skills.rag import create_skill skill = create_skill(db_path=temp_db_path) assert skill.metadata.name == "rag" assert skill.metadata.description assert skill.instructions def test_create_skill_has_expected_tools(self, temp_db_path): from haiku.rag.skills.rag import create_skill skill = create_skill(db_path=temp_db_path) tool_names = {getattr(t, "__name__") for t in skill.tools if callable(t)} assert tool_names == { "search", "list_documents", "get_document", "ask", "analyze", "research", } def test_create_skill_has_state(self, temp_db_path): from haiku.rag.skills.rag import RAGState, create_skill skill = create_skill(db_path=temp_db_path) assert skill._state_type is RAGState assert skill._state_namespace == "rag" def test_create_skill_from_env(self, monkeypatch, temp_db_path): monkeypatch.setenv("HAIKU_RAG_DB", str(temp_db_path)) from haiku.rag.skills.rag import create_skill skill = create_skill() assert skill.metadata.name == "rag" class TestSearchTool: async def test_search_returns_formatted_string(self, rag_db): from haiku.rag.skills.rag import create_skill skill = create_skill(db_path=rag_db) search = _get_tool(skill, "search") ctx = _make_ctx() result = await search(ctx, query="artificial intelligence") assert isinstance(result, str) assert len(result) > 0 async def test_search_updates_state(self, rag_db): from haiku.rag.skills.rag import RAGState, create_skill skill = create_skill(db_path=rag_db) search = _get_tool(skill, "search") state = RAGState() ctx = _make_ctx(state) await search(ctx, query="artificial intelligence") assert "artificial intelligence" in state.searches results = state.searches["artificial intelligence"] assert len(results) > 0 assert isinstance(results[0], SearchResult) async def test_search_applies_document_filter_from_state(self, rag_db): from haiku.rag.skills.rag import RAGState, create_skill skill = create_skill(db_path=rag_db) search = _get_tool(skill, "search") state = RAGState(document_filter="title = 'AI Overview'") ctx = _make_ctx(state) result = await search(ctx, query="artificial intelligence") assert "AI Overview" in result assert "ML Basics" not in result async def test_search_without_state(self, rag_db): from haiku.rag.skills.rag import create_skill skill = create_skill(db_path=rag_db) search = _get_tool(skill, "search") ctx = _make_ctx(state=None) result = await search(ctx, query="artificial intelligence") assert isinstance(result, str) class TestListDocumentsTool: async def test_list_documents_returns_results(self, rag_db): from haiku.rag.skills.rag import create_skill skill = create_skill(db_path=rag_db) list_docs = _get_tool(skill, "list_documents") ctx = _make_ctx() results = await list_docs(ctx) assert isinstance(results, list) assert len(results) == 2 async def test_list_documents_updates_state(self, rag_db): from haiku.rag.skills.rag import RAGState, create_skill skill = create_skill(db_path=rag_db) list_docs = _get_tool(skill, "list_documents") state = RAGState() ctx = _make_ctx(state) await list_docs(ctx) assert len(state.documents) == 2 assert isinstance(state.documents[0], DocumentInfo) assert state.documents[0].id is not None async def test_list_documents_no_duplicates_in_state(self, rag_db): from haiku.rag.skills.rag import RAGState, create_skill skill = create_skill(db_path=rag_db) list_docs = _get_tool(skill, "list_documents") state = RAGState() ctx = _make_ctx(state) await list_docs(ctx) await list_docs(ctx) assert len(state.documents) == 2 class TestGetDocumentTool: async def test_get_document_by_title(self, rag_db): from haiku.rag.skills.rag import create_skill skill = create_skill(db_path=rag_db) get_doc = _get_tool(skill, "get_document") ctx = _make_ctx() result = await get_doc(ctx, query="AI Overview") assert result is not None assert result["title"] == "AI Overview" async def test_get_document_updates_state(self, rag_db): from haiku.rag.skills.rag import RAGState, create_skill skill = create_skill(db_path=rag_db) get_doc = _get_tool(skill, "get_document") state = RAGState() ctx = _make_ctx(state) await get_doc(ctx, query="AI Overview") assert len(state.documents) == 1 assert isinstance(state.documents[0], DocumentInfo) assert state.documents[0].title == "AI Overview" async def test_get_document_not_found(self, rag_db): from haiku.rag.skills.rag import create_skill skill = create_skill(db_path=rag_db) get_doc = _get_tool(skill, "get_document") ctx = _make_ctx() result = await get_doc(ctx, query="nonexistent document xyz") assert result is None class TestAskTool: async def test_ask_returns_answer_with_citations(self, rag_db, monkeypatch): from haiku.rag.skills.rag import create_skill citations = [ Citation( document_id="d1", chunk_id="c1", document_uri="test://ai-overview", document_title="AI Overview", content="AI is transforming industries.", ) ] monkeypatch.setattr( HaikuRAG, "ask", AsyncMock(return_value=("AI transforms industries worldwide.", citations)), ) skill = create_skill(db_path=rag_db) ask = _get_tool(skill, "ask") ctx = _make_ctx() result = await ask(ctx, question="What is AI?") assert isinstance(result, str) assert "AI transforms industries" in result async def test_ask_updates_state(self, rag_db, monkeypatch): from haiku.rag.skills.rag import RAGState, create_skill citations = [ Citation( document_id="d1", chunk_id="c1", document_uri="test://ai-overview", content="AI content", ) ] monkeypatch.setattr( HaikuRAG, "ask", AsyncMock(return_value=("AI transforms industries.", citations)), ) skill = create_skill(db_path=rag_db) ask = _get_tool(skill, "ask") state = RAGState() ctx = _make_ctx(state) await ask(ctx, question="What is AI?") assert len(state.citations) == 1 assert len(state.qa_history) == 1 assert isinstance(state.qa_history[0], QAHistoryEntry) assert state.qa_history[0].question == "What is AI?" async def test_ask_assigns_citation_indices(self, rag_db, monkeypatch): from haiku.rag.skills.rag import RAGState, create_skill first_citations = [ Citation( document_id="d1", chunk_id="c1", document_uri="test://doc1", content="First.", ), Citation( document_id="d2", chunk_id="c2", document_uri="test://doc2", content="Second.", ), ] second_citations = [ Citation( document_id="d3", chunk_id="c3", document_uri="test://doc3", content="Third.", ), ] call_count = 0 async def mock_ask(self, question, **kwargs): nonlocal call_count call_count += 1 if call_count == 1: return ("Answer 1", first_citations) return ("Answer 2", second_citations) monkeypatch.setattr(HaikuRAG, "ask", mock_ask) skill = create_skill(db_path=rag_db) ask = _get_tool(skill, "ask") state = RAGState() ctx = _make_ctx(state) await ask(ctx, question="First question") assert state.citations[0].index == 1 assert state.citations[1].index == 2 await ask(ctx, question="Second question") assert state.citations[2].index == 3 async def test_ask_applies_document_filter_from_state(self, rag_db, monkeypatch): from haiku.rag.skills.rag import RAGState, create_skill captured_kwargs = {} async def mock_ask(self, question, **kwargs): captured_kwargs.update(kwargs) return ("Answer.", []) monkeypatch.setattr(HaikuRAG, "ask", mock_ask) skill = create_skill(db_path=rag_db) ask = _get_tool(skill, "ask") state = RAGState(document_filter="title = 'AI Overview'") ctx = _make_ctx(state) await ask(ctx, question="What is AI?") assert captured_kwargs.get("filter") == "title = 'AI Overview'" async def test_ask_includes_prior_qa_context(self, rag_db, monkeypatch): import random from haiku.rag.skills.rag import RAGState, create_skill from tests.skills.conftest import VECTOR_DIM captured_questions = [] async def mock_ask(self, question, **kwargs): captured_questions.append(question) return ("Answer about AI.", []) monkeypatch.setattr(HaikuRAG, "ask", mock_ask) skill = create_skill(db_path=rag_db) ask = _get_tool(skill, "ask") # Pre-compute the embedding the fake embedder will produce for "Tell me about AI" query_text = "Tell me about AI" random.seed(hash(query_text) % (2**32)) query_embedding = [random.random() for _ in range(VECTOR_DIM)] prior_citations = [ Citation( document_id="d1", chunk_id="c1", document_uri="test://ai-overview", document_title="AI Overview", content="AI content from source.", ) ] state = RAGState( qa_history=[ QAHistoryEntry( question="What is artificial intelligence?", answer="AI is the simulation of human intelligence by machines.", question_embedding=query_embedding, citations=prior_citations, ), ] ) ctx = _make_ctx(state) await ask(ctx, question=query_text) # rag.ask() should receive augmented question with prior context assert len(captured_questions) == 1 augmented = captured_questions[0] assert "Context from prior questions" in augmented assert "What is artificial intelligence?" in augmented assert "AI is the simulation" in augmented assert "AI Overview" in augmented assert query_text in augmented # State should store the original question, not the augmented one assert state.qa_history[-1].question == query_text async def test_ask_embeds_prior_qa_on_demand(self, rag_db, monkeypatch): from haiku.rag.skills.rag import RAGState, create_skill from tests.skills.conftest import VECTOR_DIM captured_questions = [] async def mock_ask(self, question, **kwargs): captured_questions.append(question) return ("Answer about AI.", []) monkeypatch.setattr(HaikuRAG, "ask", mock_ask) skill = create_skill(db_path=rag_db) ask = _get_tool(skill, "ask") # Use the same question text for the prior QA entry and query so # their fake embeddings are identical (cosine similarity = 1.0). prior_question = "Tell me about AI" query_text = prior_question # Leave question_embedding=None to exercise the lazy embedding path state = RAGState( qa_history=[ QAHistoryEntry( question=prior_question, answer="AI is the simulation of human intelligence by machines.", question_embedding=None, ), ] ) ctx = _make_ctx(state) await ask(ctx, question=query_text) # The lazy embedding should have populated question_embedding assert state.qa_history[0].question_embedding is not None assert len(state.qa_history[0].question_embedding) == VECTOR_DIM # The augmented question should include prior context assert len(captured_questions) == 1 assert "Context from prior questions" in captured_questions[0] assert prior_question in captured_questions[0] async def test_ask_no_prior_qa_context_when_irrelevant(self, rag_db, monkeypatch): from haiku.rag.skills.rag import RAGState, create_skill from tests.skills.conftest import VECTOR_DIM captured_questions = [] async def mock_ask(self, question, **kwargs): captured_questions.append(question) return ("Answer.", []) monkeypatch.setattr(HaikuRAG, "ask", mock_ask) skill = create_skill(db_path=rag_db) ask = _get_tool(skill, "ask") # Use orthogonal embedding — won't match the fake embedder's output orthogonal = [1.0 if i % 2 == 0 else -1.0 for i in range(VECTOR_DIM)] state = RAGState( qa_history=[ QAHistoryEntry( question="What is the weather?", answer="It is sunny today.", question_embedding=orthogonal, ), ] ) ctx = _make_ctx(state) await ask(ctx, question="Explain quantum computing") # rag.ask() should receive the original question unchanged assert len(captured_questions) == 1 assert captured_questions[0] == "Explain quantum computing" class TestAnalyzeTool: async def test_analyze_returns_result(self, rag_db, monkeypatch): from haiku.rag.skills.rag import create_skill monkeypatch.setattr( HaikuRAG, "rlm", AsyncMock(return_value=RLMResult(answer="42", program="print(42)")), ) skill = create_skill(db_path=rag_db) analyze = _get_tool(skill, "analyze") ctx = _make_ctx() result = await analyze(ctx, question="How many documents?") assert isinstance(result, str) assert "42" in result async def test_analyze_updates_state(self, rag_db, monkeypatch): from haiku.rag.skills.rag import RAGState, create_skill monkeypatch.setattr( HaikuRAG, "rlm", AsyncMock(return_value=RLMResult(answer="42", program="print(42)")), ) skill = create_skill(db_path=rag_db) analyze = _get_tool(skill, "analyze") state = RAGState() ctx = _make_ctx(state) await analyze(ctx, question="How many documents?") assert len(state.qa_history) == 1 assert state.qa_history[0].question == "How many documents?" class TestResearchTool: async def test_research_returns_report(self, rag_db, monkeypatch): from haiku.rag.skills.rag import create_skill report = ResearchReport( title="AI Research", executive_summary="AI is transforming industries.", main_findings=["Finding 1"], conclusions=["Conclusion 1"], sources_summary="Multiple sources consulted.", ) monkeypatch.setattr(HaikuRAG, "research", AsyncMock(return_value=report)) skill = create_skill(db_path=rag_db) research = _get_tool(skill, "research") ctx = _make_ctx() result = await research(ctx, question="What is AI?") assert isinstance(result, str) assert "AI Research" in result async def test_research_updates_state(self, rag_db, monkeypatch): from haiku.rag.skills.rag import RAGState, create_skill report = ResearchReport( title="AI Research", executive_summary="AI is transforming industries.", main_findings=["Finding 1"], conclusions=["Conclusion 1"], sources_summary="Multiple sources consulted.", ) monkeypatch.setattr(HaikuRAG, "research", AsyncMock(return_value=report)) skill = create_skill(db_path=rag_db) research = _get_tool(skill, "research") state = RAGState() ctx = _make_ctx(state) await research(ctx, question="What is AI?") assert len(state.reports) == 1 assert state.reports[0].question == "What is AI?" assert len(state.qa_history) == 1 assert state.qa_history[0].question == "What is AI?" assert state.qa_history[0].answer == "AI is transforming industries." async def test_research_applies_document_filter_from_state( self, rag_db, monkeypatch ): from haiku.rag.skills.rag import RAGState, create_skill captured_kwargs = {} report = ResearchReport( title="AI Research", executive_summary="Summary.", main_findings=["Finding"], conclusions=["Conclusion"], sources_summary="Sources.", ) async def mock_research(self, question, **kwargs): captured_kwargs.update(kwargs) return report monkeypatch.setattr(HaikuRAG, "research", mock_research) skill = create_skill(db_path=rag_db) research = _get_tool(skill, "research") state = RAGState(document_filter="title = 'AI Overview'") ctx = _make_ctx(state) await research(ctx, question="What is AI?") assert captured_kwargs.get("filter") == "title = 'AI Overview'" async def test_research_without_state(self, rag_db, monkeypatch): from haiku.rag.skills.rag import create_skill report = ResearchReport( title="AI Research", executive_summary="Summary.", main_findings=["Finding"], conclusions=["Conclusion"], sources_summary="Sources.", ) monkeypatch.setattr(HaikuRAG, "research", AsyncMock(return_value=report)) skill = create_skill(db_path=rag_db) research = _get_tool(skill, "research") ctx = _make_ctx(state=None) result = await research(ctx, question="What is AI?") assert isinstance(result, str)