0d7810c78a
`format_for_agent` named the database whenever one was named, so a search over a single named database carried a line with nothing to distinguish. It now takes `include_collection` from the caller, which decides from the search selection rather than from the hits: a search that could have drawn on two collections names them even when everything came back from one. `Collection:` at the model boundary, database in configuration and administration. `source` on results, documents, citations and analysis dictionaries is unchanged.
61 lines
1.9 KiB
Python
61 lines
1.9 KiB
Python
from collections.abc import Iterable
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.store.models.chunk import SearchResult, qualified_id
|
|
|
|
|
|
class CodeExecutionEntry(BaseModel):
|
|
code: str
|
|
stdout: str
|
|
stderr: str = ""
|
|
success: bool = True
|
|
|
|
|
|
async def search_corpus(
|
|
rag: HaikuRAG,
|
|
query: str,
|
|
limit: int | None = None,
|
|
document_filter: str | None = None,
|
|
sources: list[str] | None = None,
|
|
) -> tuple[str, list[SearchResult]]:
|
|
"""Search and context-expand results for a capability tool."""
|
|
results = await rag.search(
|
|
query, limit=limit, filter=document_filter, sources=sources
|
|
)
|
|
results = await rag.expand_context(results)
|
|
# Named from the selection, not the hits: a search that could have drawn on
|
|
# two collections names them even when everything came back from one.
|
|
selected = rag.source_names if sources is None else sources
|
|
include_collection = len(set(selected)) > 1
|
|
formatted = "\n\n---\n\n".join(
|
|
result.format_for_agent(
|
|
rank=index + 1, total=len(results), include_collection=include_collection
|
|
)
|
|
for index, result in enumerate(results)
|
|
)
|
|
return formatted or "No results found.", list(results)
|
|
|
|
|
|
def merge_results(
|
|
existing: list[SearchResult], incoming: Iterable[SearchResult]
|
|
) -> None:
|
|
"""Add the results not already held.
|
|
|
|
Identity is the database and the chunk id: results built by hand carry
|
|
neither and cannot be told apart, so they collapse to the first.
|
|
"""
|
|
seen = {qualified_id(result.source, result.chunk_id) for result in existing}
|
|
for result in incoming:
|
|
key = qualified_id(result.source, result.chunk_id)
|
|
if key not in seen:
|
|
existing.append(result)
|
|
seen.add(key)
|
|
|
|
|
|
__all__ = [
|
|
"CodeExecutionEntry",
|
|
"merge_results",
|
|
"search_corpus",
|
|
]
|