danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Simplify tools/ — remove AG-UI state machinery, keep core toolsets

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-19 16:31:35 +02:00
parent 524647c501
commit ed89ff0fc9
No known key found for this signature in database
41 changed files with 204 additions and 11138 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.github/workflows/test.yml vendored
View file

@ -68,6 +68,8 @@ jobs:
run: uv run python -c "from transformers import AutoTokenizer; AutoTokenizer.from_pretrained('Qwen/Qwen3-Embedding-0.6B')"
- name: Run tests with coverage
run: uv run pytest -m "not integration" --cov=haiku --cov-report=xml
env:
HF_HUB_OFFLINE: "1"
- name: Upload coverage to Codecov
uses: codecov/codecov-action@v5
with:

18
haiku_rag_slim/haiku/rag/agents/qa/agent.py
View file

@ -11,15 +11,14 @@ from haiku.rag.agents.research.models import (
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig, ModelConfig
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.search import SEARCH_NAMESPACE, SearchState, create_search_toolset
from haiku.rag.store.models import SearchResult
from haiku.rag.tools.search import create_search_toolset
from haiku.rag.utils import get_model
@dataclass
class _QARunDeps:
client: HaikuRAG
tool_context: ToolContext | None = None
class QuestionAnswerAgent:
@ -47,11 +46,12 @@ class QuestionAnswerAgent:
Returns:
Tuple of (answer text, list of resolved citations)
"""
context = ToolContext()
accumulated_results: list[SearchResult] = []
search_toolset = create_search_toolset(
self._config,
base_filter=filter,
tool_name="search_documents",
on_results=accumulated_results.extend,
)
# Agent created per-call: toolset varies with filter, and Agent
@ -66,15 +66,9 @@ class QuestionAnswerAgent:
retries=3,
)
deps = _QARunDeps(client=self._client, tool_context=context)
deps = _QARunDeps(client=self._client)
result = await agent.run(question, deps=deps)
output = result.output
# Get search results from context for citation resolution
search_state = context.get(SEARCH_NAMESPACE)
search_results = (
search_state.results if isinstance(search_state, SearchState) else []
)
citations = resolve_citations(output.cited_chunks, search_results)
citations = resolve_citations(output.cited_chunks, accumulated_results)
return output.answer, citations

39
haiku_rag_slim/haiku/rag/tools/__init__.py
View file

@ -1,52 +1,23 @@
from haiku.rag.tools.analysis import create_analysis_toolset
from haiku.rag.tools.context import (
RAGDeps,
ToolContext,
ToolContextCache,
prepare_context,
)
from haiku.rag.tools.deps import AgentDeps
from haiku.rag.tools.analysis import AnalysisResult, create_analysis_toolset
from haiku.rag.tools.context import RAGDeps
from haiku.rag.tools.document import create_document_toolset
from haiku.rag.tools.filters import (
build_document_filter,
build_multi_document_filter,
combine_filters,
get_session_filter,
)
from haiku.rag.tools.models import AnalysisResult, QAResult
from haiku.rag.tools.prompts import build_tools_prompt
from haiku.rag.tools.qa import create_qa_toolset
from haiku.rag.tools.qa import PRIOR_ANSWER_RELEVANCE_THRESHOLD, QAHistoryEntry
from haiku.rag.tools.search import create_search_toolset
from haiku.rag.tools.toolkit import (
FEATURE_ANALYSIS,
FEATURE_DOCUMENTS,
FEATURE_QA,
FEATURE_SEARCH,
Toolkit,
build_toolkit,
)
__all__ = [
"AgentDeps",
"AnalysisResult",
"FEATURE_ANALYSIS",
"FEATURE_DOCUMENTS",
"FEATURE_QA",
"FEATURE_SEARCH",
"QAResult",
"PRIOR_ANSWER_RELEVANCE_THRESHOLD",
"QAHistoryEntry",
"RAGDeps",
"ToolContext",
"ToolContextCache",
"Toolkit",
"build_document_filter",
"build_multi_document_filter",
"build_toolkit",
"build_tools_prompt",
"combine_filters",
"create_analysis_toolset",
"create_document_toolset",
"create_qa_toolset",
"create_search_toolset",
"get_session_filter",
"prepare_context",
]

18
haiku_rag_slim/haiku/rag/tools/analysis.py
View file

@ -1,3 +1,4 @@
from pydantic import BaseModel, Field
from pydantic_ai import FunctionToolset, RunContext
from haiku.rag.agents.rlm.agent import create_rlm_agent
@ -8,9 +9,17 @@ from haiku.rag.tools.context import RAGDeps
from haiku.rag.tools.filters import (
build_document_filter,
combine_filters,
get_session_filter,
)
from haiku.rag.tools.models import AnalysisResult
class AnalysisResult(BaseModel):
"""Result from the analysis toolset (RLM execution)."""
answer: str = Field(description="The answer produced by analysis")
code_executed: bool = Field(
default=True,
description="Whether code was executed to produce this answer",
)
def create_analysis_toolset(
@ -47,12 +56,9 @@ def create_analysis_toolset(
AnalysisResult with answer and execution metadata.
"""
client = ctx.deps.client
tool_context = ctx.deps.tool_context
doc_filter = build_document_filter(document_name) if document_name else None
effective_filter = combine_filters(
get_session_filter(tool_context, base_filter), doc_filter
)
effective_filter = combine_filters(base_filter, doc_filter)
rlm_context = RLMContext(filter=effective_filter)

257
haiku_rag_slim/haiku/rag/tools/context.py
View file

@ -1,13 +1,8 @@
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Protocol, TypeVar, overload, runtime_checkable
from pydantic import BaseModel, PrivateAttr
from typing import TYPE_CHECKING, Protocol, runtime_checkable
if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
T = TypeVar("T", bound=BaseModel)
@runtime_checkable
class RAGDeps(Protocol):
@ -18,253 +13,3 @@ class RAGDeps(Protocol):
"""
client: "HaikuRAG"
tool_context: "ToolContext | None"
class ToolContext(BaseModel):
"""Generic state container for haiku.rag toolsets.
Toolsets register their own Pydantic model state under namespaces.
Multiple toolsets can share state by registering under the same namespace.
All registered states must be Pydantic BaseModel subclasses, making
the entire context serializable via model_dump()/model_validate().
Example:
# Define toolset-specific state
class SearchState(BaseModel):
results: list[SearchResult] = []
filter: str | None = None
SEARCH_NAMESPACE = "haiku.rag.search"
# In toolset factory
def create_search_toolset(config):
async def search(ctx: RunContext[RAGDeps], query: str):
tool_context = ctx.deps.tool_context
if tool_context:
state = tool_context.get_or_create(SEARCH_NAMESPACE, SearchState)
...
# Usage
search_tools = create_search_toolset(config)
agent = Agent(..., toolsets=[search_tools])
await agent.run("...", deps=my_deps)
# Access accumulated state
search_state = context.get(SEARCH_NAMESPACE)
for result in search_state.results:
print(f"{result.document_title}")
# Serialize entire context
ns_data = context.dump_namespaces()
"""
state_key: str | None = None
_namespaces: dict[str, BaseModel] = PrivateAttr(default_factory=dict)
_client_snapshot: dict[str, Any] | None = PrivateAttr(default=None)
def register(self, namespace: str, state: BaseModel) -> None:
"""Register state for a namespace.
Args:
namespace: Unique identifier for the toolset (e.g., "haiku.rag.search")
state: A Pydantic BaseModel instance to store
Overwrites any existing state for the namespace.
"""
self._namespaces[namespace] = state
@overload
def get(self, namespace: str) -> BaseModel | None: ...
@overload
def get(self, namespace: str, state_type: type[T]) -> T | None: ...
def get(
self, namespace: str, state_type: type[T] | None = None
) -> BaseModel | T | None:
"""Get state for a namespace, or None if not registered.
When state_type is provided, returns the state only if it matches
the expected type, otherwise returns None.
"""
state = self._namespaces.get(namespace)
if state_type is not None:
return state if isinstance(state, state_type) else None
return state
def get_or_create(self, namespace: str, state_type: type[T]) -> T:
"""Get state for a namespace, creating it if not registered.
Args:
namespace: The namespace to get or create state for.
state_type: A Pydantic BaseModel subclass to instantiate if needed.
Returns:
The state for the namespace.
"""
if namespace not in self._namespaces:
self._namespaces[namespace] = state_type()
return self._namespaces[namespace] # type: ignore[return-value]
def clear_namespace(self, namespace: str) -> None:
"""Clear state for a specific namespace."""
if namespace in self._namespaces:
del self._namespaces[namespace]
def clear_all(self) -> None:
"""Clear all namespaces."""
self._namespaces.clear()
@property
def namespaces(self) -> list[str]:
"""List all registered namespaces."""
return list(self._namespaces.keys())
@property
def client_snapshot(self) -> dict[str, Any] | None:
"""Snapshot captured after the last restore_state_snapshot call.
Represents what the client has, before any server-side overrides.
Tools use this as the baseline for delta computation so that
server-side changes (e.g. background summarization) are included.
"""
return self._client_snapshot
def dump_namespaces(self) -> dict[str, dict[str, Any]]:
"""Serialize all namespace states to a dictionary.
Returns:
Dict mapping namespace -> serialized state dict.
"""
return {ns: state.model_dump() for ns, state in self._namespaces.items()}
def build_state_snapshot(self) -> dict[str, Any]:
"""Build a flat snapshot of all namespace states for AG-UI.
Merges model_dump(mode="json") from every registered namespace
into a single flat dict.
Returns:
Combined dict of all namespace fields.
"""
snapshot: dict[str, Any] = {}
for state in self._namespaces.values():
snapshot.update(state.model_dump(mode="json"))
return snapshot
def restore_state_snapshot(self, data: dict[str, Any]) -> None:
"""Restore namespace states from a flat snapshot dict.
For each registered namespace, finds matching fields in *data*,
validates them via the namespace model, and updates the state
in place. Fields not present in *data* are left unchanged.
After restoring, captures a snapshot as ``client_snapshot`` so
tools can compute deltas against what the client actually has.
Args:
data: Flat dict as produced by build_state_snapshot().
"""
for state in self._namespaces.values():
model_fields = state.model_fields
matching = {k: v for k, v in data.items() if k in model_fields}
if matching:
# Fill in current values for fields not in data
current = state.model_dump()
current.update(matching)
updated = state.model_validate(current)
for field_name in matching:
setattr(state, field_name, getattr(updated, field_name))
self._client_snapshot = self.build_state_snapshot()
def load_namespace(self, namespace: str, state_type: type[T], data: dict) -> T:
"""Deserialize and register state for a namespace.
Args:
namespace: The namespace to register the state under.
state_type: The Pydantic model class to deserialize into.
data: The serialized state data.
Returns:
The deserialized and registered state.
"""
state = state_type.model_validate(data)
self._namespaces[namespace] = state
return state
def prepare_context(
context: ToolContext,
features: list[str] | None = None,
state_key: str | None = None,
) -> None:
"""Register required namespaces in a ToolContext based on feature flags.
Idempotent — safe to call multiple times on the same context.
Args:
context: ToolContext to prepare.
features: List of enabled features. Defaults to ["search", "documents"].
state_key: Optional AG-UI state key to set on the context.
"""
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
if features is None:
features = ["search", "documents"]
if any(f in features for f in ("search", "qa", "analysis")):
context.get_or_create(SESSION_NAMESPACE, SessionState)
if "qa" in features:
context.get_or_create(QA_SESSION_NAMESPACE, QASessionState)
if state_key is not None:
context.state_key = state_key
class ToolContextCache:
"""In-memory cache for ToolContext instances, keyed by external session/thread ID."""
def __init__(self, ttl: timedelta = timedelta(hours=1)) -> None:
self._cache: dict[str, ToolContext] = {}
self._timestamps: dict[str, datetime] = {}
self._ttl = ttl
def get_or_create(self, key: str) -> tuple[ToolContext, bool]:
"""Get an existing context or create a new one.
Returns:
Tuple of (context, is_new) where is_new is True if a new context was created.
"""
self._cleanup()
if key in self._cache:
self._timestamps[key] = datetime.now()
return self._cache[key], False
context = ToolContext()
self._cache[key] = context
self._timestamps[key] = datetime.now()
return context, True
def remove(self, key: str) -> None:
"""Remove a specific key from the cache."""
self._cache.pop(key, None)
self._timestamps.pop(key, None)
def clear(self) -> None:
"""Clear all entries."""
self._cache.clear()
self._timestamps.clear()
def _cleanup(self) -> None:
"""Remove entries older than TTL."""
now = datetime.now()
expired = [
key for key, ts in self._timestamps.items() if (now - ts) >= self._ttl
]
for key in expired:
self._cache.pop(key, None)
self._timestamps.pop(key, None)

42
haiku_rag_slim/haiku/rag/tools/deps.py
View file

@ -1,42 +0,0 @@
from dataclasses import dataclass
from typing import Any
from haiku.rag.client import HaikuRAG
from haiku.rag.tools.context import ToolContext
@dataclass
class AgentDeps:
"""Generic dependencies for agents using haiku.rag toolsets.
Implements RAGDeps protocol and AG-UI state protocol.
"""
client: HaikuRAG
tool_context: ToolContext
@property
def state(self) -> dict[str, Any]:
"""Get current state for AG-UI protocol."""
snapshot = self.tool_context.build_state_snapshot()
state_key = self.tool_context.state_key
if state_key:
return {state_key: snapshot}
return snapshot
@state.setter
def state(self, value: dict[str, Any] | None) -> None:
"""Set state from AG-UI protocol."""
if value is None:
return
data = self._extract_state_data(value)
self.tool_context.restore_state_snapshot(data)
def _extract_state_data(self, value: dict[str, Any]) -> dict[str, Any]:
"""Extract flat state dict, unwrapping state_key if present."""
state_key = self.tool_context.state_key
if state_key and state_key in value:
nested = value[state_key]
if isinstance(nested, dict):
return nested
return value

8
haiku_rag_slim/haiku/rag/tools/document.py
View file

@ -4,7 +4,6 @@ from pydantic_ai import Agent, FunctionToolset, RunContext
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.tools.context import RAGDeps
from haiku.rag.tools.filters import get_session_filter
from haiku.rag.utils import get_model
DOCUMENT_SUMMARY_PROMPT = """Generate a summary of the document content provided below.
@ -92,17 +91,14 @@ def create_document_toolset(
Paginated list of documents with metadata.
"""
client = ctx.deps.client
tool_context = ctx.deps.tool_context
page_size = 50
offset = (page - 1) * page_size
effective_filter = get_session_filter(tool_context, base_filter)
docs = await client.list_documents(
limit=page_size, offset=offset, filter=effective_filter
limit=page_size, offset=offset, filter=base_filter
)
total = await client.count_documents(filter=effective_filter)
total = await client.count_documents(filter=base_filter)
total_pages = (total + page_size - 1) // page_size if total > 0 else 1
return DocumentListResponse(

35
haiku_rag_slim/haiku/rag/tools/filters.py
View file

@ -1,9 +1,3 @@
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from haiku.rag.tools.context import ToolContext
def build_document_filter(document_name: str) -> str:
"""Build SQL filter for document name matching.
@ -31,35 +25,6 @@ def build_multi_document_filter(document_names: list[str]) -> str | None:
return " OR ".join(f"({f})" for f in filters)
def get_session_filter(
context: "ToolContext | None",
base_filter: str | None = None,
) -> str | None:
"""Build effective filter from session state document filter and base filter.
Checks the ToolContext for a registered SessionState. If it has a
document_filter, builds a SQL filter from it and combines with base_filter.
Args:
context: Optional ToolContext that may contain a SessionState.
base_filter: Optional base SQL WHERE clause to combine with.
Returns:
Combined filter string, or None if no filters apply.
"""
if context is None:
return base_filter
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
session_state = context.get(SESSION_NAMESPACE, SessionState)
if session_state is None or not session_state.document_filter:
return base_filter
session_filter = build_multi_document_filter(session_state.document_filter)
return combine_filters(base_filter, session_filter)
def combine_filters(filter1: str | None, filter2: str | None) -> str | None:
"""Combine two SQL filters with AND logic.

37
haiku_rag_slim/haiku/rag/tools/models.py
View file

@ -1,37 +0,0 @@
from pydantic import BaseModel, Field
from haiku.rag.agents.research.models import Citation
class QAResult(BaseModel):
"""Result from the QA toolset."""
question: str = Field(description="The question that was answered")
answer: str = Field(description="The answer to the question")
confidence: float = Field(
default=1.0,
description="Confidence score for this answer (0-1)",
ge=0.0,
le=1.0,
)
citations: list[Citation] = Field(
default_factory=list,
description="Citations supporting the answer",
)
@property
def sources(self) -> list[str]:
"""Source names for display."""
return list(
dict.fromkeys(c.document_title or c.document_uri for c in self.citations)
)
class AnalysisResult(BaseModel):
"""Result from the analysis toolset (RLM execution)."""
answer: str = Field(description="The answer produced by analysis")
code_executed: bool = Field(
default=True,
description="Whether code was executed to produce this answer",
)

71
haiku_rag_slim/haiku/rag/tools/prompts.py
View file

@ -1,71 +0,0 @@
_TOOL_HEADER = """
How to decide which tool to use:"""
_TOOL_DOCUMENTS = """
- "list_documents" - Use when the user wants to browse or see what documents are available (e.g., "what documents are available?", "show me the documents", "list available docs").
- "summarize_document" - Use when the user wants an overview or summary of a specific document (e.g., "summarize document X", "what does Y cover?", "give me an overview of Z").
- "get_document" - Use when the user wants the FULL content of a specific document (e.g., "get the paper about Y", "fetch 2412.00566", "show me the full document")."""
_TOOL_QA = """
- "ask" - Use for questions about topics in the knowledge base. Searches across documents and returns answers with citations. Prior answers are recalled to avoid redundant work."""
_TOOL_SEARCH = """
- "search" - Use when the user explicitly asks to search, find, or explore documents. Handles multi-query expansion internally and returns matching passages with surrounding context."""
_TOOL_ANALYSIS = """
- "analyze" - Use when the user asks for computation, data analysis, or quantitative tasks that require code execution (e.g., "calculate the average", "compare the numbers", "plot the data"). Runs Python code in a sandbox to produce results."""
_DOCUMENT_NAME_HEADER = """
IMPORTANT - When user mentions a document in search/ask:
- If user says "search in <doc>", "find in <doc>", "answer from <doc>", or "<topic> in <doc>":
- Extract the TOPIC as `query`/`question`
- Extract the DOCUMENT NAME as `document_name`"""
_DOCUMENT_NAME_SEARCH_EXAMPLES = """
- Examples for search:
- "search for embeddings in the ML paper" → query="embeddings", document_name="ML paper"
- "find transformer architecture in 2412.00566" → query="transformer architecture", document_name="2412.00566" """
_DOCUMENT_NAME_QA_EXAMPLES = """
- Examples for ask:
- "what does the ML paper say about embeddings?" → question="what are the embedding methods?", document_name="ML paper"
- "answer from 2412.00566 about model training" → question="how is the model trained?", document_name="2412.00566" """
_FEATURE_TOOLS: dict[str, str] = {
"documents": _TOOL_DOCUMENTS,
"qa": _TOOL_QA,
"search": _TOOL_SEARCH,
"analysis": _TOOL_ANALYSIS,
}
def build_tools_prompt(features: list[str]) -> str:
"""Build tool guidance for the given features.
Returns prompt text describing when and how to use each tool.
Designed to be spliced into a custom agent's system prompt.
Args:
features: List of feature names (e.g., ["search", "documents", "qa"]).
Returns:
Tool guidance prompt text.
"""
parts: list[str] = []
tool_sections = [_FEATURE_TOOLS[f] for f in features if f in _FEATURE_TOOLS]
if tool_sections:
parts.append(_TOOL_HEADER)
parts.extend(tool_sections)
if "search" in features or "qa" in features:
parts.append(_DOCUMENT_NAME_HEADER)
if "search" in features:
parts.append(_DOCUMENT_NAME_SEARCH_EXAMPLES)
if "qa" in features:
parts.append(_DOCUMENT_NAME_QA_EXAMPLES)
return "".join(parts)

244
haiku_rag_slim/haiku/rag/tools/qa.py
View file

@ -1,29 +1,6 @@
from collections.abc import Callable
from pydantic import BaseModel, Field
from pydantic_ai import FunctionToolset, RunContext, ToolReturn
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.models import Citation, SearchAnswer
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.embeddings import get_embedder
from haiku.rag.tools.context import RAGDeps, ToolContext
from haiku.rag.tools.filters import (
build_document_filter,
combine_filters,
get_session_filter,
)
from haiku.rag.tools.models import QAResult
from haiku.rag.tools.session import (
SESSION_NAMESPACE,
SessionContext,
SessionState,
compute_combined_state_delta,
)
from haiku.rag.utils import cosine_similarity
PRIOR_ANSWER_RELEVANCE_THRESHOLD = 0.7
@ -53,224 +30,3 @@ class QAHistoryEntry(BaseModel):
cited_chunks=[c.chunk_id for c in self.citations],
citations=self.citations,
)
class QASessionState(BaseModel):
"""Extended session state for QA with embedding cache."""
qa_history: list[QAHistoryEntry] = []
session_context: SessionContext | None = None
QA_SESSION_NAMESPACE = "haiku.rag.qa_session"
MAX_QA_HISTORY = 50
async def run_qa_core(
client: HaikuRAG,
config: AppConfig,
question: str,
document_name: str | None = None,
*,
context: ToolContext | None = None,
base_filter: str | None = None,
session_context: str | None = None,
prior_answers: list[SearchAnswer] | None = None,
on_qa_complete: Callable[[QASessionState, AppConfig], None] | None = None,
) -> QAResult:
"""Run the QA flow and return a QAResult.
This is the core QA implementation shared by toolsets and client APIs.
It updates session state and QA history when context is provided.
"""
session_state: SessionState | None = None
qa_session_state: QASessionState | None = None
if context is not None:
session_state = context.get(SESSION_NAMESPACE, SessionState)
qa_session_state = context.get(QA_SESSION_NAMESPACE, QASessionState)
doc_filter = build_document_filter(document_name) if document_name else None
effective_filter = combine_filters(
get_session_filter(context, base_filter), doc_filter
)
effective_session_context = session_context
if qa_session_state is not None and qa_session_state.session_context is not None:
effective_session_context = qa_session_state.session_context.summary
effective_prior_answers = prior_answers or []
if qa_session_state is not None and qa_session_state.qa_history:
embedder = get_embedder(config)
question_embedding = await embedder.embed_query(question)
to_embed = []
to_embed_indices = []
for i, qa in enumerate(qa_session_state.qa_history):
if qa.question_embedding is None:
to_embed.append(qa.question)
to_embed_indices.append(i)
if to_embed:
new_embeddings = await embedder.embed_documents(to_embed)
for i, idx in enumerate(to_embed_indices):
qa_session_state.qa_history[idx].question_embedding = new_embeddings[i]
matched_answers = []
for qa in qa_session_state.qa_history:
if qa.question_embedding is not None:
similarity = cosine_similarity(
question_embedding, qa.question_embedding
)
if similarity >= PRIOR_ANSWER_RELEVANCE_THRESHOLD:
matched_answers.append(qa.to_search_answer())
if matched_answers:
effective_prior_answers = matched_answers
graph = build_research_graph(config=config, output_mode="conversational")
research_context = ResearchContext(
original_question=question,
session_context=effective_session_context,
qa_responses=effective_prior_answers,
)
research_state = ResearchState(
context=research_context,
max_iterations=1,
search_filter=effective_filter,
max_concurrency=config.research.max_concurrency,
)
deps = ResearchDeps(client=client)
result = await graph.run(state=research_state, deps=deps)
# Build citations with stable indices from session state
citations = []
for i, c in enumerate(result.citations):
if session_state is not None:
index = session_state.get_or_assign_index(c.chunk_id)
else:
index = i + 1
citations.append(
Citation(
index=index,
document_id=c.document_id,
chunk_id=c.chunk_id,
document_uri=c.document_uri,
document_title=c.document_title,
page_numbers=c.page_numbers,
headings=c.headings,
content=c.content,
)
)
qa_result = QAResult(
question=question,
answer=result.answer,
confidence=result.confidence,
citations=citations,
)
if session_state is not None:
session_state.citations = citations
session_state.citations_history.append(citations)
if qa_session_state is not None:
qa_session_state.qa_history.append(
QAHistoryEntry(
question=question,
answer=result.answer,
confidence=result.confidence,
citations=citations,
)
)
# Enforce FIFO limit
if len(qa_session_state.qa_history) > MAX_QA_HISTORY:
qa_session_state.qa_history = qa_session_state.qa_history[-MAX_QA_HISTORY:]
if on_qa_complete is not None:
on_qa_complete(qa_session_state, config)
return qa_result
def create_qa_toolset(
config: AppConfig,
base_filter: str | None = None,
tool_name: str = "ask",
on_ask_complete: Callable[[QASessionState, AppConfig], None] | None = None,
) -> FunctionToolset[RAGDeps]:
"""Create a toolset with Q&A capabilities using research graph.
Args:
config: Application configuration.
base_filter: Optional base SQL WHERE clause applied to searches.
tool_name: Name for the ask tool. Defaults to "ask".
on_ask_complete: Optional callback invoked after each QA cycle with
the updated QASessionState and config. Use this to trigger
background summarization or other post-processing.
Returns:
FunctionToolset with an ask tool.
"""
async def ask(
ctx: RunContext[RAGDeps],
question: str,
document_name: str | None = None,
) -> ToolReturn | QAResult:
"""Answer a question using the knowledge base.
Uses a research graph for searching and synthesizing answers.
Args:
question: The question to answer.
document_name: Optional document name/title to search within.
Returns:
QAResult with answer, confidence, and citations.
"""
client = ctx.deps.client
tool_context = ctx.deps.tool_context
state_key: str | None = None
client_snapshot: dict | None = None
if tool_context is not None:
state_key = tool_context.state_key
if tool_context.namespaces:
client_snapshot = (
tool_context.client_snapshot or tool_context.build_state_snapshot()
)
qa_result = await run_qa_core(
client=client,
config=config,
question=question,
document_name=document_name,
context=tool_context,
base_filter=base_filter,
on_qa_complete=on_ask_complete,
)
if client_snapshot is not None and tool_context is not None:
new_snapshot = tool_context.build_state_snapshot()
state_event = compute_combined_state_delta(
client_snapshot, new_snapshot, state_key=state_key
)
if state_event is not None:
answer_text = qa_result.answer
if qa_result.citations:
citation_refs = " ".join(
f"[{c.index}]" for c in qa_result.citations
)
answer_text = f"{answer_text}\n\nSources: {citation_refs}"
return ToolReturn(return_value=answer_text, metadata=[state_event])
return qa_result
toolset: FunctionToolset[RAGDeps] = FunctionToolset()
toolset.add_function(ask, name=tool_name)
return toolset

111
haiku_rag_slim/haiku/rag/tools/search.py
View file

@ -1,23 +1,11 @@
from pydantic import BaseModel
from pydantic_ai import FunctionToolset, RunContext, ToolReturn
from collections.abc import Callable
from pydantic_ai import FunctionToolset, RunContext
from haiku.rag.agents.research.models import Citation
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models import SearchResult
from haiku.rag.tools.context import RAGDeps
from haiku.rag.tools.filters import combine_filters, get_session_filter
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState, compute_state_delta
SEARCH_NAMESPACE = "haiku.rag.search"
class SearchState(BaseModel):
"""State for search toolset.
Accumulates search results across tool invocations.
"""
results: list[SearchResult] = []
from haiku.rag.tools.filters import combine_filters
def create_search_toolset(
@ -25,6 +13,7 @@ def create_search_toolset(
expand_context: bool = True,
base_filter: str | None = None,
tool_name: str = "search",
on_results: Callable[[list[SearchResult]], None] | None = None,
) -> FunctionToolset[RAGDeps]:
"""Create a toolset with search capabilities.
@ -35,6 +24,8 @@ def create_search_toolset(
base_filter: Optional base SQL WHERE clause applied to all searches.
Combined with any filter passed to the search tool.
tool_name: Name for the search tool. Defaults to "search".
on_results: Optional callback invoked with search results after each search.
Useful for accumulating results externally (e.g., for citation resolution).
Returns:
FunctionToolset with a search tool.
@ -45,7 +36,7 @@ def create_search_toolset(
query: str,
limit: int | None = None,
filter: str | None = None,
) -> ToolReturn | str:
) -> str:
"""Search the knowledge base for relevant documents.
Args:
@ -57,26 +48,8 @@ def create_search_toolset(
Formatted search results with content and metadata.
"""
client = ctx.deps.client
tool_context = ctx.deps.tool_context
search_state: SearchState | None = None
if tool_context is not None:
search_state = tool_context.get_or_create(SEARCH_NAMESPACE, SearchState)
session_state: SessionState | None = None
old_session_state: SessionState | None = None
state_key: str | None = None
if tool_context is not None:
session_state = tool_context.get(SESSION_NAMESPACE, SessionState)
state_key = tool_context.state_key
if session_state is not None:
old_session_state = session_state.model_copy(deep=True)
# Combine all filters: base_filter AND session_filter AND tool filter
effective_filter = combine_filters(
get_session_filter(tool_context, base_filter), filter
)
effective_filter = combine_filters(base_filter, filter)
effective_limit = limit or config.search.limit
results = await client.search(
query, limit=effective_limit, filter=effective_filter
@ -85,68 +58,18 @@ def create_search_toolset(
if expand_context:
results = await client.expand_context(results)
if search_state is not None:
search_state.results.extend(results)
results_list = list(results)
if not results:
if on_results:
on_results(results_list)
if not results_list:
return "No results found."
if session_state is not None:
citations = []
for r in results:
chunk_id = r.chunk_id or ""
if chunk_id:
index = session_state.get_or_assign_index(chunk_id)
else: # pragma: no cover
index = len(session_state.citation_registry) + 1
citations.append(
Citation(
index=index,
document_id=r.document_id or "",
chunk_id=chunk_id,
document_uri=r.document_uri or "",
document_title=r.document_title,
page_numbers=r.page_numbers or [],
headings=r.headings,
content=r.content,
)
)
session_state.citations = citations
session_state.citations_history.append(citations)
result_lines = []
for c in citations:
title = c.document_title or c.document_uri or "Unknown"
snippet = c.content[:300].replace("\n", " ").strip()
if len(c.content) > 300:
snippet += "..."
line = f"[{c.index}] **{title}**"
if c.page_numbers: # pragma: no cover
line += f" (pages {', '.join(map(str, c.page_numbers))})"
line += f"\n {snippet}"
result_lines.append(line)
formatted = f"Found {len(results)} results:\n\n" + "\n\n".join(result_lines)
if old_session_state is not None:
state_event = compute_state_delta(
old_session_state,
session_state,
state_key=state_key,
)
if state_event is not None:
return ToolReturn(
return_value=formatted,
metadata=[state_event],
)
return formatted # pragma: no cover
# Format results without citation indexing (standalone use)
total = len(results)
total = len(results_list)
formatted = [
r.format_for_agent(rank=i + 1, total=total) for i, r in enumerate(results)
r.format_for_agent(rank=i + 1, total=total)
for i, r in enumerate(results_list)
]
return "\n\n".join(formatted)

94
haiku_rag_slim/haiku/rag/tools/session.py
View file

@ -1,94 +0,0 @@
from datetime import datetime
from typing import Any
import jsonpatch
from ag_ui.core import EventType, StateDeltaEvent
from pydantic import BaseModel
from haiku.rag.agents.research.models import Citation
SESSION_NAMESPACE = "haiku.rag.session"
class SessionContext(BaseModel):
"""Compressed summary of conversation history for research graph."""
summary: str = ""
last_updated: datetime | None = None
class SessionState(BaseModel):
"""Session-level state for AG-UI integration.
This state is shared across toolsets and enables:
- Dynamic document filtering
- Stable citation indices across tool calls
- AG-UI state synchronization
"""
document_filter: list[str] = []
citation_registry: dict[str, int] = {}
citations: list[Citation] = []
citations_history: list[list[Citation]] = []
def get_or_assign_index(self, chunk_id: str) -> int:
"""Get or assign a stable citation index for a chunk_id.
Citation indices persist across tool calls within a session.
The first chunk gets index 1, subsequent new chunks get incrementing indices.
Same chunk_id always returns the same index.
"""
if chunk_id in self.citation_registry:
return self.citation_registry[chunk_id]
new_index = len(self.citation_registry) + 1
self.citation_registry[chunk_id] = new_index
return new_index
def compute_state_delta(
old_state: SessionState,
new_state: SessionState,
state_key: str | None = None,
) -> StateDeltaEvent | None:
"""Compute state delta between old and new session state.
Returns a StateDeltaEvent if there are changes, None otherwise.
"""
return compute_combined_state_delta(
old_state.model_dump(mode="json"),
new_state.model_dump(mode="json"),
state_key=state_key,
)
def compute_combined_state_delta(
old_snapshot: dict[str, Any],
new_snapshot: dict[str, Any],
state_key: str | None = None,
) -> StateDeltaEvent | None:
"""Compute state delta between old and new combined state snapshots.
This function computes delta for the combined chat state that includes
both SessionState and QASessionState fields.
Args:
old_snapshot: Previous state dict (e.g., from ChatDeps.state format).
new_snapshot: New state dict.
state_key: Optional namespace key for the state (e.g., "haiku.rag.chat").
Returns:
StateDeltaEvent if there are changes, None otherwise.
"""
wrapped_old = {state_key: old_snapshot} if state_key else old_snapshot
wrapped_new = {state_key: new_snapshot} if state_key else new_snapshot
patch = jsonpatch.make_patch(wrapped_old, wrapped_new)
if not patch.patch:
return None
return StateDeltaEvent(
type=EventType.STATE_DELTA,
delta=patch.patch,
)

108
haiku_rag_slim/haiku/rag/tools/toolkit.py
View file

@ -1,108 +0,0 @@
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any
from pydantic_ai import FunctionToolset
from haiku.rag.config.models import AppConfig
from haiku.rag.tools.context import ToolContext, prepare_context
from haiku.rag.tools.prompts import build_tools_prompt
FEATURE_SEARCH = "search"
FEATURE_DOCUMENTS = "documents"
FEATURE_QA = "qa"
FEATURE_ANALYSIS = "analysis"
@dataclass(frozen=True)
class Toolkit:
"""Bundled toolsets, prompt, and context factory for haiku.rag agents.
Created via build_toolkit(). Provides everything needed to compose
an agent with haiku.rag toolsets and create matching ToolContexts.
"""
toolsets: list[FunctionToolset[Any]] = field(default_factory=list)
prompt: str = ""
features: list[str] = field(default_factory=list)
def create_context(self, state_key: str | None = None) -> ToolContext:
"""Create a ToolContext with namespaces matching this toolkit's features.
Args:
state_key: Optional AG-UI state key to set on the context.
Returns:
A prepared ToolContext.
"""
context = ToolContext()
prepare_context(context, features=self.features, state_key=state_key)
return context
def prepare(self, context: ToolContext, state_key: str | None = None) -> None:
"""Register namespaces on an existing ToolContext for this toolkit's features.
Idempotent — safe to call multiple times on the same context.
Args:
context: ToolContext to prepare.
state_key: Optional AG-UI state key to set on the context.
"""
prepare_context(context, features=self.features, state_key=state_key)
def build_toolkit(
config: AppConfig,
features: list[str] | None = None,
base_filter: str | None = None,
expand_context: bool = True,
on_qa_complete: Callable | None = None,
) -> Toolkit:
"""Build a Toolkit with toolsets, prompt, and context factory for the given features.
Args:
config: Application configuration.
features: List of features to enable. Defaults to ["search", "documents"].
base_filter: Optional base SQL WHERE clause applied to all toolset factories.
expand_context: Whether to expand search results with surrounding context.
on_qa_complete: Optional callback invoked after each QA cycle.
Returns:
A Toolkit ready for agent composition.
"""
if features is None:
features = [FEATURE_SEARCH, FEATURE_DOCUMENTS]
toolsets: list[FunctionToolset[Any]] = []
if FEATURE_SEARCH in features:
from haiku.rag.tools.search import create_search_toolset
toolsets.append(
create_search_toolset(
config, expand_context=expand_context, base_filter=base_filter
)
)
if FEATURE_DOCUMENTS in features:
from haiku.rag.tools.document import create_document_toolset
toolsets.append(create_document_toolset(config, base_filter=base_filter))
if FEATURE_QA in features:
from haiku.rag.tools.qa import create_qa_toolset
toolsets.append(
create_qa_toolset(
config, base_filter=base_filter, on_ask_complete=on_qa_complete
)
)
if FEATURE_ANALYSIS in features:
from haiku.rag.tools.analysis import create_analysis_toolset
toolsets.append(create_analysis_toolset(config, base_filter=base_filter))
prompt = build_tools_prompt(features)
return Toolkit(toolsets=toolsets, prompt=prompt, features=features)

840
tests/cassettes/test_qa_tools/TestAskTool.test_ask_with_tool_context_returns_tool_return.yaml
View file

File diff suppressed because one or more lines are too long

1083
tests/cassettes/test_qa_tools/TestAskTool.test_ask_without_tool_context.yaml
View file

File diff suppressed because one or more lines are too long

1133
tests/cassettes/test_qa_tools/TestRunQACore.test_run_qa_core_fifo_limit.yaml
View file

File diff suppressed because one or more lines are too long

843
tests/cassettes/test_qa_tools/TestRunQACore.test_run_qa_core_on_qa_complete_callback.yaml
View file

File diff suppressed because one or more lines are too long

844
tests/cassettes/test_qa_tools/TestRunQACore.test_run_qa_core_populates_citations_history.yaml
View file

File diff suppressed because one or more lines are too long

844
tests/cassettes/test_qa_tools/TestRunQACore.test_run_qa_core_with_session_state.yaml
View file

File diff suppressed because one or more lines are too long

1092
tests/cassettes/test_qa_tools/TestRunQACore.test_run_qa_core_without_context.yaml
View file

File diff suppressed because one or more lines are too long

1113
tests/cassettes/test_qa_tools/TestRunQACoreWithPriorAnswers.test_run_qa_core_matches_prior_answers.yaml
View file

File diff suppressed because one or more lines are too long

0
tests/cassettes/test_search_tools/TestSearchToolExecution.test_search_multiple_accumulates.yaml → tests/cassettes/test_search_tools/TestSearchToolExecution.test_search_on_results_accumulates_across_calls.yaml
View file

0
tests/cassettes/test_search_tools/TestSearchToolExecution.test_search_accumulates_in_state.yaml → tests/cassettes/test_search_tools/TestSearchToolExecution.test_search_on_results_callback.yaml
View file

0
tests/cassettes/test_search_tools/TestSearchToolExecution.test_search_without_context.yaml → tests/cassettes/test_search_tools/TestSearchToolExecution.test_search_without_on_results.yaml
View file

162
tests/cassettes/test_search_tools/TestSearchWithSessionState.test_search_citations_accumulate_across_calls.yaml
View file

File diff suppressed because one or more lines are too long

162
tests/cassettes/test_search_tools/TestSearchWithSessionState.test_search_multiple_appends_separate_entries.yaml
View file

File diff suppressed because one or more lines are too long

122
tests/cassettes/test_search_tools/TestSearchWithSessionState.test_search_populates_citations_history.yaml
View file

File diff suppressed because one or more lines are too long

122
tests/cassettes/test_search_tools/TestSearchWithSessionState.test_search_with_session_state_formatted_output.yaml
View file

File diff suppressed because one or more lines are too long

122
tests/cassettes/test_search_tools/TestSearchWithSessionState.test_search_with_session_state_populates_citations.yaml
View file

File diff suppressed because one or more lines are too long

122
tests/cassettes/test_search_tools/TestSearchWithSessionState.test_search_with_session_state_returns_tool_return.yaml
View file

File diff suppressed because one or more lines are too long

520
tests/tools/test_context.py
View file

@ -1,512 +1,26 @@
from pydantic import BaseModel
from dataclasses import dataclass
from unittest.mock import MagicMock
from haiku.rag.tools.context import ToolContext, prepare_context
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
from haiku.rag.tools.context import RAGDeps
class TestState(BaseModel):
value: int = 0
def test_ragdeps_protocol_satisfied_by_dataclass():
"""A dataclass with a client attribute satisfies RAGDeps."""
@dataclass
class MyDeps:
client: MagicMock
class TestStateWithList(BaseModel):
items: list[str] = []
deps = MyDeps(client=MagicMock())
assert isinstance(deps, RAGDeps)
def test_tool_context_defaults():
"""Test ToolContext has sensible defaults."""
ctx = ToolContext()
assert ctx._namespaces == {}
def test_ragdeps_protocol_not_satisfied_without_client():
"""An object without client does not satisfy RAGDeps."""
@dataclass
class NoDeps:
other: str
def test_register_and_get():
"""Test register and get state for a namespace."""
ctx = ToolContext()
state = TestState(value=42)
ctx.register("test.namespace", state)
retrieved = ctx.get("test.namespace")
assert retrieved is state
assert retrieved.value == 42
def test_get_nonexistent_namespace():
"""Test get returns None for unregistered namespace."""
ctx = ToolContext()
assert ctx.get("nonexistent") is None
def test_get_or_create_creates_new():
"""Test get_or_create creates state when namespace doesn't exist."""
ctx = ToolContext()
state = ctx.get_or_create("test.namespace", TestStateWithList)
assert isinstance(state, TestStateWithList)
assert state.items == []
def test_get_or_create_returns_existing():
"""Test get_or_create returns existing state."""
ctx = ToolContext()
state1 = ctx.get_or_create("test.namespace", TestStateWithList)
state1.items.append("item1")
state2 = ctx.get_or_create("test.namespace", TestStateWithList)
assert state2 is state1
assert state2.items == ["item1"]
def test_clear_namespace():
"""Test clear_namespace removes only the specified namespace."""
ctx = ToolContext()
ctx.register("ns1", TestState(value=1))
ctx.register("ns2", TestState(value=2))
ctx.clear_namespace("ns1")
assert ctx.get("ns1") is None
ns2 = ctx.get("ns2")
assert isinstance(ns2, TestState)
assert ns2.value == 2
def test_clear_namespace_nonexistent():
"""Test clear_namespace handles nonexistent namespace gracefully."""
ctx = ToolContext()
ctx.clear_namespace("nonexistent") # Should not raise
def test_clear_all():
"""Test clear_all clears all namespaces."""
ctx = ToolContext()
ctx.register("ns1", TestState(value=1))
ctx.register("ns2", TestState(value=2))
ctx.clear_all()
assert ctx.get("ns1") is None
assert ctx.get("ns2") is None
def test_namespaces_property():
"""Test namespaces property lists all registered namespaces."""
ctx = ToolContext()
assert ctx.namespaces == []
ctx.register("ns1", TestState())
ctx.register("ns2", TestState())
assert set(ctx.namespaces) == {"ns1", "ns2"}
def test_shared_namespace_between_toolsets():
"""Test that toolsets can share state via the same namespace."""
class SharedState(BaseModel):
citations: dict[str, int] = {}
SHARED_NAMESPACE = "haiku.rag.citations"
ctx = ToolContext()
# First toolset registers the shared state
state1 = ctx.get_or_create(SHARED_NAMESPACE, SharedState)
state1.citations["chunk-a"] = 1
# Second toolset gets the same state
state2 = ctx.get_or_create(SHARED_NAMESPACE, SharedState)
assert state2 is state1
assert state2.citations == {"chunk-a": 1}
# Both see updates
state2.citations["chunk-b"] = 2
assert state1.citations == {"chunk-a": 1, "chunk-b": 2}
def test_dump_namespaces():
"""Test dump_namespaces serializes all registered states."""
ctx = ToolContext()
ctx.register("ns1", TestState(value=10))
ctx.register("ns2", TestStateWithList(items=["a", "b"]))
data = ctx.dump_namespaces()
assert data == {
"ns1": {"value": 10},
"ns2": {"items": ["a", "b"]},
}
def test_load_namespace():
"""Test load_namespace deserializes and registers state."""
ctx = ToolContext()
state = ctx.load_namespace("ns1", TestState, {"value": 42})
assert isinstance(state, TestState)
assert state.value == 42
assert ctx.get("ns1") is state
def test_serialization_roundtrip():
"""Test full serialization/deserialization roundtrip."""
# Create and populate context
original = ToolContext()
original.register("search", TestStateWithList(items=["result1", "result2"]))
original.register("qa", TestState(value=99))
# Serialize
ns_data = original.dump_namespaces()
# Deserialize
restored = ToolContext()
restored.load_namespace("search", TestStateWithList, ns_data["search"])
restored.load_namespace("qa", TestState, ns_data["qa"])
# Verify
search_state = restored.get("search")
assert isinstance(search_state, TestStateWithList)
assert search_state.items == ["result1", "result2"]
qa_state = restored.get("qa")
assert isinstance(qa_state, TestState)
assert qa_state.value == 99
def test_get_with_type_match():
"""Test get with state_type returns typed state when type matches."""
ctx = ToolContext()
state = TestState(value=42)
ctx.register("ns", state)
result = ctx.get("ns", TestState)
assert result is state
assert result.value == 42
def test_get_with_type_mismatch():
"""Test get with state_type returns None when type doesn't match."""
ctx = ToolContext()
ctx.register("ns", TestState(value=42))
result = ctx.get("ns", TestStateWithList)
assert result is None
def test_get_without_type():
"""Test get without state_type returns BaseModel (unchanged behavior)."""
ctx = ToolContext()
state = TestState(value=42)
ctx.register("ns", state)
result = ctx.get("ns")
assert result is state
def test_tool_context_state_key_default_none():
"""Test ToolContext state_key defaults to None."""
ctx = ToolContext()
assert ctx.state_key is None
def test_tool_context_state_key_set():
"""Test ToolContext state_key can be set."""
ctx = ToolContext()
ctx.state_key = "haiku.rag.chat"
assert ctx.state_key == "haiku.rag.chat"
# =============================================================================
# ToolContextCache Tests
# =============================================================================
def test_tool_context_cache_get_or_create_new():
"""Test get_or_create returns a new context with is_new=True."""
from haiku.rag.tools.context import ToolContextCache
cache = ToolContextCache()
context, is_new = cache.get_or_create("thread-1")
assert isinstance(context, ToolContext)
assert is_new is True
def test_tool_context_cache_get_or_create_existing():
"""Test get_or_create returns existing context with is_new=False."""
from haiku.rag.tools.context import ToolContextCache
cache = ToolContextCache()
ctx1, is_new1 = cache.get_or_create("thread-1")
ctx2, is_new2 = cache.get_or_create("thread-1")
assert ctx2 is ctx1
assert is_new1 is True
assert is_new2 is False
def test_tool_context_cache_different_keys():
"""Test get_or_create returns different contexts for different keys."""
from haiku.rag.tools.context import ToolContextCache
cache = ToolContextCache()
ctx1, _ = cache.get_or_create("thread-1")
ctx2, _ = cache.get_or_create("thread-2")
assert ctx1 is not ctx2
def test_tool_context_cache_ttl_expiry():
"""Test that contexts are evicted after TTL expires."""
from datetime import timedelta
from haiku.rag.tools.context import ToolContextCache
cache = ToolContextCache(ttl=timedelta(seconds=0))
ctx1, _ = cache.get_or_create("thread-1")
# With zero TTL, next access should create a new context
ctx2, is_new = cache.get_or_create("thread-1")
assert ctx2 is not ctx1
assert is_new is True
def test_tool_context_cache_remove():
"""Test remove deletes a specific key."""
from haiku.rag.tools.context import ToolContextCache
cache = ToolContextCache()
cache.get_or_create("thread-1")
cache.get_or_create("thread-2")
cache.remove("thread-1")
ctx, is_new = cache.get_or_create("thread-1")
assert is_new is True
# thread-2 should still exist
ctx2, is_new2 = cache.get_or_create("thread-2")
assert is_new2 is False
def test_tool_context_cache_remove_nonexistent():
"""Test remove handles nonexistent key gracefully."""
from haiku.rag.tools.context import ToolContextCache
cache = ToolContextCache()
cache.remove("nonexistent") # Should not raise
def test_tool_context_cache_clear():
"""Test clear removes all entries."""
from haiku.rag.tools.context import ToolContextCache
cache = ToolContextCache()
cache.get_or_create("thread-1")
cache.get_or_create("thread-2")
cache.clear()
ctx1, is_new1 = cache.get_or_create("thread-1")
ctx2, is_new2 = cache.get_or_create("thread-2")
assert is_new1 is True
assert is_new2 is True
# =============================================================================
# build_state_snapshot / restore_state_snapshot Tests
# =============================================================================
class NestedModel(BaseModel):
name: str = ""
count: int = 0
class StateWithNested(BaseModel):
nested: NestedModel | None = None
tags: list[str] = []
def test_build_state_snapshot_empty():
"""build_state_snapshot on empty context returns empty dict."""
ctx = ToolContext()
assert ctx.build_state_snapshot() == {}
def test_build_state_snapshot_single_namespace():
"""build_state_snapshot with one namespace returns its fields."""
ctx = ToolContext()
ctx.register("ns1", TestState(value=42))
snapshot = ctx.build_state_snapshot()
assert snapshot == {"value": 42}
def test_build_state_snapshot_multiple_namespaces():
"""build_state_snapshot merges fields from all namespaces."""
ctx = ToolContext()
ctx.register("ns1", TestState(value=42))
ctx.register("ns2", TestStateWithList(items=["a", "b"]))
snapshot = ctx.build_state_snapshot()
assert snapshot == {"value": 42, "items": ["a", "b"]}
def test_build_state_snapshot_nested_model():
"""build_state_snapshot serializes nested models with mode='json'."""
from datetime import datetime
class TimestampState(BaseModel):
ts: datetime | None = None
ctx = ToolContext()
ctx.register("ns", TimestampState(ts=datetime(2025, 1, 27, 12, 0, 0)))
snapshot = ctx.build_state_snapshot()
assert isinstance(snapshot["ts"], str)
assert snapshot["ts"] == "2025-01-27T12:00:00"
def test_restore_state_snapshot_empty_context():
"""restore_state_snapshot on empty context is a no-op."""
ctx = ToolContext()
ctx.restore_state_snapshot({"value": 42})
assert ctx.namespaces == []
def test_restore_state_snapshot_single_namespace():
"""restore_state_snapshot updates matching fields in registered namespaces."""
ctx = ToolContext()
ctx.register("ns1", TestState(value=0))
ctx.restore_state_snapshot({"value": 99})
state = ctx.get("ns1", TestState)
assert state is not None
assert state.value == 99
def test_restore_state_snapshot_partial_update():
"""restore_state_snapshot only touches fields present in data."""
ctx = ToolContext()
ctx.register("ns1", TestState(value=42))
ctx.register("ns2", TestStateWithList(items=["original"]))
# Only update ns2's items, not ns1's value
ctx.restore_state_snapshot({"items": ["updated"]})
ns1 = ctx.get("ns1", TestState)
ns2 = ctx.get("ns2", TestStateWithList)
assert ns1 is not None
assert ns2 is not None
assert ns1.value == 42
assert ns2.items == ["updated"]
def test_restore_state_snapshot_nested_model():
"""restore_state_snapshot deserializes nested models from dicts."""
ctx = ToolContext()
ctx.register("ns", StateWithNested())
ctx.restore_state_snapshot({"nested": {"name": "foo", "count": 5}, "tags": ["x"]})
state = ctx.get("ns", StateWithNested)
assert state is not None
assert state.nested is not None
assert state.nested.name == "foo"
assert state.nested.count == 5
assert state.tags == ["x"]
def test_state_snapshot_roundtrip():
"""build then restore produces equivalent state."""
ctx = ToolContext()
ctx.register("ns1", TestState(value=42))
ctx.register("ns2", TestStateWithList(items=["a", "b"]))
snapshot = ctx.build_state_snapshot()
ctx2 = ToolContext()
ctx2.register("ns1", TestState())
ctx2.register("ns2", TestStateWithList())
ctx2.restore_state_snapshot(snapshot)
ns1 = ctx2.get("ns1", TestState)
ns2 = ctx2.get("ns2", TestStateWithList)
assert ns1 is not None
assert ns2 is not None
assert ns1.value == 42
assert ns2.items == ["a", "b"]
def test_restore_state_snapshot_ignores_unknown_fields():
"""restore_state_snapshot ignores fields not in any registered namespace."""
ctx = ToolContext()
ctx.register("ns1", TestState(value=0))
ctx.restore_state_snapshot({"value": 10, "unknown_field": "ignored"})
ns1 = ctx.get("ns1", TestState)
assert ns1 is not None
assert ns1.value == 10
def test_restore_state_snapshot_captures_client_snapshot():
"""restore_state_snapshot stores the restored state as client_snapshot.
This baseline is used by tools to compute deltas against what the
client actually has, so server-side changes (e.g. background
summarization) appear in the delta.
"""
ctx = ToolContext()
ctx.register("ns1", TestState(value=0))
ctx.register("ns2", TestStateWithList(items=[]))
assert ctx.client_snapshot is None
ctx.restore_state_snapshot({"value": 10, "items": ["a"]})
assert ctx.client_snapshot == {"value": 10, "items": ["a"]}
# Mutating state after restore doesn't affect the captured snapshot
ns1 = ctx.get("ns1", TestState)
assert ns1 is not None
ns1.value = 99
assert ctx.client_snapshot == {"value": 10, "items": ["a"]}
# --- prepare_context tests ---
def test_prepare_context_default_features():
"""Default features register SessionState only."""
ctx = ToolContext()
prepare_context(ctx)
assert ctx.get(SESSION_NAMESPACE, SessionState) is not None
assert ctx.get(QA_SESSION_NAMESPACE, QASessionState) is None
def test_prepare_context_with_qa():
"""QA feature registers both SessionState and QASessionState."""
ctx = ToolContext()
prepare_context(ctx, features=["search", "qa"])
assert ctx.get(SESSION_NAMESPACE, SessionState) is not None
assert ctx.get(QA_SESSION_NAMESPACE, QASessionState) is not None
def test_prepare_context_sets_state_key():
"""state_key is set on context when provided."""
ctx = ToolContext()
prepare_context(ctx, state_key="my_app")
assert ctx.state_key == "my_app"
def test_prepare_context_no_state_key_by_default():
"""state_key is not set when not provided."""
ctx = ToolContext()
prepare_context(ctx)
assert ctx.state_key is None
def test_prepare_context_idempotent():
"""Calling prepare_context twice doesn't create duplicate state."""
ctx = ToolContext()
prepare_context(ctx, features=["search", "qa"])
session1 = ctx.get(SESSION_NAMESPACE, SessionState)
qa1 = ctx.get(QA_SESSION_NAMESPACE, QASessionState)
prepare_context(ctx, features=["search", "qa"])
assert ctx.get(SESSION_NAMESPACE, SessionState) is session1
assert ctx.get(QA_SESSION_NAMESPACE, QASessionState) is qa1
deps = NoDeps(other="x")
assert not isinstance(deps, RAGDeps)

106
tests/tools/test_deps.py
View file

@ -1,106 +0,0 @@
from unittest.mock import MagicMock
import pytest
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.deps import AgentDeps
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
@pytest.fixture
def mock_client():
return MagicMock()
def test_agent_deps_state_getter_empty(mock_client):
"""state returns empty dict when no namespaces are registered."""
ctx = ToolContext()
deps = AgentDeps(client=mock_client, tool_context=ctx)
assert deps.state == {}
def test_agent_deps_state_getter_with_session(mock_client):
"""state returns flat snapshot of registered namespaces."""
ctx = ToolContext()
ctx.register(SESSION_NAMESPACE, SessionState())
deps = AgentDeps(client=mock_client, tool_context=ctx)
snapshot = deps.state
assert "citations" in snapshot
assert "citation_registry" in snapshot
def test_agent_deps_state_getter_with_state_key(mock_client):
"""state wraps snapshot under state_key when set on context."""
ctx = ToolContext()
ctx.state_key = "my_app"
ctx.register(SESSION_NAMESPACE, SessionState())
deps = AgentDeps(client=mock_client, tool_context=ctx)
snapshot = deps.state
assert "my_app" in snapshot
assert "citations" in snapshot["my_app"]
def test_agent_deps_state_setter_restores(mock_client):
"""state setter restores namespace fields from flat dict."""
ctx = ToolContext()
ctx.register(SESSION_NAMESPACE, SessionState())
deps = AgentDeps(client=mock_client, tool_context=ctx)
deps.state = {"document_filter": ["doc1", "doc2"]}
session = ctx.get(SESSION_NAMESPACE, SessionState)
assert session is not None
assert session.document_filter == ["doc1", "doc2"]
def test_agent_deps_state_setter_with_state_key(mock_client):
"""state setter extracts data from namespaced key."""
ctx = ToolContext()
ctx.state_key = "my_app"
ctx.register(SESSION_NAMESPACE, SessionState())
deps = AgentDeps(client=mock_client, tool_context=ctx)
deps.state = {"my_app": {"document_filter": ["doc1"]}}
session = ctx.get(SESSION_NAMESPACE, SessionState)
assert session is not None
assert session.document_filter == ["doc1"]
def test_agent_deps_state_setter_ignores_none(mock_client):
"""state setter is a no-op when value is None."""
ctx = ToolContext()
ctx.register(SESSION_NAMESPACE, SessionState())
deps = AgentDeps(client=mock_client, tool_context=ctx)
deps.state = None
session = ctx.get(SESSION_NAMESPACE, SessionState)
assert session is not None
assert session.document_filter == []
def test_agent_deps_state_roundtrip(mock_client):
"""Build snapshot then restore produces equivalent state."""
ctx = ToolContext()
ctx.state_key = "app"
ctx.register(SESSION_NAMESPACE, SessionState(document_filter=["doc1"]))
ctx.register(QA_SESSION_NAMESPACE, QASessionState())
deps = AgentDeps(client=mock_client, tool_context=ctx)
snapshot = deps.state
ctx2 = ToolContext()
ctx2.state_key = "app"
ctx2.register(SESSION_NAMESPACE, SessionState())
ctx2.register(QA_SESSION_NAMESPACE, QASessionState())
deps2 = AgentDeps(client=mock_client, tool_context=ctx2)
deps2.state = snapshot
session = ctx2.get(SESSION_NAMESPACE, SessionState)
assert session is not None
assert session.document_filter == ["doc1"]
def test_agent_deps_satisfies_rag_deps_protocol(mock_client):
"""AgentDeps satisfies the RAGDeps protocol."""
from haiku.rag.tools.context import RAGDeps
ctx = ToolContext()
deps = AgentDeps(client=mock_client, tool_context=ctx)
assert isinstance(deps, RAGDeps)

4
tests/tools/test_document.py
View file

@ -15,9 +15,9 @@ def vcr_cassette_dir():
return str(Path(__file__).parent.parent / "cassettes" / "test_document_tools")
def make_ctx(client, context=None):
def make_ctx(client):
"""Create a lightweight RunContext-like object for direct tool function calls."""
return SimpleNamespace(deps=SimpleNamespace(client=client, tool_context=context))
return SimpleNamespace(deps=SimpleNamespace(client=client))
class TestDocumentModels:

37
tests/tools/test_filters.py
View file

@ -1,11 +1,8 @@
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.filters import (
build_document_filter,
build_multi_document_filter,
combine_filters,
get_session_filter,
)
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
def test_build_document_filter_simple():
@ -80,37 +77,3 @@ def test_combine_filters_both():
"""Test combine_filters combines with AND."""
result = combine_filters("uri = 'test'", "title = 'doc'")
assert result == "(uri = 'test') AND (title = 'doc')"
def test_get_session_filter_no_context():
"""Returns base_filter as-is when context is None."""
assert get_session_filter(None) is None
assert get_session_filter(None, "uri = 'test'") == "uri = 'test'"
def test_get_session_filter_no_document_filter():
"""Returns base_filter when SessionState has no document_filter."""
context = ToolContext()
context.register(SESSION_NAMESPACE, SessionState())
assert get_session_filter(context) is None
assert get_session_filter(context, "uri = 'test'") == "uri = 'test'"
def test_get_session_filter_with_document_filter():
"""Builds filter from SessionState.document_filter."""
context = ToolContext()
context.register(SESSION_NAMESPACE, SessionState(document_filter=["mytest"]))
result = get_session_filter(context)
assert result is not None
assert "mytest" in result
def test_get_session_filter_combines_with_base_filter():
"""Combines session filter with base_filter using AND."""
context = ToolContext()
context.register(SESSION_NAMESPACE, SessionState(document_filter=["mytest"]))
result = get_session_filter(context, "uri = 'base'")
assert result is not None
assert "uri = 'base'" in result
assert "mytest" in result
assert "AND" in result

81
tests/tools/test_models.py
View file

@ -1,83 +1,4 @@
from haiku.rag.agents.research.models import Citation
from haiku.rag.tools.models import AnalysisResult, QAResult
def test_qa_result_defaults():
"""Test QAResult has sensible defaults."""
result = QAResult(question="What is X?", answer="X is Y.")
assert result.confidence == 1.0
assert result.citations == []
def test_qa_result_with_citations():
"""Test QAResult with citations."""
citation = Citation(
index=1,
document_id="doc-1",
chunk_id="chunk-1",
document_uri="test.md",
document_title="Test Doc",
content="Citation content",
)
result = QAResult(
question="What is X?",
answer="X is Y.",
confidence=0.9,
citations=[citation],
)
assert result.confidence == 0.9
assert len(result.citations) == 1
assert result.citations[0].document_title == "Test Doc"
def test_qa_result_sources_property():
"""Test QAResult.sources returns unique source names."""
citations = [
Citation(
document_id="doc-1",
chunk_id="chunk-1",
document_uri="doc1.md",
document_title="Document One",
content="Content 1",
),
Citation(
document_id="doc-1",
chunk_id="chunk-2",
document_uri="doc1.md",
document_title="Document One",
content="Content 2",
),
Citation(
document_id="doc-2",
chunk_id="chunk-3",
document_uri="doc2.md",
document_title="Document Two",
content="Content 3",
),
]
result = QAResult(question="Q", answer="A", citations=citations)
sources = result.sources
assert len(sources) == 2
assert "Document One" in sources
assert "Document Two" in sources
def test_qa_result_sources_uses_uri_as_fallback():
"""Test QAResult.sources uses uri when title is None."""
citations = [
Citation(
document_id="doc-1",
chunk_id="chunk-1",
document_uri="test.md",
document_title=None,
content="Content",
),
]
result = QAResult(question="Q", answer="A", citations=citations)
sources = result.sources
assert sources == ["test.md"]
from haiku.rag.tools.analysis import AnalysisResult
def test_analysis_result_defaults():

53
tests/tools/test_prompts.py
View file

@ -1,53 +0,0 @@
from haiku.rag.tools.prompts import build_tools_prompt
def test_empty_features():
result = build_tools_prompt([])
assert result == ""
def test_single_feature_search():
result = build_tools_prompt(["search"])
assert "search" in result
assert "document_name" in result.lower()
def test_single_feature_documents():
result = build_tools_prompt(["documents"])
assert "list_documents" in result
assert "summarize_document" in result
assert "get_document" in result
def test_single_feature_qa():
result = build_tools_prompt(["qa"])
assert "ask" in result
assert "document_name" in result.lower()
def test_single_feature_analysis():
result = build_tools_prompt(["analysis"])
assert "analyze" in result
def test_multiple_features():
result = build_tools_prompt(["search", "qa", "documents"])
assert "search" in result
assert "ask" in result
assert "list_documents" in result
def test_search_and_qa_both_add_document_name_examples():
result = build_tools_prompt(["search", "qa"])
assert "search for embeddings" in result.lower() or "embeddings" in result
assert "what does the ML paper say" in result or "ML paper" in result
def test_unknown_features_ignored():
result = build_tools_prompt(["nonexistent", "also_fake"])
assert result == ""
def test_unknown_mixed_with_known():
result = build_tools_prompt(["nonexistent", "search"])
assert "search" in result

347
tests/tools/test_qa.py
View file

@ -1,274 +1,95 @@
from pathlib import Path
from types import SimpleNamespace
import pytest
from pydantic_ai import ToolReturn
from haiku.rag.tools import ToolContext, prepare_context
from haiku.rag.tools.models import QAResult
from haiku.rag.tools.qa import (
MAX_QA_HISTORY,
QA_SESSION_NAMESPACE,
QASessionState,
create_qa_toolset,
run_qa_core,
)
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
from haiku.rag.agents.research.models import Citation
from haiku.rag.tools.qa import PRIOR_ANSWER_RELEVANCE_THRESHOLD, QAHistoryEntry
@pytest.fixture(scope="module")
def vcr_cassette_dir():
return str(Path(__file__).parent.parent / "cassettes" / "test_qa_tools")
class TestQAHistoryEntry:
"""Tests for QAHistoryEntry model."""
def test_defaults(self):
"""QAHistoryEntry has sensible defaults."""
entry = QAHistoryEntry(question="What is X?", answer="X is Y.")
assert entry.confidence == 0.9
assert entry.citations == []
assert entry.question_embedding is None
def make_ctx(client, context=None):
"""Create a lightweight RunContext-like object for direct tool function calls."""
return SimpleNamespace(deps=SimpleNamespace(client=client, tool_context=context))
class TestQAToolset:
"""Tests for create_qa_toolset."""
def test_create_qa_toolset_returns_function_toolset(self, qa_config):
"""create_qa_toolset returns a FunctionToolset."""
from pydantic_ai import FunctionToolset
toolset = create_qa_toolset(qa_config)
assert isinstance(toolset, FunctionToolset)
def test_qa_toolset_has_ask_tool(self, qa_config):
"""The toolset includes an 'ask' tool."""
toolset = create_qa_toolset(qa_config)
assert "ask" in toolset.tools
def test_qa_toolset_custom_tool_name(self, qa_config):
"""Toolset supports custom tool name."""
toolset = create_qa_toolset(qa_config, tool_name="answer_question")
assert "answer_question" in toolset.tools
assert "ask" not in toolset.tools
@pytest.mark.vcr()
class TestRunQACore:
"""Tests for run_qa_core."""
@pytest.mark.asyncio
async def test_run_qa_core_with_session_state(
self, allow_model_requests, qa_client, qa_config
):
"""run_qa_core with SessionState assigns citation indices via registry."""
context = ToolContext()
prepare_context(context, features=["qa"])
result = await run_qa_core(
client=qa_client,
config=qa_config,
question="What is Python?",
context=context,
)
assert isinstance(result, QAResult)
assert result.answer
session_state = context.get(SESSION_NAMESPACE, SessionState)
assert session_state is not None
# If citations were returned, they should use registry indices
if result.citations:
assert len(session_state.citation_registry) > 0
@pytest.mark.asyncio
async def test_run_qa_core_populates_citations_history(
self, allow_model_requests, qa_client, qa_config
):
"""run_qa_core appends to SessionState.citations_history."""
context = ToolContext()
prepare_context(context, features=["qa"])
await run_qa_core(
client=qa_client,
config=qa_config,
question="What is Python?",
context=context,
)
session_state = context.get(SESSION_NAMESPACE, SessionState)
assert session_state is not None
assert len(session_state.citations_history) == 1
assert session_state.citations_history[0] == session_state.citations
@pytest.mark.asyncio
async def test_run_qa_core_without_context(
self, allow_model_requests, qa_client, qa_config
):
"""run_qa_core without context uses sequential fallback indices."""
result = await run_qa_core(
client=qa_client,
config=qa_config,
question="What is Python?",
context=None,
)
assert isinstance(result, QAResult)
assert result.answer
# Without context, citation indices are i+1
for i, c in enumerate(result.citations):
assert c.index == i + 1
@pytest.mark.asyncio
async def test_run_qa_core_on_qa_complete_callback(
self, allow_model_requests, qa_client, qa_config
):
"""run_qa_core invokes on_qa_complete callback when context is provided."""
context = ToolContext()
prepare_context(context, features=["qa"])
callback_calls: list[tuple] = []
def on_complete(qa_session_state, config):
callback_calls.append((qa_session_state, config))
await run_qa_core(
client=qa_client,
config=qa_config,
question="What is Python?",
context=context,
on_qa_complete=on_complete,
)
assert len(callback_calls) == 1
assert isinstance(callback_calls[0][0], QASessionState)
@pytest.mark.asyncio
async def test_run_qa_core_fifo_limit(
self, allow_model_requests, qa_client, qa_config
):
"""run_qa_core trims qa_history beyond MAX_QA_HISTORY."""
context = ToolContext()
prepare_context(context, features=["qa"])
qa_session_state = context.get(QA_SESSION_NAMESPACE, QASessionState)
assert qa_session_state is not None
# Pre-fill with MAX_QA_HISTORY entries
from haiku.rag.tools.qa import QAHistoryEntry
qa_session_state.qa_history = [
QAHistoryEntry(question=f"Q{i}", answer=f"A{i}", confidence=0.9)
for i in range(MAX_QA_HISTORY)
def test_sources_property(self):
"""sources returns unique document titles."""
citations = [
Citation(
document_id="d1",
chunk_id="c1",
document_uri="doc1.md",
document_title="Document One",
content="Content 1",
),
Citation(
document_id="d1",
chunk_id="c2",
document_uri="doc1.md",
document_title="Document One",
content="Content 2",
),
Citation(
document_id="d2",
chunk_id="c3",
document_uri="doc2.md",
document_title="Document Two",
content="Content 3",
),
]
entry = QAHistoryEntry(question="Q", answer="A", citations=citations)
sources = entry.sources
assert len(sources) == 2
assert "Document One" in sources
assert "Document Two" in sources
await run_qa_core(
client=qa_client,
config=qa_config,
question="One more question?",
context=context,
)
# After adding one more, FIFO should trim to MAX_QA_HISTORY
assert len(qa_session_state.qa_history) == MAX_QA_HISTORY
# The oldest entry (Q0) should have been trimmed
assert qa_session_state.qa_history[0].question != "Q0"
@pytest.mark.vcr()
class TestRunQACoreWithPriorAnswers:
"""Tests for run_qa_core prior answer matching."""
@pytest.mark.asyncio
async def test_run_qa_core_matches_prior_answers(
self, allow_model_requests, qa_client, qa_config
):
"""run_qa_core matches prior answers when embedding similarity is high."""
from unittest.mock import AsyncMock, patch
from haiku.rag.tools.qa import QAHistoryEntry
context = ToolContext()
prepare_context(context, features=["qa"])
qa_session_state = context.get(QA_SESSION_NAMESPACE, QASessionState)
assert qa_session_state is not None
# Pre-populate with a prior answer that has a known embedding
prior_embedding = [0.5] * 2560
qa_session_state.qa_history = [
QAHistoryEntry(
question="What is Python?",
answer="A programming language.",
confidence=0.9,
question_embedding=prior_embedding,
)
def test_sources_uses_uri_as_fallback(self):
"""sources uses uri when title is None."""
citations = [
Citation(
document_id="d1",
chunk_id="c1",
document_uri="test.md",
document_title=None,
content="Content",
),
]
entry = QAHistoryEntry(question="Q", answer="A", citations=citations)
assert entry.sources == ["test.md"]
# Mock the embedder to return a near-identical embedding for the new question
mock_embedder = AsyncMock()
mock_embedder.embed_query = AsyncMock(return_value=[0.5] * 2560)
with patch("haiku.rag.tools.qa.get_embedder", return_value=mock_embedder):
result = await run_qa_core(
client=qa_client,
config=qa_config,
question="Tell me about Python",
context=context,
)
assert isinstance(result, QAResult)
assert result.answer
@pytest.mark.vcr()
class TestAskTool:
"""Tests for the ask tool in create_qa_toolset."""
@pytest.mark.asyncio
async def test_ask_without_tool_context(
self, allow_model_requests, qa_client, qa_config
):
"""ask tool without tool context returns raw QAResult."""
toolset = create_qa_toolset(qa_config)
ask_tool = toolset.tools["ask"]
ctx = make_ctx(qa_client, None)
result = await ask_tool.function(ctx, "What is Python?")
assert isinstance(result, QAResult)
assert result.answer
@pytest.mark.asyncio
async def test_ask_with_tool_context_returns_tool_return(
self, allow_model_requests, qa_client, qa_config
):
"""ask tool with tool context returns ToolReturn with state snapshot."""
context = ToolContext()
prepare_context(context, features=["qa"])
toolset = create_qa_toolset(qa_config)
ask_tool = toolset.tools["ask"]
ctx = make_ctx(qa_client, context)
result = await ask_tool.function(ctx, "What is Python?")
assert isinstance(result, ToolReturn)
assert result.metadata is not None
assert len(result.metadata) > 0
@pytest.fixture
async def qa_client(temp_db_path):
"""Create a HaikuRAG client with test documents for QA tests."""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as rag:
await rag.create_document(
"Python is a programming language. It is widely used for web development.",
uri="test://python",
title="Python Guide",
def test_to_search_answer(self):
"""to_search_answer converts to SearchAnswer."""
citation = Citation(
document_id="d1",
chunk_id="c1",
document_uri="doc1.md",
document_title="Doc",
content="Content",
)
yield rag
entry = QAHistoryEntry(
question="What is X?",
answer="X is Y.",
confidence=0.85,
citations=[citation],
)
sa = entry.to_search_answer()
assert sa.query == "What is X?"
assert sa.answer == "X is Y."
assert sa.confidence == 0.85
assert sa.cited_chunks == ["c1"]
assert len(sa.citations) == 1
def test_question_embedding_excluded_from_serialization(self):
"""question_embedding is excluded from model_dump."""
entry = QAHistoryEntry(
question="Q",
answer="A",
question_embedding=[0.1, 0.2],
)
data = entry.model_dump()
assert "question_embedding" not in data
@pytest.fixture
def qa_config():
"""Default AppConfig for QA tests."""
from haiku.rag.config import Config
return Config
def test_prior_answer_relevance_threshold():
"""PRIOR_ANSWER_RELEVANCE_THRESHOLD is a sensible value."""
assert 0 < PRIOR_ANSWER_RELEVANCE_THRESHOLD < 1

292
tests/tools/test_search.py
View file

@ -2,11 +2,9 @@ from pathlib import Path
from types import SimpleNamespace
import pytest
from pydantic_ai import ToolReturn
from haiku.rag.tools import ToolContext, prepare_context
from haiku.rag.tools.search import SEARCH_NAMESPACE, SearchState, create_search_toolset
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
from haiku.rag.store.models import SearchResult
from haiku.rag.tools.search import create_search_toolset
@pytest.fixture(scope="module")
@ -14,52 +12,9 @@ def vcr_cassette_dir():
return str(Path(__file__).parent.parent / "cassettes" / "test_search_tools")
def make_ctx(client, context=None):
def make_ctx(client):
"""Create a lightweight RunContext-like object for direct tool function calls."""
return SimpleNamespace(deps=SimpleNamespace(client=client, tool_context=context))
class TestSearchState:
"""Tests for SearchState model."""
def test_search_state_defaults(self):
"""SearchState initializes with empty results."""
state = SearchState()
assert state.results == []
def test_search_state_add_results(self):
"""Can add results to SearchState."""
from haiku.rag.store.models import SearchResult
state = SearchState()
result = SearchResult(content="test content", score=0.9, chunk_id="chunk1")
state.results.append(result)
assert len(state.results) == 1
assert state.results[0].chunk_id == "chunk1"
def test_search_state_serialization(self):
"""SearchState serializes and deserializes correctly."""
from haiku.rag.store.models import SearchResult
state = SearchState()
state.results.append(
SearchResult(
content="test",
score=0.8,
chunk_id="c1",
document_title="Doc Title",
)
)
# Serialize
data = state.model_dump()
assert "results" in data
assert len(data["results"]) == 1
# Deserialize
restored = SearchState.model_validate(data)
assert len(restored.results) == 1
assert restored.results[0].chunk_id == "c1"
return SimpleNamespace(deps=SimpleNamespace(client=client))
@pytest.mark.vcr()
@ -88,34 +43,15 @@ class TestSearchToolExecution:
@pytest.mark.asyncio
async def test_search_returns_formatted_results(self, search_client, search_config):
"""Search tool returns formatted results."""
context = ToolContext()
toolset = create_search_toolset(search_config)
# Get the search function
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client, context)
ctx = make_ctx(search_client)
result = await search_tool.function(ctx, "Python")
assert "Python" in result or "programming" in result
assert "No results found" not in result
@pytest.mark.asyncio
async def test_search_accumulates_in_state(self, search_client, search_config):
"""Search tool accumulates results in SearchState."""
context = ToolContext()
toolset = create_search_toolset(search_config)
# Run search
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client, context)
await search_tool.function(ctx, "Python")
# Check state was updated
state = context.get(SEARCH_NAMESPACE)
assert isinstance(state, SearchState)
assert len(state.results) > 0
assert any("Python" in r.content for r in state.results)
@pytest.mark.asyncio
async def test_search_with_no_results(self, temp_db_path, search_config):
"""Search tool returns appropriate message when no results."""
@ -123,11 +59,10 @@ class TestSearchToolExecution:
# Use empty database
async with HaikuRAG(temp_db_path, create=True) as empty_client:
context = ToolContext()
toolset = create_search_toolset(search_config)
search_tool = toolset.tools["search"]
ctx = make_ctx(empty_client, context)
ctx = make_ctx(empty_client)
result = await search_tool.function(ctx, "anything")
assert result == "No results found."
@ -135,73 +70,74 @@ class TestSearchToolExecution:
@pytest.mark.asyncio
async def test_search_with_filter(self, search_client, search_config):
"""Search tool respects filter parameter."""
context = ToolContext()
toolset = create_search_toolset(search_config)
accumulated: list[SearchResult] = []
toolset = create_search_toolset(search_config, on_results=accumulated.extend)
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client, context)
# Filter to only Python documents
ctx = make_ctx(search_client)
await search_tool.function(ctx, "programming", filter="title LIKE '%Python%'")
# Should find Python but not JavaScript
state = context.get(SEARCH_NAMESPACE)
assert isinstance(state, SearchState)
for r in state.results:
for r in accumulated:
assert "JavaScript" not in (r.document_title or "")
@pytest.mark.asyncio
async def test_search_without_context(self, search_client, search_config):
"""Search tool works without ToolContext."""
toolset = create_search_toolset(search_config)
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client, None)
result = await search_tool.function(ctx, "Python")
# Should still return results
assert "Python" in result or "programming" in result
@pytest.mark.asyncio
async def test_search_multiple_accumulates(self, search_client, search_config):
"""Multiple searches accumulate results in state."""
context = ToolContext()
toolset = create_search_toolset(search_config)
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client, context)
await search_tool.function(ctx, "Python")
state = context.get(SEARCH_NAMESPACE)
assert isinstance(state, SearchState)
first_count = len(state.results)
await search_tool.function(ctx, "JavaScript")
state = context.get(SEARCH_NAMESPACE)
assert isinstance(state, SearchState)
second_count = len(state.results)
assert second_count > first_count
@pytest.mark.asyncio
async def test_search_with_base_filter(self, search_client, search_config):
"""Search toolset respects base_filter parameter."""
context = ToolContext()
# Create toolset with base_filter for Python documents only
accumulated: list[SearchResult] = []
toolset = create_search_toolset(
search_config,
base_filter="title LIKE '%Python%'",
on_results=accumulated.extend,
)
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client, context)
ctx = make_ctx(search_client)
await search_tool.function(ctx, "programming")
# Should only find Python documents
state = context.get(SEARCH_NAMESPACE)
assert isinstance(state, SearchState)
assert len(state.results) > 0
for r in state.results:
assert len(accumulated) > 0
for r in accumulated:
assert "JavaScript" not in (r.document_title or "")
@pytest.mark.asyncio
async def test_search_on_results_callback(self, search_client, search_config):
"""on_results callback receives search results."""
accumulated: list[SearchResult] = []
toolset = create_search_toolset(search_config, on_results=accumulated.extend)
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client)
await search_tool.function(ctx, "Python")
assert len(accumulated) > 0
assert any("Python" in r.content for r in accumulated)
@pytest.mark.asyncio
async def test_search_on_results_accumulates_across_calls(
self, search_client, search_config
):
"""Multiple searches accumulate results via on_results callback."""
accumulated: list[SearchResult] = []
toolset = create_search_toolset(search_config, on_results=accumulated.extend)
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client)
await search_tool.function(ctx, "Python")
first_count = len(accumulated)
await search_tool.function(ctx, "JavaScript")
assert len(accumulated) > first_count
@pytest.mark.asyncio
async def test_search_without_on_results(self, search_client, search_config):
"""Search works without on_results callback."""
toolset = create_search_toolset(search_config)
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client)
result = await search_tool.function(ctx, "Python")
assert "Python" in result or "programming" in result
@pytest.fixture
async def search_client(temp_db_path):
@ -222,124 +158,6 @@ async def search_client(temp_db_path):
yield rag
@pytest.mark.vcr()
class TestSearchWithSessionState:
"""Tests for search tool with session state (citation indexing path)."""
@pytest.mark.asyncio
async def test_search_with_session_state_returns_tool_return(
self, search_client, search_config
):
"""Search with SessionState returns ToolReturn with StateDeltaEvent."""
context = ToolContext()
prepare_context(context, features=["search"])
toolset = create_search_toolset(search_config)
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client, context)
result = await search_tool.function(ctx, "Python")
assert isinstance(result, ToolReturn)
assert result.metadata is not None
assert len(result.metadata) > 0
@pytest.mark.asyncio
async def test_search_with_session_state_populates_citations(
self, search_client, search_config
):
"""Search populates SessionState.citation_registry and citations."""
context = ToolContext()
prepare_context(context, features=["search"])
toolset = create_search_toolset(search_config)
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client, context)
await search_tool.function(ctx, "Python")
session_state = context.get(SESSION_NAMESPACE, SessionState)
assert session_state is not None
assert len(session_state.citation_registry) > 0
assert len(session_state.citations) > 0
@pytest.mark.asyncio
async def test_search_with_session_state_formatted_output(
self, search_client, search_config
):
"""Search with SessionState formats results with [index] **Title**."""
context = ToolContext()
prepare_context(context, features=["search"])
toolset = create_search_toolset(search_config)
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client, context)
result = await search_tool.function(ctx, "Python")
assert isinstance(result, ToolReturn)
output = result.return_value
assert "Found" in output
assert "[1]" in output
assert "**" in output
@pytest.mark.asyncio
async def test_search_citations_accumulate_across_calls(
self, search_client, search_config
):
"""Multiple searches accumulate citation indices across calls."""
context = ToolContext()
prepare_context(context, features=["search"])
toolset = create_search_toolset(search_config)
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client, context)
await search_tool.function(ctx, "Python")
session_state = context.get(SESSION_NAMESPACE, SessionState)
assert session_state is not None
first_count = len(session_state.citation_registry)
await search_tool.function(ctx, "JavaScript")
# New chunks should get higher indices
assert len(session_state.citation_registry) >= first_count
@pytest.mark.asyncio
async def test_search_populates_citations_history(
self, search_client, search_config
):
"""Search appends to SessionState.citations_history."""
context = ToolContext()
prepare_context(context, features=["search"])
toolset = create_search_toolset(search_config)
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client, context)
await search_tool.function(ctx, "Python")
session_state = context.get(SESSION_NAMESPACE, SessionState)
assert session_state is not None
assert len(session_state.citations_history) == 1
assert session_state.citations_history[0] == session_state.citations
@pytest.mark.asyncio
async def test_search_multiple_appends_separate_entries(
self, search_client, search_config
):
"""Multiple searches append separate entries to citations_history."""
context = ToolContext()
prepare_context(context, features=["search"])
toolset = create_search_toolset(search_config)
search_tool = toolset.tools["search"]
ctx = make_ctx(search_client, context)
await search_tool.function(ctx, "Python")
await search_tool.function(ctx, "JavaScript")
session_state = context.get(SESSION_NAMESPACE, SessionState)
assert session_state is not None
assert len(session_state.citations_history) == 2
# Latest citations should match the last entry
assert session_state.citations_history[1] == session_state.citations
@pytest.fixture
def search_config():
"""Default AppConfig for search tests."""

