Make search agent more like a qa agent

src/haiku/rag/research/dependencies.py

@@ -20,6 +20,9 @@ class ResearchContext(BaseModel):
     search_results: list[dict[str, Any]] = Field(
         default_factory=list, description="Accumulated search results"
     )
+    qa_responses: list[dict[str, Any]] = Field(
+        default_factory=list, description="Question-answer pairs with sources"
+    )
     insights: list[str] = Field(
         default_factory=list, description="Key insights discovered"
     )
@@ -39,6 +42,18 @@ class ResearchContext(BaseModel):
             }
         )
 
+    def add_qa_response(
+        self, question: str, answer: str, sources: list["SearchResult"]
+    ) -> None:
+        """Add a QA response with its source documents."""
+        self.qa_responses.append(
+            {
+                "question": question,
+                "answer": answer,
+                "sources": sources,
+            }
+        )
+
     def add_insight(self, insight: str) -> None:
         """Add a key insight."""
         if insight not in self.insights:

src/haiku/rag/research/orchestrator.py

@@ -202,30 +202,40 @@ class ResearchOrchestrator(BaseResearchAgent):
             if context.follow_up_questions:
                 # Use follow-up questions from previous clarification
                 search_target = context.follow_up_questions[:3]  # Take top 3 follow-ups
-                search_prompt = f"Search for: {', '.join(search_target)}"
+                search_prompt = ", ".join(search_target)
             elif iteration < len(context.sub_questions):
                 # Use pre-planned sub-questions
-                search_prompt = f"Search for: {context.sub_questions[iteration]}"
+                search_prompt = context.sub_questions[iteration]
             else:
-                # Fall back to original question with variation
+                # Fall back to original question
-                search_prompt = f"Additional search for: {question}"
+                search_prompt = question
 
             # Search phase - directly call the search agent
             if console:
-                console.print(f"\n[bold cyan]🔍 Searching:[/bold cyan] {search_prompt}")
+                console.print(
+                    f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {search_prompt}"
+                )
 
             await self.search_agent.run(search_prompt, deps=deps)
 
-            if console and context.search_results:
+            if console:
-                latest_results = context.search_results[-1]
+                # Show documents found
-                console.print(
+                if context.search_results:
-                    f"   Found [green]{len(latest_results.get('results', []))} documents[/green]"
+                    latest_results = context.search_results[-1]
-                )
-                for i, result in enumerate(latest_results.get("results", [])[:3], 1):
                     console.print(
-                        f"   {i}. Score: [yellow]{result.score:.3f}[/yellow] - {result.document_uri}"
+                        f"   Found [green]{len(latest_results.get('results', []))} documents[/green]"
                     )
 
+                # Show the answer generated
+                if context.qa_responses:
+                    latest_qa = context.qa_responses[-1]
+                    answer_preview = (
+                        latest_qa["answer"][:200] + "..."
+                        if len(latest_qa["answer"]) > 200
+                        else latest_qa["answer"]
+                    )
+                    console.print(f"   [bold]Answer:[/bold] {answer_preview}")
+
             # Analysis phase (only if we have results)
             if context.search_results:
                 if console:

src/haiku/rag/research/search_agent.py

@@ -1,47 +1,69 @@
-"""Search specialist agent for document retrieval."""
+"""Search specialist agent that answers questions using RAG."""
+
+from typing import Any
 
 from pydantic_ai import RunContext
 
 from haiku.rag.research.base import BaseResearchAgent
 from haiku.rag.research.dependencies import ResearchDependencies
-from haiku.rag.store.models.chunk import Chunk
 
 
 class SearchSpecialistAgent(BaseResearchAgent):
-    """Agent specialized in document search and retrieval."""
+    """Agent specialized in answering questions using RAG search."""
 
     def __init__(self, provider: str, model: str):
-        # No specific output type needed - the tool handles everything
+        # Output is a string answer, like the QA agent
-        super().__init__(provider, model)
+        super().__init__(provider, model, output_type=str)
+
+    async def run(self, prompt: str, deps: ResearchDependencies, **kwargs) -> Any:
+        """Execute the agent and store QA response in context."""
+        # Run the base agent
+        result = await super().run(prompt, deps, **kwargs)
+
+        # Store the QA response if we got an answer
+        if hasattr(result, "output") and result.output:
+            # Get the sources from the last search (which the tool just stored)
+            if deps.context.search_results:
+                last_search = deps.context.search_results[-1]
+                sources = last_search.get("results", [])
+                deps.context.add_qa_response(prompt, result.output, sources)
+
+        return result
 
     def get_system_prompt(self) -> str:
-        return """You are a search specialist agent focused on document retrieval from a knowledge base that uses hybrid (semantic and full-text search) search.
+        return """You are a search and question-answering specialist.
-        Your role is to:
-        1. Understand the search query and context
-        2. Execute targeted searches to find relevant documents
 
