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

Merge pull request #263 from ggozad/feat/state-deltas

Browse source 
Send delta snapshots instead of full snapshots in AGUI
This commit is contained in:
Yiorgis Gozadinos 2026-01-28 17:44:55 +02:00 committed by GitHub
commit 2afb10f6a4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
24 changed files with 6140 additions and 6297 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
CHANGELOG.md
View file

@ -17,10 +17,11 @@
### Changed
- **Selective Citation Filtering**: Synthesis steps now select only relevant citations instead of including all
- LLM receives `<available_citations>` with chunk IDs and content previews
- LLM populates `cited_chunks` with only chunks that directly support the answer
- `ResearchReport` now has `cited_chunks` and `citations` fields; removed `sources_summary`
- **AG-UI State Delta Updates**: Web application now sends `StateDeltaEvent` (JSON Patch RFC 6902) instead of full `StateSnapshotEvent` for state updates
- Reduces bandwidth when state grows large (e.g., 50 Q&As with citations)
- First request still sends full snapshot; subsequent requests send only changes
- Backend logging shows incoming/outgoing state events for debugging
## [0.27.1] - 2026-01-27

75
app/backend/main.py
View file

@ -79,38 +79,71 @@ async def stream_chat(request: Request) -> Response:
accept = request.headers.get("accept", SSE_CONTENT_TYPE)
run_input = AGUIAdapter.build_run_input(body)
# Restore session state from incoming AG-UI state (look under namespaced key)
initial_qa_history: list[QAResponse] = []
session_id: str | None = None
document_filter: list[str] = []
initial_context: str | None = None
# 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:
chat_state = state.get(AGUI_STATE_KEY, state)
if "qa_history" in chat_state:
initial_qa_history = [
QAResponse(**qa) for qa in chat_state.get("qa_history", [])
]
session_id = chat_state.get("session_id")
document_filter = chat_state.get("document_filter", [])
initial_context = chat_state.get("initial_context")
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")
deps = ChatDeps(
client=get_client(db_path),
config=Config,
session_state=ChatSessionState(
qa_history=initial_qa_history,
document_filter=document_filter,
initial_context=initial_context,
**({"session_id": session_id} if session_id else {}),
),
session_state=session_state,
state_key=AGUI_STATE_KEY,
)
# Use AGUIAdapter for streaming
adapter = AGUIAdapter(agent=chat_agent, run_input=run_input, accept=accept)
event_stream = adapter.run_stream(deps=deps)
sse_event_stream = adapter.encode_stream(event_stream)
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())
return StreamingResponse(
sse_event_stream,

6
app/frontend/components/Chat.tsx
View file

@ -47,6 +47,7 @@ interface ChatSessionState {
qa_history: QAResponse[];
session_context: SessionContext | null;
document_filter: string[];
citation_registry: Record<string, number>;
}
// AG-UI state is namespaced under AGUI_STATE_KEY
@ -404,6 +405,7 @@ function ChatContentInner() {
qa_history: [],
session_context: null,
document_filter: [],
citation_registry: {},
},
},
},
@ -429,6 +431,8 @@ function ChatContentInner() {
qa_history: agentState?.[AGUI_STATE_KEY]?.qa_history ?? [],
session_context: agentState?.[AGUI_STATE_KEY]?.session_context ?? null,
document_filter: selected,
citation_registry:
agentState?.[AGUI_STATE_KEY]?.citation_registry ?? {},
},
});
};
@ -445,6 +449,8 @@ function ChatContentInner() {
qa_history: agentState?.[AGUI_STATE_KEY]?.qa_history ?? [],
session_context: agentState?.[AGUI_STATE_KEY]?.session_context ?? null,
document_filter: agentState?.[AGUI_STATE_KEY]?.document_filter ?? [],
citation_registry:
agentState?.[AGUI_STATE_KEY]?.citation_registry ?? {},
},
});
};

5
docs/configuration/prompts.md
View file

