Simplify SearchAgent. Omit search results, but keep QA results.

src/haiku/rag/config.py

@@ -31,7 +31,7 @@ class AppConfig(BaseModel):
 
     # Research defaults (fallback to QA if not provided via env)
     RESEARCH_PROVIDER: str = "ollama"
-    RESEARCH_MODEL: str = "qwen3"
+    RESEARCH_MODEL: str = "gpt-oss"
 
     CHUNK_SIZE: int = 256
     CONTEXT_CHUNK_RADIUS: int = 0

src/haiku/rag/research/__init__.py

@@ -1,6 +1,11 @@
 """Multi-agent research workflow for advanced RAG queries."""
 
-from haiku.rag.research.base import BaseResearchAgent, ResearchOutput, SearchResult
+from haiku.rag.research.base import (
+    BaseResearchAgent,
+    ResearchOutput,
+    SearchAnswer,
+    SearchResult,
+)
 from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
 from haiku.rag.research.evaluation_agent import (
     AnalysisEvaluationAgent,
@@ -18,6 +23,7 @@ __all__ = [
     "SearchResult",
     "ResearchOutput",
     # Specialized agents
+    "SearchAnswer",
     "SearchSpecialistAgent",
     "AnalysisEvaluationAgent",
     "EvaluationResult",

src/haiku/rag/research/base.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
-from typing import Any, Generic, TypeVar
+from typing import TYPE_CHECKING, Any
 
 from pydantic import BaseModel, Field
 from pydantic_ai import Agent
@@ -9,12 +11,12 @@ from pydantic_ai.providers.openai import OpenAIProvider
 from pydantic_ai.run import AgentRunResult
 
 from haiku.rag.config import Config
-from haiku.rag.research.dependencies import ResearchDependencies
 
-T = TypeVar("T")
+if TYPE_CHECKING:
+    from haiku.rag.research.dependencies import ResearchDependencies
 
 
-class BaseResearchAgent(ABC, Generic[T]):
+class BaseResearchAgent[T](ABC):
     """Base class for all research agents."""
 
     def __init__(
@@ -29,6 +31,9 @@ class BaseResearchAgent(ABC, Generic[T]):
 
         model_obj = self._get_model(provider, model)
 
+        # Import deps type lazily to avoid circular import during module load
+        from haiku.rag.research.dependencies import ResearchDependencies
+
         self._agent = Agent(
             model=model_obj,
             deps_type=ResearchDependencies,
@@ -75,7 +80,7 @@ class BaseResearchAgent(ABC, Generic[T]):
         return await self._agent.run(prompt, deps=deps, **kwargs)
 
     @property
-    def agent(self) -> Agent[ResearchDependencies, T]:
+    def agent(self) -> Agent[Any, T]:
         """Access the underlying Pydantic AI agent."""
         return self._agent
 
@@ -96,3 +101,23 @@ class ResearchOutput(BaseModel):
     detailed_findings: list[str]
     sources: list[str]
     confidence: float
+
+
+class SearchAnswer(BaseModel):
+    """Structured output for the SearchSpecialist agent."""
+
+    query: str = Field(description="The search query that was performed")
+    answer: str = Field(description="The answer generated based on the context")
+    context: list[str] = Field(
+        description=(
+            "Only the minimal set of relevant snippets (verbatim) that directly "
+            "support the answer"
+        )
+    )
+    sources: list[str] = Field(
+        description=(
+            "Document URIs corresponding to the snippets actually used in the"
+            " answer (one URI per snippet; omit if none)"
+        ),
+        default_factory=list,
+    )

src/haiku/rag/research/dependencies.py

@@ -1,11 +1,7 @@
-from typing import TYPE_CHECKING, Any
-
 from pydantic import BaseModel, Field
 
 from haiku.rag.client import HaikuRAG
-
+from haiku.rag.research.base import SearchAnswer
-if TYPE_CHECKING:
-    from haiku.rag.research.base import SearchResult
 
 
 class ResearchContext(BaseModel):
@@ -15,11 +11,8 @@ class ResearchContext(BaseModel):
     sub_questions: list[str] = Field(
         default_factory=list, description="Decomposed sub-questions"
     )
-    search_results: list[dict[str, Any]] = Field(
+    qa_responses: list["SearchAnswer"] = Field(
-        default_factory=list, description="Accumulated search results"
+        default_factory=list, description="Structured QA pairs used during research"
-    )
-    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"
@@ -28,26 +21,9 @@ class ResearchContext(BaseModel):
         default_factory=list, description="Identified information gaps"
     )
 
-    def add_search_result(self, query: str, results: list["SearchResult"]) -> None:
+    def add_qa_response(self, qa: "SearchAnswer") -> None:
-        """Add search results to context."""
+        """Add a structured QA response (minimal context already included)."""
-        self.search_results.append(
+        self.qa_responses.append(qa)
-            {
-                "query": query,
-                "results": results,
-            }
-        )
-
-    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."""

src/haiku/rag/research/orchestrator.py

@@ -61,7 +61,7 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
             "original_question": context.original_question,
             "unanswered_questions": context.sub_questions,
             "qa_responses": [
-                {"question": qa["question"], "answer": qa["answer"]}
+                {"question": qa.query, "answer": qa.answer}
                 for qa in context.qa_responses
             ],
             "insights": context.insights,
@@ -149,18 +149,23 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
             # Run searches for all questions and remove answered ones
             answered_questions = []
             for search_question in questions_to_search:
-                await self.search_agent.run(search_question, deps=deps)
+                try:
-
+                    await self.search_agent.run(search_question, deps=deps)
-                # Mark this question as answered
+                except Exception as e:  # pragma: no cover - defensive
-                answered_questions.append(search_question)
+                    if console:
+                        console.print(
+                            f"\n   [red]×[/red] Omitting failed question: {search_question} ({e})"
+                        )
+                finally:
+                    answered_questions.append(search_question)
 
                 if console and context.qa_responses:
                     # Show the last QA response (which should be for this question)
                     latest_qa = context.qa_responses[-1]
                     answer_preview = (
-                        latest_qa["answer"][:150] + "..."
+                        latest_qa.answer[:150] + "..."
-                        if len(latest_qa["answer"]) > 150
+                        if len(latest_qa.answer) > 150
-                        else latest_qa["answer"]
+                        else latest_qa.answer
                     )
                     console.print(
                         f"\n   [green]✓[/green] {search_question[:50]}..."

src/haiku/rag/research/prompts.py

@@ -11,24 +11,42 @@ Create a research plan that:
 - Breaks down the question into at most 3 focused sub-questions
 - Each sub-question should target a specific aspect of the research
 - Prioritize the most important aspects to investigate
-- Ensure comprehensive coverage within the 3-question limit"""
+- Ensure comprehensive coverage within the 3-question limit
+- IMPORTANT: Make each sub-question a standalone, self-contained query that can
+  be executed without additional context. Include necessary entities, scope,
+  timeframe, and qualifiers. Avoid pronouns like "it/they/this"; write queries
+  that make sense in isolation."""
 
 SEARCH_AGENT_PROMPT = """You are a search and question-answering specialist.
 
 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
+3. Provide an accurate answer strictly grounded in the retrieved context
-4. Base your answer strictly on the information found
 
-Use the search_and_answer tool to retrieve relevant documents and formulate your response.
+Output format:
+- You must return a SearchAnswer model with fields:
+  - query: the question being answered (echo the user query)
+  - answer: your final answer based only on the provided context
+  - context: list[str] of only the minimal set of verbatim snippet texts you
+    used to justify the answer (do not include unrelated text; do not invent)
+  - sources: list[str] of document_uri values corresponding to the snippets you
+    actually used in the answer (one URI per context snippet, order aligned)
+
+Tool usage:
+- Always call the search_and_answer tool before drafting any answer.
+- The tool returns XML containing only a list of snippets, where each snippet
+  has the verbatim `text`, a `score` indicating relevance, and the
+  `document_uri` it came from.
+- Use scores to prioritize evidence, but include only the minimal subset of
+  snippet texts (verbatim) in SearchAnswer.context.
+- Set SearchAnswer.sources to the matching document_uris for the snippets you
+  used (one URI per snippet, aligned by order). Context must be text-only.
+- If no relevant information is found, say so and return an empty context list.
 
 Important:
-- Always call the search_and_answer tool before drafting any answer.
+- Do not include any content in the answer that is not supported by the context.
-- Never answer without first using the tool at least once.
+- Keep context snippets short (just the necessary lines), verbatim, and focused."""
-- If no relevant information is found, state that clearly and avoid speculation.
-
-Be thorough and specific in your answers, citing relevant information from the sources."""
 
 EVALUATION_AGENT_PROMPT = """You are an analysis and evaluation specialist for research workflows.
 
@@ -67,7 +85,10 @@ Generate new sub-questions that:
 - Explore important edge cases or exceptions
 - Are focused and actionable (max 3)
 - Do NOT repeat or rephrase questions that have already been answered (see qa_responses)
-- Should be genuinely new areas to explore"""
+- Should be genuinely new areas to explore
+- Must be standalone, self-contained queries: include entities, scope, and any
+  needed qualifiers (e.g., timeframe, region), and avoid ambiguous pronouns so
+  they can be executed independently."""
 
 SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist agent focused on creating comprehensive research reports.
 

src/haiku/rag/research/search_agent.py

@@ -1,32 +1,30 @@
 from pydantic_ai import RunContext
+from pydantic_ai.format_prompt import format_as_xml
 from pydantic_ai.run import AgentRunResult
 
-from haiku.rag.research.base import BaseResearchAgent
+from haiku.rag.research.base import BaseResearchAgent, SearchAnswer
 from haiku.rag.research.dependencies import ResearchDependencies
 from haiku.rag.research.prompts import SEARCH_AGENT_PROMPT
 
 
-class SearchSpecialistAgent(BaseResearchAgent[str]):
+class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]):
     """Agent specialized in answering questions using RAG search."""
 
     def __init__(self, provider: str, model: str) -> None:
-        # Output is a string answer, like the QA agent
+        super().__init__(provider, model, output_type=SearchAnswer)
-        super().__init__(provider, model, output_type=str)
 
     async def run(
         self, prompt: str, deps: ResearchDependencies, **kwargs
-    ) -> AgentRunResult[str]:
+    ) -> AgentRunResult[SearchAnswer]:
-        """Execute the agent and store QA response in context."""
+        """Execute the agent and persist the QA pair in shared context.
-        # Run the base agent
+
+        Pydantic AI enforces `SearchAnswer` as the output model; we just store
+        the QA response with the last search results as sources.
+        """
         result = await super().run(prompt, deps, **kwargs)
 
-        # Store the QA response if we got an answer
         if result.output:
-            # Get the sources from the last search (which the tool just stored)
+            deps.context.add_qa_response(result.output)
-            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
 
@@ -42,37 +40,24 @@ class SearchSpecialistAgent(BaseResearchAgent[str]):
             query: str,
             limit: int = 5,
         ) -> str:
-            """Search for information and provide context for answering the question."""
+            """Search the KB and return a concise context pack."""
-            # Use the default hybrid search
+            # Remove quotes from queries as this requires positional indexing in lancedb
+            query = query.replace('"', "")
             search_results = await ctx.deps.client.search(query, limit=limit)
-
-            # Expand context for better relevance
             expanded = await ctx.deps.client.expand_context(search_results)
 
-            # Convert to SearchResult for context storage
+            snippet_entries = [
-            from haiku.rag.research.base import SearchResult
+                {
+                    "text": chunk.content,
+                    "score": score,
+                    "document_uri": (chunk.document_uri or ""),
+                }
+                for chunk, score in expanded
+            ]
 
-            results_for_context = []
+            # Return an XML-formatted payload with the question and snippets.
-            context_texts = []
+            if snippet_entries:
-
+                return format_as_xml(snippet_entries, root_tag="snippets")
-            for chunk, score in expanded:
-                results_for_context.append(
-                    SearchResult(
-                        content=chunk.content,
-                        score=score,
-                        document_uri=chunk.document_uri or "",
-                        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)
-
-            # Format context for the LLM to answer the question
-            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

@@ -1,4 +1,4 @@
-from haiku.rag.research.search_agent import SearchSpecialistAgent
+from haiku.rag.research import SearchAnswer, SearchSpecialistAgent
 
 
 class TestSearchSpecialistAgent:
@@ -8,4 +8,4 @@ class TestSearchSpecialistAgent:
         agent = SearchSpecialistAgent(provider="openai", model="gpt-4")
         assert agent.provider == "openai"
         assert agent.model == "gpt-4"
-        assert agent.output_type is str
+        assert agent.output_type is SearchAnswer