86 lines
2.9 KiB
Python
86 lines
2.9 KiB
Python
from pydantic_ai import Agent, RunContext
|
|
|
|
from haiku.rag.agents.chat.prompts import SEARCH_SYSTEM_PROMPT
|
|
from haiku.rag.agents.chat.state import SearchDeps
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.store.models import SearchResult
|
|
from haiku.rag.utils import get_model
|
|
|
|
|
|
class SearchAgent:
|
|
"""Agent that generates multiple queries and consolidates results."""
|
|
|
|
def __init__(self, client: HaikuRAG, config: AppConfig):
|
|
self._client = client
|
|
self._config = config
|
|
|
|
model = get_model(config.qa.model, config)
|
|
self._agent: Agent[SearchDeps, str] = Agent(
|
|
model,
|
|
deps_type=SearchDeps,
|
|
output_type=str,
|
|
instructions=SEARCH_SYSTEM_PROMPT,
|
|
retries=3,
|
|
)
|
|
|
|
@self._agent.tool
|
|
async def run_search(
|
|
ctx: RunContext[SearchDeps],
|
|
query: str,
|
|
limit: int | None = None,
|
|
) -> str:
|
|
"""Run a single search query against the knowledge base.
|
|
|
|
Args:
|
|
query: The search query
|
|
limit: Number of results to fetch (default: 5)
|
|
"""
|
|
effective_limit = limit or 5
|
|
results = await ctx.deps.client.search(
|
|
query, limit=effective_limit, filter=ctx.deps.filter
|
|
)
|
|
results = await ctx.deps.client.expand_context(results)
|
|
ctx.deps.search_results.extend(results)
|
|
|
|
if not results:
|
|
return f"No results for: {query}"
|
|
return f"Found {len(results)} results for: {query}"
|
|
|
|
async def search(
|
|
self,
|
|
query: str,
|
|
context: str | None = None,
|
|
filter: str | None = None,
|
|
limit: int | None = None,
|
|
) -> list[SearchResult]:
|
|
"""Execute search with query expansion and deduplication.
|
|
|
|
Args:
|
|
query: The user's search request
|
|
context: Optional conversation context
|
|
filter: Optional SQL WHERE clause to filter documents
|
|
limit: Maximum number of results to return (default: config limit)
|
|
|
|
Returns:
|
|
Deduplicated list of SearchResult sorted by score
|
|
"""
|
|
prompt = query
|
|
if context:
|
|
prompt = f"Context: {context}\n\nSearch request: {query}"
|
|
|
|
deps = SearchDeps(client=self._client, config=self._config, filter=filter)
|
|
await self._agent.run(prompt, deps=deps)
|
|
|
|
# Deduplicate by chunk_id, keeping highest score
|
|
seen: dict[str, SearchResult] = {}
|
|
for result in deps.search_results:
|
|
chunk_id = result.chunk_id or ""
|
|
if chunk_id not in seen or result.score > seen[chunk_id].score:
|
|
seen[chunk_id] = result
|
|
|
|
# Sort by score descending and apply limit
|
|
effective_limit = limit or self._config.search.limit
|
|
return sorted(seen.values(), key=lambda r: r.score, reverse=True)[
|
|
:effective_limit
|
|
]
|