115
tests/tools/test_session.py
View file

@ -1,115 +0,0 @@
from ag_ui.core import EventType, StateDeltaEvent
from haiku.rag.agents.research.models import Citation
from haiku.rag.tools.session import (
SessionState,
compute_combined_state_delta,
compute_state_delta,
)
class TestSessionState:
"""Tests for SessionState model."""
def test_citations_history_defaults_to_empty(self):
"""SessionState.citations_history defaults to empty list."""
state = SessionState()
assert state.citations_history == []
def test_citations_history_serialization_roundtrip(self):
"""citations_history survives serialize/deserialize."""
citation = Citation(
index=1,
document_id="d1",
chunk_id="c1",
document_uri="test://doc",
document_title="Doc",
page_numbers=[],
headings=None,
content="some content",
)
state = SessionState(citations_history=[[citation]])
data = state.model_dump(mode="json")
restored = SessionState.model_validate(data)
assert len(restored.citations_history) == 1
assert len(restored.citations_history[0]) == 1
assert restored.citations_history[0][0].chunk_id == "c1"
class TestComputeStateDelta:
"""Tests for compute_state_delta."""
def test_returns_delta_on_change(self):
"""compute_state_delta returns StateDeltaEvent when state changed."""
old = SessionState()
new = SessionState(citation_registry={"chunk-a": 1})
result = compute_state_delta(old, new)
assert isinstance(result, StateDeltaEvent)
assert result.type == EventType.STATE_DELTA
assert len(result.delta) > 0
def test_returns_none_on_no_change(self):
"""compute_state_delta returns None when states are identical."""
state = SessionState(document_filter=["doc1"])
result = compute_state_delta(state, state.model_copy(deep=True))
assert result is None
def test_with_state_key(self):
"""compute_state_delta wraps delta under state_key."""
old = SessionState()
new = SessionState(document_filter=["doc1"])
result = compute_state_delta(old, new, state_key="my.key")
assert isinstance(result, StateDeltaEvent)
# The delta paths should be prefixed with /my.key/
paths = [op["path"] for op in result.delta]
assert all(p.startswith("/my.key/") for p in paths)
class TestComputeCombinedStateDelta:
"""Tests for compute_combined_state_delta."""
def test_returns_delta_on_change(self):
"""compute_combined_state_delta returns StateDeltaEvent when snapshots differ."""
old = {"citations": []}
new = {"citations": [{"index": 1, "chunk_id": "c1"}]}
result = compute_combined_state_delta(old, new)
assert isinstance(result, StateDeltaEvent)
assert result.type == EventType.STATE_DELTA
def test_returns_none_on_no_change(self):
"""compute_combined_state_delta returns None when snapshots are identical."""
snapshot = {"citations": [], "document_filter": []}
result = compute_combined_state_delta(snapshot, snapshot.copy())
assert result is None
def test_with_state_key_wraps(self):
"""compute_combined_state_delta wraps under state_key."""
old = {"value": 1}
new = {"value": 2}
result = compute_combined_state_delta(old, new, state_key="ns")
assert isinstance(result, StateDeltaEvent)
paths = [op["path"] for op in result.delta]
assert all(p.startswith("/ns/") for p in paths)
def test_without_state_key(self):
"""compute_combined_state_delta works without state_key."""
old = {"value": 1}
new = {"value": 2}
result = compute_combined_state_delta(old, new)
assert isinstance(result, StateDeltaEvent)
paths = [op["path"] for op in result.delta]
assert any(p == "/value" for p in paths)

