391 lines
14 KiB
Python
391 lines
14 KiB
Python
from collections.abc import AsyncIterator
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel
|
|
from pydantic_ai import ModelRetry, RunContext
|
|
from pydantic_ai.messages import ToolReturn
|
|
|
|
from haiku.rag.client import HaikuRAG
|
|
from haiku.rag.config.models import AppConfig, ModelConfig
|
|
from haiku.rag.skills._deps import AnalysisRunDeps, RAGRunDeps
|
|
from haiku.rag.store.models.chunk import SearchResult
|
|
from haiku.rag.store.models.citation import Citation
|
|
from haiku.rag.tools.search import build_binary_parts_from_results
|
|
|
|
|
|
class CodeExecutionEntry(BaseModel):
|
|
code: str
|
|
stdout: str
|
|
stderr: str = ""
|
|
success: bool = True
|
|
|
|
|
|
async def skill_search(
|
|
rag: HaikuRAG,
|
|
query: str,
|
|
limit: int | None = None,
|
|
document_filter: str | None = None,
|
|
) -> tuple[str, list[SearchResult]]:
|
|
results = await rag.search(query, limit=limit, filter=document_filter)
|
|
results = await rag.expand_context(results)
|
|
formatted = "\n\n---\n\n".join(
|
|
r.format_for_agent(rank=i + 1, total=len(results))
|
|
for i, r in enumerate(results)
|
|
)
|
|
return formatted, list(results)
|
|
|
|
|
|
async def skill_list_documents(
|
|
rag: HaikuRAG,
|
|
filter: str | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
documents = await rag.list_documents(filter=filter)
|
|
return [
|
|
{
|
|
"id": doc.id,
|
|
"title": doc.title,
|
|
"uri": doc.uri,
|
|
"metadata": doc.metadata,
|
|
"created_at": str(doc.created_at),
|
|
"updated_at": str(doc.updated_at),
|
|
}
|
|
for doc in documents
|
|
]
|
|
|
|
|
|
async def skill_get_document(
|
|
rag: HaikuRAG,
|
|
query: str,
|
|
) -> dict[str, Any] | None:
|
|
document = await rag.resolve_document(query)
|
|
if document is None:
|
|
return None
|
|
return {
|
|
"id": document.id,
|
|
"content": document.content,
|
|
"title": document.title,
|
|
"uri": document.uri,
|
|
"metadata": document.metadata,
|
|
"created_at": str(document.created_at),
|
|
"updated_at": str(document.updated_at),
|
|
}
|
|
|
|
|
|
def _get_state(ctx: RunContext[RAGRunDeps], state_type: type[BaseModel]) -> Any:
|
|
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, state_type):
|
|
return ctx.deps.state
|
|
return None
|
|
|
|
|
|
def _require_rag(ctx: RunContext[RAGRunDeps]) -> HaikuRAG:
|
|
assert ctx.deps is not None and ctx.deps.rag is not None, (
|
|
"RAGRunDeps.rag is not set — skill lifespan must run before tools."
|
|
)
|
|
return ctx.deps.rag
|
|
|
|
|
|
@asynccontextmanager
|
|
async def _serialized(ctx: RunContext[RAGRunDeps]) -> AsyncIterator[None]:
|
|
"""Serialize access to the shared connection through the run's lock.
|
|
|
|
pydantic-ai runs a turn's tool calls concurrently and LanceDB's
|
|
per-connection state cannot take two in-flight operations at once. The lock
|
|
is always present (``RAGRunDeps`` creates one by default); the no-op branch
|
|
is a guard for a missing deps/lock.
|
|
"""
|
|
lock = ctx.deps.rag_lock if ctx.deps is not None else None
|
|
if lock is None:
|
|
yield
|
|
else:
|
|
async with lock:
|
|
yield
|
|
|
|
|
|
def _register_citations(state: Any, citations: "list[Citation]") -> None:
|
|
"""Add citations to the index and record cited chunk IDs for this invocation."""
|
|
next_index = len(state.citation_index) + 1
|
|
for citation in citations:
|
|
cid = citation.chunk_id
|
|
if cid not in state.citation_index:
|
|
citation.index = next_index
|
|
next_index += 1
|
|
state.citation_index[cid] = citation
|
|
if cid not in state.citations:
|
|
state.citations.append(cid)
|
|
|
|
|
|
def create_skill_extras(
|
|
db_path: Path,
|
|
config: AppConfig,
|
|
) -> dict[str, Any]:
|
|
"""Create non-tool utility functions bound to a specific database.
|
|
|
|
Returns a dict of values that can be attached to a Skill's extras:
|
|
|
|
Keys:
|
|
- 'db_path': path to the LanceDB used to configure the skill
|
|
- 'config': config passed to (or derived for) the skill
|
|
- 'list_documents': returns info for documents in the database
|
|
- 'visualize_chunk': returns visualizations for chunks in the database
|
|
"""
|
|
|
|
async def visualize_chunk(chunk_id: str) -> list:
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
|
|
chunk = await rag.get_chunk_by_id(chunk_id)
|
|
if chunk is None:
|
|
return []
|
|
return await rag.visualize_chunk(chunk)
|
|
|
|
async def list_documents(
|
|
limit: int | None = None,
|
|
offset: int | None = None,
|
|
filter: str | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
from haiku.rag.client import HaikuRAG
|
|
|
|
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
|
|
documents = await rag.list_documents(limit, offset, filter=filter)
|
|
return [
|
|
{
|
|
"id": doc.id,
|
|
"title": doc.title,
|
|
"uri": doc.uri,
|
|
"metadata": doc.metadata,
|
|
"created_at": str(doc.created_at),
|
|
"updated_at": str(doc.updated_at),
|
|
}
|
|
for doc in documents
|
|
]
|
|
|
|
return {
|
|
"db_path": db_path,
|
|
"config": config,
|
|
"visualize_chunk": visualize_chunk,
|
|
"list_documents": list_documents,
|
|
}
|
|
|
|
|
|
def create_skill_tools(
|
|
db_path: Path,
|
|
config: AppConfig,
|
|
state_type: type[BaseModel],
|
|
tool_names: list[str],
|
|
model: ModelConfig,
|
|
) -> dict[str, Any]:
|
|
"""Create tool closures for a skill.
|
|
|
|
Returns a dict mapping tool name to async callable.
|
|
Each tool extracts state from RunContext, calls the shared implementation,
|
|
and updates state. ``model`` is the driving model for the skill (e.g.
|
|
``config.qa.model`` for the RAG skill, or
|
|
``config.analysis.model or config.qa.model`` for the analysis skill,
|
|
which defaults to ``None`` and inherits QA's model when unconfigured);
|
|
its ``vision`` flag gates picture-bytes attachment on the ``search``
|
|
tool.
|
|
"""
|
|
tools: dict[str, Any] = {}
|
|
|
|
if "search" in tool_names:
|
|
max_searches = config.qa.max_searches
|
|
|
|
async def search(
|
|
ctx: RunContext[RAGRunDeps], query: str, limit: int | None = None
|
|
) -> str | ToolReturn:
|
|
"""Search the knowledge base using hybrid search (vector + full-text).
|
|
|
|
Returns ranked results with content and metadata. When picture
|
|
content is in the result set and the driving skill model is
|
|
vision-capable, picture bytes are attached as ``BinaryContent``
|
|
parts so the model sees figures alongside text.
|
|
|
|
Args:
|
|
query: The search query.
|
|
limit: Maximum number of results.
|
|
"""
|
|
ctx.deps.search_count += 1
|
|
if ctx.deps.search_count > max_searches:
|
|
return (
|
|
"Search limit reached. Answer the question using "
|
|
"the results you already have."
|
|
)
|
|
|
|
state = _get_state(ctx, state_type)
|
|
async with _serialized(ctx):
|
|
formatted, results = await skill_search(
|
|
_require_rag(ctx),
|
|
query,
|
|
limit=limit,
|
|
document_filter=state.document_filter if state else None,
|
|
)
|
|
if state:
|
|
state.searches[query] = results
|
|
|
|
if not model.vision:
|
|
return formatted
|
|
|
|
binary_parts = build_binary_parts_from_results(results)
|
|
if binary_parts:
|
|
return ToolReturn(return_value=formatted, content=binary_parts)
|
|
return formatted
|
|
|
|
tools["search"] = search
|
|
|
|
if "list_documents" in tool_names:
|
|
|
|
async def list_documents(
|
|
ctx: RunContext[RAGRunDeps],
|
|
) -> list[dict[str, Any]]:
|
|
"""List all documents in the knowledge base."""
|
|
state = _get_state(ctx, state_type)
|
|
async with _serialized(ctx):
|
|
return await skill_list_documents(
|
|
_require_rag(ctx),
|
|
filter=state.document_filter if state else None,
|
|
)
|
|
|
|
tools["list_documents"] = list_documents
|
|
|
|
if "get_document" in tool_names:
|
|
|
|
async def get_document(
|
|
ctx: RunContext[RAGRunDeps], query: str
|
|
) -> dict[str, Any] | None:
|
|
"""Retrieve a document by ID, title, or URI.
|
|
|
|
Args:
|
|
query: Document ID, title, or URI to look up.
|
|
"""
|
|
async with _serialized(ctx):
|
|
return await skill_get_document(_require_rag(ctx), query)
|
|
|
|
tools["get_document"] = get_document
|
|
|
|
if "execute_code" in tool_names:
|
|
max_executions = config.analysis.max_executions
|
|
|
|
async def execute_code(ctx: RunContext[AnalysisRunDeps], code: str) -> str:
|
|
"""Execute Python code in a sandboxed interpreter.
|
|
|
|
The code has access to search() and list_documents() functions
|
|
and a virtual filesystem at /documents/ with document content
|
|
and structure (metadata.json, content.txt, items.jsonl, toc.json
|
|
per document).
|
|
|
|
Use print() to output results. Variables persist between calls
|
|
within the same skill invocation.
|
|
|
|
Args:
|
|
code: Python code to execute.
|
|
"""
|
|
ctx.deps.execute_count += 1
|
|
if ctx.deps.execute_count > max_executions:
|
|
return (
|
|
"Code-execution limit reached. Give your final answer now "
|
|
"from what you already have; do not call execute_code again."
|
|
)
|
|
|
|
assert ctx.deps is not None and ctx.deps.sandbox is not None, (
|
|
"AnalysisRunDeps.sandbox is not set — skill lifespan must run before execute_code."
|
|
)
|
|
sandbox = ctx.deps.sandbox
|
|
result = await sandbox.execute(code)
|
|
|
|
state = _get_state(ctx, state_type)
|
|
if state and sandbox._search_results:
|
|
existing = state.searches.get("_sandbox", [])
|
|
seen = {r.chunk_id for r in existing}
|
|
for sr in sandbox._search_results:
|
|
if sr.chunk_id not in seen:
|
|
existing.append(sr)
|
|
seen.add(sr.chunk_id)
|
|
state.searches["_sandbox"] = existing
|
|
|
|
if state:
|
|
state.executions.append(
|
|
CodeExecutionEntry(
|
|
code=code,
|
|
stdout=result.stdout,
|
|
stderr=result.stderr,
|
|
success=result.success,
|
|
)
|
|
)
|
|
|
|
if result.success:
|
|
return result.stdout if result.stdout else "No output."
|
|
return f"Error: {result.stderr}\n\nOutput: {result.stdout}"
|
|
|
|
tools["execute_code"] = execute_code
|
|
|
|
if "cite" in tool_names:
|
|
|
|
async def cite(ctx: RunContext[RAGRunDeps], chunk_ids: list[str]) -> str:
|
|
"""Register chunk IDs as citations for your answer.
|
|
|
|
Accepts chunk_ids from search results AND from direct file reads
|
|
(items.jsonl, toc.json). Verbatim copies only — chunk_ids that
|
|
don't exist in the database trigger a retry.
|
|
|
|
Args:
|
|
chunk_ids: List of chunk_id values from search results or VFS reads.
|
|
"""
|
|
from haiku.rag.store.models.chunk import SearchResult
|
|
from haiku.rag.store.models.citation import resolve_citations
|
|
|
|
state = _get_state(ctx, state_type)
|
|
if not state:
|
|
return "No state available."
|
|
|
|
if not chunk_ids:
|
|
return "Registered 0 citations (empty chunk_ids)."
|
|
|
|
all_results: list[SearchResult] = []
|
|
for results_list in state.searches.values():
|
|
all_results.extend(results_list)
|
|
|
|
citations = resolve_citations(chunk_ids, all_results)
|
|
resolved_ids = {c.chunk_id for c in citations}
|
|
missing = [
|
|
cid.strip("[]")
|
|
for cid in chunk_ids
|
|
if cid.strip("[]") not in resolved_ids
|
|
]
|
|
|
|
if missing:
|
|
async with _serialized(ctx):
|
|
rag = _require_rag(ctx)
|
|
synthetic: list[SearchResult] = []
|
|
doc_cache: dict[str, Any] = {}
|
|
for cid in missing:
|
|
chunk = await rag.get_chunk_by_id(cid)
|
|
if chunk is None or not chunk.document_id:
|
|
continue
|
|
did = chunk.document_id
|
|
if did in doc_cache:
|
|
doc = doc_cache[did]
|
|
else:
|
|
doc = await rag.get_document_by_id(did)
|
|
doc_cache[did] = doc
|
|
chunk.document_uri = doc.uri if doc else None
|
|
chunk.document_title = doc.title if doc else None
|
|
synthetic.append(SearchResult.from_chunk(chunk, score=1.0))
|
|
if synthetic:
|
|
citations.extend(resolve_citations(missing, synthetic))
|
|
|
|
if citations:
|
|
_register_citations(state, citations)
|
|
return f"Registered {len(citations)} citation(s)."
|
|
|
|
raise ModelRetry(
|
|
f"None of the supplied chunk_ids {list(chunk_ids)} could be "
|
|
"resolved. Copy chunk_ids verbatim from `search` results or "
|
|
"from the `chunk_ids` field on items.jsonl / toc.json rows — "
|
|
"never reconstruct, abbreviate, or paraphrase them."
|
|
)
|
|
|
|
tools["cite"] = cite
|
|
|
|
return tools
|