From 05e08c48e4787d95caf07eb41ca9d3c5e17c62b8 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 17 Sep 2025 10:11:05 +0300 Subject: [PATCH] Simplify SearchAgent. Omit search results, but keep QA results. --- src/haiku/rag/config.py | 2 +- src/haiku/rag/research/__init__.py | 8 +++- src/haiku/rag/research/base.py | 35 ++++++++++++-- src/haiku/rag/research/dependencies.py | 36 +++----------- src/haiku/rag/research/orchestrator.py | 21 +++++---- src/haiku/rag/research/prompts.py | 41 ++++++++++++---- src/haiku/rag/research/search_agent.py | 65 ++++++++++---------------- tests/research/test_search_agent.py | 4 +- 8 files changed, 115 insertions(+), 97 deletions(-) diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index a51b0b4b..2f5366c9 100644 --- a/src/haiku/rag/config.py +++ b/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 diff --git a/src/haiku/rag/research/__init__.py b/src/haiku/rag/research/__init__.py index d947df6e..e0716ee8 100644 --- a/src/haiku/rag/research/__init__.py +++ b/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", diff --git a/src/haiku/rag/research/base.py b/src/haiku/rag/research/base.py index e547a50f..517bf26d 100644 --- a/src/haiku/rag/research/base.py +++ b/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, + ) diff --git a/src/haiku/rag/research/dependencies.py b/src/haiku/rag/research/dependencies.py index bb64bf1b..3438a796 100644 --- a/src/haiku/rag/research/dependencies.py +++ b/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 - -if TYPE_CHECKING: - from haiku.rag.research.base import SearchResult +from haiku.rag.research.base import SearchAnswer 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( - default_factory=list, description="Accumulated search results" - ) - qa_responses: list[dict[str, Any]] = Field( - default_factory=list, description="Question-answer pairs with sources" + qa_responses: list["SearchAnswer"] = Field( + default_factory=list, description="Structured QA pairs used during research" ) 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: - """Add search results to context.""" - self.search_results.append( - { - "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_qa_response(self, qa: "SearchAnswer") -> None: + """Add a structured QA response (minimal context already included).""" + self.qa_responses.append(qa) def add_insight(self, insight: str) -> None: """Add a key insight.""" diff --git a/src/haiku/rag/research/orchestrator.py b/src/haiku/rag/research/orchestrator.py index 2e43c8d7..ffde5736 100644 --- a/src/haiku/rag/research/orchestrator.py +++ b/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) - - # Mark this question as answered - answered_questions.append(search_question) + try: + await self.search_agent.run(search_question, deps=deps) + except Exception as e: # pragma: no cover - defensive + 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] + "..." - if len(latest_qa["answer"]) > 150 - else latest_qa["answer"] + latest_qa.answer[:150] + "..." + if len(latest_qa.answer) > 150 + else latest_qa.answer ) console.print( f"\n [green]✓[/green] {search_question[:50]}..." diff --git a/src/haiku/rag/research/prompts.py b/src/haiku/rag/research/prompts.py index 05ba7be3..ca540f0f 100644 --- a/src/haiku/rag/research/prompts.py +++ b/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 -4. Base your answer strictly on the information found +3. Provide an accurate answer strictly grounded in the retrieved context -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. -- Never answer without first using the tool at least once. -- 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.""" +- Do not include any content in the answer that is not supported by the context. +- Keep context snippets short (just the necessary lines), verbatim, and focused.""" 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. diff --git a/src/haiku/rag/research/search_agent.py b/src/haiku/rag/research/search_agent.py index 82e937c4..200f619a 100644 --- a/src/haiku/rag/research/search_agent.py +++ b/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=str) + super().__init__(provider, model, output_type=SearchAnswer) async def run( self, prompt: str, deps: ResearchDependencies, **kwargs - ) -> AgentRunResult[str]: - """Execute the agent and store QA response in context.""" - # Run the base agent + ) -> AgentRunResult[SearchAnswer]: + """Execute the agent and persist the QA pair in shared context. + + 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) - 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) + deps.context.add_qa_response(result.output) 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.""" - # Use the default hybrid search + """Search the KB and return a concise context pack.""" + # 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 - from haiku.rag.research.base import SearchResult + snippet_entries = [ + { + "text": chunk.content, + "score": score, + "document_uri": (chunk.document_uri or ""), + } + for chunk, score in expanded + ] - results_for_context = [] - context_texts = [] - - 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}" + # Return an XML-formatted payload with the question and snippets. + if snippet_entries: + return format_as_xml(snippet_entries, root_tag="snippets") else: return ( f"No relevant information found in the knowledge base for: {query}" diff --git a/tests/research/test_search_agent.py b/tests/research/test_search_agent.py index d6d18b1e..a7a9062c 100644 --- a/tests/research/test_search_agent.py +++ b/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