fa319596cc
Search results and capsule entries name the collection they came from; the images attached beside them carried only the chunk id and reference. Two collections can return the same picture of the same document, so the two labels were identical and the model could place neither. The decision is the one already made for the text: `covers_multiple` at the generic search tool, and the flag `search_corpus` computed for the capability tools, which it now returns.
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], bool]:
|
|
"""Search and context-expand results, and whether they name their collection."""
|
|
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), include_collection
|
|
|
|
|
|
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",
|
|
]
|