Simplify SearchAgent. Omit search results, but keep QA results.
This commit is contained in:
Yiorgis Gozadinos
parent
99a737e5b8
commit
05e08c48e4
8 changed files with 115 additions and 97 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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]}..."
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue