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

Remove session_id from state layer, remove module-level caches

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-11 14:55:53 +02:00
parent d88f2f003a
commit 9e63ff1ff7
No known key found for this signature in database
12 changed files with 131 additions and 460 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
app/backend/main.py
View file

@ -20,10 +20,13 @@ from haiku.rag.agents.chat import (
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config import load_yaml_config from haiku.rag.config import load_yaml_config
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
from haiku.rag.tools import ToolContext from haiku.rag.tools.context import ToolContextCache
load_dotenv(find_dotenv(usecwd=True)) load_dotenv(find_dotenv(usecwd=True))
# Cache ToolContext instances by thread_id across requests
context_cache = ToolContextCache()
# Configure logfire (only sends data if LOGFIRE_TOKEN is present) # Configure logfire (only sends data if LOGFIRE_TOKEN is present)
try: try:
import logfire import logfire
@ -68,30 +71,24 @@ def get_client() -> HaikuRAG:
async def stream_chat(request: Request) -> Response: async def stream_chat(request: Request) -> Response:
"""Chat streaming endpoint with AG-UI protocol. """Chat streaming endpoint with AG-UI protocol.
This endpoint is stateless - all state flows via AG-UI protocol: Uses ToolContextCache to maintain state across requests for the same thread.
- Fresh ToolContext created per request AGUIAdapter restores client-sent state via ChatDeps.state setter.
- 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() body = await request.body()
accept = request.headers.get("accept", SSE_CONTENT_TYPE) accept = request.headers.get("accept", SSE_CONTENT_TYPE)
run_input = AGUIAdapter.build_run_input(body) run_input = AGUIAdapter.build_run_input(body)
# Fresh context per request - state restored by AGUIAdapter via ChatDeps.state setter thread_id = getattr(run_input, "thread_id", None) or "default"
context = ToolContext() context, is_new = context_cache.get_or_create(thread_id)
agent = create_chat_agent(Config, get_client(), context) agent = create_chat_agent(Config, get_client(), context)
deps = ChatDeps( deps = ChatDeps(
config=Config, config=Config,
tool_context=context, tool_context=context,
is_new=is_new,
state_key=AGUI_STATE_KEY, state_key=AGUI_STATE_KEY,
) )
# Use AGUIAdapter for streaming
# 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) adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
event_stream = adapter.run_stream(deps=deps) event_stream = adapter.run_stream(deps=deps)
sse_event_stream = adapter.encode_stream(event_stream) sse_event_stream = adapter.encode_stream(event_stream)

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

@ -1,12 +1,8 @@
import uuid
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any from typing import Any
from pydantic_ai import Agent from pydantic_ai import Agent
from haiku.rag.agents.chat.context import (
get_cached_session_context,
)
from haiku.rag.agents.chat.context import ( from haiku.rag.agents.chat.context import (
trigger_background_summarization as _trigger_summarization, trigger_background_summarization as _trigger_summarization,
) )
@ -47,7 +43,7 @@ class ChatDeps:
config: AppConfig config: AppConfig
tool_context: ToolContext tool_context: ToolContext
session_id: str = "" is_new: bool = True
state_key: str | None = None state_key: str | None = None
@property @property
@ -78,59 +74,57 @@ class ChatDeps:
state_data = nested state_data = nested
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) session_state = self.tool_context.get(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 = [ if self.is_new:
Citation(**c) if isinstance(c, dict) else c # First request for this context: fully populate from client state
for c in state_data.get("citations", []) 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
# Restore session_id from client or generate one session_state.citations = [
client_session_id = state_data.get("session_id", "") Citation(**c) if isinstance(c, dict) else c
if client_session_id: for c in state_data.get("citations", [])
self.session_id = client_session_id ]
elif not self.session_id:
self.session_id = str(uuid.uuid4())
if session_state is not None: qa_session_state = self.tool_context.get(
session_state.session_id = self.session_id 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 = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState) qa_session_state.qa_history = [
if qa_session_state is not None: QAHistoryEntry(**qa) if isinstance(qa, dict) else qa
if "qa_history" in state_data: for qa in state_data.get("qa_history", [])
from haiku.rag.tools.qa import QAHistoryEntry ]
qa_session_state.qa_history = [ # Restore session_context from client
QAHistoryEntry(**qa) if isinstance(qa, dict) else qa session_context = state_data.get("session_context")
for qa in state_data.get("qa_history", []) if isinstance(session_context, dict):
] qa_session_state.session_context = SessionContext(
**session_context
).summary
elif session_context is None:
qa_session_state.session_context = None
# Restore session_context from client # Handle initial_context -> session_context for first message
session_context = state_data.get("session_context") if "initial_context" in state_data:
if isinstance(session_context, dict): initial = state_data.get("initial_context")
qa_session_state.session_context = SessionContext( if initial and not qa_session_state.session_context:
**session_context qa_session_state.session_context = initial
).summary else:
elif session_context is None: # Returning request: only merge client-controlled fields
qa_session_state.session_context = None if session_state is not None:
if "document_filter" in state_data:
# Check cache for fresher session_context from background summarization session_state.document_filter = state_data.get(
if self.session_id: "document_filter", []
cached = get_cached_session_context(self.session_id) )
if cached and cached.summary:
qa_session_state.session_context = cached.summary
# Handle initial_context -> session_context for first message
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( def create_chat_agent(
@ -204,22 +198,16 @@ def trigger_background_summarization(deps: ChatDeps) -> None:
Call this after agent.run() or agent.run_stream() completes to update Call this after agent.run() or agent.run_stream() completes to update
the session context summary in the background. 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: Args:
deps: Chat dependencies with tool_context containing QASessionState. deps: Chat dependencies with tool_context containing QASessionState.
""" """
qa_session_state = deps.tool_context.get(QA_SESSION_NAMESPACE, QASessionState) qa_session_state = deps.tool_context.get(QA_SESSION_NAMESPACE, QASessionState)
if qa_session_state is None or not qa_session_state.qa_history: if qa_session_state is None or not qa_session_state.qa_history:
return return
if not deps.session_id:
return
_trigger_summarization( _trigger_summarization(
qa_session_state=qa_session_state, qa_session_state=qa_session_state,
config=deps.config, config=deps.config,
session_id=deps.session_id,
) )

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

@ -1,6 +1,5 @@
import asyncio import asyncio
from dataclasses import dataclass, field from datetime import datetime
from datetime import datetime, timedelta
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from pydantic_ai import Agent from pydantic_ai import Agent
@ -14,68 +13,8 @@ if TYPE_CHECKING:
from haiku.rag.tools.qa import QAHistoryEntry, QASessionState from haiku.rag.tools.qa import QAHistoryEntry, QASessionState
@dataclass # Track summarization tasks to allow cancellation
class SessionCache: _summarization_tasks: dict[int, asyncio.Task[None]] = {}
"""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_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."""
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()
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( async def summarize_session(
@ -115,15 +54,13 @@ async def summarize_session(
async def update_session_context( async def update_session_context(
qa_history: list["QAHistoryEntry"], qa_history: list["QAHistoryEntry"],
config: AppConfig, config: AppConfig,
session_id: str = "",
current_context: str | None = None, current_context: str | None = None,
) -> SessionContext: ) -> SessionContext:
"""Summarize qa_history and cache the resulting session context. """Summarize qa_history and return the resulting session context.
Args: Args:
qa_history: List of Q&A pairs from the conversation. qa_history: List of Q&A pairs from the conversation.
config: AppConfig for model selection. config: AppConfig for model selection.
session_id: Session ID for caching. If empty, result is not cached.
current_context: Previous summary to incorporate. current_context: Previous summary to incorporate.
Returns: Returns:
@ -132,13 +69,10 @@ async def update_session_context(
summary = await summarize_session( summary = await summarize_session(
qa_history, config, current_context=current_context qa_history, config, current_context=current_context
) )
context = SessionContext( return SessionContext(
summary=summary, summary=summary,
last_updated=datetime.now(), last_updated=datetime.now(),
) )
if session_id:
cache_session_context(session_id, context)
return context
def _format_qa_history(qa_history: list["QAHistoryEntry"]) -> str: def _format_qa_history(qa_history: list["QAHistoryEntry"]) -> str:
@ -159,14 +93,12 @@ def _format_qa_history(qa_history: list["QAHistoryEntry"]) -> str:
async def _update_context_background( async def _update_context_background(
qa_session_state: "QASessionState", qa_session_state: "QASessionState",
config: AppConfig, config: AppConfig,
session_id: str,
) -> None: ) -> None:
"""Background task to update session context after an ask.""" """Background task to update session context after an ask."""
try: try:
result = await update_session_context( result = await update_session_context(
qa_history=list(qa_session_state.qa_history), qa_history=list(qa_session_state.qa_history),
config=config, config=config,
session_id=session_id,
current_context=qa_session_state.session_context, current_context=qa_session_state.session_context,
) )
@ -184,31 +116,28 @@ async def _update_context_background(
def trigger_background_summarization( def trigger_background_summarization(
qa_session_state: "QASessionState", qa_session_state: "QASessionState",
config: AppConfig, config: AppConfig,
session_id: str,
) -> None: ) -> None:
"""Trigger background session summarization if qa_history has entries. """Trigger background session summarization if qa_history has entries.
Args: Args:
qa_session_state: QASessionState with qa_history to summarize. qa_session_state: QASessionState with qa_history to summarize.
config: AppConfig for model selection. config: AppConfig for model selection.
session_id: Session ID for caching results.
""" """
if not qa_session_state.qa_history or not session_id: if not qa_session_state.qa_history:
return return
# Cancel any existing summarization task for this session key = id(qa_session_state)
if session_id in _summarization_tasks:
_summarization_tasks[session_id].cancel() # Cancel any existing summarization task for this state
if key in _summarization_tasks:
_summarization_tasks[key].cancel()
# Spawn background task # Spawn background task
task = asyncio.create_task( task = asyncio.create_task(
_update_context_background( _update_context_background(
qa_session_state=qa_session_state, qa_session_state=qa_session_state,
config=config, config=config,
session_id=session_id,
) )
) )
_summarization_tasks[session_id] = task _summarization_tasks[key] = task
task.add_done_callback( task.add_done_callback(lambda _t, k=key: _summarization_tasks.pop(k, None))
lambda _t, sid=session_id: _summarization_tasks.pop(sid, None)
)

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

@ -24,7 +24,6 @@ class SessionContext(BaseModel):
class ChatSessionState(BaseModel): class ChatSessionState(BaseModel):
"""State shared between frontend and agent via AG-UI.""" """State shared between frontend and agent via AG-UI."""
session_id: str = ""
initial_context: str | None = None initial_context: str | None = None
citations: list[Citation] = [] citations: list[Citation] = []
qa_history: list["QAHistoryEntry"] = [] qa_history: list["QAHistoryEntry"] = []
@ -56,12 +55,11 @@ def build_chat_state_snapshot(
Returns: Returns:
Snapshot dict. Snapshot dict.
""" """
snapshot: dict[str, Any] = {"session_id": ""} snapshot: dict[str, Any] = {}
if session_state is not None: if session_state is not None:
snapshot.update( snapshot.update(
{ {
"session_id": session_state.session_id,
"document_filter": session_state.document_filter.copy(), "document_filter": session_state.document_filter.copy(),
"citation_registry": session_state.citation_registry.copy(), "citation_registry": session_state.citation_registry.copy(),
"citations": [c.model_dump() for c in session_state.citations], "citations": [c.model_dump() for c in session_state.citations],

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

@ -169,7 +169,6 @@ class ChatApp(App):
# Keep ChatSessionState for UI state sync (used by _sync_session_state) # Keep ChatSessionState for UI state sync (used by _sync_session_state)
self.session_state = ChatSessionState( self.session_state = ChatSessionState(
session_id=str(uuid.uuid4()),
initial_context=self._initial_context, initial_context=self._initial_context,
document_filter=self._document_filter, document_filter=self._document_filter,
) )
@ -191,8 +190,6 @@ class ChatApp(App):
from haiku.rag.agents.research.models import Citation from haiku.rag.agents.research.models import Citation
# Update specific fields rather than replacing the entire state # 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: if "document_filter" in chat_state:
self.session_state.document_filter = chat_state["document_filter"] self.session_state.document_filter = chat_state["document_filter"]
if "citation_registry" in chat_state: if "citation_registry" in chat_state:
@ -322,7 +319,6 @@ class ChatApp(App):
# Sync session state to tool context before running # Sync session state to tool context before running
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState)
if session_state is not None: 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.document_filter = self.session_state.document_filter
session_state.citation_registry = self.session_state.citation_registry session_state.citation_registry = self.session_state.citation_registry
session_state.citations = list(self.session_state.citations) session_state.citations = list(self.session_state.citations)
@ -345,7 +341,7 @@ class ChatApp(App):
deps = ChatDeps( deps = ChatDeps(
config=self.config, config=self.config,
tool_context=self.tool_context, tool_context=self.tool_context,
session_id=self.session_state.session_id, is_new=False,
state_key=AGUI_STATE_KEY, state_key=AGUI_STATE_KEY,
) )
@ -414,7 +410,6 @@ class ChatApp(App):
# Reset context lock and session state (reset to CLI value) # Reset context lock and session state (reset to CLI value)
self._context_locked = False self._context_locked = False
self.session_state = ChatSessionState( self.session_state = ChatSessionState(
session_id=str(uuid.uuid4()),
initial_context=self._initial_context, initial_context=self._initial_context,
document_filter=self._document_filter, document_filter=self._document_filter,
) )

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

@ -3,11 +3,7 @@ import math
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from pydantic_ai import FunctionToolset, ToolReturn from pydantic_ai import FunctionToolset, ToolReturn
from haiku.rag.agents.chat.context import ( from haiku.rag.agents.chat.context import trigger_background_summarization
cache_question_embedding,
get_cached_embedding,
trigger_background_summarization,
)
from haiku.rag.agents.chat.state import ( from haiku.rag.agents.chat.state import (
build_chat_state_delta, build_chat_state_delta,
build_chat_state_snapshot, build_chat_state_snapshot,
@ -121,7 +117,6 @@ async def run_qa_core(
effective_session_context = qa_session_state.session_context effective_session_context = qa_session_state.session_context
effective_prior_answers = prior_answers or [] 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: if qa_session_state is not None and qa_session_state.qa_history:
embedder = get_embedder(config) embedder = get_embedder(config)
question_embedding = await embedder.embed_query(question) question_embedding = await embedder.embed_query(question)
@ -130,25 +125,13 @@ async def run_qa_core(
to_embed_indices = [] to_embed_indices = []
for i, qa in enumerate(qa_session_state.qa_history): for i, qa in enumerate(qa_session_state.qa_history):
if qa.question_embedding is None: if qa.question_embedding is None:
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.append(qa.question)
to_embed_indices.append(i) to_embed_indices.append(i)
if to_embed: if to_embed:
new_embeddings = await embedder.embed_documents(to_embed) new_embeddings = await embedder.embed_documents(to_embed)
for i, idx in enumerate(to_embed_indices): for i, idx in enumerate(to_embed_indices):
embedding = new_embeddings[i] qa_session_state.qa_history[idx].question_embedding = new_embeddings[i]
qa_session_state.qa_history[idx].question_embedding = embedding
if session_id:
cache_question_embedding(
session_id,
qa_session_state.qa_history[idx].question,
embedding,
)
matched_answers = [] matched_answers = []
for qa in qa_session_state.qa_history: for qa in qa_session_state.qa_history:
@ -225,7 +208,6 @@ async def run_qa_core(
trigger_background_summarization( trigger_background_summarization(
qa_session_state=qa_session_state, qa_session_state=qa_session_state,
config=config, config=config,
session_id=session_id,
) )
return qa_result return qa_result

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

@ -13,13 +13,11 @@ class SessionState(BaseModel):
"""Session-level state for AG-UI integration. """Session-level state for AG-UI integration.
This state is shared across toolsets and enables: This state is shared across toolsets and enables:
- Session identification
- Dynamic document filtering - Dynamic document filtering
- Stable citation indices across tool calls - Stable citation indices across tool calls
- AG-UI state synchronization - AG-UI state synchronization
""" """
session_id: str = ""
document_filter: list[str] = [] document_filter: list[str] = []
citation_registry: dict[str, int] = {} citation_registry: dict[str, int] = {}
citations: list[Citation] = [] citations: list[Citation] = []

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

@ -9,10 +9,7 @@ from haiku.rag.agents.chat import (
ChatSessionState, ChatSessionState,
create_chat_agent, create_chat_agent,
) )
from haiku.rag.agents.chat.context import ( from haiku.rag.agents.chat.context import _summarization_tasks
_summarization_tasks,
get_cached_session_context,
)
from haiku.rag.agents.research.models import Citation from haiku.rag.agents.research.models import Citation
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config from haiku.rag.config import Config
@ -38,9 +35,7 @@ def extract_state_from_result(result, state_key: str = AGUI_STATE_KEY) -> dict |
elif isinstance(meta, StateDeltaEvent): elif isinstance(meta, StateDeltaEvent):
# Apply delta to empty state to get final state # Apply delta to empty state to get final state
empty_state = { empty_state = {
state_key: ChatSessionState(session_id="").model_dump( state_key: ChatSessionState().model_dump(mode="json")
mode="json"
)
} }
patched = jsonpatch.apply_patch(empty_state, meta.delta) patched = jsonpatch.apply_patch(empty_state, meta.delta)
return patched.get(state_key) return patched.get(state_key)
@ -69,7 +64,7 @@ def test_chat_deps_initialization(temp_db_path):
assert deps.config is Config assert deps.config is Config
assert deps.tool_context is context assert deps.tool_context is context
assert deps.session_id == "" assert deps.is_new is True
assert deps.state_key is None assert deps.state_key is None
@ -109,7 +104,6 @@ def test_chat_deps_state_setter_handles_initial_context():
# Client sends initial_context with no session_context # Client sends initial_context with no session_context
incoming_state = { incoming_state = {
AGUI_STATE_KEY: { AGUI_STATE_KEY: {
"session_id": "",
"initial_context": "Background info about the project", "initial_context": "Background info about the project",
"session_context": None, "session_context": None,
"qa_history": [], "qa_history": [],
@ -140,7 +134,6 @@ def test_chat_deps_state_setter_parses_session_context_dict():
# Client sends session_context as a dict (as it comes from JSON) # Client sends session_context as a dict (as it comes from JSON)
incoming_state = { incoming_state = {
AGUI_STATE_KEY: { AGUI_STATE_KEY: {
"session_id": "test-session",
"session_context": { "session_context": {
"summary": "Previous conversation summary", "summary": "Previous conversation summary",
"last_updated": "2025-01-27T12:00:00", "last_updated": "2025-01-27T12:00:00",
@ -160,44 +153,9 @@ def test_chat_deps_state_setter_parses_session_context_dict():
assert qa_session_state.session_context == "Previous conversation summary" assert qa_session_state.session_context == "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(): def test_chat_session_state():
"""Test ChatSessionState model.""" """Test ChatSessionState model."""
state = ChatSessionState(session_id="test-session") state = ChatSessionState()
assert state.session_id == "test-session"
assert state.citations == [] assert state.citations == []
assert state.qa_history == [] assert state.qa_history == []
@ -356,7 +314,6 @@ async def test_chat_agent_search_tool(allow_model_requests, temp_db_path):
deps = ChatDeps( deps = ChatDeps(
config=Config, config=Config,
tool_context=context, tool_context=context,
session_id="test-search",
) )
# Ask something that should trigger the search tool # Ask something that should trigger the search tool
@ -391,7 +348,6 @@ async def test_chat_agent_search_tool_with_filter(allow_model_requests, temp_db_
deps = ChatDeps( deps = ChatDeps(
config=Config, config=Config,
tool_context=context, tool_context=context,
session_id="test-search-filter",
) )
# Ask to search within a specific document # Ask to search within a specific document
@ -503,10 +459,16 @@ async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path)
async def test_chat_agent_ask_triggers_background_summarization( async def test_chat_agent_ask_triggers_background_summarization(
allow_model_requests, temp_db_path allow_model_requests, temp_db_path
): ):
"""Test that the ask tool triggers background session context summarization.""" """Test that the ask tool triggers background session context summarization.
import asyncio
from haiku.rag.agents.chat.agent import run_chat_agent Patches the internal trigger in run_qa_core to avoid concurrent HTTP calls
that break VCR cassette replay ordering. Triggers summarization explicitly
after the agent run completes.
"""
from unittest.mock import patch
from haiku.rag.agents.chat.agent import trigger_background_summarization
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
async with HaikuRAG(temp_db_path, create=True) as client: async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document( await client.create_document(
@ -520,33 +482,31 @@ async def test_chat_agent_ask_triggers_background_summarization(
deps = ChatDeps( deps = ChatDeps(
config=Config, config=Config,
tool_context=context, tool_context=context,
session_id="test-summarization",
state_key=AGUI_STATE_KEY, state_key=AGUI_STATE_KEY,
) )
# Ask a question using run_chat_agent to trigger background summarization # Patch internal trigger to avoid concurrent HTTP calls during VCR
result = await run_chat_agent( with patch("haiku.rag.tools.qa.trigger_background_summarization"):
agent, result = await agent.run(
deps, "What is the highest count class in the DocLayNet dataset?",
"What is the highest count class in the DocLayNet dataset?", deps=deps,
) )
assert result is not None assert result.output is not None
# Trigger summarization explicitly (sequential, deterministic)
trigger_background_summarization(deps)
# Wait for background task to complete # Wait for background task to complete
# The task caches session_context server-side qa_session_state = context.get(QA_SESSION_NAMESPACE, QASessionState)
session_id = deps.session_id assert qa_session_state is not None
cached_context = None key = id(qa_session_state)
for _ in range(50): # Wait up to 5 seconds if key in _summarization_tasks:
cached_context = get_cached_session_context(session_id) await _summarization_tasks[key]
if cached_context is not None:
break
await asyncio.sleep(0.1)
# Verify session_context was populated by background task # Verify session_context was populated by background task
assert cached_context is not None assert qa_session_state.session_context is not None
assert cached_context.summary != "" assert qa_session_state.session_context != ""
assert cached_context.last_updated is not None
@pytest.mark.asyncio @pytest.mark.asyncio
@ -592,7 +552,6 @@ async def test_chat_agent_multi_turn_with_context(allow_model_requests, temp_db_
# Set initial state with initial_context (mimicking AG-UI client) # Set initial state with initial_context (mimicking AG-UI client)
deps.state = { deps.state = {
AGUI_STATE_KEY: { AGUI_STATE_KEY: {
"session_id": "",
"initial_context": "The user is researching the DocLayNet dataset for a paper on document layout analysis.", "initial_context": "The user is researching the DocLayNet dataset for a paper on document layout analysis.",
"session_context": None, "session_context": None,
"qa_history": [], "qa_history": [],
@ -602,10 +561,6 @@ async def test_chat_agent_multi_turn_with_context(allow_model_requests, temp_db_
} }
} }
# session_id should be auto-generated
assert deps.session_id != ""
session_id = deps.session_id
# initial_context should be transferred to QASessionState # initial_context should be transferred to QASessionState
qa_session = context.get(QA_SESSION_NAMESPACE, QASessionState) qa_session = context.get(QA_SESSION_NAMESPACE, QASessionState)
assert qa_session is not None assert qa_session is not None
@ -628,12 +583,12 @@ async def test_chat_agent_multi_turn_with_context(allow_model_requests, temp_db_
# Trigger summarization explicitly (sequential, no concurrency) # Trigger summarization explicitly (sequential, no concurrency)
trigger_background_summarization(deps) trigger_background_summarization(deps)
if session_id in _summarization_tasks: key = id(qa_session)
await _summarization_tasks[session_id] if key in _summarization_tasks:
await _summarization_tasks[key]
cached_context = get_cached_session_context(session_id) assert qa_session.session_context is not None
assert cached_context is not None assert qa_session.session_context != ""
assert cached_context.summary != ""
# qa_history should have one entry # qa_history should have one entry
qa_session = context.get(QA_SESSION_NAMESPACE, QASessionState) qa_session = context.get(QA_SESSION_NAMESPACE, QASessionState)
@ -653,18 +608,18 @@ async def test_chat_agent_multi_turn_with_context(allow_model_requests, temp_db_
# Trigger summarization explicitly # Trigger summarization explicitly
trigger_background_summarization(deps) trigger_background_summarization(deps)
if session_id in _summarization_tasks:
await _summarization_tasks[session_id]
# qa_history should have two entries
qa_session = context.get(QA_SESSION_NAMESPACE, QASessionState) qa_session = context.get(QA_SESSION_NAMESPACE, QASessionState)
assert qa_session is not None assert qa_session is not None
key = id(qa_session)
if key in _summarization_tasks:
await _summarization_tasks[key]
# qa_history should have two entries
assert len(qa_session.qa_history) >= 2 assert len(qa_session.qa_history) >= 2
# Session context should be updated with newer summary # Session context should be updated with newer summary
updated = get_cached_session_context(session_id) assert qa_session.session_context is not None
assert updated is not None assert qa_session.session_context != ""
assert updated.summary != ""
@pytest.mark.asyncio @pytest.mark.asyncio
@ -740,7 +695,6 @@ def test_fifo_limit_enforcement():
] ]
session_state = ChatSessionState( session_state = ChatSessionState(
session_id="test-fifo",
qa_history=qa_history, qa_history=qa_history,
) )
@ -759,7 +713,6 @@ def test_fifo_limit_enforcement():
def test_chat_session_state_document_filter(): def test_chat_session_state_document_filter():
"""Test ChatSessionState with document_filter.""" """Test ChatSessionState with document_filter."""
state = ChatSessionState( state = ChatSessionState(
session_id="test-filter",
document_filter=["doc1.pdf", "doc2.pdf"], document_filter=["doc1.pdf", "doc2.pdf"],
) )
assert state.document_filter == ["doc1.pdf", "doc2.pdf"] assert state.document_filter == ["doc1.pdf", "doc2.pdf"]
@ -767,7 +720,7 @@ def test_chat_session_state_document_filter():
def test_chat_session_state_document_filter_default_empty(): def test_chat_session_state_document_filter_default_empty():
"""Test ChatSessionState document_filter defaults to empty list.""" """Test ChatSessionState document_filter defaults to empty list."""
state = ChatSessionState(session_id="test") state = ChatSessionState()
assert state.document_filter == [] assert state.document_filter == []
@ -801,7 +754,6 @@ async def test_chat_agent_search_with_session_filter(
deps = ChatDeps( deps = ChatDeps(
config=Config, config=Config,
tool_context=context, tool_context=context,
session_id="test-session-filter",
) )
# Search should only return results from the filtered document # Search should only return results from the filtered document
@ -1032,7 +984,7 @@ def test_qa_response_embedding_default_none():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_summarization_task_cancellation(): async def test_summarization_task_cancellation():
"""Test that new summarization tasks cancel previous ones for same session.""" """Test that new summarization tasks cancel previous ones for same state object."""
import asyncio import asyncio
from haiku.rag.agents.chat.context import _summarization_tasks from haiku.rag.agents.chat.context import _summarization_tasks
@ -1040,7 +992,7 @@ async def test_summarization_task_cancellation():
# Clear any existing tasks # Clear any existing tasks
_summarization_tasks.clear() _summarization_tasks.clear()
session_id = "test-cancel-session" key = 12345 # Simulates id(qa_session_state)
# Create a slow task that simulates summarization # Create a slow task that simulates summarization
async def slow_task(): async def slow_task():
@ -1048,18 +1000,18 @@ async def test_summarization_task_cancellation():
# Start first task # Start first task
task1 = asyncio.create_task(slow_task()) task1 = asyncio.create_task(slow_task())
_summarization_tasks[session_id] = task1 _summarization_tasks[key] = task1
# Simulate what happens when second ask comes in - cancel first task # Simulate what happens when second ask comes in - cancel first task
if session_id in _summarization_tasks: if key in _summarization_tasks:
_summarization_tasks[session_id].cancel() _summarization_tasks[key].cancel()
# Yield to let cancellation propagate # Yield to let cancellation propagate
await asyncio.sleep(0) await asyncio.sleep(0)
# Start second task # Start second task
task2 = asyncio.create_task(slow_task()) task2 = asyncio.create_task(slow_task())
_summarization_tasks[session_id] = task2 _summarization_tasks[key] = task2
# First task should be cancelled # First task should be cancelled
assert task1.cancelled() or task1.done() assert task1.cancelled() or task1.done()
@ -1143,7 +1095,6 @@ async def test_list_documents_with_session_filter(allow_model_requests, temp_db_
deps = ChatDeps( deps = ChatDeps(
config=Config, config=Config,
tool_context=context, tool_context=context,
session_id="test-list-filter",
) )
# Ask to list documents - should only show filtered documents # Ask to list documents - should only show filtered documents
@ -1320,18 +1271,18 @@ async def test_summarization_task_cleanup_on_completion():
_summarization_tasks.clear() _summarization_tasks.clear()
session_id = "test-cleanup-session" key = 67890 # Simulates id(qa_session_state)
# Create a fast task # Create a fast task
async def fast_task(): async def fast_task():
await asyncio.sleep(0.01) await asyncio.sleep(0.01)
task = asyncio.create_task(fast_task()) task = asyncio.create_task(fast_task())
_summarization_tasks[session_id] = task _summarization_tasks[key] = task
task.add_done_callback(lambda t: _summarization_tasks.pop(session_id, None)) task.add_done_callback(lambda t: _summarization_tasks.pop(key, None))
# Wait for completion # Wait for completion
await task await task
# Task should be cleaned up # Task should be cleaned up
assert session_id not in _summarization_tasks assert key not in _summarization_tasks

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

@ -212,7 +212,6 @@ class TestUpdateSessionContext:
result = await update_session_context( result = await update_session_context(
qa_history=qa_history, qa_history=qa_history,
config=Config, config=Config,
session_id="test-session",
) )
assert result.summary != "" assert result.summary != ""
@ -231,129 +230,8 @@ class TestUpdateSessionContext:
assert result.summary == "" assert result.summary == ""
class TestSessionContextCache: class TestUpdateSessionContextPassesCurrentContext:
"""Tests for server-side session context caching.""" """Tests for update_session_context current_context forwarding."""
def test_cache_and_retrieve_session_context(self):
"""Test caching and retrieving a session context."""
from haiku.rag.agents.chat.context import (
_session_cache,
cache_session_context,
get_cached_session_context,
)
# Clear cache
_session_cache.clear()
now = datetime.now()
ctx = SessionContext(summary="Test summary", last_updated=now)
cache_session_context("session-1", ctx)
result = get_cached_session_context("session-1")
assert result is not None
assert result.summary == "Test summary"
assert result.last_updated == now
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_cache,
get_cached_session_context,
)
_session_cache.clear()
result = get_cached_session_context("nonexistent-session")
assert result is None
def test_cache_ttl_cleanup_removes_stale_entries(self):
"""Test that stale cache entries are cleaned up."""
from datetime import timedelta
from haiku.rag.agents.chat.context import (
_CACHE_TTL,
_cache_timestamps,
_session_cache,
cache_session_context,
get_cached_session_context,
)
_session_cache.clear()
_cache_timestamps.clear()
# Add an entry
ctx = SessionContext(summary="Old summary", last_updated=datetime.now())
cache_session_context("stale-session", ctx)
# Make the entry stale by backdating its timestamp
_cache_timestamps["stale-session"] = (
datetime.now() - _CACHE_TTL - timedelta(seconds=1)
)
# Getting session context should trigger cleanup
result = get_cached_session_context("stale-session")
# Should be None because the entry was cleaned up
assert result is None
assert "stale-session" not in _session_cache
@pytest.mark.asyncio
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_cache,
get_cached_session_context,
update_session_context,
)
_session_cache.clear()
qa_history = [
QAHistoryEntry(
question="What is Python?",
answer="A programming language.",
confidence=0.95,
)
]
with patch(
"haiku.rag.agents.chat.context.summarize_session",
new=AsyncMock(return_value="Mocked summary"),
):
result = await update_session_context(
qa_history=qa_history,
config=Config,
session_id="cache-test-session",
)
assert result.summary == "Mocked summary"
cached = get_cached_session_context("cache-test-session")
assert cached is not None
assert cached.summary == "Mocked summary"
@pytest.mark.asyncio
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_cache,
get_cached_session_context,
update_session_context,
)
_session_cache.clear()
await update_session_context(
qa_history=[],
config=Config,
)
# Nothing should be cached (no session_id)
cached = get_cached_session_context("")
assert cached is None
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_update_session_context_passes_current_context(self): async def test_update_session_context_passes_current_context(self):

3
tests/agents/chat/test_features.py
View file

@ -107,10 +107,9 @@ def test_chat_deps_state_without_qa(temp_db_path):
context = ToolContext() context = ToolContext()
create_chat_agent(Config, client, context, features=[FEATURE_SEARCH]) create_chat_agent(Config, client, context, features=[FEATURE_SEARCH])
deps = ChatDeps(config=Config, tool_context=context, session_id="test") deps = ChatDeps(config=Config, tool_context=context)
state = deps.state state = deps.state
assert "session_id" in state
# SessionState fields should be present # SessionState fields should be present
assert "document_filter" in state assert "document_filter" in state
assert "citation_registry" in state assert "citation_registry" in state

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

@ -57,7 +57,7 @@ def test_citation_registry_stability():
def test_citation_registry_serialization_roundtrip(): def test_citation_registry_serialization_roundtrip():
"""Test citation_registry serializes and deserializes correctly for AG-UI state.""" """Test citation_registry serializes and deserializes correctly for AG-UI state."""
# Create state and assign indices # Create state and assign indices
original = ChatSessionState(session_id="test") original = ChatSessionState()
original.citation_registry = {"chunk-a": 1, "chunk-b": 2} original.citation_registry = {"chunk-a": 1, "chunk-b": 2}
# Serialize # Serialize
@ -70,22 +70,6 @@ def test_citation_registry_serialization_roundtrip():
assert restored.citation_registry == {"chunk-a": 1, "chunk-b": 2} assert restored.citation_registry == {"chunk-a": 1, "chunk-b": 2}
def test_chat_session_state_defaults_to_empty_session_id():
"""New ChatSessionState should default to empty session_id.
Tools in agent.py detect the empty string and assign a UUID,
which then appears in the state delta so clients receive it.
"""
state = ChatSessionState()
assert state.session_id == ""
def test_chat_session_state_preserves_explicit_session_id():
"""Explicit session_id should be preserved."""
state = ChatSessionState(session_id="my-custom-id")
assert state.session_id == "my-custom-id"
def test_chat_session_state_initial_context_default_none(): def test_chat_session_state_initial_context_default_none():
"""Initial context should default to None.""" """Initial context should default to None."""
state = ChatSessionState() state = ChatSessionState()
@ -101,7 +85,6 @@ def test_chat_session_state_initial_context_preserved():
def test_chat_session_state_initial_context_serialization(): def test_chat_session_state_initial_context_serialization():
"""initial_context should serialize and deserialize correctly.""" """initial_context should serialize and deserialize correctly."""
state = ChatSessionState( state = ChatSessionState(
session_id="test-123",
initial_context="User is working on authentication", initial_context="User is working on authentication",
) )
state_dict = state.model_dump() state_dict = state.model_dump()
@ -121,7 +104,6 @@ def test_chat_session_state_model_dump_json_serializes_datetime():
from datetime import datetime from datetime import datetime
session_state = ChatSessionState( session_state = ChatSessionState(
session_id="test",
session_context=SessionContext( session_context=SessionContext(
summary="Test summary", summary="Test summary",
last_updated=datetime(2025, 1, 27, 12, 0, 0), last_updated=datetime(2025, 1, 27, 12, 0, 0),

30
tests/chat/test_chat_app.py
View file

@ -245,25 +245,9 @@ async def test_chat_history_thinking_indicator(temp_db_path: Path):
assert len(list(thinking)) == 0 assert len(list(thinking)) == 0
@pytest.mark.asyncio
async def test_chat_app_generates_session_id(temp_db_path: Path):
"""Test that ChatApp generates a UUID session_id on mount."""
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():
assert app.session_state.session_id != ""
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_clear_chat_resets_session(temp_db_path: Path): async def test_clear_chat_resets_session(temp_db_path: Path):
"""Test that clearing chat resets the session state with a new session_id.""" """Test that clearing chat resets the session state."""
from haiku.rag.chat.app import ChatApp from haiku.rag.chat.app import ChatApp
from haiku.rag.chat.widgets.chat_history import ChatHistory from haiku.rag.chat.widgets.chat_history import ChatHistory
@ -282,10 +266,6 @@ async def test_clear_chat_resets_session(temp_db_path: Path):
await chat_history.add_message("assistant", "Hi there") await chat_history.add_message("assistant", "Hi there")
assert len(chat_history.messages) == 2 assert len(chat_history.messages) == 2
# Record the original session_id
original_session_id = app.session_state.session_id
assert original_session_id != ""
# Clear chat via action (available through command palette) # Clear chat via action (available through command palette)
await app.action_clear_chat() await app.action_clear_chat()
await pilot.pause() await pilot.pause()
@ -293,10 +273,8 @@ async def test_clear_chat_resets_session(temp_db_path: Path):
# Verify messages cleared # Verify messages cleared
assert len(chat_history.messages) == 0 assert len(chat_history.messages) == 0
# Verify session state reset with a new session_id # Verify session state reset
assert app.session_state is not None assert app.session_state is not None
assert app.session_state.session_id != ""
assert app.session_state.session_id != original_session_id
assert app.session_state.qa_history == [] assert app.session_state.qa_history == []
assert app.session_state.citations == [] assert app.session_state.citations == []
@ -326,7 +304,6 @@ async def test_handle_stream_event_extracts_citations_from_state_snapshot(
type=EventType.STATE_SNAPSHOT, type=EventType.STATE_SNAPSHOT,
snapshot={ snapshot={
AGUI_STATE_KEY: { AGUI_STATE_KEY: {
"session_id": "test",
"citations": [ "citations": [
{ {
"index": 1, "index": 1,
@ -392,7 +369,6 @@ async def test_handle_stream_event_extracts_citations_from_state_delta(
type=EventType.STATE_SNAPSHOT, type=EventType.STATE_SNAPSHOT,
snapshot={ snapshot={
AGUI_STATE_KEY: { AGUI_STATE_KEY: {
"session_id": "test",
"citations": [], "citations": [],
"qa_history": [], "qa_history": [],
"citation_registry": {}, "citation_registry": {},
@ -479,7 +455,6 @@ async def test_handle_stream_event_delta_with_preinitialized_state(
# Pre-initialize _agui_state_snapshot (simulating what _run_agent does) # Pre-initialize _agui_state_snapshot (simulating what _run_agent does)
app._agui_state_snapshot = { app._agui_state_snapshot = {
AGUI_STATE_KEY: { AGUI_STATE_KEY: {
"session_id": "test",
"citations": [], "citations": [],
"qa_history": [], "qa_history": [],
"citation_registry": {}, "citation_registry": {},
@ -552,7 +527,6 @@ async def test_handle_stream_event_syncs_session_context(temp_db_path: Path):
# Pre-initialize state # Pre-initialize state
app._agui_state_snapshot = { app._agui_state_snapshot = {
AGUI_STATE_KEY: { AGUI_STATE_KEY: {
"session_id": "test",
"citations": [], "citations": [],
"qa_history": [], "qa_history": [],
"citation_registry": {}, "citation_registry": {},