85594a1fd1
`CapabilityEvidenceRecord` holds the relationships a transcript cannot express: which chunks a capability retrieved, which it cited, in which questions, and at which point in the conversation. RAG and analysis each own one in their own state namespace. Nothing is co-written: the host's state is JSON storage, so a shared record would be overwritten by whichever capability synced last, and merging happens in transient per-request views instead. Both clocks are derived from the conversation rather than counted locally, so every participant computes the same values without sharing a counter. Question identity is the message count when the question arrived; epoch is the message count at an outcome. Epochs are therefore globally comparable, which is what lets `citation_status` require a declaration to follow the newest evidence of every capability, and what makes equal epochs mean one request. A declaration is written only after `resolve_citations` succeeds, so a call naming only unresolvable ids is not a citation. Status is derived, never stored, so refs and status cannot contradict. Resuming a question requires the host to carry the capability state from the run being resumed. Without it the identity of the question in progress is unknowable, and adopting the current message count would relabel that question as a new one and judge every declaration in it against the wrong identity. Nothing reads the records yet and no wire behaviour changes.
112 lines
3.6 KiB
Python
112 lines
3.6 KiB
Python
from dataclasses import dataclass
|
|
from functools import cache
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel, Field
|
|
from pydantic_ai import RunContext
|
|
from pydantic_ai.messages import ToolReturn
|
|
from pydantic_ai.toolsets import FunctionToolset
|
|
|
|
from haiku.rag.capabilities._base import (
|
|
RAGCapabilityBase,
|
|
resolve_db_path,
|
|
)
|
|
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
|
|
from haiku.rag.config.models import AppConfig
|
|
from haiku.rag.store.models.chunk import SearchResult
|
|
from haiku.rag.store.models.citation import Citation
|
|
|
|
AGENT_PREAMBLE = """You are a helpful research assistant powered by haiku.rag, a knowledge base system.
|
|
|
|
CRITICAL RULES:
|
|
1. For greetings or casual chat: respond directly WITHOUT using any tools
|
|
2. NEVER make up information - always use capabilities to get facts from the knowledge base
|
|
3. When a capability returns citations, always include them in your response
|
|
"""
|
|
|
|
STATE_NAMESPACE = "rag"
|
|
_CAPABILITY_ID = "haiku-rag"
|
|
_TOOL_NAMES = frozenset({"rag_search", "rag_cite"})
|
|
_instructions_path = Path(__file__).parent / "instructions" / "rag.md"
|
|
|
|
|
|
class RAGState(BaseModel):
|
|
citation_index: dict[str, Citation] = Field(default_factory=dict)
|
|
citations: list[str] = Field(default_factory=list)
|
|
evidence: CapabilityEvidenceRecord = Field(default_factory=CapabilityEvidenceRecord)
|
|
document_filter: str | None = None
|
|
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
|
|
|
|
|
|
@cache
|
|
def instructions() -> str:
|
|
return _instructions_path.read_text().strip()
|
|
|
|
|
|
@dataclass
|
|
class RAGCapability(RAGCapabilityBase[RAGState]):
|
|
"""Deferred, native Pydantic AI capability for grounded RAG queries."""
|
|
|
|
def get_toolset(self) -> FunctionToolset[Any]:
|
|
async def rag_search(
|
|
ctx: RunContext[Any], query: str, limit: int | None = None
|
|
) -> str | ToolReturn:
|
|
"""Search the knowledge base using hybrid vector and full-text search."""
|
|
return await self._with_state(self._search(query, limit))
|
|
|
|
async def rag_cite(ctx: RunContext[Any], chunk_ids: list[str]) -> Any:
|
|
"""Register exact search-result chunk IDs as citations for the answer."""
|
|
return await self._with_state(self._cite(chunk_ids))
|
|
|
|
return FunctionToolset(
|
|
[rag_search, rag_cite],
|
|
id=_CAPABILITY_ID,
|
|
max_retries=3,
|
|
sequential=True,
|
|
)
|
|
|
|
|
|
def create_capability(
|
|
db_path: Path | None = None,
|
|
config: AppConfig | None = None,
|
|
*,
|
|
defer_loading: bool = True,
|
|
request_limit: int | None = 20,
|
|
vision: bool | None = None,
|
|
) -> RAGCapability:
|
|
"""Create a native Pydantic AI RAG capability.
|
|
|
|
``vision`` gates whether picture chunks are attached to search results as
|
|
images, and should reflect the model the hosting agent actually runs.
|
|
Defaults to ``config.qa.model.vision``.
|
|
"""
|
|
if config is None:
|
|
from haiku.rag.config import get_config
|
|
|
|
config = get_config()
|
|
return RAGCapability(
|
|
db_path=resolve_db_path(db_path, config),
|
|
config=config,
|
|
state_type=RAGState,
|
|
state_namespace=STATE_NAMESPACE,
|
|
instruction_text=instructions(),
|
|
vision=config.qa.model.vision if vision is None else vision,
|
|
tool_names=_TOOL_NAMES,
|
|
request_limit=request_limit,
|
|
id=_CAPABILITY_ID,
|
|
description=(
|
|
"Search the haiku.rag knowledge base and cite evidence for grounded answers."
|
|
),
|
|
defer_loading=defer_loading,
|
|
)
|
|
|
|
|
|
__all__ = [
|
|
"AGENT_PREAMBLE",
|
|
"RAGCapability",
|
|
"RAGState",
|
|
"STATE_NAMESPACE",
|
|
"create_capability",
|
|
"instructions",
|
|
]
|