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

Simplify search agent

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-09-15 15:28:50 +03:00
parent 4df5d31b18
commit 53905c2d69
No known key found for this signature in database
3 changed files with 91 additions and 97 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -22,18 +22,18 @@ class ResearchPlan(BaseModel):
sub_questions: list[str] = Field(
description="Decomposed sub-questions to investigate"
)
search_strategies: list[str] = Field(
description="Different search approaches to use"
)
success_criteria: list[str] = Field(description="Criteria for successful research")
class ResearchOrchestrator(BaseResearchAgent):
"""Orchestrator agent that coordinates the research workflow."""
def __init__(
self, provider: str = Config.RERANK_PROVIDER, model: str = Config.RERANK_MODEL
self, provider: str | None = Config.RESEARCH_PROVIDER, model: str | None = None
):
# Use provided values or fall back to config defaults
provider = provider or Config.RESEARCH_PROVIDER or Config.QA_PROVIDER
model = model or Config.RESEARCH_MODEL or Config.QA_MODEL
super().__init__(provider, model, output_type=ResearchPlan)
self.search_agent = SearchSpecialistAgent(provider, model)
@ -55,7 +55,8 @@ class ResearchOrchestrator(BaseResearchAgent):
- Breaks down complex questions into manageable parts
- Identifies multiple search strategies
- Defines clear success criteria
- Ensures thorough investigation"""
- Ensures thorough investigation
/no_think"""
def register_tools(self) -> None:
"""Register orchestration tools."""
@ -63,13 +64,21 @@ class ResearchOrchestrator(BaseResearchAgent):
@self.agent.tool
async def delegate_search(
ctx: RunContext[ResearchDependencies], queries: list[str], limit: int = 5
) -> Any:
"""Delegate search to the search specialist agent."""
# Pass the context to maintain usage tracking
result = await self.search_agent.run(
f"Search for: {', '.join(queries)}", deps=ctx.deps, usage=ctx.usage
)
return result
) -> list[Any]:
"""Delegate search to the search specialist agent for multiple queries."""
all_results = []
# Search for each query
# The search agent will automatically store results in context
for query in queries:
result = await self.search_agent.run(
f"Search for: {query} with limit {limit}",
deps=ctx.deps,
usage=ctx.usage,
)
all_results.append(result)
return all_results
@self.agent.tool
async def delegate_analysis(
@ -158,10 +167,8 @@ class ResearchOrchestrator(BaseResearchAgent):
f"Create a research plan for: {question}", deps=deps
)
if hasattr(plan_result, "output") and isinstance(
plan_result.output, ResearchPlan
):
context.sub_questions = plan_result.output.sub_questions
assert plan_result.output and isinstance(plan_result.output, ResearchPlan)
context.sub_questions = plan_result.output.sub_questions
# Execute research iterations
for iteration in range(max_iterations):
@ -176,16 +183,18 @@ class ResearchOrchestrator(BaseResearchAgent):
else:
# Fall back to original question with variation
search_prompt = f"Additional search for: {question}"
# Search phase - directly call the search agent
# Search phase
await self.run(search_prompt, deps=deps)
await self.search_agent.run(search_prompt, deps=deps)
# Analysis phase (only if we have results)
if context.search_results:
await self.run("Analyze the gathered information", deps=deps)
await self.analysis_agent.run(
"Analyze the gathered information", deps=deps
)
# Clarification phase - evaluate completeness
clarification_result = await self.run(
clarification_result = await self.clarification_agent.run(
f"Evaluate the completeness of research for: {question}. "
f"Consider all information gathered so far and determine if we have sufficient "
f"information to provide a comprehensive answer.",
@ -202,7 +211,9 @@ class ResearchOrchestrator(BaseResearchAgent):
break
# Generate final report
report_result = await self.run("Generate the final research report", deps=deps)
report_result = await self.synthesis_agent.run(
"Generate the final research report", deps=deps
)
return (
report_result.output if hasattr(report_result, "output") else report_result
)

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

@ -1,26 +1,26 @@
"""Search specialist agent for advanced document retrieval."""
"""Search specialist agent for document retrieval."""
from pydantic_ai import RunContext
from haiku.rag.research.base import BaseResearchAgent, SearchResult
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 advanced document search and retrieval."""
"""Agent specialized in document search and retrieval."""
def __init__(self, provider: str, model: str):
super().__init__(provider, model, output_type=list[SearchResult])
# No specific output type needed - the tool handles everything
super().__init__(provider, model)
def get_system_prompt(self) -> str:
return """You are a search specialist agent focused on document retrieval.
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. Generate multiple search queries from different perspectives
2. Identify key terms and synonyms for comprehensive search
3. Execute searches and rank results by relevance
4. Return the most relevant documents for the research question
1. Understand the search query and context
2. Execute targeted searches to find relevant documents
Use the search tools to explore the knowledge base thoroughly."""
Use the search tool to perform the searches on the knowledge base."""
def register_tools(self) -> None:
"""Register search-specific tools."""
@ -28,46 +28,30 @@ class SearchSpecialistAgent(BaseResearchAgent):
@self.agent.tool
async def search(
ctx: RunContext[ResearchDependencies],
queries: str | list[str],
query: str,
limit: int = 5,
) -> list[SearchResult]:
"""Execute search with single or multiple query variants."""
# Normalize to list
query_list = [queries] if isinstance(queries, str) else queries
) -> list[tuple[Chunk, float]]:
"""Execute search and return raw results from client."""
# Use the default hybrid search
search_results = await ctx.deps.client.search(query, limit=limit)
all_results = []
seen_chunk_ids = set()
# Expand context for better relevance
expanded = await ctx.deps.client.expand_context(search_results)
for query in query_list:
# Use the default hybrid search
search_results = await ctx.deps.client.search(query, limit=limit)
# Store in context (convert to SearchResult for context storage)
from haiku.rag.research.base import SearchResult
# Expand context for better relevance
expanded = await ctx.deps.client.expand_context(search_results)
results_for_context = []
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 {},
)
)
ctx.deps.context.add_search_result(query, results_for_context)
for chunk, score in expanded:
# Avoid duplicates based on chunk ID
if chunk.id and chunk.id not in seen_chunk_ids:
seen_chunk_ids.add(chunk.id)
all_results.append(
SearchResult(
content=chunk.content,
score=score,
document_uri=chunk.document_uri or "",
metadata={"query": query, "chunk_id": chunk.id},
)
)
# Sort by score and limit results
all_results.sort(key=lambda x: x.score, reverse=True)
final_results = all_results[: limit * len(query_list)]
# Store in context
query_summary = (
query_list[0]
if len(query_list) == 1
else f"Multi-query: {len(query_list)} variants"
)
ctx.deps.context.add_search_result(query_summary, final_results)
return final_results
# Return raw chunk, score tuples
return expanded

