162 lines
6 KiB
Python
162 lines
6 KiB
Python
from pydantic import BaseModel
|
|
from pydantic_ai import FunctionToolset, ToolReturn
|
|
|
|
from haiku.rag.agents.research.models import Citation
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.store.models import SearchResult
|
|
from haiku.rag.tools.context import ToolContext
|
|
from haiku.rag.tools.filters import build_multi_document_filter, combine_filters
|
|
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState, compute_state_delta
|
|
|
|
SEARCH_NAMESPACE = "haiku.rag.search"
|
|
|
|
|
|
class SearchState(BaseModel):
|
|
"""State for search toolset.
|
|
|
|
Accumulates search results across tool invocations.
|
|
"""
|
|
|
|
results: list[SearchResult] = []
|
|
|
|
|
|
def create_search_toolset(
|
|
client: HaikuRAG,
|
|
config: AppConfig,
|
|
context: ToolContext | None = None,
|
|
expand_context: bool = True,
|
|
base_filter: str | None = None,
|
|
tool_name: str = "search",
|
|
) -> FunctionToolset:
|
|
"""Create a toolset with search capabilities.
|
|
|
|
Args:
|
|
client: HaikuRAG client for search operations.
|
|
config: Application configuration.
|
|
context: Optional ToolContext for state accumulation.
|
|
If provided, search results are accumulated in SearchState.
|
|
If SessionState is registered, it will be used for dynamic
|
|
document filtering and citation indexing.
|
|
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".
|
|
|
|
Returns:
|
|
FunctionToolset with a search tool.
|
|
"""
|
|
# Get or create search state if context provided
|
|
search_state: SearchState | None = None
|
|
if context is not None:
|
|
search_state = context.get_or_create(SEARCH_NAMESPACE, SearchState)
|
|
|
|
async def search(
|
|
query: str,
|
|
limit: int | None = None,
|
|
filter: str | None = None,
|
|
) -> ToolReturn | 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).
|
|
filter: Optional SQL WHERE clause to filter documents.
|
|
|
|
Returns:
|
|
Formatted search results with content and metadata.
|
|
"""
|
|
# Get session state for dynamic filters and citation indexing
|
|
session_state: SessionState | None = None
|
|
old_session_state: SessionState | None = None
|
|
if context is not None:
|
|
session_state = context.get_typed(SESSION_NAMESPACE, SessionState)
|
|
if session_state is not None:
|
|
old_session_state = session_state.model_copy(deep=True)
|
|
|
|
# Build session filter from session state's document_filter
|
|
session_filter = None
|
|
if session_state is not None and session_state.document_filter:
|
|
session_filter = build_multi_document_filter(session_state.document_filter)
|
|
|
|
# Combine all filters: base_filter AND session_filter AND tool filter
|
|
effective_filter = combine_filters(
|
|
combine_filters(base_filter, session_filter), 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)
|
|
|
|
# Accumulate results in search state if context provided
|
|
if search_state is not None:
|
|
search_state.results.extend(results)
|
|
|
|
if not results:
|
|
return "No results found."
|
|
|
|
# Build citations if session state is available
|
|
if session_state is not None:
|
|
citations = []
|
|
for r in results:
|
|
chunk_id = r.chunk_id or ""
|
|
if chunk_id:
|
|
index = session_state.get_or_assign_index(chunk_id)
|
|
else:
|
|
index = len(session_state.citation_registry) + 1
|
|
citations.append(
|
|
Citation(
|
|
index=index,
|
|
document_id=r.document_id or "",
|
|
chunk_id=chunk_id,
|
|
document_uri=r.document_uri or "",
|
|
document_title=r.document_title,
|
|
page_numbers=r.page_numbers or [],
|
|
headings=r.headings,
|
|
content=r.content,
|
|
)
|
|
)
|
|
session_state.citations = citations
|
|
|
|
# Format results with citation indices
|
|
result_lines = []
|
|
for c in citations:
|
|
title = c.document_title or c.document_uri or "Unknown"
|
|
snippet = c.content[:300].replace("\n", " ").strip()
|
|
if len(c.content) > 300:
|
|
snippet += "..."
|
|
|
|
line = f"[{c.index}] **{title}**"
|
|
if c.page_numbers:
|
|
line += f" (pages {', '.join(map(str, c.page_numbers))})"
|
|
line += f"\n {snippet}"
|
|
result_lines.append(line)
|
|
|
|
formatted = f"Found {len(results)} results:\n\n" + "\n\n".join(result_lines)
|
|
|
|
# Compute state delta if session state changed
|
|
if old_session_state is not None:
|
|
state_event = compute_state_delta(old_session_state, session_state)
|
|
if state_event is not None:
|
|
return ToolReturn(
|
|
return_value=formatted,
|
|
metadata=[state_event],
|
|
)
|
|
|
|
return formatted
|
|
|
|
# Format results without citation indexing (standalone use)
|
|
total = len(results)
|
|
formatted = [
|
|
r.format_for_agent(rank=i + 1, total=total) for i, r in enumerate(results)
|
|
]
|
|
return "\n\n".join(formatted)
|
|
|
|
toolset = FunctionToolset()
|
|
toolset.add_function(search, name=tool_name)
|
|
return toolset
|