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

Use the tools inside the chat agent.

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-05 12:48:35 +01:00
parent 5bcbcc3928
commit 3ed9cbb7d3
No known key found for this signature in database
20 changed files with 3140 additions and 1292 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

110
app/backend/main.py
View file

@ -15,8 +15,7 @@ from starlette.routing import Route
from haiku.rag.agents.chat import (
AGUI_STATE_KEY,
ChatDeps,
ChatSessionState,
QAResponse,
ToolContext,
create_chat_agent,
)
from haiku.rag.client import HaikuRAG
@ -54,96 +53,48 @@ db_path = Path(db_path_str)
logger.info(f"Database path: {db_path}")
logger.info(f"QA Provider: {Config.qa.model.provider}, Model: {Config.qa.model.name}")
# Create the chat agent
chat_agent = create_chat_agent(Config)
# Client cache for proper lifecycle
_client_cache: dict[str, HaikuRAG] = {}
# Only HaikuRAG client is a singleton (expensive to create)
_client: HaikuRAG | None = None
def get_client(effective_db_path: Path) -> HaikuRAG:
def get_client() -> HaikuRAG:
"""Get or create cached client."""
path_key = str(effective_db_path)
if path_key not in _client_cache:
_client_cache[path_key] = HaikuRAG(
db_path=effective_db_path, config=Config, create=True
)
return _client_cache[path_key]
global _client
if _client is None:
_client = HaikuRAG(db_path=db_path, config=Config, create=True)
return _client
async def stream_chat(request: Request) -> Response:
"""Chat streaming endpoint with AG-UI protocol."""
body = await request.body()
"""Chat streaming endpoint with AG-UI protocol.
# Parse request to build run_input
This endpoint is stateless - all state flows via AG-UI protocol:
- Fresh ToolContext created per request
- AGUIAdapter restores state via ChatDeps.state setter
- ChatDeps generates session_id if not provided
- Ask tool triggers background summarization internally
- ChatDeps.state getter emits final state in response
"""
body = await request.body()
accept = request.headers.get("accept", SSE_CONTENT_TYPE)
run_input = AGUIAdapter.build_run_input(body)
# Restore session state from incoming AG-UI state
session_state = ChatSessionState(session_id="") # New session: empty session_id
state = getattr(run_input, "state", None)
if state and AGUI_STATE_KEY in state:
chat_state = state[AGUI_STATE_KEY]
if chat_state and chat_state.get("session_id"):
session_state = ChatSessionState(
session_id=chat_state["session_id"],
qa_history=[
QAResponse(**qa) for qa in chat_state.get("qa_history", [])
],
document_filter=chat_state.get("document_filter", []),
initial_context=chat_state.get("initial_context"),
citation_registry=chat_state.get("citation_registry", {}),
)
logger.info(
f"Incoming state: session={session_state.session_id[:8]}, "
f"qa_history={len(session_state.qa_history)}, "
f"citations={len(session_state.citation_registry)}"
)
else:
logger.info("Incoming state: new session")
else:
logger.info("Incoming state: new session")
# Fresh context per request - state restored by AGUIAdapter via ChatDeps.state setter
context = ToolContext()
agent = create_chat_agent(Config, get_client(), context)
deps = ChatDeps(
client=get_client(db_path),
config=Config,
session_state=session_state,
tool_context=context,
state_key=AGUI_STATE_KEY,
)
# Use AGUIAdapter for streaming
adapter = AGUIAdapter(agent=chat_agent, run_input=run_input, accept=accept)
# State restoration happens automatically via ChatDeps.state setter
# Background summarization triggered by ask() tool internally
adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
event_stream = adapter.run_stream(deps=deps)
async def logged_event_stream():
async for event in event_stream:
event_type = getattr(event, "type", None)
if event_type and "state" in str(event_type).lower():
delta = getattr(event, "delta", None)
snapshot = getattr(event, "snapshot", None)
if delta is not None:
logger.info(f"Outgoing StateDeltaEvent: {len(delta)} ops")
for op in delta:
# Extract key from path like /haiku.rag.chat/qa_history/0
parts = op["path"].split("/")
key = "/".join(parts[2:]) if len(parts) > 2 else op["path"]
logger.info(f" {op['op']} {key}")
elif snapshot is not None:
chat_state = snapshot.get(AGUI_STATE_KEY, {})
sid = chat_state.get("session_id", "")[:8] if chat_state else ""
qa_len = len(chat_state.get("qa_history", [])) if chat_state else 0
reg_len = (
len(chat_state.get("citation_registry", {}))
if chat_state
else 0
)
logger.info(
f"Outgoing StateSnapshotEvent: session={sid}, "
f"qa={qa_len}, keys={reg_len}"
)
yield event
sse_event_stream = adapter.encode_stream(logged_event_stream())
sse_event_stream = adapter.encode_stream(event_stream)
return StreamingResponse(
sse_event_stream,
@ -158,10 +109,13 @@ async def stream_chat(request: Request) -> Response:
async def health_check(_: Request) -> JSONResponse:
"""Health check endpoint."""
# Create a temporary agent just for health check
context = ToolContext()
agent = create_chat_agent(Config, get_client(), context)
return JSONResponse(
{
"status": "healthy",
"agent_model": str(chat_agent.model),
"agent_model": str(agent.model),
"qa_provider": Config.qa.model.provider,
"qa_model": Config.qa.model.name,
"db_path": str(db_path),
@ -175,7 +129,7 @@ async def list_documents(_: Request) -> JSONResponse:
if not db_path.exists():
return JSONResponse({"documents": [], "error": "Database not found"})
client = get_client(db_path)
client = get_client()
docs = await client.document_repository.list_all()
return JSONResponse(
{
@ -198,7 +152,7 @@ async def db_info(_: Request) -> JSONResponse:
}
)
client = get_client(db_path)
client = get_client()
stats = client.store.get_stats()
return JSONResponse(
@ -224,7 +178,7 @@ async def visualize_chunk(request: Request) -> JSONResponse:
if not db_path.exists():
return JSONResponse({"error": "Database not found"}, status_code=404)
client = get_client(db_path)
client = get_client()
chunk = await client.chunk_repository.get_by_id(chunk_id)
if not chunk:

12
haiku_rag_slim/haiku/rag/agents/chat/__init__.py
View file

@ -1,4 +1,9 @@
from haiku.rag.agents.chat.agent import create_chat_agent
from haiku.rag.agents.chat.agent import (
ChatDeps,
create_chat_agent,
run_chat_agent,
trigger_background_summarization,
)
from haiku.rag.agents.chat.context import (
summarize_session,
update_session_context,
@ -6,7 +11,6 @@ from haiku.rag.agents.chat.context import (
from haiku.rag.agents.chat.search import SearchAgent
from haiku.rag.agents.chat.state import (
AGUI_STATE_KEY,
ChatDeps,
ChatSessionState,
DocumentInfo,
DocumentListResponse,
@ -14,10 +18,13 @@ from haiku.rag.agents.chat.state import (
SearchDeps,
SessionContext,
)
from haiku.rag.tools.context import ToolContext
__all__ = [
"AGUI_STATE_KEY",
"create_chat_agent",
"run_chat_agent",
"trigger_background_summarization",
"SearchAgent",
"ChatDeps",
"ChatSessionState",
@ -26,6 +33,7 @@ __all__ = [
"QAResponse",
"SearchDeps",
"SessionContext",
"ToolContext",
"summarize_session",
"update_session_context",
]

675
haiku_rag_slim/haiku/rag/agents/chat/agent.py
View file

@ -1,452 +1,285 @@
import asyncio
import math
import uuid
from dataclasses import dataclass
from typing import Any, cast
from pydantic_ai import Agent, RunContext, ToolReturn
from pydantic_ai import Agent
from haiku.rag.agents.chat.context import (
get_cached_session_context,
update_session_context,
)
from haiku.rag.agents.chat.prompts import CHAT_SYSTEM_PROMPT, DOCUMENT_SUMMARY_PROMPT
from haiku.rag.agents.chat.search import SearchAgent
from haiku.rag.agents.chat.context import (
trigger_background_summarization as _trigger_summarization,
)
from haiku.rag.agents.chat.prompts import CHAT_SYSTEM_PROMPT
from haiku.rag.agents.chat.state import (
MAX_QA_HISTORY,
ChatDeps,
AGUI_STATE_KEY,
ChatSessionState,
DocumentInfo,
DocumentListResponse,
QAResponse,
SearchDeps,
SessionContext,
emit_state_event,
)
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
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.filters import (
build_document_filter,
build_multi_document_filter,
combine_filters,
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.document import DocumentListResponse, create_document_toolset
from haiku.rag.tools.qa import (
QA_SESSION_NAMESPACE,
QASessionState,
create_qa_toolset,
)
from haiku.rag.tools.search import create_search_toolset
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
from haiku.rag.utils import get_model
# Similarity threshold for finding relevant prior answers
PRIOR_ANSWER_RELEVANCE_THRESHOLD = 0.7
@dataclass
class ChatDeps:
"""Dependencies for chat agent.
def _cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
"""Compute cosine similarity between two vectors."""
dot_product = sum(a * b for a, b in zip(vec1, vec2))
norm1 = math.sqrt(sum(a * a for a in vec1))
norm2 = math.sqrt(sum(b * b for b in vec2))
if norm1 == 0 or norm2 == 0:
return 0.0
return dot_product / (norm1 * norm2)
Implements StateHandler protocol for AG-UI state management.
"""
config: AppConfig
tool_context: ToolContext
session_id: str = ""
state_key: str | None = None
# Track summarization tasks per session to allow cancellation
_summarization_tasks: dict[str, asyncio.Task[None]] = {}
@property
def state(self) -> dict[str, Any]:
"""Get current state for AG-UI protocol.
Combines SessionState and QASessionState into a single state dict
matching the ChatSessionState schema expected by AG-UI clients.
"""
snapshot: dict[str, Any] = {"session_id": self.session_id}
async def _update_context_background(
qa_history: list[QAResponse],
config: AppConfig,
session_state: ChatSessionState,
) -> None:
"""Background task to update session context after an ask."""
try:
await update_session_context(
qa_history=qa_history,
config=config,
session_state=session_state,
# Add SessionState fields
session_state = self.tool_context.get_typed(SESSION_NAMESPACE, SessionState)
if session_state is not None:
snapshot["document_filter"] = session_state.document_filter
snapshot["citation_registry"] = session_state.citation_registry
snapshot["citations"] = [c.model_dump() for c in session_state.citations]
# Add QASessionState fields
qa_session_state = self.tool_context.get_typed(
QA_SESSION_NAMESPACE, QASessionState
)
except asyncio.CancelledError:
pass
if qa_session_state is not None:
snapshot["qa_history"] = [
qa.model_dump() for qa in qa_session_state.qa_history
]
# Convert string to SessionContext model for frontend
if qa_session_state.session_context:
snapshot["session_context"] = SessionContext(
summary=qa_session_state.session_context
).model_dump(mode="json")
else:
snapshot["session_context"] = None
if self.state_key:
return {self.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
# Extract from namespaced key if present
state_data: dict[str, Any] = value
if self.state_key and self.state_key in value:
nested = value[self.state_key]
if isinstance(nested, dict):
state_data = nested
# Update SessionState from incoming state
session_state = self.tool_context.get_typed(SESSION_NAMESPACE, SessionState)
if session_state is not None:
if "document_filter" in state_data:
session_state.document_filter = state_data.get("document_filter", [])
if "citation_registry" in state_data:
session_state.citation_registry = state_data["citation_registry"]
if "citations" in state_data:
from haiku.rag.agents.research.models import Citation
session_state.citations = [
Citation(**c) if isinstance(c, dict) else c
for c in state_data.get("citations", [])
]
# Extract session_id if present, or generate one
# Track what the client sent (for delta computation)
incoming_session_id = state_data.get("session_id", "")
if incoming_session_id:
self.session_id = incoming_session_id
elif not self.session_id:
# Generate session_id now so ask() tool can use it
self.session_id = str(uuid.uuid4())
# Sync session_id to SessionState (track incoming for delta computation)
if session_state is not None:
session_state.session_id = self.session_id
session_state.incoming_session_id = incoming_session_id
# Update QASessionState from incoming state
qa_session_state = self.tool_context.get_typed(
QA_SESSION_NAMESPACE, QASessionState
)
if qa_session_state is not None:
if "qa_history" in state_data:
from haiku.rag.tools.qa import QAHistoryEntry
qa_session_state.qa_history = [
QAHistoryEntry(**qa) if isinstance(qa, dict) else qa
for qa in state_data.get("qa_history", [])
]
# Track what client sent for delta computation
incoming_session_context = state_data.get("session_context")
if isinstance(incoming_session_context, dict):
qa_session_state.incoming_session_context = SessionContext(
**incoming_session_context
)
qa_session_state.session_context = (
qa_session_state.incoming_session_context.summary
)
elif incoming_session_context is None:
qa_session_state.incoming_session_context = None
qa_session_state.session_context = None
# Check cache for fresher session_context from background summarization
# Cache is authoritative - always use it if available
if self.session_id:
cached = get_cached_session_context(self.session_id)
if cached and cached.summary:
qa_session_state.session_context = cached.render_markdown()
# Handle initial_context -> session_context for first message
# Only applies if session_context is still empty after restoring and cache check
if "initial_context" in state_data:
initial = state_data.get("initial_context")
if initial and not qa_session_state.session_context:
qa_session_state.session_context = initial
def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
"""Create the chat agent with search and ask tools."""
def create_chat_agent(
config: AppConfig,
client: HaikuRAG,
context: ToolContext,
) -> Agent[ChatDeps, str]:
"""Create the chat agent with composed toolsets.
Args:
config: Application configuration.
client: HaikuRAG client for database operations.
context: ToolContext for shared state across toolsets.
Should have SessionState and QASessionState registered
(will be auto-registered if not present).
Returns:
The configured chat agent.
Example:
async with HaikuRAG(db_path, create=True) as client:
context = ToolContext()
agent = create_chat_agent(config, client, context)
deps = ChatDeps(config=config, tool_context=context)
result = await agent.run("Search for X", deps=deps)
"""
# Ensure session states are registered with proper AG-UI state key
existing = context.get_typed(SESSION_NAMESPACE, SessionState)
if existing is None:
context.register(SESSION_NAMESPACE, SessionState(state_key=AGUI_STATE_KEY))
elif existing.state_key is None:
existing.state_key = AGUI_STATE_KEY
if context.get_typed(QA_SESSION_NAMESPACE, QASessionState) is None:
context.register(QA_SESSION_NAMESPACE, QASessionState())
# Create toolsets - these capture client, config, and context in closures
search_toolset = create_search_toolset(client, config, context=context)
document_toolset = create_document_toolset(client, config, context=context)
qa_toolset = create_qa_toolset(client, config, context=context)
# Create the agent with composed toolsets
model = get_model(config.qa.model, config)
agent: Agent[ChatDeps, str] = Agent(
model,
deps_type=ChatDeps,
output_type=str,
instructions=CHAT_SYSTEM_PROMPT,
retries=3,
agent = cast(
Agent[ChatDeps, str],
Agent(
model,
deps_type=ChatDeps,
output_type=str,
instructions=CHAT_SYSTEM_PROMPT,
toolsets=[search_toolset, document_toolset, qa_toolset], # type: ignore[arg-type]
retries=3,
),
)
@agent.tool
async def search(
ctx: RunContext[ChatDeps],
query: str,
document_name: str | None = None,
limit: int | None = None,
) -> ToolReturn:
"""Search the knowledge base for relevant documents.
Use this when you need to find documents or explore the knowledge base.
Results are displayed to the user - just list the titles found.
Args:
query: The search query (what to search for)
document_name: Optional document name/title to search within
limit: Number of results to return (default: 5)
"""
# Build session filter from document_filter
session_filter = build_multi_document_filter(
ctx.deps.session_state.document_filter
)
# Build tool filter from document_name parameter
tool_filter = build_document_filter(document_name) if document_name else None
# Combine filters: session AND tool
doc_filter = combine_filters(session_filter, tool_filter)
# Use search agent for query expansion and deduplication
search_agent = SearchAgent(ctx.deps.client, ctx.deps.config)
results = await search_agent.search(query, filter=doc_filter, limit=limit)
# Store for potential citation resolution
ctx.deps.search_results = results
if not results:
return ToolReturn(return_value="No results found.")
new_state = ctx.deps.session_state.model_copy(deep=True)
if not new_state.session_id:
new_state.session_id = str(uuid.uuid4())
# Build citation infos using the copy's registry
citation_infos = []
for r in results:
chunk_id = r.chunk_id or ""
if chunk_id:
index = new_state.get_or_assign_index(chunk_id)
else:
index = len(citation_infos) + 1
citation_infos.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,
)
)
# Update new_state with citations and fresh session_context
new_state.citations = citation_infos
if new_state.session_id:
new_state.session_context = get_cached_session_context(new_state.session_id)
# Return detailed results for the agent to present
result_lines = []
for c in citation_infos:
title = c.document_title or c.document_uri or "Unknown"
# Truncate content for display
snippet = c.content[:300].replace("\n", " ").strip()
if len(c.content) > 300:
snippet += "..."
line = f"[{c.index}] **{title}**"
if c.page_numbers:
line += f" (pages {', '.join(map(str, c.page_numbers))})"
line += f"\n {snippet}"
result_lines.append(line)
state_event = emit_state_event(
ctx.deps.session_state, new_state, ctx.deps.state_key
)
return ToolReturn(
return_value=f"Found {len(results)} results:\n\n"
+ "\n\n".join(result_lines),
metadata=[state_event] if state_event else None,
)
@agent.tool
async def ask(
ctx: RunContext[ChatDeps],
question: str,
document_name: str | None = None,
) -> ToolReturn:
"""Answer a specific question using the knowledge base.
Use this for direct questions that need a focused answer with citations.
Uses a research graph for planning, searching, and synthesis.
Args:
question: The question to answer
document_name: Optional document name/title to search within (e.g., "tbmed593", "army manual")
"""
# Build session filter from document_filter
session_filter = build_multi_document_filter(
ctx.deps.session_state.document_filter
)
# Build tool filter from document_name parameter
tool_filter = build_document_filter(document_name) if document_name else None
# Combine filters: session AND tool
doc_filter = combine_filters(session_filter, tool_filter)
# Build and run the conversational research graph
graph = build_research_graph(
config=ctx.deps.config, output_mode="conversational"
)
session_id = ctx.deps.session_state.session_id
# Get session context from server cache for planning, fallback to initial_context
cached_context = get_cached_session_context(session_id)
session_context = (
cached_context.render_markdown()
if cached_context and cached_context.summary
else ctx.deps.session_state.initial_context
)
# Find relevant prior answers from qa_history
prior_answers = []
if ctx.deps.session_state.qa_history:
embedder = get_embedder(ctx.deps.config)
question_embedding = await embedder.embed_query(question)
# Collect questions that need embedding (not cached)
to_embed = []
to_embed_indices = []
for i, qa in enumerate(ctx.deps.session_state.qa_history):
if qa.question_embedding is None:
to_embed.append(qa.question)
to_embed_indices.append(i)
# Batch embed uncached questions
if to_embed:
new_embeddings = await embedder.embed_documents(to_embed)
for i, idx in enumerate(to_embed_indices):
ctx.deps.session_state.qa_history[
idx
].question_embedding = new_embeddings[i]
# Compare against all questions and collect relevant prior answers
for qa in ctx.deps.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:
prior_answers.append(qa.to_search_answer())
context = ResearchContext(
original_question=question,
session_context=session_context,
qa_responses=prior_answers,
)
state = ResearchState(
context=context,
max_iterations=1,
search_filter=doc_filter,
max_concurrency=ctx.deps.config.research.max_concurrency,
)
deps = ResearchDeps(
client=ctx.deps.client,
)
result = await graph.run(state=state, deps=deps)
new_state = ctx.deps.session_state.model_copy(deep=True)
if not new_state.session_id:
new_state.session_id = str(uuid.uuid4())
# Build citation infos using the copy's registry
citation_infos = []
for c in result.citations:
index = new_state.get_or_assign_index(c.chunk_id)
citation_infos.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,
)
)
# Add Q&A to the copy's history
qa_response = QAResponse(
question=question,
answer=result.answer,
confidence=result.confidence,
citations=citation_infos,
)
new_state.qa_history.append(qa_response)
# Enforce FIFO limit
if len(new_state.qa_history) > MAX_QA_HISTORY:
new_state.qa_history = new_state.qa_history[-MAX_QA_HISTORY:]
# Update citations and session_context
new_state.citations = citation_infos
if new_state.session_id:
new_state.session_context = get_cached_session_context(new_state.session_id)
# Spawn background task to update session context
if new_state.session_id in _summarization_tasks:
_summarization_tasks[new_state.session_id].cancel()
task = asyncio.create_task(
_update_context_background(
qa_history=list(new_state.qa_history),
config=ctx.deps.config,
session_state=new_state,
)
)
_summarization_tasks[new_state.session_id] = task
task.add_done_callback(
lambda t, sid=new_state.session_id: _summarization_tasks.pop(sid, None)
)
# Format answer with citation references using stable indices
answer_text = result.answer
if citation_infos:
citation_refs = " ".join(f"[{c.index}]" for c in citation_infos)
answer_text = f"{answer_text}\n\nSources: {citation_refs}"
state_event = emit_state_event(
ctx.deps.session_state, new_state, ctx.deps.state_key
)
return ToolReturn(
return_value=answer_text,
metadata=[state_event] if state_event else None,
)
@agent.tool
async def list_documents(
ctx: RunContext[ChatDeps],
page: int = 1,
) -> DocumentListResponse:
"""List available documents in the knowledge base.
Use this when the user wants to browse or see what documents are available.
Args:
page: Page number (default: 1, 50 documents per page)
"""
page_size = 50
offset = (page - 1) * page_size
doc_filter = build_multi_document_filter(ctx.deps.session_state.document_filter)
docs = await ctx.deps.client.list_documents(
limit=page_size, offset=offset, filter=doc_filter
)
total = await ctx.deps.client.count_documents(filter=doc_filter)
total_pages = (total + page_size - 1) // page_size if total > 0 else 1
return DocumentListResponse(
documents=[
DocumentInfo(
title=doc.title or "Untitled",
uri=doc.uri or "",
created=doc.created_at.strftime("%Y-%m-%d"),
)
for doc in docs
],
page=page,
total_pages=total_pages,
total_documents=total,
)
async def _find_document(client: HaikuRAG, query: str):
"""Find a document by exact URI, partial URI, or partial title match."""
# Try exact URI match first
doc = await client.get_document_by_uri(query)
if doc is not None:
return doc
escaped_query = query.replace("'", "''")
# Also try without spaces for matching "TB MED 593" to "tbmed593"
no_spaces = escaped_query.replace(" ", "")
# Try partial URI match (with and without spaces)
docs = await client.list_documents(
limit=1,
filter=f"LOWER(uri) LIKE LOWER('%{escaped_query}%') OR LOWER(uri) LIKE LOWER('%{no_spaces}%')",
)
if docs:
return docs[0]
# Try partial title match (with and without spaces)
docs = await client.list_documents(
limit=1,
filter=f"LOWER(title) LIKE LOWER('%{escaped_query}%') OR LOWER(title) LIKE LOWER('%{no_spaces}%')",
)
if docs:
return docs[0]
return None
@agent.tool
async def get_document(
ctx: RunContext[ChatDeps],
query: str,
) -> str:
"""Retrieve a specific document by title or URI.
Use this when the user wants to fetch/get/retrieve a specific document.
Args:
query: The document title or URI to look up
"""
doc = await _find_document(ctx.deps.client, query)
if doc is None:
return f"Document not found: {query}"
return (
f"**{doc.title or 'Untitled'}**\n\n"
f"- ID: {doc.id}\n"
f"- URI: {doc.uri}\n"
f"- Created: {doc.created_at.strftime('%Y-%m-%d %H:%M')}\n\n"
f"**Content:**\n{doc.content}"
)
@agent.tool
async def summarize_document(
ctx: RunContext[ChatDeps],
query: str,
) -> str:
"""Generate a summary of a specific document.
Use this when the user wants an overview or summary of a document's content.
Args:
query: The document title or URI to summarize
"""
doc = await _find_document(ctx.deps.client, query)
if doc is None:
return f"Document not found: {query}"
# Use LLM to generate summary
summary_model = get_model(ctx.deps.config.qa.model, ctx.deps.config)
summary_agent: Agent[None, str] = Agent(
summary_model,
output_type=str,
)
result = await summary_agent.run(
DOCUMENT_SUMMARY_PROMPT.format(content=doc.content or "")
)
return f"**Summary of {doc.title or doc.uri}:**\n\n{result.output}"
return agent
def trigger_background_summarization(deps: ChatDeps) -> None:
"""Trigger background session summarization if qa_history has entries.
Call this after agent.run() or agent.run_stream() completes to update
the session context summary in the background.
Note: The ask() tool now triggers summarization internally, so this
function is primarily for explicit triggering when needed.
Args:
deps: Chat dependencies with tool_context containing QASessionState.
"""
qa_session_state = deps.tool_context.get_typed(QA_SESSION_NAMESPACE, QASessionState)
if qa_session_state is None or not qa_session_state.qa_history:
return
if not deps.session_id:
return
_trigger_summarization(
qa_session_state=qa_session_state,
config=deps.config,
session_id=deps.session_id,
)
async def run_chat_agent(
agent: Agent[ChatDeps, str],
deps: ChatDeps,
message: str,
) -> str:
"""Run the chat agent and trigger background summarization.
This wrapper handles post-processing like background summarization.
Args:
agent: The chat agent.
deps: Chat dependencies.
message: User message.
Returns:
Agent response.
"""
result = await agent.run(message, deps=deps)
trigger_background_summarization(deps)
return result.output
__all__ = [
"create_chat_agent",
"run_chat_agent",
"trigger_background_summarization",
"ChatDeps",
"ChatSessionState",
"DocumentInfo",
"DocumentListResponse",
"QAResponse",
"SearchDeps",
"SessionContext",
"emit_state_event",
"AGUI_STATE_KEY",
]

133
haiku_rag_slim/haiku/rag/agents/chat/context.py
View file

@ -1,4 +1,7 @@
import asyncio
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import TYPE_CHECKING
from pydantic_ai import Agent
@ -7,33 +10,72 @@ from haiku.rag.agents.chat.state import ChatSessionState, QAResponse, SessionCon
from haiku.rag.config.models import AppConfig
from haiku.rag.utils import get_model
# Cache for session contexts (session_id -> SessionContext)
# Used to persist async summarization results between requests
_session_context_cache: dict[str, SessionContext] = {}
if TYPE_CHECKING:
from haiku.rag.tools.qa import QASessionState
@dataclass
class SessionCache:
"""Per-session cache for context and embeddings."""
context: SessionContext | None = None
embeddings: dict[str, list[float]] = field(default_factory=dict)
# Cache for session data (session_id -> SessionCache)
# Used to persist async summarization results and embeddings between requests
_session_cache: dict[str, SessionCache] = {}
_cache_timestamps: dict[str, datetime] = {}
_CACHE_TTL = timedelta(hours=1)
# Track summarization tasks per session to allow cancellation
_summarization_tasks: dict[str, asyncio.Task[None]] = {}
def _cleanup_stale_cache() -> None:
"""Remove cache entries older than TTL."""
now = datetime.now()
stale = [sid for sid, ts in _cache_timestamps.items() if now - ts > _CACHE_TTL]
for sid in stale:
_session_context_cache.pop(sid, None)
_session_cache.pop(sid, None)
_cache_timestamps.pop(sid, None)
def _get_or_create_session_cache(session_id: str) -> SessionCache:
"""Get or create session cache for a given session_id."""
_cleanup_stale_cache()
if session_id not in _session_cache:
_session_cache[session_id] = SessionCache()
_cache_timestamps[session_id] = datetime.now()
return _session_cache[session_id]
def cache_session_context(session_id: str, context: SessionContext) -> None:
"""Store session context in cache."""
_cleanup_stale_cache()
_session_context_cache[session_id] = context
_cache_timestamps[session_id] = datetime.now()
cache = _get_or_create_session_cache(session_id)
cache.context = context
def get_cached_session_context(session_id: str) -> SessionContext | None:
"""Get session context from server cache."""
_cleanup_stale_cache()
return _session_context_cache.get(session_id)
cache = _session_cache.get(session_id)
return cache.context if cache else None
def cache_question_embedding(
session_id: str, question: str, embedding: list[float]
) -> None:
"""Store question embedding in session cache."""
cache = _get_or_create_session_cache(session_id)
cache.embeddings[question] = embedding
def get_cached_embedding(session_id: str, question: str) -> list[float] | None:
"""Get cached embedding for a question in this session."""
_cleanup_stale_cache()
cache = _session_cache.get(session_id)
return cache.embeddings.get(question) if cache else None
async def summarize_session(
@ -114,3 +156,78 @@ def _format_qa_history(qa_history: list[QAResponse]) -> str:
lines.append("")
return "\n".join(lines)
async def _update_context_background(
qa_session_state: "QASessionState",
config: AppConfig,
session_id: str,
) -> None:
"""Background task to update session context after an ask."""
try:
# Convert QAHistoryEntry to QAResponse format for update_session_context
qa_history = [
QAResponse(
question=entry.question,
answer=entry.answer,
confidence=entry.confidence,
citations=list(entry.citations),
)
for entry in qa_session_state.qa_history
]
session_state = ChatSessionState(
session_id=session_id,
qa_history=qa_history,
)
await update_session_context(
qa_history=qa_history,
config=config,
session_state=session_state,
)
# Update the QASessionState with the new context
cached = get_cached_session_context(session_id)
if cached and cached.summary:
qa_session_state.session_context = cached.render_markdown()
except asyncio.CancelledError:
pass
except Exception as e:
import logging
logging.getLogger(__name__).exception(f"Background summarization failed: {e}")
def trigger_background_summarization(
qa_session_state: "QASessionState",
config: AppConfig,
session_id: str,
) -> None:
"""Trigger background session summarization if qa_history has entries.
Args:
qa_session_state: QASessionState with qa_history to summarize.
config: AppConfig for model selection.
session_id: Session ID for caching results.
"""
if not qa_session_state.qa_history or not session_id:
return
# Cancel any existing summarization task for this session
if session_id in _summarization_tasks:
_summarization_tasks[session_id].cancel()
# Spawn background task
task = asyncio.create_task(
_update_context_background(
qa_session_state=qa_session_state,
config=config,
session_id=session_id,
)
)
_summarization_tasks[session_id] = task
task.add_done_callback(
lambda _t, sid=session_id: _summarization_tasks.pop(sid, None)
)

116
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -17,14 +17,19 @@ from pydantic_ai import (
)
from pydantic_ai.messages import ModelMessage
from haiku.rag.agents.chat.agent import create_chat_agent
from haiku.rag.agents.chat.agent import (
ChatDeps,
create_chat_agent,
trigger_background_summarization,
)
from haiku.rag.agents.chat.state import (
AGUI_STATE_KEY,
ChatDeps,
ChatSessionState,
)
from haiku.rag.client import HaikuRAG
from haiku.rag.config import get_config
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
if TYPE_CHECKING:
from textual.app import ComposeResult
@ -153,8 +158,16 @@ class ChatApp(App):
)
await self.client.__aenter__()
# Create agent and session state
self.agent = create_chat_agent(self.config)
# Create tool context and agent
self.tool_context = ToolContext()
self.agent = create_chat_agent(self.config, self.client, self.tool_context)
# Initialize session state in tool context
session_state = self.tool_context.get_typed(SESSION_NAMESPACE, SessionState)
if session_state is not None:
session_state.document_filter = self._document_filter
# Keep ChatSessionState for UI state sync (used by _sync_session_state)
self.session_state = ChatSessionState(
session_id=str(uuid.uuid4()),
initial_context=self._initial_context,
@ -170,8 +183,40 @@ class ChatApp(App):
await self.client.__aexit__(None, None, None)
def _sync_session_state(self, chat_state: dict[str, Any]) -> None:
"""Sync session_state from AG-UI state."""
self.session_state = ChatSessionState.model_validate(chat_state)
"""Sync session_state from AG-UI state.
Updates existing session_state with fields from incoming state,
preserving fields not present in the update (e.g., initial_context).
"""
from haiku.rag.agents.research.models import Citation
# Update specific fields rather than replacing the entire state
if "session_id" in chat_state:
self.session_state.session_id = chat_state["session_id"]
if "document_filter" in chat_state:
self.session_state.document_filter = chat_state["document_filter"]
if "citation_registry" in chat_state:
self.session_state.citation_registry = chat_state["citation_registry"]
if "citations" in chat_state:
self.session_state.citations = [
Citation(**c) if isinstance(c, dict) else c
for c in chat_state["citations"]
]
if "qa_history" in chat_state:
from haiku.rag.agents.chat.state import QAResponse
self.session_state.qa_history = [
QAResponse(**qa) if isinstance(qa, dict) else qa
for qa in chat_state["qa_history"]
]
if "session_context" in chat_state:
from haiku.rag.agents.chat.state import SessionContext
ctx = chat_state["session_context"]
if ctx is not None:
self.session_state.session_context = (
SessionContext(**ctx) if isinstance(ctx, dict) else ctx
)
async def _handle_stream_event(self, event: AgentStreamEvent) -> None:
"""Handle streaming events from the agent."""
@ -274,10 +319,33 @@ class ChatApp(App):
AGUI_STATE_KEY: self.session_state.model_dump(mode="json")
}
# Sync session state to tool context before running
session_state = self.tool_context.get_typed(SESSION_NAMESPACE, SessionState)
if session_state is not None:
session_state.session_id = self.session_state.session_id
session_state.document_filter = self.session_state.document_filter
session_state.citation_registry = self.session_state.citation_registry
session_state.citations = list(self.session_state.citations)
# Sync initial_context to QA session state
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
qa_session_state = self.tool_context.get_typed(
QA_SESSION_NAMESPACE, QASessionState
)
if qa_session_state is not None:
if (
not qa_session_state.session_context
and self.session_state.initial_context
):
qa_session_state.session_context = (
self.session_state.initial_context
)
deps = ChatDeps(
client=self.client,
config=self.config,
session_state=self.session_state,
tool_context=self.tool_context,
session_id=self.session_state.session_id,
state_key=AGUI_STATE_KEY,
)
@ -307,6 +375,23 @@ class ChatApp(App):
if self.session_state.citations:
await chat_history.add_citations(self.session_state.citations)
# Trigger background summarization
trigger_background_summarization(deps)
# Sync session context from QASessionState to ChatSessionState for modal
qa_session_state = self.tool_context.get_typed(
QA_SESSION_NAMESPACE, QASessionState
)
if qa_session_state is not None and qa_session_state.session_context:
from datetime import datetime
from haiku.rag.agents.chat.state import SessionContext
self.session_state.session_context = SessionContext(
summary=qa_session_state.session_context,
last_updated=datetime.now(),
)
except asyncio.CancelledError:
chat_history.hide_thinking()
await chat_history.add_message("assistant", "*Cancelled*")
@ -383,6 +468,21 @@ class ChatApp(App):
async def action_show_context(self) -> None:
"""Show context modal (edit initial context or view session context)."""
from haiku.rag.chat.widgets.context_modal import ContextModal
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
# Sync session context from QASessionState before showing modal
qa_session_state = self.tool_context.get_typed(
QA_SESSION_NAMESPACE, QASessionState
)
if qa_session_state is not None and qa_session_state.session_context:
from datetime import datetime
from haiku.rag.agents.chat.state import SessionContext
self.session_state.session_context = SessionContext(
summary=qa_session_state.session_context,
last_updated=datetime.now(),
)
await self.push_screen(
ContextModal(self.session_state, is_locked=self._context_locked)

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

@ -18,8 +18,20 @@ from haiku.rag.tools.filters import (
combine_filters,
)
from haiku.rag.tools.models import AnalysisResult, QAResult
from haiku.rag.tools.qa import QA_NAMESPACE, QAState, create_qa_toolset
from haiku.rag.tools.qa import (
QA_NAMESPACE,
QA_SESSION_NAMESPACE,
QAHistoryEntry,
QASessionState,
QAState,
create_qa_toolset,
)
from haiku.rag.tools.search import SEARCH_NAMESPACE, SearchState, create_search_toolset
from haiku.rag.tools.session import (
SESSION_NAMESPACE,
SessionState,
compute_state_delta,
)
__all__ = [
"ToolContext",
@ -38,9 +50,15 @@ __all__ = [
"create_document_toolset",
"find_document",
"QA_NAMESPACE",
"QA_SESSION_NAMESPACE",
"QAState",
"QASessionState",
"QAHistoryEntry",
"create_qa_toolset",
"ANALYSIS_NAMESPACE",
"AnalysisState",
"create_analysis_toolset",
"SESSION_NAMESPACE",
"SessionState",
"compute_state_delta",
]

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

@ -7,8 +7,13 @@ from haiku.rag.agents.rlm.models import CodeExecution
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.filters import build_document_filter, combine_filters
from haiku.rag.tools.filters import (
build_document_filter,
build_multi_document_filter,
combine_filters,
)
from haiku.rag.tools.models import AnalysisResult
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
ANALYSIS_NAMESPACE = "haiku.rag.analysis"
@ -36,6 +41,8 @@ def create_analysis_toolset(
config: Application configuration.
context: Optional ToolContext for state accumulation.
If provided, code executions are tracked in AnalysisState.
If SessionState is registered, it will be used for dynamic
document filtering.
base_filter: Optional base SQL WHERE clause applied to searches.
tool_name: Name for the analyze tool. Defaults to "analyze".
@ -63,9 +70,20 @@ def create_analysis_toolset(
Returns:
AnalysisResult with answer and execution metadata.
"""
# Build filter from base_filter and document_name
# Get session filter from session state
session_filter = None
if context is not None:
session_state = context.get_typed(SESSION_NAMESPACE, SessionState)
if session_state is not None and session_state.document_filter:
session_filter = build_multi_document_filter(
session_state.document_filter
)
# Build filter from base_filter, session_filter, and document_name
doc_filter = build_document_filter(document_name) if document_name else None
effective_filter = combine_filters(base_filter, doc_filter)
effective_filter = combine_filters(
combine_filters(base_filter, session_filter), doc_filter
)
# Create RLM context and deps
rlm_context = RLMContext(filter=effective_filter)

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

@ -62,6 +62,16 @@ class ToolContext(BaseModel):
"""Get state for a namespace, or None if not registered."""
return self._namespaces.get(namespace)
def get_typed(self, namespace: str, expected_type: type[T]) -> T | None:
"""Get state for a namespace with type checking.
Returns the state cast to expected_type if it matches, None otherwise.
"""
state = self._namespaces.get(namespace)
if isinstance(state, expected_type):
return state
return None
def get_or_create(self, namespace: str, factory: Callable[[], T]) -> T:
"""Get state for a namespace, creating it if not registered.

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

@ -4,6 +4,8 @@ from pydantic_ai import Agent, FunctionToolset
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.filters import build_multi_document_filter, combine_filters
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
from haiku.rag.utils import get_model
DOCUMENT_NAMESPACE = "haiku.rag.document"
@ -91,6 +93,8 @@ def create_document_toolset(
config: Application configuration (used for summarization LLM).
context: Optional ToolContext for state tracking.
If provided, accessed documents are tracked in DocumentState.
If SessionState is registered, it will be used for dynamic
document filtering.
base_filter: Optional base SQL WHERE clause applied to list operations.
Returns:
@ -113,10 +117,21 @@ def create_document_toolset(
page_size = 50
offset = (page - 1) * page_size
# Get session filter from session state
session_filter = None
if context is not None:
session_state = context.get_typed(SESSION_NAMESPACE, SessionState)
if session_state is not None and session_state.document_filter:
session_filter = build_multi_document_filter(
session_state.document_filter
)
effective_filter = combine_filters(base_filter, session_filter)
docs = await client.list_documents(
limit=page_size, offset=offset, filter=base_filter
limit=page_size, offset=offset, filter=effective_filter
)
total = await client.count_documents(filter=base_filter)
total = await client.count_documents(filter=effective_filter)
total_pages = (total + page_size - 1) // page_size if total > 0 else 1
return DocumentListResponse(

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

@ -1,18 +1,64 @@
from pydantic import BaseModel
from pydantic_ai import FunctionToolset
import math
from pydantic import BaseModel, Field
from pydantic_ai import FunctionToolset, ToolReturn
from haiku.rag.agents.chat.context import (
cache_question_embedding,
get_cached_embedding,
trigger_background_summarization,
)
from haiku.rag.agents.chat.state import SessionContext
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 ToolContext
from haiku.rag.tools.filters import build_document_filter, combine_filters
from haiku.rag.tools.models import QAResult
from haiku.rag.tools.session import (
SESSION_NAMESPACE,
SessionState,
compute_combined_state_delta,
)
QA_NAMESPACE = "haiku.rag.qa"
PRIOR_ANSWER_RELEVANCE_THRESHOLD = 0.7
def _cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
"""Compute cosine similarity between two vectors."""
dot_product = sum(a * b for a, b in zip(vec1, vec2))
norm1 = math.sqrt(sum(a * a for a in vec1))
norm2 = math.sqrt(sum(b * b for b in vec2))
if norm1 == 0 or norm2 == 0:
return 0.0
return dot_product / (norm1 * norm2)
class QAHistoryEntry(BaseModel):
"""A Q&A pair with optional cached embedding for similarity matching."""
question: str
answer: str
confidence: float = 0.9
citations: list[Citation] = []
question_embedding: list[float] | None = Field(default=None, exclude=True)
def to_search_answer(self) -> SearchAnswer:
"""Convert to SearchAnswer for research graph context."""
return SearchAnswer(
query=self.question,
answer=self.answer,
confidence=self.confidence,
cited_chunks=[c.chunk_id for c in self.citations],
citations=self.citations,
)
class QAState(BaseModel):
"""State for QA toolset.
@ -23,6 +69,20 @@ class QAState(BaseModel):
history: list[QAResult] = []
class QASessionState(BaseModel):
"""Extended session state for QA with embedding cache."""
qa_history: list[QAHistoryEntry] = []
session_context: str | None = None
incoming_session_context: SessionContext | None = Field(
default=None, exclude=True
) # Track what client sent
QA_SESSION_NAMESPACE = "haiku.rag.qa_session"
MAX_QA_HISTORY = 50
def create_qa_toolset(
client: HaikuRAG,
config: AppConfig,
@ -39,10 +99,14 @@ def create_qa_toolset(
config: Application configuration.
context: Optional ToolContext for state accumulation.
If provided, Q&A results are accumulated in QAState.
If SessionState is registered, it will be used for dynamic
document filtering and citation indexing.
base_filter: Optional base SQL WHERE clause applied to searches.
tool_name: Name for the ask tool. Defaults to "ask".
session_context: Optional session context for the research graph.
Overridden by QASessionState.session_context if available.
prior_answers: Optional list of prior answers for context.
Overridden by similarity-matched answers from QASessionState if available.
Returns:
FunctionToolset with an ask tool.
@ -55,7 +119,7 @@ def create_qa_toolset(
async def ask(
question: str,
document_name: str | None = None,
) -> QAResult:
) -> ToolReturn | QAResult:
"""Answer a question using the knowledge base.
Uses a research graph for searching and synthesizing answers.
@ -67,17 +131,110 @@ def create_qa_toolset(
Returns:
QAResult with answer, confidence, and citations.
"""
# Build filter from base_filter and document_name
# Get session states
session_state: SessionState | None = None
qa_session_state: QASessionState | None = None
old_state_snapshot: dict | None = None
if context is not None:
session_state = context.get_typed(SESSION_NAMESPACE, SessionState)
qa_session_state = context.get_typed(QA_SESSION_NAMESPACE, QASessionState)
# Capture combined state snapshot before changes
# Use incoming values (what client sent) so delta shows server-side updates
if session_state is not None:
old_state_snapshot = {
"session_id": session_state.incoming_session_id,
"document_filter": session_state.document_filter.copy(),
"citation_registry": session_state.citation_registry.copy(),
"citations": [c.model_dump() for c in session_state.citations],
}
if qa_session_state is not None:
old_state_snapshot["qa_history"] = [
qa.model_dump() for qa in qa_session_state.qa_history
]
# Use incoming_session_context so delta shows what client sent
if qa_session_state.incoming_session_context is not None:
old_state_snapshot["session_context"] = (
qa_session_state.incoming_session_context.model_dump(
mode="json"
)
)
else:
old_state_snapshot["session_context"] = None
# Build filter from session state, base_filter, and document_name
session_filter = None
if session_state is not None and session_state.document_filter:
from haiku.rag.tools.filters import build_multi_document_filter
session_filter = build_multi_document_filter(session_state.document_filter)
doc_filter = build_document_filter(document_name) if document_name else None
effective_filter = combine_filters(base_filter, doc_filter)
effective_filter = combine_filters(
combine_filters(base_filter, session_filter), doc_filter
)
# Determine session context
effective_session_context = session_context
if qa_session_state is not None and qa_session_state.session_context:
effective_session_context = qa_session_state.session_context
# Find relevant prior answers via similarity matching
effective_prior_answers = prior_answers or []
session_id = session_state.session_id if session_state is not None else ""
if qa_session_state is not None and qa_session_state.qa_history:
embedder = get_embedder(config)
question_embedding = await embedder.embed_query(question)
# Collect questions that need embedding
to_embed = []
to_embed_indices = []
for i, qa in enumerate(qa_session_state.qa_history):
if qa.question_embedding is None:
# Check per-session cache first
if session_id:
cached = get_cached_embedding(session_id, qa.question)
if cached:
qa.question_embedding = cached
continue
to_embed.append(qa.question)
to_embed_indices.append(i)
# Batch embed uncached questions
if to_embed:
new_embeddings = await embedder.embed_documents(to_embed)
for i, idx in enumerate(to_embed_indices):
embedding = new_embeddings[i]
qa_session_state.qa_history[idx].question_embedding = embedding
# Cache per-session for next request
if session_id:
cache_question_embedding(
session_id,
qa_session_state.qa_history[idx].question,
embedding,
)
# Find similar prior answers
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
# Build and run the research graph
graph = build_research_graph(config=config, output_mode="conversational")
research_context = ResearchContext(
original_question=question,
session_context=session_context,
qa_responses=prior_answers or [],
session_context=effective_session_context,
qa_responses=effective_prior_answers,
)
research_state = ResearchState(
context=research_context,
@ -89,20 +246,26 @@ def create_qa_toolset(
result = await graph.run(state=research_state, deps=deps)
# Convert to QAResult
citations = [
Citation(
index=i + 1,
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,
# 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,
)
)
for i, c in enumerate(result.citations)
]
qa_result = QAResult(
question=question,
@ -111,10 +274,71 @@ def create_qa_toolset(
citations=citations,
)
# Accumulate in state if context provided
# Accumulate in QA state if context provided
if state is not None:
state.history.append(qa_result)
# Update session state with citations
if session_state is not None:
session_state.citations = citations
# Update QA session state with history entry
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:
]
# Trigger background summarization
trigger_background_summarization(
qa_session_state=qa_session_state,
config=config,
session_id=session_id,
)
# Compute and return state delta if session state changed
if session_state is not None and old_state_snapshot is not None:
# Build new combined state snapshot
new_state_snapshot = {
"session_id": session_state.session_id,
"document_filter": session_state.document_filter,
"citation_registry": session_state.citation_registry,
"citations": [c.model_dump() for c in session_state.citations],
}
if qa_session_state is not None:
new_state_snapshot["qa_history"] = [
qa.model_dump() for qa in qa_session_state.qa_history
]
if qa_session_state.session_context:
new_state_snapshot["session_context"] = SessionContext(
summary=qa_session_state.session_context
).model_dump(mode="json")
else:
new_state_snapshot["session_context"] = None
state_event = compute_combined_state_delta(
old_state_snapshot,
new_state_snapshot,
state_key=session_state.state_key,
)
# Format answer with citation references
answer_text = result.answer
if citations:
citation_refs = " ".join(f"[{c.index}]" for c in citations)
answer_text = f"{answer_text}\n\nSources: {citation_refs}"
metadata = [state_event] if state_event is not None else None
return ToolReturn(return_value=answer_text, metadata=metadata)
return qa_result
toolset = FunctionToolset()

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

@ -1,11 +1,13 @@
from pydantic import BaseModel
from pydantic_ai import FunctionToolset
from pydantic_ai import FunctionToolset, ToolReturn
from haiku.rag.agents.research.models import Citation
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models import SearchResult
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.filters import combine_filters
from haiku.rag.tools.filters import build_multi_document_filter, combine_filters
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState, compute_state_delta
SEARCH_NAMESPACE = "haiku.rag.search"
@ -34,6 +36,8 @@ def create_search_toolset(
config: Application configuration.
context: Optional ToolContext for state accumulation.
If provided, search results are accumulated in SearchState.
If SessionState is registered, it will be used for dynamic
document filtering and citation indexing.
expand_context: Whether to expand search results with surrounding context.
Defaults to True.
base_filter: Optional base SQL WHERE clause applied to all searches.
@ -44,15 +48,15 @@ def create_search_toolset(
FunctionToolset with a search tool.
"""
# Get or create search state if context provided
state: SearchState | None = None
search_state: SearchState | None = None
if context is not None:
state = context.get_or_create(SEARCH_NAMESPACE, SearchState)
search_state = context.get_or_create(SEARCH_NAMESPACE, SearchState)
async def search(
query: str,
limit: int | None = None,
filter: str | None = None,
) -> str:
) -> ToolReturn | str:
"""Search the knowledge base for relevant documents.
Args:
@ -63,8 +67,25 @@ def create_search_toolset(
Returns:
Formatted search results with content and metadata.
"""
# Get session state for dynamic filters and citation indexing
session_state: SessionState | None = None
old_session_state: SessionState | None = None
if context is not None:
session_state = context.get_typed(SESSION_NAMESPACE, SessionState)
if session_state is not None:
old_session_state = session_state.model_copy(deep=True)
# Build session filter from session state's document_filter
session_filter = None
if session_state is not None and session_state.document_filter:
session_filter = build_multi_document_filter(session_state.document_filter)
# Combine all filters: base_filter AND session_filter AND tool filter
effective_filter = combine_filters(
combine_filters(base_filter, session_filter), filter
)
effective_limit = limit or config.search.limit
effective_filter = combine_filters(base_filter, filter)
results = await client.search(
query, limit=effective_limit, filter=effective_filter
)
@ -72,14 +93,64 @@ def create_search_toolset(
if expand_context:
results = await client.expand_context(results)
# Accumulate results in state if context provided
if state is not None:
state.results.extend(results)
# Accumulate results in search state if context provided
if search_state is not None:
search_state.results.extend(results)
if not results:
return "No results found."
# Format results for agent context
# Build citations if session state is available
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:
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
# Format results with citation indices
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:
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)
# Compute state delta if session state changed
if old_session_state is not None:
state_event = compute_state_delta(old_session_state, session_state)
if state_event is not None:
return ToolReturn(
return_value=formatted,
metadata=[state_event],
)
return formatted
# Format results without citation indexing (standalone use)
total = len(results)
formatted = [
r.format_for_agent(rank=i + 1, total=total) for i, r in enumerate(results)

89
haiku_rag_slim/haiku/rag/tools/session.py Normal file
View file

@ -0,0 +1,89 @@
from typing import Any
import jsonpatch
from ag_ui.core import EventType, StateDeltaEvent
from pydantic import BaseModel, Field
from haiku.rag.agents.research.models import Citation
SESSION_NAMESPACE = "haiku.rag.session"
class SessionState(BaseModel):
"""Session-level state for AG-UI integration.
This state is shared across toolsets and enables:
- Session identification
- Dynamic document filtering
- Stable citation indices across tool calls
- AG-UI state synchronization
"""
session_id: str = ""
incoming_session_id: str = Field(default="", exclude=True) # Track what client sent
document_filter: list[str] = []
citation_registry: dict[str, int] = {}
citations: list[Citation] = []
state_key: str | None = Field(default=None, exclude=True)
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,
) -> StateDeltaEvent | None:
"""Compute state delta between old and new session state.
Returns a StateDeltaEvent if there are changes, None otherwise.
The state_key from new_state is used for namespacing.
"""
return compute_combined_state_delta(
old_state.model_dump(mode="json"),
new_state.model_dump(mode="json"),
state_key=new_state.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,
)

356
tests/agents/chat/test_chat_agent.py
View file

@ -9,6 +9,7 @@ from haiku.rag.agents.chat import (
ChatSessionState,
QAResponse,
SearchAgent,
ToolContext,
create_chat_agent,
)
from haiku.rag.agents.chat.context import get_cached_session_context
@ -16,6 +17,7 @@ from haiku.rag.agents.chat.state import MAX_QA_HISTORY
from haiku.rag.agents.research.models import Citation
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
def extract_state_from_result(result, state_key: str = AGUI_STATE_KEY) -> dict | None:
@ -49,26 +51,25 @@ def vcr_cassette_dir():
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_chat_agent")
def test_create_chat_agent():
def test_create_chat_agent(temp_db_path):
"""Test that create_chat_agent returns a properly configured agent."""
agent = create_chat_agent(Config)
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
agent = create_chat_agent(Config, client, context)
assert agent is not None
assert agent.name == "chat_agent" or agent.name is None
client.close()
def test_chat_deps_initialization(temp_db_path):
"""Test ChatDeps can be initialized with required fields."""
client = HaikuRAG(temp_db_path, create=True)
deps = ChatDeps(client=client, config=Config)
context = ToolContext()
deps = ChatDeps(config=Config, tool_context=context)
assert deps.client is client
assert deps.config is Config
assert deps.search_results is None
assert deps.session_state is not None
assert deps.session_state.qa_history == []
assert deps.session_state.citations == []
client.close()
assert deps.tool_context is context
assert deps.session_id == ""
assert deps.state_key is None
def test_agui_state_key_constant():
@ -76,26 +77,126 @@ def test_agui_state_key_constant():
assert AGUI_STATE_KEY == "haiku.rag.chat"
def test_chat_deps_with_state_key(temp_db_path):
def test_chat_deps_with_state_key():
"""Test ChatDeps can be initialized with state_key for keyed state emission."""
client = HaikuRAG(temp_db_path, create=True)
deps = ChatDeps(client=client, config=Config, state_key="my_state")
context = ToolContext()
deps = ChatDeps(config=Config, tool_context=context, state_key="my_state")
assert deps.client is client
assert deps.config is Config
assert deps.state_key == "my_state"
client.close()
def test_chat_deps_state_key_default_none(temp_db_path):
def test_chat_deps_state_key_default_none():
"""Test ChatDeps state_key defaults to None."""
client = HaikuRAG(temp_db_path, create=True)
deps = ChatDeps(client=client, config=Config)
context = ToolContext()
deps = ChatDeps(config=Config, tool_context=context)
assert deps.state_key is None
client.close()
def test_chat_deps_state_setter_handles_initial_context():
"""Test ChatDeps.state setter transfers initial_context to qa_session_state."""
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
context = ToolContext()
# Register QASessionState (normally done by create_chat_agent)
context.register(QA_SESSION_NAMESPACE, QASessionState())
context.register(SESSION_NAMESPACE, SessionState())
deps = ChatDeps(config=Config, tool_context=context, state_key=AGUI_STATE_KEY)
# Client sends initial_context with no session_context
incoming_state = {
AGUI_STATE_KEY: {
"session_id": "",
"initial_context": "Background info about the project",
"session_context": None,
"qa_history": [],
"citations": [],
"document_filter": [],
"citation_registry": {},
}
}
deps.state = incoming_state
# initial_context should be copied to qa_session_state.session_context
qa_session_state = context.get(QA_SESSION_NAMESPACE)
assert isinstance(qa_session_state, QASessionState)
assert qa_session_state.session_context == "Background info about the project"
def test_chat_deps_state_setter_parses_session_context_dict():
"""Test ChatDeps.state setter parses session_context dict into SessionContext model."""
from haiku.rag.agents.chat.state import SessionContext
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
context = ToolContext()
context.register(QA_SESSION_NAMESPACE, QASessionState())
context.register(SESSION_NAMESPACE, SessionState())
deps = ChatDeps(config=Config, tool_context=context, state_key=AGUI_STATE_KEY)
# Client sends session_context as a dict (as it comes from JSON)
incoming_state = {
AGUI_STATE_KEY: {
"session_id": "test-session",
"session_context": {
"summary": "Previous conversation summary",
"last_updated": "2025-01-27T12:00:00",
},
"qa_history": [],
"citations": [],
"document_filter": [],
"citation_registry": {},
}
}
deps.state = incoming_state
# session_context dict should be parsed into SessionContext model
qa_session_state = context.get(QA_SESSION_NAMESPACE)
assert isinstance(qa_session_state, QASessionState)
assert qa_session_state.session_context == "Previous conversation summary"
assert isinstance(qa_session_state.incoming_session_context, SessionContext)
assert (
qa_session_state.incoming_session_context.summary
== "Previous conversation summary"
)
def test_chat_deps_state_setter_generates_session_id():
"""Test ChatDeps.state setter generates session_id if client sends empty."""
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
context = ToolContext()
context.register(QA_SESSION_NAMESPACE, QASessionState())
context.register(SESSION_NAMESPACE, SessionState())
deps = ChatDeps(config=Config, tool_context=context, state_key=AGUI_STATE_KEY)
# Client sends empty session_id
incoming_state = {
AGUI_STATE_KEY: {
"session_id": "",
"session_context": None,
"qa_history": [],
"citations": [],
"document_filter": [],
"citation_registry": {},
}
}
deps.state = incoming_state
# session_id should be generated (UUID format)
assert deps.session_id != ""
assert len(deps.session_id) == 36 # UUID length with dashes
# Should also be synced to SessionState
session_state = context.get(SESSION_NAMESPACE)
assert isinstance(session_state, SessionState)
assert session_state.session_id == deps.session_id
def test_chat_session_state():
@ -263,12 +364,12 @@ async def test_chat_agent_search_tool(allow_model_requests, temp_db_path):
title="DocLayNet Annotation",
)
agent = create_chat_agent(Config)
session_state = ChatSessionState(session_id="test-search")
context = ToolContext()
agent = create_chat_agent(Config, client, context)
deps = ChatDeps(
client=client,
config=Config,
session_state=session_state,
tool_context=context,
session_id="test-search",
)
# Ask something that should trigger the search tool
@ -298,12 +399,12 @@ async def test_chat_agent_search_tool_with_filter(allow_model_requests, temp_db_
title="DocLayNet Sources",
)
agent = create_chat_agent(Config)
session_state = ChatSessionState(session_id="test-search-filter")
context = ToolContext()
agent = create_chat_agent(Config, client, context)
deps = ChatDeps(
client=client,
config=Config,
session_state=session_state,
tool_context=context,
session_id="test-search-filter",
)
# Ask to search within a specific document
@ -327,10 +428,11 @@ async def test_chat_agent_get_document_tool(allow_model_requests, temp_db_path):
title="DocLayNet Class Labels",
)
agent = create_chat_agent(Config)
context = ToolContext()
agent = create_chat_agent(Config, client, context)
deps = ChatDeps(
client=client,
config=Config,
tool_context=context,
)
# Ask to get a specific document
@ -349,10 +451,11 @@ async def test_chat_agent_get_document_tool(allow_model_requests, temp_db_path):
async def test_chat_agent_get_document_not_found(allow_model_requests, temp_db_path):
"""Test the chat agent's get_document tool when document is not found."""
async with HaikuRAG(temp_db_path, create=True) as client:
agent = create_chat_agent(Config)
context = ToolContext()
agent = create_chat_agent(Config, client, context)
deps = ChatDeps(
client=client,
config=Config,
tool_context=context,
)
# Ask for a document that doesn't exist
@ -469,7 +572,7 @@ async def test_search_agent_no_results(allow_model_requests, temp_db_path):
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path):
"""Test that the ask tool adds citations to the response."""
"""Test that the ask tool is called and can add citations to session state."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add a document with specific content
await client.create_document(
@ -478,14 +581,15 @@ async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path)
title="DocLayNet Class Labels",
)
agent = create_chat_agent(Config)
context = ToolContext()
agent = create_chat_agent(Config, client, context)
deps = ChatDeps(
client=client,
config=Config,
tool_context=context,
state_key=AGUI_STATE_KEY,
)
# Ask a question that should use the ask tool with citations
# Ask a question that should use the ask tool
result = await agent.run(
"What is the highest count class in the DocLayNet dataset?",
deps=deps,
@ -493,11 +597,20 @@ async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path)
assert result.output is not None
# Extract emitted state from result metadata
emitted_state = extract_state_from_result(result)
assert emitted_state is not None
# The qa_history should have been updated with the new Q&A
assert len(emitted_state.get("qa_history", [])) >= 1
# Verify the agent used the ask tool by checking for tool calls
tool_calls = [
part
for msg in result.all_messages()
if hasattr(msg, "parts")
for part in msg.parts
if hasattr(part, "tool_name") and part.tool_name == "ask"
]
assert len(tool_calls) >= 1, "Expected ask tool to be called"
# Session state should be registered (citations may or may not be present
# depending on whether the research graph found relevant evidence)
session_state = context.get(SESSION_NAMESPACE)
assert isinstance(session_state, SessionState)
@pytest.mark.asyncio
@ -508,6 +621,8 @@ async def test_chat_agent_ask_triggers_background_summarization(
"""Test that the ask tool triggers background session context summarization."""
import asyncio
from haiku.rag.agents.chat.agent import run_chat_agent
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
@ -515,30 +630,27 @@ async def test_chat_agent_ask_triggers_background_summarization(
title="DocLayNet Class Labels",
)
agent = create_chat_agent(Config)
context = ToolContext()
agent = create_chat_agent(Config, client, context)
deps = ChatDeps(
client=client,
config=Config,
tool_context=context,
session_id="test-summarization",
state_key=AGUI_STATE_KEY,
)
# Ask a question
result = await agent.run(
# Ask a question using run_chat_agent to trigger background summarization
result = await run_chat_agent(
agent,
deps,
"What is the highest count class in the DocLayNet dataset?",
deps=deps,
)
assert result.output is not None
# Extract emitted state to get the session_id
emitted_state = extract_state_from_result(result)
assert emitted_state is not None
session_id = emitted_state.get("session_id")
assert session_id is not None
assert len(emitted_state.get("qa_history", [])) >= 1
assert result is not None
# Wait for background task to complete
# The task caches session_context server-side
session_id = deps.session_id
cached_context = None
for _ in range(50): # Wait up to 5 seconds
cached_context = get_cached_session_context(session_id)
@ -559,12 +671,10 @@ async def test_chat_agent_ask_with_prior_answer_retrieval(
):
"""Test that ask tool retrieves relevant prior answers from qa_history.
This exercises the prior answer retrieval logic (agent.py lines 231-257):
This exercises the prior answer retrieval logic:
1. First ask populates qa_history with question_embedding
2. Second similar ask should find the prior answer via embedding similarity
"""
import asyncio
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
@ -572,10 +682,11 @@ async def test_chat_agent_ask_with_prior_answer_retrieval(
title="DocLayNet Class Labels",
)
agent = create_chat_agent(Config)
context = ToolContext()
agent = create_chat_agent(Config, client, context)
deps1 = ChatDeps(
client=client,
config=Config,
tool_context=context,
state_key=AGUI_STATE_KEY,
)
@ -586,47 +697,26 @@ async def test_chat_agent_ask_with_prior_answer_retrieval(
)
assert result1.output is not None
# Extract emitted state from first call
state1 = extract_state_from_result(result1)
assert state1 is not None
assert len(state1.get("qa_history", [])) == 1
session_id = state1.get("session_id")
assert session_id is not None
# Wait for background summarization to complete
for _ in range(50):
if get_cached_session_context(session_id) is not None:
break
await asyncio.sleep(0.1)
# Create new session state from emitted state for second call
# (simulating client sending state back to server)
session_state2 = ChatSessionState(
session_id=session_id,
qa_history=[QAResponse(**qa) for qa in state1.get("qa_history", [])],
citation_registry=state1.get("citation_registry", {}),
)
deps2 = ChatDeps(
client=client,
config=Config,
session_state=session_state2,
state_key=AGUI_STATE_KEY,
)
# Check that session state has citations after first call
session_state = context.get(SESSION_NAMESPACE)
assert isinstance(session_state, SessionState)
# Citations might be 0 if the answer came from prior context
assert len(session_state.citations) >= 0
# Second ask - similar question triggers prior answer retrieval
# This will embed the first question and compare similarity
result2 = await agent.run(
"Tell me about DocLayNet class labels",
deps=deps2,
deps=deps1,
)
assert result2.output is not None
# The important thing is that prior answer retrieval happened
# We can verify this by checking session_state2 was used (embedding added)
assert session_state2.qa_history[0].question_embedding is not None
# Embedding should be a list of floats
assert isinstance(session_state2.qa_history[0].question_embedding, list)
assert len(session_state2.qa_history[0].question_embedding) > 0
# The QA session state should have history entries
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
qa_session = context.get(QA_SESSION_NAMESPACE)
assert isinstance(qa_session, QASessionState)
# After two asks, we should have entries in qa_history
assert len(qa_session.qa_history) >= 1
def test_fifo_limit_enforcement():
@ -697,16 +787,18 @@ async def test_chat_agent_search_with_session_filter(
title="DocLayNet Sources",
)
agent = create_chat_agent(Config)
context = ToolContext()
agent = create_chat_agent(Config, client, context)
# Set session filter to only include the labels document
session_state = ChatSessionState(
session_id="test-session-filter",
document_filter=["DocLayNet Class Labels"],
)
session_state = context.get(SESSION_NAMESPACE)
assert isinstance(session_state, SessionState)
session_state.document_filter = ["DocLayNet Class Labels"]
deps = ChatDeps(
client=client,
config=Config,
session_state=session_state,
tool_context=context,
session_id="test-session-filter",
)
# Search should only return results from the filtered document
@ -716,8 +808,15 @@ async def test_chat_agent_search_with_session_filter(
)
assert result.output is not None
# Results should only reference the Labels document, not Sources
assert "Labels" in result.output or "class" in result.output.lower()
# Check that citations in context are only from the filtered document
session_state = context.get(SESSION_NAMESPACE)
assert isinstance(session_state, SessionState)
# If citations were added, they should only be from the labels document
for citation in session_state.citations:
assert "labels" in citation.document_uri.lower() or "Labels" in (
citation.document_title or ""
)
@pytest.mark.asyncio
@ -861,7 +960,7 @@ def test_search_tool_citation_registry_logic():
def test_cosine_similarity_identical_vectors():
"""Test cosine similarity returns 1.0 for identical vectors."""
from haiku.rag.agents.chat.agent import _cosine_similarity
from haiku.rag.tools.qa import _cosine_similarity
vec = [1.0, 2.0, 3.0]
assert _cosine_similarity(vec, vec) == pytest.approx(1.0)
@ -869,7 +968,7 @@ def test_cosine_similarity_identical_vectors():
def test_cosine_similarity_orthogonal_vectors():
"""Test cosine similarity returns 0.0 for orthogonal vectors."""
from haiku.rag.agents.chat.agent import _cosine_similarity
from haiku.rag.tools.qa import _cosine_similarity
vec1 = [1.0, 0.0, 0.0]
vec2 = [0.0, 1.0, 0.0]
@ -878,7 +977,7 @@ def test_cosine_similarity_orthogonal_vectors():
def test_cosine_similarity_opposite_vectors():
"""Test cosine similarity returns -1.0 for opposite vectors."""
from haiku.rag.agents.chat.agent import _cosine_similarity
from haiku.rag.tools.qa import _cosine_similarity
vec1 = [1.0, 2.0, 3.0]
vec2 = [-1.0, -2.0, -3.0]
@ -887,7 +986,7 @@ def test_cosine_similarity_opposite_vectors():
def test_cosine_similarity_zero_vector():
"""Test cosine similarity handles zero vectors gracefully."""
from haiku.rag.agents.chat.agent import _cosine_similarity
from haiku.rag.tools.qa import _cosine_similarity
vec = [1.0, 2.0, 3.0]
zero = [0.0, 0.0, 0.0]
@ -898,14 +997,14 @@ def test_cosine_similarity_zero_vector():
def test_prior_answer_relevance_threshold_constant():
"""Test PRIOR_ANSWER_RELEVANCE_THRESHOLD is set to expected value."""
from haiku.rag.agents.chat.agent import PRIOR_ANSWER_RELEVANCE_THRESHOLD
from haiku.rag.tools.qa import PRIOR_ANSWER_RELEVANCE_THRESHOLD
assert PRIOR_ANSWER_RELEVANCE_THRESHOLD == 0.7
def test_prior_answer_matching_above_threshold():
"""Test that similar questions (above threshold) are matched."""
from haiku.rag.agents.chat.agent import (
from haiku.rag.tools.qa import (
PRIOR_ANSWER_RELEVANCE_THRESHOLD,
_cosine_similarity,
)
@ -920,7 +1019,7 @@ def test_prior_answer_matching_above_threshold():
def test_prior_answer_matching_below_threshold():
"""Test that dissimilar questions (below threshold) are not matched."""
from haiku.rag.agents.chat.agent import (
from haiku.rag.tools.qa import (
PRIOR_ANSWER_RELEVANCE_THRESHOLD,
_cosine_similarity,
)
@ -970,7 +1069,7 @@ async def test_summarization_task_cancellation():
"""Test that new summarization tasks cancel previous ones for same session."""
import asyncio
from haiku.rag.agents.chat.agent import _summarization_tasks
from haiku.rag.agents.chat.context import _summarization_tasks
# Clear any existing tasks
_summarization_tasks.clear()
@ -1033,10 +1132,11 @@ async def test_list_documents_basic(allow_model_requests, temp_db_path):
title="DocLayNet Annotation",
)
agent = create_chat_agent(Config)
context = ToolContext()
agent = create_chat_agent(Config, client, context)
deps = ChatDeps(
client=client,
config=Config,
tool_context=context,
)
# Ask to list documents
@ -1067,16 +1167,17 @@ async def test_list_documents_with_session_filter(allow_model_requests, temp_db_
title="DocLayNet Sources",
)
agent = create_chat_agent(Config)
# Set session filter to only include the labels document
session_state = ChatSessionState(
session_id="test-list-filter",
document_filter=["DocLayNet Class Labels"],
context = ToolContext()
context.register(
SESSION_NAMESPACE,
SessionState(document_filter=["DocLayNet Class Labels"]),
)
agent = create_chat_agent(Config, client, context)
deps = ChatDeps(
client=client,
config=Config,
session_state=session_state,
tool_context=context,
session_id="test-list-filter",
)
# Ask to list documents - should only show filtered documents
@ -1112,10 +1213,11 @@ async def test_list_documents_pagination(allow_model_requests, temp_db_path):
title="DocLayNet Sources",
)
agent = create_chat_agent(Config)
context = ToolContext()
agent = create_chat_agent(Config, client, context)
deps = ChatDeps(
client=client,
config=Config,
tool_context=context,
)
# Ask to list first 2 documents
@ -1144,10 +1246,11 @@ async def test_summarize_document_found(allow_model_requests, temp_db_path):
title="DocLayNet Class Labels",
)
agent = create_chat_agent(Config)
context = ToolContext()
agent = create_chat_agent(Config, client, context)
deps = ChatDeps(
client=client,
config=Config,
tool_context=context,
)
# Ask to summarize a specific document
@ -1166,10 +1269,11 @@ async def test_summarize_document_found(allow_model_requests, temp_db_path):
async def test_summarize_document_not_found(allow_model_requests, temp_db_path):
"""Test that summarize_document handles not found documents gracefully."""
async with HaikuRAG(temp_db_path, create=True) as client:
agent = create_chat_agent(Config)
context = ToolContext()
agent = create_chat_agent(Config, client, context)
deps = ChatDeps(
client=client,
config=Config,
tool_context=context,
)
# Ask to summarize a document that doesn't exist
@ -1246,7 +1350,7 @@ async def test_summarization_task_cleanup_on_completion():
"""Test that completed tasks are cleaned up from _summarization_tasks."""
import asyncio
from haiku.rag.agents.chat.agent import _summarization_tasks
from haiku.rag.agents.chat.context import _summarization_tasks
_summarization_tasks.clear()

32
tests/agents/chat/test_context.py
View file

@ -261,13 +261,13 @@ class TestSessionContextCache:
def test_cache_and_retrieve_session_context(self):
"""Test caching and retrieving a session context."""
from haiku.rag.agents.chat.context import (
_session_context_cache,
_session_cache,
cache_session_context,
get_cached_session_context,
)
# Clear cache
_session_context_cache.clear()
_session_cache.clear()
now = datetime.now()
ctx = SessionContext(summary="Test summary", last_updated=now)
@ -282,11 +282,11 @@ class TestSessionContextCache:
def test_get_cached_session_context_returns_none_when_not_cached(self):
"""Test get_cached_session_context returns None when nothing cached."""
from haiku.rag.agents.chat.context import (
_session_context_cache,
_session_cache,
get_cached_session_context,
)
_session_context_cache.clear()
_session_cache.clear()
result = get_cached_session_context("nonexistent-session")
assert result is None
@ -298,12 +298,12 @@ class TestSessionContextCache:
from haiku.rag.agents.chat.context import (
_CACHE_TTL,
_cache_timestamps,
_session_context_cache,
_session_cache,
cache_session_context,
get_cached_session_context,
)
_session_context_cache.clear()
_session_cache.clear()
_cache_timestamps.clear()
# Add an entry
@ -320,21 +320,21 @@ class TestSessionContextCache:
# Should be None because the entry was cleaned up
assert result is None
assert "stale-session" not in _session_context_cache
assert "stale-session" not in _session_cache
@pytest.mark.asyncio
async def test_update_session_context_caches_result(self):
async def test_update_session_caches_result(self):
"""Test update_session_context stores result in cache."""
from unittest.mock import AsyncMock, patch
from haiku.rag.agents.chat.context import (
_session_context_cache,
_session_cache,
get_cached_session_context,
update_session_context,
)
from haiku.rag.agents.chat.state import ChatSessionState
_session_context_cache.clear()
_session_cache.clear()
session_state = ChatSessionState(session_id="cache-test-session")
@ -368,13 +368,13 @@ class TestSessionContextCache:
async def test_update_session_context_no_cache_without_session_id(self):
"""Test update_session_context doesn't cache without session_id."""
from haiku.rag.agents.chat.context import (
_session_context_cache,
_session_cache,
get_cached_session_context,
update_session_context,
)
from haiku.rag.agents.chat.state import ChatSessionState
_session_context_cache.clear()
_session_cache.clear()
# No session_id
session_state = ChatSessionState()
@ -395,12 +395,12 @@ class TestSessionContextCache:
from unittest.mock import patch
from haiku.rag.agents.chat.context import (
_session_context_cache,
_session_cache,
update_session_context,
)
from haiku.rag.agents.chat.state import ChatSessionState
_session_context_cache.clear()
_session_cache.clear()
# Create session_state with initial_context but no session_context
session_state = ChatSessionState(
@ -446,12 +446,12 @@ class TestSessionContextCache:
from unittest.mock import patch
from haiku.rag.agents.chat.context import (
_session_context_cache,
_session_cache,
update_session_context,
)
from haiku.rag.agents.chat.state import ChatSessionState, SessionContext
_session_context_cache.clear()
_session_cache.clear()
# Create session_state with BOTH initial_context and session_context
session_state = ChatSessionState(

766
tests/cassettes/test_chat_agent/test_chat_agent_ask_adds_citations.yaml
View file

File diff suppressed because one or more lines are too long

1665
tests/cassettes/test_chat_agent/test_chat_agent_ask_emits_qa_state.yaml Normal file
View file

File diff suppressed because one or more lines are too long

12
tests/tools/test_analysis.py
View file

@ -72,18 +72,12 @@ class TestAnalysisToolset:
@pytest.fixture
def analysis_client(temp_db_path):
async def analysis_client(temp_db_path):
"""Create a HaikuRAG client for analysis tests."""
import asyncio
from haiku.rag.client import HaikuRAG
async def setup():
rag = HaikuRAG(temp_db_path, create=True)
await rag.__aenter__()
return rag
return asyncio.get_event_loop().run_until_complete(setup())
async with HaikuRAG(temp_db_path, create=True) as rag:
yield rag
@pytest.fixture

12
tests/tools/test_document.py
View file

@ -162,15 +162,11 @@ class TestDocumentToolExecution:
@pytest.fixture
def doc_client(temp_db_path):
async def doc_client(temp_db_path):
"""Create a HaikuRAG client with test documents."""
import asyncio
from haiku.rag.client import HaikuRAG
async def setup():
rag = HaikuRAG(temp_db_path, create=True)
await rag.__aenter__()
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",
@ -181,9 +177,7 @@ def doc_client(temp_db_path):
uri="test://javascript",
title="JavaScript Guide",
)
return rag
return asyncio.get_event_loop().run_until_complete(setup())
yield rag
@pytest.fixture

12
tests/tools/test_qa.py
View file

@ -73,18 +73,12 @@ class TestQAToolset:
@pytest.fixture
def qa_client_simple(temp_db_path):
async def qa_client_simple(temp_db_path):
"""Create a HaikuRAG client without documents for basic tests."""
import asyncio
from haiku.rag.client import HaikuRAG
async def setup():
rag = HaikuRAG(temp_db_path, create=True)
await rag.__aenter__()
return rag
return asyncio.get_event_loop().run_until_complete(setup())
async with HaikuRAG(temp_db_path, create=True) as rag:
yield rag
@pytest.fixture

12
tests/tools/test_search.py
View file

@ -213,15 +213,11 @@ class TestSearchToolExecution:
@pytest.fixture
def search_client(temp_db_path):
async def search_client(temp_db_path):
"""Create a HaikuRAG client with test data for search tests."""
import asyncio
from haiku.rag.client import HaikuRAG
async def setup():
rag = HaikuRAG(temp_db_path, create=True)
await rag.__aenter__()
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",
@ -232,9 +228,7 @@ def search_client(temp_db_path):
uri="test://javascript",
title="JavaScript Guide",
)
return rag
return asyncio.get_event_loop().run_until_complete(setup())
yield rag
@pytest.fixture