99
tests/tools/test_toolkit.py
View file

@ -1,99 +0,0 @@
import pytest
from haiku.rag.config import Config
from haiku.rag.tools.prompts import build_tools_prompt
from haiku.rag.tools.toolkit import (
FEATURE_ANALYSIS,
FEATURE_DOCUMENTS,
FEATURE_QA,
FEATURE_SEARCH,
build_toolkit,
)
def test_build_toolkit_default_features():
"""Defaults to ["search", "documents"], producing 2 toolsets."""
toolkit = build_toolkit(Config)
assert len(toolkit.toolsets) == 2
assert toolkit.features == [FEATURE_SEARCH, FEATURE_DOCUMENTS]
def test_build_toolkit_all_features():
"""All 4 features produce 4 toolsets."""
features = [FEATURE_SEARCH, FEATURE_DOCUMENTS, FEATURE_QA, FEATURE_ANALYSIS]
toolkit = build_toolkit(Config, features=features)
assert len(toolkit.toolsets) == 4
assert toolkit.features == features
def test_build_toolkit_single_feature():
"""Single feature produces 1 toolset."""
toolkit = build_toolkit(Config, features=[FEATURE_SEARCH])
assert len(toolkit.toolsets) == 1
assert toolkit.features == [FEATURE_SEARCH]
def test_build_toolkit_prompt_matches_features():
"""Toolkit prompt matches build_tools_prompt for the same features."""
features = [FEATURE_SEARCH, FEATURE_DOCUMENTS, FEATURE_QA]
toolkit = build_toolkit(Config, features=features)
expected = build_tools_prompt(features)
assert toolkit.prompt == expected
def test_toolkit_create_context_registers_namespaces():
"""create_context registers correct namespaces for the features."""
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
toolkit = build_toolkit(Config, features=[FEATURE_SEARCH, FEATURE_QA])
context = toolkit.create_context()
assert context.get(SESSION_NAMESPACE, SessionState) is not None
assert context.get(QA_SESSION_NAMESPACE, QASessionState) is not None
def test_toolkit_create_context_no_qa_skips_qa_state():
"""create_context without QA feature skips QASessionState."""
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
toolkit = build_toolkit(Config, features=[FEATURE_SEARCH, FEATURE_DOCUMENTS])
context = toolkit.create_context()
assert context.get(SESSION_NAMESPACE, SessionState) is not None
assert context.get(QA_SESSION_NAMESPACE, QASessionState) is None
def test_toolkit_create_context_sets_state_key():
"""create_context propagates state_key to the ToolContext."""
toolkit = build_toolkit(Config)
context = toolkit.create_context(state_key="my.state.key")
assert context.state_key == "my.state.key"
def test_toolkit_create_context_no_state_key():
"""create_context without state_key leaves it None."""
toolkit = build_toolkit(Config)
context = toolkit.create_context()
assert context.state_key is None
def test_toolkit_prepare_existing_context():
"""prepare registers namespaces on an existing ToolContext."""
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
toolkit = build_toolkit(Config, features=[FEATURE_SEARCH])
context = ToolContext()
toolkit.prepare(context, state_key="test.key")
assert context.get(SESSION_NAMESPACE, SessionState) is not None
assert context.state_key == "test.key"
def test_toolkit_frozen():
"""Toolkit is immutable after creation."""
toolkit = build_toolkit(Config)
with pytest.raises(AttributeError):
toolkit.features = ["search"] # type: ignore[misc]