51 lines
1.4 KiB
Python
51 lines
1.4 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
|
|
|
|
|
|
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,
|
|
) -> tuple[str, list[SearchResult]]:
|
|
"""Search and context-expand results for a capability tool."""
|
|
results = await rag.search(query, limit=limit, filter=document_filter)
|
|
results = await rag.expand_context(results)
|
|
formatted = "\n\n---\n\n".join(
|
|
result.format_for_agent(rank=index + 1, total=len(results))
|
|
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 chunk id, which every stored chunk carries; results built by
|
|
hand without one cannot be told apart and collapse to the first.
|
|
"""
|
|
seen = {result.chunk_id for result in existing}
|
|
for result in incoming:
|
|
if result.chunk_id not in seen:
|
|
existing.append(result)
|
|
seen.add(result.chunk_id)
|
|
|
|
|
|
__all__ = [
|
|
"CodeExecutionEntry",
|
|
"merge_results",
|
|
"search_corpus",
|
|
]
|