53
tests/research/test_search_agent.py
View file

@ -101,12 +101,14 @@ async def test_search_single_query(mock_client, research_deps):
# Test the tool
ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage())
results = await search_tool(ctx, queries="climate change")
results = await search_tool(ctx, query="climate change")
# Verify results
# Verify results - should be list of (Chunk, float) tuples
assert isinstance(results, list)
assert len(results) == 2
assert results[0].content == "Climate change is a global phenomenon"
assert results[0].metadata["chunk_id"] == "chunk1"
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 mock was called
mock_client.search.assert_called_once_with("climate change", limit=5)
@ -114,20 +116,19 @@ async def test_search_single_query(mock_client, research_deps):
@pytest.mark.asyncio
async def test_search_multiple_queries_deduplication(mock_client, research_deps):
"""Test deduplication when searching with multiple queries."""
# Create chunks with duplicate IDs across queries
chunks_q1 = [
async def test_search_with_limit(mock_client, research_deps):
"""Test that search respects the limit parameter."""
# Create more chunks than the limit
mock_chunks = [
create_mock_chunk("chunk1", "Content 1", 0.9),
create_mock_chunk("chunk2", "Content 2", 0.7),
]
chunks_q2 = [
create_mock_chunk("chunk1", "Content 1", 0.9), # Duplicate
create_mock_chunk("chunk3", "Content 3", 0.8),
create_mock_chunk("chunk2", "Content 2", 0.8),
create_mock_chunk("chunk3", "Content 3", 0.7),
create_mock_chunk("chunk4", "Content 4", 0.6),
create_mock_chunk("chunk5", "Content 5", 0.5),
]
mock_client.search.side_effect = [[chunks_q1[0]], [chunks_q2[0]]]
mock_client.expand_context.side_effect = [chunks_q1, chunks_q2]
mock_client.search.return_value = mock_chunks[:3]
mock_client.expand_context.return_value = mock_chunks[:3]
agent = SearchSpecialistAgent(provider="openai", model="gpt-4")
@ -135,20 +136,18 @@ async def test_search_multiple_queries_deduplication(mock_client, research_deps)
search_tool = get_agent_tool(agent, "search")
assert search_tool is not None
# Test the tool
# Test the tool with limit
ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage())
results = await search_tool(ctx, queries=["query1", "query2"])
results = await search_tool(ctx, query="test query", limit=3)
# Check deduplication - chunk1 should appear only once
chunk_ids = [r.metadata["chunk_id"] for r in results]
assert chunk_ids.count("chunk1") == 1
assert "chunk2" in chunk_ids
assert "chunk3" in chunk_ids
# 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 sorting by score
assert all(
results[i].score >= results[i + 1].score for i in range(len(results) - 1)
)
# Verify mock was called with correct limit
mock_client.search.assert_called_once_with("test query", limit=3)
@pytest.mark.asyncio
@ -167,7 +166,7 @@ async def test_search_updates_context(mock_client, research_deps):
# Test the tool
ctx = RunContext(deps=research_deps, model=TestModel(), usage=RunUsage())
await search_tool(ctx, queries="test query")
await search_tool(ctx, query="test query")
# Verify context was updated
assert len(research_deps.context.search_results) == 1