@ -71,7 +71,7 @@ prompts:
Replace the research report synthesis prompt by setting `prompts.synthesis`. This controls how the multi-agent research workflow generates its final report.
The prompt should produce a `ResearchReport` with: `title`, `executive_summary`, `main_findings`, `conclusions`, `recommendations`, `limitations`, and `cited_chunks`.
The prompt should produce a `ResearchReport` with: `title`, `executive_summary`, `main_findings`, `conclusions`, `recommendations`, `limitations`, and `sources_summary`.
**Example:**
@ -87,13 +87,12 @@ prompts:
- conclusions: 2-4 bullet points
- recommendations: 2-5 actionable recommendations
- limitations: 1-3 limitations or gaps
- cited_chunks: List of chunk IDs that directly support the report
- sources_summary: Brief description of sources used
Guidelines:
- Base all content strictly on collected evidence
- Be specific and objective
- Avoid meta-commentary like "This report covers..."
- Only include chunks in cited_chunks that directly support claims in the report
```
## Picture Description Prompt

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

@ -1,7 +1,7 @@
import asyncio
import math
import uuid
from ag_ui.core import EventType, StateSnapshotEvent
from pydantic_ai import Agent, RunContext, ToolReturn
from haiku.rag.agents.chat.context import (
@ -20,6 +20,7 @@ from haiku.rag.agents.chat.state import (
build_document_filter,
build_multi_document_filter,
combine_filters,
emit_state_event,
)
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_conversational_graph
@ -94,11 +95,9 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
limit: Number of results to return (default: 5)
"""
# Build session filter from document_filter
session_filter = None
if ctx.deps.session_state and ctx.deps.session_state.document_filter:
session_filter = build_multi_document_filter(
ctx.deps.session_state.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
@ -116,12 +115,16 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
if not results:
return ToolReturn(return_value="No results found.")
# Build citation infos using stable registry indices
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 ctx.deps.session_state is not None and chunk_id:
index = ctx.deps.session_state.get_or_assign_index(chunk_id)
if chunk_id:
index = new_state.get_or_assign_index(chunk_id)
else:
index = len(citation_infos) + 1
citation_infos.append(
@ -137,26 +140,10 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
)
)
# Build new state with citations and registry
session_id = ctx.deps.session_state.session_id if ctx.deps.session_state else ""
new_state = ChatSessionState(
session_id=session_id,
citations=citation_infos,
qa_history=(
ctx.deps.session_state.qa_history if ctx.deps.session_state else []
),
session_context=get_cached_session_context(session_id)
if session_id
else None,
document_filter=(
ctx.deps.session_state.document_filter if ctx.deps.session_state else []
),
citation_registry=(
ctx.deps.session_state.citation_registry
if ctx.deps.session_state
else {}
),
)
# 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 = []
@ -173,19 +160,14 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
line += f"\n {snippet}"
result_lines.append(line)
snapshot = new_state.model_dump(mode="json")
if ctx.deps.state_key:
snapshot = {ctx.deps.state_key: snapshot}
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=[
StateSnapshotEvent(
type=EventType.STATE_SNAPSHOT,
snapshot=snapshot,
)
],
metadata=[state_event] if state_event else None,
)
@agent.tool
@ -204,11 +186,9 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
document_name: Optional document name/title to search within (e.g., "tbmed593", "army manual")
"""
# Build session filter from document_filter
session_filter = None
if ctx.deps.session_state and ctx.deps.session_state.document_filter:
session_filter = build_multi_document_filter(
ctx.deps.session_state.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
@ -218,23 +198,19 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
# Build and run the conversational research graph
graph = build_conversational_graph(config=ctx.deps.config)
session_id = ctx.deps.session_state.session_id if ctx.deps.session_state else ""
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) if session_id else None
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
if ctx.deps.session_state
else None
)
else ctx.deps.session_state.initial_context
)
# Find relevant prior answers from qa_history
prior_answers = []
if ctx.deps.session_state and ctx.deps.session_state.qa_history:
if ctx.deps.session_state.qa_history:
embedder = get_embedder(ctx.deps.config)
question_embedding = await embedder.embed_query(question)
@ -281,14 +257,14 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
result = await graph.run(state=state, deps=deps)
# Build citation infos using stable registry indices
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:
# Use registry for stable indices across calls
if ctx.deps.session_state is not None:
index = ctx.deps.session_state.get_or_assign_index(c.chunk_id)
else:
index = len(citation_infos) + 1
index = new_state.get_or_assign_index(c.chunk_id)
citation_infos.append(
Citation(
index=index,
@ -302,54 +278,37 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
)
)
# Accumulate Q&A in session state with full citation metadata
if ctx.deps.session_state is not None:
qa_response = QAResponse(
question=question,
answer=result.answer,
confidence=result.confidence,
citations=citation_infos,
)
ctx.deps.session_state.qa_history.append(qa_response)
# Enforce FIFO limit
if len(ctx.deps.session_state.qa_history) > MAX_QA_HISTORY:
ctx.deps.session_state.qa_history = ctx.deps.session_state.qa_history[
-MAX_QA_HISTORY:
]
# Spawn background task to update session context
# Cancel any previous summarization for this session
if session_id in _summarization_tasks:
_summarization_tasks[session_id].cancel()
task = asyncio.create_task(
_update_context_background(
qa_history=list(ctx.deps.session_state.qa_history),
config=ctx.deps.config,
session_state=ctx.deps.session_state,
)
)
_summarization_tasks[session_id] = task
task.add_done_callback(lambda t: _summarization_tasks.pop(session_id, None))
# Build new state with citations, qa_history, and registry
new_state = ChatSessionState(
session_id=session_id,
# Add Q&A to the copy's history
qa_response = QAResponse(
question=question,
answer=result.answer,
confidence=result.confidence,
citations=citation_infos,
qa_history=(
ctx.deps.session_state.qa_history if ctx.deps.session_state else []
),
session_context=get_cached_session_context(session_id)
if session_id
else None,
document_filter=(
ctx.deps.session_state.document_filter if ctx.deps.session_state else []
),
citation_registry=(
ctx.deps.session_state.citation_registry
if ctx.deps.session_state
else {}
),
)
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
@ -358,18 +317,13 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
citation_refs = " ".join(f"[{c.index}]" for c in citation_infos)
answer_text = f"{answer_text}\n\nSources: {citation_refs}"
snapshot = new_state.model_dump(mode="json")
if ctx.deps.state_key:
snapshot = {ctx.deps.state_key: snapshot}
state_event = emit_state_event(
ctx.deps.session_state, new_state, ctx.deps.state_key
)
return ToolReturn(
return_value=answer_text,
metadata=[
StateSnapshotEvent(
type=EventType.STATE_SNAPSHOT,
snapshot=snapshot,
)
],
metadata=[state_event] if state_event else None,
)
@agent.tool
@ -387,12 +341,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
page_size = 50
offset = (page - 1) * page_size
# Build session filter from document_filter
doc_filter = None
if ctx.deps.session_state and ctx.deps.session_state.document_filter:
doc_filter = build_multi_document_filter(
ctx.deps.session_state.document_filter
)
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

74
haiku_rag_slim/haiku/rag/agents/chat/state.py
View file

@ -3,6 +3,8 @@ from dataclasses import dataclass, field
from datetime import datetime
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, SearchAnswer
@ -106,14 +108,14 @@ class ChatDeps:
client: HaikuRAG
config: AppConfig
search_results: list[SearchResult] | None = None
session_state: ChatSessionState | None = None
session_state: ChatSessionState = field(
default_factory=lambda: ChatSessionState(session_id="")
)
state_key: str | None = None
@property
def state(self) -> dict[str, Any] | None:
def state(self) -> dict[str, Any]:
"""Get current state for AG-UI protocol."""
if self.session_state is None:
return None
snapshot = self.session_state.model_dump()
if self.state_key:
return {self.state_key: snapshot}
@ -131,29 +133,24 @@ class ChatDeps:
if isinstance(nested, dict):
state_data = nested
# Update session_state from incoming state
if self.session_state is not None:
if "qa_history" in state_data:
self.session_state.qa_history = [
QAResponse(**qa) if isinstance(qa, dict) else qa
for qa in state_data.get("qa_history", [])
]
if "citations" in state_data:
self.session_state.citations = [
Citation(**c) if isinstance(c, dict) else c
for c in state_data.get("citations", [])
]
if state_data.get("session_id"):
self.session_state.session_id = state_data["session_id"]
if "document_filter" in state_data:
self.session_state.document_filter = state_data.get(
"document_filter", []
)
if "citation_registry" in state_data:
self.session_state.citation_registry = state_data["citation_registry"]
if "initial_context" in state_data:
self.session_state.initial_context = state_data.get("initial_context")
# NOTE: session_context is server-managed; we don't accept it from the client
# to maintain server-side ownership of conversation summarization
if "qa_history" in state_data:
self.session_state.qa_history = [
QAResponse(**qa) if isinstance(qa, dict) else qa
for qa in state_data.get("qa_history", [])
]
if "citations" in state_data:
self.session_state.citations = [
Citation(**c) if isinstance(c, dict) else c
for c in state_data.get("citations", [])
]
if state_data.get("session_id"):
self.session_state.session_id = state_data["session_id"]
if "document_filter" in state_data:
self.session_state.document_filter = state_data.get("document_filter", [])
if "citation_registry" in state_data:
self.session_state.citation_registry = state_data["citation_registry"]
if "initial_context" in state_data:
self.session_state.initial_context = state_data.get("initial_context")
@dataclass
@ -194,3 +191,26 @@ def combine_filters(filter1: str | None, filter2: str | None) -> str | None:
if len(filters) == 1:
return filters[0]
return f"({filters[0]}) AND ({filters[1]})"
def emit_state_event(
current_state: ChatSessionState,
new_state: ChatSessionState,
state_key: str | None = None,
) -> StateDeltaEvent | None:
"""Emit state delta against current state, or None if no changes."""
new_snapshot = new_state.model_dump(mode="json")
wrapped_new = {state_key: new_snapshot} if state_key else new_snapshot
current_snapshot = current_state.model_dump(mode="json")
wrapped_current = {state_key: current_snapshot} if state_key else current_snapshot
patch = jsonpatch.make_patch(wrapped_current, wrapped_new)
if not patch.patch:
return None
return StateDeltaEvent(
type=EventType.STATE_DELTA,
delta=patch.patch,
)

6
haiku_rag_slim/haiku/rag/agents/qa/prompts.py
View file

@ -26,11 +26,7 @@ Each result includes:
- Type: content type like paragraph, table, code, list_item (when available)
- Content: the actual text
Citation guidelines:
- In cited_chunks, include ONLY chunk IDs that directly support your answer.
- Do NOT cite chunks that are merely related or that you reviewed but did not use.
- Quality over quantity: fewer relevant citations are better than many marginal ones.
- Use the EXACT, COMPLETE chunk IDs (full UUIDs).
In your response, include the chunk IDs you used in cited_chunks.
Guidelines:
- Base answers strictly on retrieved content - do not use external knowledge

81
haiku_rag_slim/haiku/rag/agents/research/graph.py
View file

@ -32,7 +32,6 @@ from haiku.rag.utils import build_prompt, get_model
def format_context_for_prompt(
context: ResearchContext,
include_pending_questions: bool = True,
include_citations: bool = False,
) -> str:
"""Format the research context as XML for prompts.
@ -40,8 +39,6 @@ def format_context_for_prompt(
context: The research context to format.
include_pending_questions: Whether to include pending sub-questions.
Set to False for synthesis prompts where pending questions aren't relevant.
include_citations: Whether to include available citations for selection.
Set to True for synthesis prompts where the LLM should select relevant citations.
"""
context_data: dict[str, object] = {}
@ -64,26 +61,6 @@ def format_context_for_prompt(
for qa in context.qa_responses
]
if include_citations and context.qa_responses:
seen_chunks: set[str] = set()
available_citations: list[dict[str, str]] = []
for qa in context.qa_responses:
for c in qa.citations:
if c.chunk_id not in seen_chunks:
seen_chunks.add(c.chunk_id)
content_preview = (
c.content[:500] + "..." if len(c.content) > 500 else c.content
)
available_citations.append(
{
"chunk_id": c.chunk_id,
"document": c.document_title or c.document_uri,
"content": content_preview,
}
)
if available_citations:
context_data["available_citations"] = available_citations
return format_as_xml(context_data, root_tag="context")
@ -110,7 +87,7 @@ async def _plan_step_logic(
else plan_prompt
)
plan_agent: Agent[ResearchDependencies, ResearchPlan] = Agent( # type: ignore[assignment]
plan_agent: Agent[ResearchDependencies, ResearchPlan] = Agent( # type: ignore[invalid-assignment]
model=get_model(model_config, config),
output_type=ResearchPlan,
instructions=effective_plan_prompt,
@ -177,7 +154,7 @@ async def _search_one_step_logic(
deps.semaphore = asyncio.Semaphore(state.max_concurrency)
async with deps.semaphore:
agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent( # type: ignore[assignment]
agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent( # type: ignore[invalid-assignment]
model=get_model(model_config, config),
output_type=ToolOutput(RawSearchAnswer, max_retries=3),
instructions=search_prompt,
@ -303,7 +280,7 @@ def build_research_graph(
state = ctx.state
deps = ctx.deps
agent: Agent[ResearchDependencies, EvaluationResult] = Agent( # type: ignore[assignment]
agent: Agent[ResearchDependencies, EvaluationResult] = Agent( # type: ignore[invalid-assignment]
model=get_model(model_config, config),
output_type=EvaluationResult,
instructions=decision_prompt,
@ -364,7 +341,7 @@ def build_research_graph(
state = ctx.state
deps = ctx.deps
agent: Agent[ResearchDependencies, ResearchReport] = Agent( # type: ignore[assignment]
agent: Agent[ResearchDependencies, ResearchReport] = Agent( # type: ignore[invalid-assignment]
model=get_model(model_config, config),
output_type=ResearchReport,
instructions=synthesis_prompt,
@ -373,10 +350,7 @@ def build_research_graph(
deps_type=ResearchDependencies,
)
# Include available citations for the LLM to select from
context_xml = format_context_for_prompt(
state.context, include_pending_questions=False, include_citations=True
)
context_xml = format_context_for_prompt(state.context)
prompt = (
"Generate a comprehensive research report based on all gathered information.\n\n"
f"{context_xml}\n\n"
@ -387,21 +361,7 @@ def build_research_graph(
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
report = result.output
citation_lookup: dict[str, Citation] = {}
for qa in state.context.qa_responses:
for c in qa.citations:
if c.chunk_id not in citation_lookup:
citation_lookup[c.chunk_id] = c
resolved_citations: list[Citation] = []
for chunk_id in report.cited_chunks:
if chunk_id in citation_lookup:
resolved_citations.append(citation_lookup[chunk_id])
report.citations = resolved_citations
return report
return result.output
# Build the graph structure
collect_answers = g.join(
@ -519,19 +479,17 @@ def build_conversational_graph(
state = ctx.state
deps = ctx.deps
# Use RawSearchAnswer so LLM can select which chunks to cite
agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent( # type: ignore[assignment]
agent: Agent[ResearchDependencies, ConversationalAnswer] = Agent( # type: ignore[invalid-assignment]
model=get_model(config.research.model, config),
output_type=RawSearchAnswer,
output_type=ConversationalAnswer,
instructions=conversational_prompt,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
# Include available citations for the LLM to select from
context_xml = format_context_for_prompt(
state.context, include_pending_questions=False, include_citations=True
state.context, include_pending_questions=False
)
prompt = f"Answer the question based on the gathered evidence.\n\n{context_xml}"
agent_deps = ResearchDependencies(
@ -539,23 +497,20 @@ def build_conversational_graph(
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
raw_answer = result.output
citation_lookup: dict[str, Citation] = {}
# Collect unique citations from qa_responses (dedupe by chunk_id)
seen_chunks: set[str] = set()
unique_citations: list[Citation] = []
for qa in state.context.qa_responses:
for c in qa.citations:
if c.chunk_id not in citation_lookup:
citation_lookup[c.chunk_id] = c
filtered_citations: list[Citation] = []
for chunk_id in raw_answer.cited_chunks:
if chunk_id in citation_lookup:
filtered_citations.append(citation_lookup[chunk_id])
if c.chunk_id not in seen_chunks:
seen_chunks.add(c.chunk_id)
unique_citations.append(c)
return ConversationalAnswer(
answer=raw_answer.answer,
citations=filtered_citations,
confidence=raw_answer.confidence,
answer=result.output.answer,
citations=unique_citations,
confidence=result.output.confidence,
)
# Build the graph structure (simplified: plan → search → synthesize)

9
haiku_rag_slim/haiku/rag/agents/research/models.py
View file

@ -163,11 +163,6 @@ class ResearchReport(BaseModel):
recommendations: list[str] = Field(
description="Actionable recommendations based on findings", default=[]
)
cited_chunks: list[str] = Field(
default_factory=list,
description="Chunk IDs selected by synthesis as directly supporting the report",
)
citations: list[Citation] = Field(
default_factory=list,
description="Resolved citations with full metadata",
sources_summary: str = Field(
description="Summary of sources used and their reliability"
)

21
haiku_rag_slim/haiku/rag/agents/research/prompts.py
View file

@ -116,7 +116,6 @@ Goals:
2. Present findings clearly and concisely.
3. Draw evidence-based conclusions and recommendations.
4. State limitations and uncertainties transparently.
5. Select only the citations that directly support your final answer.
Report guidelines (map to output fields):
- title: concise (5-12 words), informative.
@ -128,17 +127,10 @@ Report guidelines (map to output fields):
- conclusions: list of plain strings, 2-4 bullets following logically from findings.
- recommendations: list of plain strings, 2-5 actionable bullets tied to findings.
- limitations: list of plain strings, 1-3 bullets describing constraints or uncertainties.
- cited_chunks: list of chunk IDs that DIRECTLY support your report.
- sources_summary: single string listing sources with document paths and page numbers.
All list fields must contain plain strings only, not objects.
Citation selection:
- Review the <available_citations> section in the context.
- Include ONLY chunk IDs whose content directly supports specific claims in your report.
- Do NOT include chunks that are merely related, tangential, or were reviewed but unused.
- Quality over quantity: fewer relevant citations are better than many marginal ones.
- Use the EXACT chunk IDs from the available_citations (full UUIDs).
Style:
- Base all content solely on the collected evidence.
- Be professional, objective, and specific.
@ -149,11 +141,9 @@ CONVERSATIONAL_SYNTHESIS_PROMPT = """Generate a direct, conversational answer
to the question based on the gathered evidence.
Output:
- query: Echo the original question being answered.
- answer: Direct, comprehensive answer with a natural, helpful tone.
Write the actual answer, not a description of what you found.
Use as many sentences as needed to fully address the question.
- cited_chunks: List of chunk IDs that DIRECTLY support your answer.
- confidence: Score from 0.0 to 1.0 indicating answer quality.
Guidelines:
@ -163,11 +153,4 @@ Guidelines:
- Use formatting (bullet points, numbered lists) when it improves clarity.
- Do NOT use meta-commentary like "Based on the research..." or "The evidence shows..."
Instead, directly state the information.
- If the evidence is incomplete, acknowledge limitations briefly.
Citation selection:
- Review the <available_citations> section in the context.
- Include ONLY chunk IDs whose content directly supports specific statements in your answer.
- Do NOT include chunks that are merely related, tangential, or were reviewed but unused.
- Quality over quantity: fewer relevant citations are better than many marginal ones.
- Use the EXACT chunk IDs from available_citations (full UUIDs)."""
- If the evidence is incomplete, acknowledge limitations briefly."""

15
haiku_rag_slim/haiku/rag/app.py
View file

@ -416,9 +416,10 @@ class HaikuRAGApp:
self.console.print("[bold cyan]Key Findings:[/bold cyan]")
for finding in report.main_findings:
self.console.print(f"• {finding}")
if report.citations:
for renderable in format_citations_rich(report.citations):
self.console.print(renderable)
if report.sources_summary:
self.console.print()
self.console.print("[bold cyan]Sources:[/bold cyan]")
self.console.print(report.sources_summary)
else:
self.console.print("[yellow]No answer generated.[/yellow]")
else:
@ -512,10 +513,10 @@ class HaikuRAGApp:
self.console.print(f"• {limitation}")
self.console.print()
# Sources
if report.citations:
for renderable in format_citations_rich(report.citations):
self.console.print(renderable)
# Sources Summary
if report.sources_summary:
self.console.print("[bold cyan]Sources:[/bold cyan]")
self.console.print(report.sources_summary)
async def rebuild(self, mode: RebuildMode = RebuildMode.FULL):
async with HaikuRAG(

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

@ -6,6 +6,7 @@ from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any
import jsonpatch
from ag_ui.core import EventType
from pydantic_ai import (
Agent,
@ -99,13 +100,14 @@ class ChatApp(App):
self.client: HaikuRAG | None = None
self.config = get_config()
self.agent: Agent[ChatDeps, str] | None = None
self.session_state: ChatSessionState | None = None
self.session_state = ChatSessionState()
self._is_processing = False
self._tool_call_widgets: dict[str, Any] = {}
self._last_citations: list[Citation] = []
self._current_worker: Worker[None] | None = None
self._message_history: list[ModelMessage] = []
self._document_filter: list[str] = []
self._agui_state_snapshot: dict[str, Any] = {}
def compose(self) -> "ComposeResult":
"""Compose the UI layout."""
@ -185,20 +187,33 @@ class ChatApp(App):
widget = self._tool_call_widgets[tool_call_id]
chat_history.mark_tool_complete(widget)
# Extract citations from StateSnapshotEvent in tool metadata
# Extract citations from state events in tool metadata
result = getattr(event, "result", None)
metadata = getattr(result, "metadata", None) if result else None
if metadata:
for meta_event in metadata:
if (
hasattr(meta_event, "type")
and meta_event.type == EventType.STATE_SNAPSHOT
):
if not hasattr(meta_event, "type"):
continue
if meta_event.type == EventType.STATE_SNAPSHOT:
snapshot = getattr(meta_event, "snapshot", {})
self._agui_state_snapshot = snapshot
chat_state = snapshot.get(AGUI_STATE_KEY, snapshot)
self._last_citations = [
Citation(**c) for c in chat_state["citations"]
]
citations = chat_state.get("citations", [])
self._last_citations = [Citation(**c) for c in citations]
elif meta_event.type == EventType.STATE_DELTA:
delta = getattr(meta_event, "delta", [])
if delta:
patch = jsonpatch.JsonPatch(delta)
self._agui_state_snapshot = patch.apply(
self._agui_state_snapshot
)
chat_state = self._agui_state_snapshot.get(
AGUI_STATE_KEY, self._agui_state_snapshot
)
citations = chat_state.get("citations", [])
self._last_citations = [Citation(**c) for c in citations]
async def _event_stream_handler(
self,
@ -252,6 +267,12 @@ class ChatApp(App):
await chat_history.show_thinking()
try:
# Initialize AGUI state snapshot from session state for delta application
if self.session_state:
self._agui_state_snapshot = {
AGUI_STATE_KEY: self.session_state.model_dump(mode="json")
}
deps = ChatDeps(
client=self.client,
config=self.config,
@ -304,6 +325,7 @@ class ChatApp(App):
await chat_history.clear_messages()
self._last_citations.clear()
self._message_history.clear()
self._agui_state_snapshot = {}
# Reset context lock and session state (reset to CLI value)
self._context_locked = False
self.session_state = ChatSessionState(

1
haiku_rag_slim/pyproject.toml
View file

@ -24,6 +24,7 @@ classifiers = [
dependencies = [
"docling-core==2.60.1",
"httpx>=0.28.1",
"jsonpatch>=1.33",
"lancedb==0.27.0",
"pathspec>=1.0.3",
"pydantic>=2.12.5",

2
pyproject.toml
View file

@ -75,7 +75,7 @@ dev = [
"pydantic-ai-slim[bedrock]",
"pydantic-ai-slim[google]",
"pydantic-ai-slim[groq]",
"ty>=0.0.12",
"ty>=0.0.14",
"pytest>=9.0.2",
"pytest-asyncio>=1.3.0",
"pytest-cov>=7.0.0",

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

@ -1,6 +1,7 @@
from pathlib import Path
import pytest
from ag_ui.core import StateDeltaEvent, StateSnapshotEvent
from haiku.rag.agents.chat import (
AGUI_STATE_KEY,
@ -10,12 +11,39 @@ from haiku.rag.agents.chat import (
SearchAgent,
create_chat_agent,
)
from haiku.rag.agents.chat.context import get_cached_session_context
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
def extract_state_from_result(result, state_key: str = AGUI_STATE_KEY) -> dict | None:
"""Extract emitted state from agent result's tool return metadata.
For deltas, applies the patch to an empty state to get the final state.
"""
import jsonpatch
for message in result.all_messages():
if hasattr(message, "parts"):
for part in message.parts:
if hasattr(part, "metadata") and part.metadata:
for meta in part.metadata:
if isinstance(meta, StateSnapshotEvent):
return meta.snapshot.get(state_key)
elif isinstance(meta, StateDeltaEvent):
# Apply delta to empty state to get final state
empty_state = {
state_key: ChatSessionState(session_id="").model_dump(
mode="json"
)
}
patched = jsonpatch.apply_patch(empty_state, meta.delta)
return patched.get(state_key)
return None
@pytest.fixture(scope="module")
def vcr_cassette_dir():
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_chat_agent")
@ -36,7 +64,9 @@ def test_chat_deps_initialization(temp_db_path):
assert deps.client is client
assert deps.config is Config
assert deps.search_results is None
assert deps.session_state is None
assert deps.session_state is not None
assert deps.session_state.qa_history == []
assert deps.session_state.citations == []
client.close()
@ -449,11 +479,10 @@ async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path)
)
agent = create_chat_agent(Config)
session_state = ChatSessionState(session_id="test-citations")
deps = ChatDeps(
client=client,
config=Config,
session_state=session_state,
state_key=AGUI_STATE_KEY,
)
# Ask a question that should use the ask tool with citations
@ -463,8 +492,12 @@ 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(session_state.qa_history) >= 1
assert len(emitted_state.get("qa_history", [])) >= 1
@pytest.mark.asyncio
@ -483,16 +516,12 @@ async def test_chat_agent_ask_triggers_background_summarization(
)
agent = create_chat_agent(Config)
session_state = ChatSessionState(session_id="test-summarization")
deps = ChatDeps(
client=client,
config=Config,
session_state=session_state,
state_key=AGUI_STATE_KEY,
)
# Initially no session_context
assert session_state.session_context is None
# Ask a question
result = await agent.run(
"What is the highest count class in the DocLayNet dataset?",
@ -500,19 +529,27 @@ async def test_chat_agent_ask_triggers_background_summarization(
)
assert result.output is not None
assert len(session_state.qa_history) >= 1
# 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
# Wait for background task to complete
# The task should update session_state.session_context
# The task caches session_context server-side
cached_context = None
for _ in range(50): # Wait up to 5 seconds
if session_state.session_context is not None:
cached_context = get_cached_session_context(session_id)
if cached_context is not None:
break
await asyncio.sleep(0.1)
# Verify session_context was populated by background task
assert session_state.session_context is not None
assert session_state.session_context.summary != ""
assert session_state.session_context.last_updated is not None
assert cached_context is not None
assert cached_context.summary != ""
assert cached_context.last_updated is not None
@pytest.mark.asyncio
@ -536,52 +573,60 @@ async def test_chat_agent_ask_with_prior_answer_retrieval(
)
agent = create_chat_agent(Config)
session_state = ChatSessionState(session_id="test-prior-answers")
deps = ChatDeps(
deps1 = ChatDeps(
client=client,
config=Config,
session_state=session_state,
state_key=AGUI_STATE_KEY,
)
# First ask - establishes qa_history
result1 = await agent.run(
"What are the class labels in DocLayNet?",
deps=deps,
deps=deps1,
)
assert result1.output is not None
assert len(session_state.qa_history) == 1
# First question should NOT have embedding yet (set lazily on next ask)
assert session_state.qa_history[0].question_embedding is 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 session_state.session_context is not None:
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,
)
# 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=deps,
deps=deps2,
)
assert result2.output is not None
# qa_history should now have 2 entries
assert len(session_state.qa_history) == 2
# First question should now have embedding (set during second ask's recall check)
assert session_state.qa_history[0].question_embedding 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_state.qa_history[0].question_embedding, list)
assert len(session_state.qa_history[0].question_embedding) > 0
# Verify prior answer was reused without new searches:
# Second answer's citations should be subset of first answer's citations
first_chunk_ids = {c.chunk_id for c in session_state.qa_history[0].citations}
second_chunk_ids = {c.chunk_id for c in session_state.qa_history[1].citations}
assert second_chunk_ids <= first_chunk_ids, (
"Second answer should reuse prior citations, not perform new searches"
)
assert isinstance(session_state2.qa_history[0].question_embedding, list)
assert len(session_state2.qa_history[0].question_embedding) > 0
def test_fifo_limit_enforcement():

112
tests/agents/chat/test_state.py
View file

@ -1,5 +1,7 @@
import uuid
from ag_ui.core import StateDeltaEvent
from haiku.rag.agents.chat.state import (
MAX_QA_HISTORY,
ChatSessionState,
@ -144,8 +146,8 @@ def test_chat_deps_state_getter_without_namespace():
assert state["session_id"] == "test-123"
def test_chat_deps_state_getter_returns_none_without_session():
"""Test ChatDeps.state getter returns None when no session_state."""
def test_chat_deps_state_getter_returns_default_state():
"""Test ChatDeps.state getter returns default state when not explicitly set."""
from unittest.mock import MagicMock
from haiku.rag.agents.chat.state import ChatDeps
@ -156,10 +158,13 @@ def test_chat_deps_state_getter_returns_none_without_session():
deps = ChatDeps(
client=mock_client,
config=mock_config,
session_state=None,
)
assert deps.state is None
state = deps.state
assert state is not None
assert "session_id" in state
assert state["qa_history"] == []
assert state["citations"] == []
def test_chat_deps_state_setter_updates_from_namespaced_state():
@ -221,8 +226,8 @@ def test_chat_deps_state_setter_handles_none():
assert deps.session_state.session_id == "original"
def test_chat_deps_state_setter_without_session_state():
"""Test ChatDeps.state setter does nothing when session_state is None."""
def test_chat_deps_state_setter_updates_default_state():
"""Test ChatDeps.state setter updates the default session_state."""
from unittest.mock import MagicMock
from haiku.rag.agents.chat.state import ChatDeps
@ -233,13 +238,15 @@ def test_chat_deps_state_setter_without_session_state():
deps = ChatDeps(
client=mock_client,
config=mock_config,
session_state=None,
)
# Should not raise even with valid incoming state
deps.state = {"session_id": "test", "qa_history": [], "citations": []}
original_session_id = deps.session_state.session_id
assert deps.session_state is None
# Update with incoming state
deps.state = {"session_id": "updated-123", "qa_history": [], "citations": []}
assert deps.session_state.session_id == "updated-123"
assert deps.session_state.session_id != original_session_id
def test_chat_deps_state_setter_with_citation_dicts():
@ -641,3 +648,88 @@ def test_chat_session_state_model_dump_json_serializes_datetime():
# datetime should be serialized as ISO string, not datetime object
assert isinstance(snapshot["session_context"]["last_updated"], str)
assert snapshot["session_context"]["last_updated"] == "2025-01-27T12:00:00"
def test_emit_state_event_returns_none_when_no_changes():
"""emit_state_event returns None when states are identical."""
from haiku.rag.agents.chat.state import emit_state_event
state = ChatSessionState(session_id="test-123", qa_history=[], citations=[])
event = emit_state_event(state, state)
assert event is None
def test_emit_state_event_returns_delta_with_changes():
"""emit_state_event returns StateDeltaEvent with JSON Patch ops for changes."""
from ag_ui.core import EventType, StateDeltaEvent
from haiku.rag.agents.chat.state import emit_state_event
current_state = ChatSessionState(session_id="test-123", qa_history=[], citations=[])
new_state = ChatSessionState(
session_id="test-123",
qa_history=[QAResponse(question="Q1", answer="A1", confidence=0.9)],
citations=[],
)
event = emit_state_event(current_state, new_state)
assert isinstance(event, StateDeltaEvent)
assert event.type == EventType.STATE_DELTA
assert len(event.delta) > 0
# Delta should contain an "add" operation for the new qa_history entry
ops = event.delta
qa_history_op = next((op for op in ops if "/qa_history" in op["path"]), None)
assert qa_history_op is not None
def test_emit_state_event_delta_with_state_key():
"""emit_state_event wraps delta paths with state_key namespace."""
from ag_ui.core import StateDeltaEvent
from haiku.rag.agents.chat.state import AGUI_STATE_KEY, emit_state_event
current_state = ChatSessionState(session_id="test-123", qa_history=[])
new_state = ChatSessionState(
session_id="test-123",
qa_history=[QAResponse(question="Q1", answer="A1", confidence=0.9)],
)
event = emit_state_event(current_state, new_state, state_key=AGUI_STATE_KEY)
assert isinstance(event, StateDeltaEvent)
# Paths should be namespaced under state_key
for op in event.delta:
assert op["path"].startswith(f"/{AGUI_STATE_KEY}")
def test_emit_state_event_delta_produces_valid_patch():
"""emit_state_event delta can be applied to reproduce new state."""
import jsonpatch
from haiku.rag.agents.chat.state import emit_state_event
current_state = ChatSessionState(
session_id="test-123",
qa_history=[QAResponse(question="Q1", answer="A1", confidence=0.9)],
citations=[],
)
new_state = ChatSessionState(
session_id="test-123",
qa_history=[
QAResponse(question="Q1", answer="A1", confidence=0.9),
QAResponse(question="Q2", answer="A2", confidence=0.8),
],
citations=[],
)
event = emit_state_event(current_state, new_state)
assert isinstance(event, StateDeltaEvent)
# Apply patch to current state and verify it produces new state
current_snapshot = current_state.model_dump(mode="json")
patched = jsonpatch.apply_patch(current_snapshot, event.delta)
new_snapshot = new_state.model_dump(mode="json")
assert patched == new_snapshot

43
tests/agents/research/test_research_graph.py
View file

@ -184,46 +184,3 @@ def test_format_context_for_prompt_with_prior_answers():
assert "<prior_answers>" in result
assert "Sub question?" in result
assert "The answer is here." in result
def test_format_context_for_prompt_with_citations():
"""Test format_context_for_prompt includes available_citations when requested."""
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import format_context_for_prompt
from haiku.rag.agents.research.models import Citation, SearchAnswer
context = ResearchContext(original_question="Main question?")
context.add_qa_response(
SearchAnswer(
query="Sub question?",
answer="The answer is here.",
confidence=0.9,
cited_chunks=["chunk-123", "chunk-456"],
citations=[
Citation(
document_id="doc-1",
chunk_id="chunk-123",
document_uri="test://doc1",
document_title="Test Document",
content="This is the chunk content.",
),
Citation(
document_id="doc-1",
chunk_id="chunk-456",
document_uri="test://doc1",
document_title="Test Document",
content="More chunk content here.",
),
],
)
)
result_without = format_context_for_prompt(context, include_citations=False)
assert "<available_citations>" not in result_without
result_with = format_context_for_prompt(context, include_citations=True)
assert "<available_citations>" in result_with
assert "chunk-123" in result_with
assert "chunk-456" in result_with
assert "Test Document" in result_with
assert "This is the chunk content." in result_with

3356
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

2790
tests/cassettes/test_chat_agent/test_chat_agent_ask_triggers_background_summarization.yaml
View file

File diff suppressed because one or more lines are too long

5211
tests/cassettes/test_chat_agent/test_chat_agent_ask_with_prior_answer_retrieval.yaml
View file

File diff suppressed because one or more lines are too long

229
tests/chat/test_chat_app.py
View file

@ -282,6 +282,235 @@ async def test_clear_chat_resets_session(temp_db_path: Path):
assert app.session_state.session_id != original_session_id
@pytest.mark.asyncio
async def test_handle_stream_event_extracts_citations_from_state_snapshot(
temp_db_path: Path,
):
"""Test that _handle_stream_event extracts citations from STATE_SNAPSHOT events."""
from ag_ui.core import EventType, StateSnapshotEvent
from pydantic_ai import FunctionToolResultEvent
from pydantic_ai.messages import ToolReturnPart
from haiku.rag.agents.chat.state import AGUI_STATE_KEY
from haiku.rag.chat.app import ChatApp
mock_client = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
app = ChatApp(temp_db_path, read_only=True)
async with app.run_test():
# Create a STATE_SNAPSHOT event with citations
snapshot_event = StateSnapshotEvent(
type=EventType.STATE_SNAPSHOT,
snapshot={
AGUI_STATE_KEY: {
"session_id": "test",
"citations": [
{
"index": 1,
"document_id": "doc1",
"chunk_id": "chunk1",
"document_uri": "test.pdf",
"document_title": "Test Doc",
"content": "Test content",
}
],
"qa_history": [],
"citation_registry": {"chunk1": 1},
}
},
)
# Create a tool result with the snapshot in metadata
tool_return = ToolReturnPart(
tool_name="search",
content="Found results",
tool_call_id="test-call-1",
metadata=[snapshot_event],
)
event = FunctionToolResultEvent(result=tool_return)
# Handle the event
await app._handle_stream_event(event)
# Citations should be extracted
assert len(app._last_citations) == 1
assert app._last_citations[0].chunk_id == "chunk1"
@pytest.mark.asyncio
async def test_handle_stream_event_extracts_citations_from_state_delta(
temp_db_path: Path,
):
"""Test that _handle_stream_event extracts citations from STATE_DELTA events.
This test demonstrates that state deltas need to be applied to extract citations.
After the first STATE_SNAPSHOT, subsequent tool calls emit STATE_DELTA events
containing JSON Patch operations.
"""
from ag_ui.core import EventType, StateDeltaEvent, StateSnapshotEvent
from pydantic_ai import FunctionToolResultEvent
from pydantic_ai.messages import ToolReturnPart
from haiku.rag.agents.chat.state import AGUI_STATE_KEY
from haiku.rag.chat.app import ChatApp
mock_client = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
app = ChatApp(temp_db_path, read_only=True)
async with app.run_test():
# First, handle a STATE_SNAPSHOT to establish initial state
initial_snapshot = StateSnapshotEvent(
type=EventType.STATE_SNAPSHOT,
snapshot={
AGUI_STATE_KEY: {
"session_id": "test",
"citations": [],
"qa_history": [],
"citation_registry": {},
}
},
)
tool_return1 = ToolReturnPart(
tool_name="search",
content="Initial search",
tool_call_id="test-call-1",
metadata=[initial_snapshot],
)
event1 = FunctionToolResultEvent(result=tool_return1)
await app._handle_stream_event(event1)
assert len(app._last_citations) == 0
# Now handle a STATE_DELTA event that adds citations
delta_event = StateDeltaEvent(
type=EventType.STATE_DELTA,
delta=[
{
"op": "replace",
"path": f"/{AGUI_STATE_KEY}/citations",
"value": [
{
"index": 1,
"document_id": "doc1",
"chunk_id": "chunk1",
"document_uri": "test.pdf",
"document_title": "Test Doc",
"content": "Test content from delta",
}
],
},
{
"op": "add",
"path": f"/{AGUI_STATE_KEY}/citation_registry/chunk1",
"value": 1,
},
],
)
tool_return2 = ToolReturnPart(
tool_name="ask",
content="Answer with citations",
tool_call_id="test-call-2",
metadata=[delta_event],
)
event2 = FunctionToolResultEvent(result=tool_return2)
# Handle the delta event
await app._handle_stream_event(event2)
# Citations should be extracted from the delta
assert len(app._last_citations) == 1
assert app._last_citations[0].chunk_id == "chunk1"
assert app._last_citations[0].content == "Test content from delta"
@pytest.mark.asyncio
async def test_handle_stream_event_delta_with_preinitialized_state(
temp_db_path: Path,
):
"""Test delta handling when _agui_state_snapshot is pre-initialized.
This is the actual TUI scenario: session_state exists from the start,
so the agent emits deltas (not snapshots) even on the first tool call.
The TUI pre-initializes _agui_state_snapshot from session_state.
"""
from ag_ui.core import EventType, StateDeltaEvent
from pydantic_ai import FunctionToolResultEvent
from pydantic_ai.messages import ToolReturnPart
from haiku.rag.agents.chat.state import AGUI_STATE_KEY
from haiku.rag.chat.app import ChatApp
mock_client = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
app = ChatApp(temp_db_path, read_only=True)
async with app.run_test():
# Pre-initialize _agui_state_snapshot (simulating what _run_agent does)
app._agui_state_snapshot = {
AGUI_STATE_KEY: {
"session_id": "test",
"citations": [],
"qa_history": [],
"citation_registry": {},
"document_filter": [],
"initial_context": None,
"session_context": None,
}
}
# Now handle a STATE_DELTA event directly (no prior snapshot event)
delta_event = StateDeltaEvent(
type=EventType.STATE_DELTA,
delta=[
{
"op": "replace",
"path": f"/{AGUI_STATE_KEY}/citations",
"value": [
{
"index": 1,
"document_id": "doc1",
"chunk_id": "chunk1",
"document_uri": "test.pdf",
"document_title": "Test Doc",
"content": "Content from first delta",
}
],
},
{
"op": "add",
"path": f"/{AGUI_STATE_KEY}/citation_registry/chunk1",
"value": 1,
},
],
)
tool_return = ToolReturnPart(
tool_name="ask",
content="Answer with citations",
tool_call_id="test-call-1",
metadata=[delta_event],
)
event = FunctionToolResultEvent(result=tool_return)
# Handle the delta event
await app._handle_stream_event(event)
# Citations should be extracted from the delta
assert len(app._last_citations) == 1
assert app._last_citations[0].chunk_id == "chunk1"
assert app._last_citations[0].content == "Content from first delta"
@pytest.mark.asyncio
async def test_citation_expand_collapse_with_enter(temp_db_path: Path):
"""Test that pressing Enter on a focused citation toggles expand/collapse."""

2
tests/test_app.py
View file

@ -349,6 +349,7 @@ async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch):
executive_summary="Deep research answer",
main_findings=["Finding 1"],
conclusions=["Conclusion 1"],
sources_summary="Sources",
)
mock_graph = AsyncMock()
@ -386,6 +387,7 @@ async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch):
executive_summary="Deep research answer",
main_findings=["Finding 1"],
conclusions=["Conclusion 1"],
sources_summary="Sources",
)
mock_graph = AsyncMock()

1
tests/test_mcp.py
View file

@ -277,6 +277,7 @@ async def test_mcp_research_question():
main_findings=["Finding 1"],
conclusions=["Conclusion 1"],
recommendations=["Recommendation 1"],
sources_summary="Sources used",
)
with (

41
uv.lock
View file

@ -1335,7 +1335,7 @@ dev = [
{ name = "pytest-cov", specifier = ">=7.0.0" },
{ name = "pytest-recording", specifier = ">=0.13.4" },
{ name = "ruff", specifier = ">=0.14.13" },
{ name = "ty", specifier = ">=0.0.12" },
{ name = "ty", specifier = ">=0.0.14" },
]
[[package]]
@ -1368,6 +1368,7 @@ source = { editable = "haiku_rag_slim" }
dependencies = [
{ name = "docling-core" },
{ name = "httpx" },
{ name = "jsonpatch" },
{ name = "lancedb" },
{ name = "pathspec" },
{ name = "pydantic" },
@ -1429,6 +1430,7 @@ requires-dist = [
{ name = "docling", marker = "extra == 'docling'", specifier = "==2.69.1" },
{ name = "docling-core", specifier = "==2.60.1" },
{ name = "httpx", specifier = ">=0.28.1" },
{ name = "jsonpatch", specifier = ">=1.33" },
{ name = "lancedb", specifier = "==0.27.0" },
{ name = "mxbai-rerank", marker = "extra == 'mxbai'", specifier = ">=0.1.6" },
{ name = "opencv-python-headless", marker = "extra == 'docling'", specifier = ">=4.13.0.90" },
@ -4984,27 +4986,26 @@ wheels = [
[[package]]
name = "ty"
version = "0.0.12"
version = "0.0.14"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/b5/78/ba1a4ad403c748fbba8be63b7e774a90e80b67192f6443d624c64fe4aaab/ty-0.0.12.tar.gz", hash = "sha256:cd01810e106c3b652a01b8f784dd21741de9fdc47bd595d02c122a7d5cefeee7", size = 4981303, upload-time = "2026-01-14T22:30:48.537Z" }
sdist = { url = "https://files.pythonhosted.org/packages/af/57/22c3d6bf95c2229120c49ffc2f0da8d9e8823755a1c3194da56e51f1cc31/ty-0.0.14.tar.gz", hash = "sha256:a691010565f59dd7f15cf324cdcd1d9065e010c77a04f887e1ea070ba34a7de2", size = 5036573, upload-time = "2026-01-27T00:57:31.427Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/7d/8f/c21314d074dda5fb13d3300fa6733fd0d8ff23ea83a721818740665b6314/ty-0.0.12-py3-none-linux_armv6l.whl", hash = "sha256:eb9da1e2c68bd754e090eab39ed65edf95168d36cbeb43ff2bd9f86b4edd56d1", size = 9614164, upload-time = "2026-01-14T22:30:44.016Z" },
{ url = "https://files.pythonhosted.org/packages/09/28/f8a4d944d13519d70c486e8f96d6fa95647ac2aa94432e97d5cfec1f42f6/ty-0.0.12-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:c181f42aa19b0ed7f1b0c2d559980b1f1d77cc09419f51c8321c7ddf67758853", size = 9542337, upload-time = "2026-01-14T22:30:05.687Z" },
{ url = "https://files.pythonhosted.org/packages/e1/9c/f576e360441de7a8201daa6dc4ebc362853bc5305e059cceeb02ebdd9a48/ty-0.0.12-py3-none-macosx_11_0_arm64.whl", hash = "sha256:1f829e1eecd39c3e1b032149db7ae6a3284f72fc36b42436e65243a9ed1173db", size = 8909582, upload-time = "2026-01-14T22:30:46.089Z" },
{ url = "https://files.pythonhosted.org/packages/d6/13/0898e494032a5d8af3060733d12929e3e7716db6c75eac63fa125730a3e7/ty-0.0.12-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f45162e7826e1789cf3374627883cdeb0d56b82473a0771923e4572928e90be3", size = 9384932, upload-time = "2026-01-14T22:30:13.769Z" },
{ url = "https://files.pythonhosted.org/packages/e4/1a/b35b6c697008a11d4cedfd34d9672db2f0a0621ec80ece109e13fca4dfef/ty-0.0.12-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:d11fec40b269bec01e751b2337d1c7ffa959a2c2090a950d7e21c2792442cccd", size = 9453140, upload-time = "2026-01-14T22:30:11.131Z" },
{ url = "https://files.pythonhosted.org/packages/dd/1e/71c9edbc79a3c88a0711324458f29c7dbf6c23452c6e760dc25725483064/ty-0.0.12-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:09d99e37e761a4d2651ad9d5a610d11235fbcbf35dc6d4bc04abf54e7cf894f1", size = 9960680, upload-time = "2026-01-14T22:30:33.621Z" },
{ url = "https://files.pythonhosted.org/packages/0e/75/39375129f62dd22f6ad5a99cd2a42fd27d8b91b235ce2db86875cdad397d/ty-0.0.12-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:d9ca0cdb17bd37397da7b16a7cd23423fc65c3f9691e453ad46c723d121225a1", size = 10904518, upload-time = "2026-01-14T22:30:08.464Z" },
{ url = "https://files.pythonhosted.org/packages/32/5e/26c6d88fafa11a9d31ca9f4d12989f57782ec61e7291d4802d685b5be118/ty-0.0.12-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:fcf2757b905e7eddb7e456140066335b18eb68b634a9f72d6f54a427ab042c64", size = 10525001, upload-time = "2026-01-14T22:30:16.454Z" },
{ url = "https://files.pythonhosted.org/packages/c2/a5/2f0b91894af13187110f9ad7ee926d86e4e6efa755c9c88a820ed7f84c85/ty-0.0.12-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:00cf34c1ebe1147efeda3021a1064baa222c18cdac114b7b050bbe42deb4ca80", size = 10307103, upload-time = "2026-01-14T22:30:41.221Z" },
{ url = "https://files.pythonhosted.org/packages/4b/77/13d0410827e4bc713ebb7fdaf6b3590b37dcb1b82e0a81717b65548f2442/ty-0.0.12-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2bb3a655bd869352e9a22938d707631ac9fbca1016242b1f6d132d78f347c851", size = 10072737, upload-time = "2026-01-14T22:30:51.783Z" },
{ url = "https://files.pythonhosted.org/packages/e1/dd/fc36d8bac806c74cf04b4ca735bca14d19967ca84d88f31e121767880df1/ty-0.0.12-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:4658e282c7cb82be304052f8f64f9925f23c3c4f90eeeb32663c74c4b095d7ba", size = 9368726, upload-time = "2026-01-14T22:30:18.683Z" },
{ url = "https://files.pythonhosted.org/packages/54/70/9e8e461647550f83e2fe54bc632ccbdc17a4909644783cdbdd17f7296059/ty-0.0.12-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:c167d838eaaa06e03bb66a517f75296b643d950fbd93c1d1686a187e5a8dbd1f", size = 9454704, upload-time = "2026-01-14T22:30:22.759Z" },
{ url = "https://files.pythonhosted.org/packages/04/9b/6292cf7c14a0efeca0539cf7d78f453beff0475cb039fbea0eb5d07d343d/ty-0.0.12-py3-none-musllinux_1_2_i686.whl", hash = "sha256:2956e0c9ab7023533b461d8a0e6b2ea7b78e01a8dde0688e8234d0fce10c4c1c", size = 9649829, upload-time = "2026-01-14T22:30:31.234Z" },
{ url = "https://files.pythonhosted.org/packages/49/bd/472a5d2013371e4870886cff791c94abdf0b92d43d305dd0f8e06b6ff719/ty-0.0.12-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:5c6a3fd7479580009f21002f3828320621d8a82d53b7ba36993234e3ccad58c8", size = 10162814, upload-time = "2026-01-14T22:30:36.174Z" },
{ url = "https://files.pythonhosted.org/packages/31/e9/2ecbe56826759845a7c21d80aa28187865ea62bc9757b056f6cbc06f78ed/ty-0.0.12-py3-none-win32.whl", hash = "sha256:a91c24fd75c0f1796d8ede9083e2c0ec96f106dbda73a09fe3135e075d31f742", size = 9140115, upload-time = "2026-01-14T22:30:38.903Z" },
{ url = "https://files.pythonhosted.org/packages/5d/6d/d9531eff35a5c0ec9dbc10231fac21f9dd6504814048e81d6ce1c84dc566/ty-0.0.12-py3-none-win_amd64.whl", hash = "sha256:df151894be55c22d47068b0f3b484aff9e638761e2267e115d515fcc9c5b4a4b", size = 9884532, upload-time = "2026-01-14T22:30:25.112Z" },
{ url = "https://files.pythonhosted.org/packages/e9/f3/20b49e75967023b123a221134548ad7000f9429f13fdcdda115b4c26305f/ty-0.0.12-py3-none-win_arm64.whl", hash = "sha256:cea99d334b05629de937ce52f43278acf155d3a316ad6a35356635f886be20ea", size = 9313974, upload-time = "2026-01-14T22:30:27.44Z" },
{ url = "https://files.pythonhosted.org/packages/99/cb/cc6d1d8de59beb17a41f9a614585f884ec2d95450306c173b3b7cc090d2e/ty-0.0.14-py3-none-linux_armv6l.whl", hash = "sha256:32cf2a7596e693094621d3ae568d7ee16707dce28c34d1762947874060fdddaa", size = 10034228, upload-time = "2026-01-27T00:57:53.133Z" },
{ url = "https://files.pythonhosted.org/packages/f3/96/dd42816a2075a8f31542296ae687483a8d047f86a6538dfba573223eaf9a/ty-0.0.14-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:f971bf9805f49ce8c0968ad53e29624d80b970b9eb597b7cbaba25d8a18ce9a2", size = 9939162, upload-time = "2026-01-27T00:57:43.857Z" },
{ url = "https://files.pythonhosted.org/packages/ff/b4/73c4859004e0f0a9eead9ecb67021438b2e8e5fdd8d03e7f5aca77623992/ty-0.0.14-py3-none-macosx_11_0_arm64.whl", hash = "sha256:45448b9e4806423523268bc15e9208c4f3f2ead7c344f615549d2e2354d6e924", size = 9418661, upload-time = "2026-01-27T00:58:03.411Z" },
{ url = "https://files.pythonhosted.org/packages/58/35/839c4551b94613db4afa20ee555dd4f33bfa7352d5da74c5fa416ffa0fd2/ty-0.0.14-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ee94a9b747ff40114085206bdb3205a631ef19a4d3fb89e302a88754cbbae54c", size = 9837872, upload-time = "2026-01-27T00:57:23.718Z" },
{ url = "https://files.pythonhosted.org/packages/41/2b/bbecf7e2faa20c04bebd35fc478668953ca50ee5847ce23e08acf20ea119/ty-0.0.14-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:6756715a3c33182e9ab8ffca2bb314d3c99b9c410b171736e145773ee0ae41c3", size = 9848819, upload-time = "2026-01-27T00:57:58.501Z" },
{ url = "https://files.pythonhosted.org/packages/be/60/3c0ba0f19c0f647ad9d2b5b5ac68c0f0b4dc899001bd53b3a7537fb247a2/ty-0.0.14-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:89d0038a2f698ba8b6fec5cf216a4e44e2f95e4a5095a8c0f57fe549f87087c2", size = 10324371, upload-time = "2026-01-27T00:57:29.291Z" },
{ url = "https://files.pythonhosted.org/packages/24/32/99d0a0b37d0397b0a989ffc2682493286aa3bc252b24004a6714368c2c3d/ty-0.0.14-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:2c64a83a2d669b77f50a4957039ca1450626fb474619f18f6f8a3eb885bf7544", size = 10865898, upload-time = "2026-01-27T00:57:33.542Z" },
{ url = "https://files.pythonhosted.org/packages/1a/88/30b583a9e0311bb474269cfa91db53350557ebec09002bfc3fb3fc364e8c/ty-0.0.14-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:242488bfb547ef080199f6fd81369ab9cb638a778bb161511d091ffd49c12129", size = 10555777, upload-time = "2026-01-27T00:58:05.853Z" },
{ url = "https://files.pythonhosted.org/packages/cd/a2/cb53fb6325dcf3d40f2b1d0457a25d55bfbae633c8e337bde8ec01a190eb/ty-0.0.14-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4790c3866f6c83a4f424fc7d09ebdb225c1f1131647ba8bdc6fcdc28f09ed0ff", size = 10412913, upload-time = "2026-01-27T00:57:38.834Z" },
{ url = "https://files.pythonhosted.org/packages/42/8f/f2f5202d725ed1e6a4e5ffaa32b190a1fe70c0b1a2503d38515da4130b4c/ty-0.0.14-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:950f320437f96d4ea9a2332bbfb5b68f1c1acd269ebfa4c09b6970cc1565bd9d", size = 9837608, upload-time = "2026-01-27T00:57:55.898Z" },
{ url = "https://files.pythonhosted.org/packages/f7/ba/59a2a0521640c489dafa2c546ae1f8465f92956fede18660653cce73b4c5/ty-0.0.14-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:4a0ec3ee70d83887f86925bbc1c56f4628bd58a0f47f6f32ddfe04e1f05466df", size = 9884324, upload-time = "2026-01-27T00:57:46.786Z" },
{ url = "https://files.pythonhosted.org/packages/03/95/8d2a49880f47b638743212f011088552ecc454dd7a665ddcbdabea25772a/ty-0.0.14-py3-none-musllinux_1_2_i686.whl", hash = "sha256:a1a4e6b6da0c58b34415955279eff754d6206b35af56a18bb70eb519d8d139ef", size = 10033537, upload-time = "2026-01-27T00:58:01.149Z" },
{ url = "https://files.pythonhosted.org/packages/e9/40/4523b36f2ce69f92ccf783855a9e0ebbbd0f0bb5cdce6211ee1737159ed3/ty-0.0.14-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:dc04384e874c5de4c5d743369c277c8aa73d1edea3c7fc646b2064b637db4db3", size = 10495910, upload-time = "2026-01-27T00:57:26.691Z" },
{ url = "https://files.pythonhosted.org/packages/08/d5/655beb51224d1bfd4f9ddc0bb209659bfe71ff141bcf05c418ab670698f0/ty-0.0.14-py3-none-win32.whl", hash = "sha256:b20e22cf54c66b3e37e87377635da412d9a552c9bf4ad9fc449fed8b2e19dad2", size = 9507626, upload-time = "2026-01-27T00:57:41.43Z" },
{ url = "https://files.pythonhosted.org/packages/b6/d9/c569c9961760e20e0a4bc008eeb1415754564304fd53997a371b7cf3f864/ty-0.0.14-py3-none-win_amd64.whl", hash = "sha256:e312ff9475522d1a33186657fe74d1ec98e4a13e016d66f5758a452c90ff6409", size = 10437980, upload-time = "2026-01-27T00:57:36.422Z" },
{ url = "https://files.pythonhosted.org/packages/ad/0c/186829654f5bfd9a028f6648e9caeb11271960a61de97484627d24443f91/ty-0.0.14-py3-none-win_arm64.whl", hash = "sha256:b6facdbe9b740cb2c15293a1d178e22ffc600653646452632541d01c36d5e378", size = 9885831, upload-time = "2026-01-27T00:57:49.747Z" },
]
[[package]]