danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Make search agent more like a qa agent

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-09-15 16:17:50 +03:00
parent 9dc2e05f3f
commit c029350d79
No known key found for this signature in database
4 changed files with 100 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
src/haiku/rag/research/dependencies.py
View file

@ -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:

34
src/haiku/rag/research/orchestrator.py
View file

@ -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
search_prompt = f"Additional search for: {question}"
# Fall back to original 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:
latest_results = context.search_results[-1]
console.print(
f" Found [green]{len(latest_results.get('results', []))} documents[/green]"
)
for i, result in enumerate(latest_results.get("results", [])[:3], 1):
if console:
# Show documents found
if context.search_results:
latest_results = context.search_results[-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:

63
src/haiku/rag/research/search_agent.py
View file

@ -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
super().__init__(provider, model)
# Output is a string answer, like the QA agent
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.
Your role is to:
1. Understand the search query and context
2. Execute targeted searches to find relevant documents
return """You are a search and question-answering specialist.
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]]:
"""Execute search and return raw results from client."""
) -> str:
"""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
return expanded
# 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}"
)

34
tests/research/test_search_agent.py
View file

@ -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
assert isinstance(results, list)
assert len(results) == 2
assert results[0][0].content == "Climate change is a global phenomenon"
assert results[0][0].id == "chunk1"
assert results[0][1] == 0.8 # score
# Verify result - should be a formatted string with context
assert isinstance(result, str)
assert "Climate change is a global phenomenon" in result
assert "Rising temperatures affect ecosystems" in result
# 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
assert isinstance(results, list)
assert len(results) == 3
assert all(isinstance(r, tuple) and len(r) == 2 for r in results)
assert all(isinstance(r[0], Chunk) and isinstance(r[1], float) for r in results)
# Verify result is a formatted string
assert isinstance(result, str)
assert "Content 1" in result
assert "Content 2" in result
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