89 lines
3 KiB
Python
89 lines
3 KiB
Python
from collections.abc import Callable
|
|
|
|
from pydantic_ai import FunctionToolset, RunContext
|
|
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.store.models import SearchResult
|
|
from haiku.rag.tools.context import RAGDeps
|
|
|
|
|
|
def create_search_toolset(
|
|
config: AppConfig,
|
|
expand_context: bool = True,
|
|
base_filter: str | None = None,
|
|
tool_name: str = "search",
|
|
on_results: Callable[[list[SearchResult]], None] | None = None,
|
|
max_searches: int | None = None,
|
|
) -> FunctionToolset[RAGDeps]:
|
|
"""Create a toolset with search capabilities.
|
|
|
|
Args:
|
|
config: Application configuration.
|
|
expand_context: Whether to expand search results with surrounding context.
|
|
Defaults to True.
|
|
base_filter: Optional base SQL WHERE clause applied to all searches.
|
|
Combined with any filter passed to the search tool.
|
|
tool_name: Name for the search tool. Defaults to "search".
|
|
on_results: Optional callback invoked with search results after each search.
|
|
Useful for accumulating results externally (e.g., for citation resolution).
|
|
max_searches: Maximum number of searches allowed. When exceeded, returns
|
|
a message directing the agent to answer with existing results.
|
|
|
|
Returns:
|
|
FunctionToolset with a search tool.
|
|
"""
|
|
# Per-run search counter keyed by run_id. Safe for concurrent runs
|
|
# and reuse across sequential agent.run() calls.
|
|
search_counts: dict[str, int] = {}
|
|
|
|
async def search(
|
|
ctx: RunContext[RAGDeps],
|
|
query: str,
|
|
limit: int | None = None,
|
|
) -> str:
|
|
"""Search the knowledge base for relevant documents.
|
|
|
|
Args:
|
|
query: The search query (what to search for).
|
|
limit: Number of results to return (default: from config).
|
|
|
|
Returns:
|
|
Formatted search results with content and metadata.
|
|
"""
|
|
rid = ctx.run_id or ""
|
|
search_counts[rid] = search_counts.get(rid, 0) + 1
|
|
if max_searches is not None and search_counts[rid] > max_searches:
|
|
return (
|
|
"Search limit reached. "
|
|
"Answer the question using the results you already have."
|
|
)
|
|
|
|
client = ctx.deps.client
|
|
|
|
effective_filter = base_filter
|
|
effective_limit = limit or config.search.limit
|
|
results = await client.search(
|
|
query, limit=effective_limit, filter=effective_filter
|
|
)
|
|
|
|
if expand_context:
|
|
results = await client.expand_context(results)
|
|
|
|
results_list = list(results)
|
|
|
|
if on_results:
|
|
on_results(results_list)
|
|
|
|
if not results_list:
|
|
return "No results found."
|
|
|
|
total = len(results_list)
|
|
formatted = [
|
|
r.format_for_agent(rank=i + 1, total=total)
|
|
for i, r in enumerate(results_list)
|
|
]
|
|
return "\n\n".join(formatted)
|
|
|
|
toolset: FunctionToolset[RAGDeps] = FunctionToolset()
|
|
toolset.add_function(search, name=tool_name, retries=3)
|
|
return toolset
|