-        Use the search tool to perform the searches on the knowledge base."""
+        Your role is to:
+        1. Search the knowledge base for relevant information
+        2. Analyze the retrieved documents
+        3. Provide a comprehensive answer to the question
+        4. Base your answer strictly on the information found
+
+        Use the search_and_answer tool to retrieve relevant documents and formulate your response.
+        Be thorough and specific in your answers, citing relevant information from the sources."""
 
     def register_tools(self) -> None:
         """Register search-specific tools."""
 
         @self.agent.tool
-        async def search(
+        async def search_and_answer(
             ctx: RunContext[ResearchDependencies],
             query: str,
             limit: int = 5,
-        ) -> list[tuple[Chunk, float]]:
+        ) -> str:
-            """Execute search and return raw results from client."""
+            """Search for information and provide context for answering the question."""
             # Use the default hybrid search
             search_results = await ctx.deps.client.search(query, limit=limit)
 
             # Expand context for better relevance
             expanded = await ctx.deps.client.expand_context(search_results)
 
-            # Store in context (convert to SearchResult for context storage)
+            # Convert to SearchResult for context storage
             from haiku.rag.research.base import SearchResult
 
             results_for_context = []
+            context_texts = []
+
             for chunk, score in expanded:
                 results_for_context.append(
                     SearchResult(
@@ -51,7 +73,16 @@ class SearchSpecialistAgent(BaseResearchAgent):
                         metadata={"chunk_id": chunk.id} if chunk.id else {},
                     )
                 )
+                context_texts.append(chunk.content)
+
+            # Store raw search results for analysis by other agents
             ctx.deps.context.add_search_result(query, results_for_context)
 
-            # Return raw chunk, score tuples
+            # Format context for the LLM to answer the question
-            return expanded
+            if context_texts:
+                context = "\n\n---\n\n".join(context_texts)
+                return f"Based on the following information from the knowledge base:\n\n{context}\n\nAnswer the question: {query}"
+            else:
+                return (
+                    f"No relevant information found in the knowledge base for: {query}"
+                )

tests/research/test_search_agent.py

@@ -77,7 +77,7 @@ async def test_search_agent_has_search_tool():
     tools = test_model.last_model_request_parameters.function_tools
     assert tools is not None
     assert len(tools) == 1
-    assert tools[0].name == "search"
+    assert tools[0].name == "search_and_answer"
 
 
 @pytest.mark.asyncio
@@ -96,21 +96,19 @@ async def test_search_single_query(mock_client, research_deps):
     agent = SearchSpecialistAgent(provider="openai", model="gpt-4")
 
     # Get the search tool
-    search_tool = get_agent_tool(agent, "search")
+    search_tool = get_agent_tool(agent, "search_and_answer")
     assert search_tool is not None
 
     # Test the tool
     ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage())
-    results = await search_tool(ctx, query="climate change")
+    result = await search_tool(ctx, query="climate change")
 
-    # Verify results - should be list of (Chunk, float) tuples
+    # Verify result - should be a formatted string with context
-    assert isinstance(results, list)
+    assert isinstance(result, str)
-    assert len(results) == 2
+    assert "Climate change is a global phenomenon" in result
-    assert results[0][0].content == "Climate change is a global phenomenon"
+    assert "Rising temperatures affect ecosystems" in result
-    assert results[0][0].id == "chunk1"
-    assert results[0][1] == 0.8  # score
 
-    # Verify mock was called
+    # Verify mock was called with default limit
     mock_client.search.assert_called_once_with("climate change", limit=5)
     mock_client.expand_context.assert_called_once()
 
@@ -133,18 +131,18 @@ async def test_search_with_limit(mock_client, research_deps):
     agent = SearchSpecialistAgent(provider="openai", model="gpt-4")
 
     # Get the search tool
-    search_tool = get_agent_tool(agent, "search")
+    search_tool = get_agent_tool(agent, "search_and_answer")
     assert search_tool is not None
 
     # Test the tool with limit
     ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage())
-    results = await search_tool(ctx, query="test query", limit=3)
+    result = await search_tool(ctx, query="test query", limit=3)
 
-    # Verify results respect limit
+    # Verify result is a formatted string
-    assert isinstance(results, list)
+    assert isinstance(result, str)
-    assert len(results) == 3
+    assert "Content 1" in result
-    assert all(isinstance(r, tuple) and len(r) == 2 for r in results)
+    assert "Content 2" in result
-    assert all(isinstance(r[0], Chunk) and isinstance(r[1], float) for r in results)
+    assert "Content 3" in result
 
     # Verify mock was called with correct limit
     mock_client.search.assert_called_once_with("test query", limit=3)
@@ -161,7 +159,7 @@ async def test_search_updates_context(mock_client, research_deps):
     agent = SearchSpecialistAgent(provider="openai", model="gpt-4")
 
     # Get the search tool
-    search_tool = get_agent_tool(agent, "search")
+    search_tool = get_agent_tool(agent, "search_and_answer")
     assert search_tool is not None
 
     # Test the tool