From 9d71ccb21390cfcdb5d4797b2ac0ac4302690636 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 23 Jan 2026 16:46:51 +0200 Subject: [PATCH] Move session context to agent, make AGUI work --- app/backend/main.py | 18 +-- .../haiku/rag/agents/chat/__init__.py | 5 +- haiku_rag_slim/haiku/rag/agents/chat/agent.py | 60 ++++---- .../haiku/rag/agents/chat/context.py | 40 +++++- haiku_rag_slim/haiku/rag/agents/chat/state.py | 14 +- haiku_rag_slim/haiku/rag/chat/app.py | 13 ++ tests/agents/chat/test_context.py | 134 ++++++++++++++++++ tests/agents/chat/test_state.py | 58 ++------ 8 files changed, 244 insertions(+), 98 deletions(-) diff --git a/app/backend/main.py b/app/backend/main.py index 3ad29691..76944324 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -1,5 +1,6 @@ import logging import os +import uuid from pathlib import Path from dotenv import find_dotenv, load_dotenv @@ -17,7 +18,6 @@ from haiku.rag.agents.chat import ( ChatDeps, ChatSessionState, QAResponse, - SessionContext, create_chat_agent, ) from haiku.rag.client import HaikuRAG @@ -75,7 +75,6 @@ def get_client(effective_db_path: Path) -> HaikuRAG: async def stream_chat(request: Request) -> Response: """Chat streaming endpoint with AG-UI protocol.""" body = await request.body() - logger.info("Received chat request") # Parse request to build run_input accept = request.headers.get("accept", SSE_CONTENT_TYPE) @@ -84,7 +83,7 @@ async def stream_chat(request: Request) -> Response: # Restore session state from incoming AG-UI state (look under namespaced key) initial_qa_history: list[QAResponse] = [] background_context: str | None = None - session_context: SessionContext | None = None + session_id: str | None = None state = getattr(run_input, "state", None) if state: chat_state = state.get(AGUI_STATE_KEY, state) @@ -93,20 +92,21 @@ async def stream_chat(request: Request) -> Response: QAResponse(**qa) for qa in chat_state.get("qa_history", []) ] background_context = chat_state.get("background_context") - ctx_data = chat_state.get("session_context") - if ctx_data and isinstance(ctx_data, dict): - session_context = SessionContext(**ctx_data) + session_id = chat_state.get("session_id") - # Build deps with session state + # Determine session_id: prefer state, fall back to thread_id, generate UUID if neither thread_id = getattr(run_input, "thread_id", None) + if not session_id: + session_id = thread_id or str(uuid.uuid4()) + deps = ChatDeps( client=get_client(db_path), config=Config, session_state=ChatSessionState( - session_id=thread_id or "", + session_id=session_id, qa_history=initial_qa_history, background_context=background_context, - session_context=session_context, + # session_context intentionally NOT set - agent will fetch from cache ), state_key=AGUI_STATE_KEY, ) diff --git a/haiku_rag_slim/haiku/rag/agents/chat/__init__.py b/haiku_rag_slim/haiku/rag/agents/chat/__init__.py index c462ebfa..5ffd5116 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/__init__.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/__init__.py @@ -1,5 +1,8 @@ from haiku.rag.agents.chat.agent import create_chat_agent -from haiku.rag.agents.chat.context import summarize_session, update_session_context +from haiku.rag.agents.chat.context import ( + summarize_session, + update_session_context, +) from haiku.rag.agents.chat.search import SearchAgent from haiku.rag.agents.chat.state import ( AGUI_STATE_KEY, diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index d0862c18..f109e727 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py @@ -1,10 +1,12 @@ import asyncio -import logging from ag_ui.core import EventType, StateSnapshotEvent from pydantic_ai import Agent, RunContext, ToolReturn -from haiku.rag.agents.chat.context import update_session_context +from haiku.rag.agents.chat.context import ( + get_cached_session_context, + update_session_context, +) from haiku.rag.agents.chat.prompts import CHAT_SYSTEM_PROMPT from haiku.rag.agents.chat.search import SearchAgent from haiku.rag.agents.chat.state import ( @@ -21,8 +23,6 @@ from haiku.rag.agents.research.state import ResearchDeps, ResearchState from haiku.rag.config.models import AppConfig from haiku.rag.utils import get_model -logger = logging.getLogger(__name__) - # Track summarization tasks per session to allow cancellation _summarization_tasks: dict[str, asyncio.Task[None]] = {} @@ -39,11 +39,10 @@ async def _update_context_background( config=config, session_state=session_state, ) - logger.debug("Session context updated") except asyncio.CancelledError: - logger.debug("Session context summarization cancelled") + pass except Exception: - logger.exception("Failed to update session context") + pass def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: @@ -111,10 +110,9 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: ] # Build new state with citations + session_id = ctx.deps.session_state.session_id if ctx.deps.session_state else "" new_state = ChatSessionState( - session_id=( - ctx.deps.session_state.session_id if ctx.deps.session_state else "" - ), + session_id=session_id, citations=citation_infos, qa_history=( ctx.deps.session_state.qa_history if ctx.deps.session_state else [] @@ -124,11 +122,9 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: if ctx.deps.session_state else None ), - session_context=( - ctx.deps.session_state.session_context - if ctx.deps.session_state - else None - ), + session_context=get_cached_session_context(session_id) + if session_id + else None, ) # Return detailed results for the agent to present @@ -183,19 +179,20 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: graph = build_conversational_graph(config=ctx.deps.config) # Determine context strategy: - # 1. If session_context exists, use compressed summary (skip raw qa_history) - # 2. Otherwise, fall back to explicit background_context + # 1. Read from server cache (ignoring client state) + # 2. Fall back to explicit background_context (first request) background_context: str | None = None + session_id = ctx.deps.session_state.session_id if ctx.deps.session_state else "" if ctx.deps.session_state: - if ( - ctx.deps.session_state.session_context - and ctx.deps.session_state.session_context.summary - ): - # Use compressed SessionContext - no need for raw qa_history - background_context = ( - ctx.deps.session_state.session_context.render_markdown() - ) + cached_context = ( + get_cached_session_context(session_id) if session_id else None + ) + + if cached_context and cached_context.summary: + # Use cached SessionContext from previous summarization + background_context = cached_context.render_markdown() elif ctx.deps.session_state.background_context: + # Fall back to explicit background_context (first request) background_context = ctx.deps.session_state.background_context context = ResearchContext( @@ -247,7 +244,6 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: # Spawn background task to update session context # Cancel any previous summarization for this session - session_id = ctx.deps.session_state.session_id if session_id in _summarization_tasks: _summarization_tasks[session_id].cancel() @@ -263,9 +259,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: # Build new state with citations AND accumulated qa_history new_state = ChatSessionState( - session_id=( - ctx.deps.session_state.session_id if ctx.deps.session_state else "" - ), + session_id=session_id, citations=citation_infos, qa_history=( ctx.deps.session_state.qa_history if ctx.deps.session_state else [] @@ -275,11 +269,9 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: if ctx.deps.session_state else None ), - session_context=( - ctx.deps.session_state.session_context - if ctx.deps.session_state - else None - ), + session_context=get_cached_session_context(session_id) + if session_id + else None, ) # Format answer with citation references and confidence diff --git a/haiku_rag_slim/haiku/rag/agents/chat/context.py b/haiku_rag_slim/haiku/rag/agents/chat/context.py index 473cd168..42b7ce3a 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/context.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/context.py @@ -1,4 +1,4 @@ -from datetime import datetime +from datetime import datetime, timedelta from pydantic_ai import Agent @@ -7,6 +7,41 @@ from haiku.rag.agents.chat.state import ChatSessionState, QAResponse, SessionCon from haiku.rag.config.models import AppConfig from haiku.rag.utils import get_model +# Cache for session contexts (session_id -> SessionContext) +# Used to persist async summarization results between requests +_session_context_cache: dict[str, SessionContext] = {} +_cache_timestamps: dict[str, datetime] = {} +_CACHE_TTL = timedelta(hours=1) + + +def _cleanup_stale_cache() -> None: + """Remove cache entries older than TTL.""" + now = datetime.now() + stale = [sid for sid, ts in _cache_timestamps.items() if now - ts > _CACHE_TTL] + for sid in stale: + _session_context_cache.pop(sid, None) + _cache_timestamps.pop(sid, None) + + +def cache_session_context(session_id: str, context: SessionContext) -> None: + """Store session context in cache.""" + _cleanup_stale_cache() + _session_context_cache[session_id] = context + _cache_timestamps[session_id] = datetime.now() + + +def get_cached_session_context(session_id: str) -> SessionContext | None: + """Get session context from server cache. + + Args: + session_id: The session identifier. + + Returns: + Cached SessionContext, or None if not cached. + """ + _cleanup_stale_cache() + return _session_context_cache.get(session_id) + async def summarize_session( qa_history: list[QAResponse], @@ -70,6 +105,9 @@ async def update_session_context( summary=summary, last_updated=datetime.now(), ) + # Also cache for next-run delivery in stateless contexts + if session_state.session_id: + cache_session_context(session_state.session_id, session_state.session_context) def _format_qa_history(qa_history: list[QAResponse]) -> str: diff --git a/haiku_rag_slim/haiku/rag/agents/chat/state.py b/haiku_rag_slim/haiku/rag/agents/chat/state.py index 320c4fd4..456fafa4 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/state.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/state.py @@ -101,16 +101,10 @@ class ChatDeps: self.session_state.background_context = state_data.get( "background_context" ) - if "session_id" in state_data: - self.session_state.session_id = state_data.get("session_id", "") - if "session_context" in state_data: - ctx_data = state_data.get("session_context") - if ctx_data is None: - self.session_state.session_context = None - elif isinstance(ctx_data, dict): - self.session_state.session_context = SessionContext(**ctx_data) - else: - self.session_state.session_context = ctx_data + if state_data.get("session_id"): + self.session_state.session_id = state_data["session_id"] + # NOTE: session_context intentionally NOT updated from client + # The agent owns this via server-side cache @dataclass diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 7b075503..338153ed 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -21,6 +21,7 @@ from haiku.rag.agents.chat.state import ( AGUI_STATE_KEY, ChatDeps, ChatSessionState, + SessionContext, ) from haiku.rag.agents.research.models import Citation from haiku.rag.client import HaikuRAG @@ -128,9 +129,15 @@ class ChatApp(App): # Create agent and session state self.agent = create_chat_agent(self.config) + initial_context = ( + SessionContext(summary=self.background_context, last_updated=datetime.now()) + if self.background_context + else None + ) self.session_state = ChatSessionState( session_id=str(uuid.uuid4()), background_context=self.background_context, + session_context=initial_context, ) # Focus the input field @@ -277,9 +284,15 @@ class ChatApp(App): self._selected_citation_idx = None self._message_history.clear() # Reset session state for fresh conversation (preserve background_context) + initial_context = ( + SessionContext(summary=self.background_context, last_updated=datetime.now()) + if self.background_context + else None + ) self.session_state = ChatSessionState( session_id=str(uuid.uuid4()), background_context=self.background_context, + session_context=initial_context, ) def action_focus_input(self) -> None: diff --git a/tests/agents/chat/test_context.py b/tests/agents/chat/test_context.py index e5f8ed71..496cac8a 100644 --- a/tests/agents/chat/test_context.py +++ b/tests/agents/chat/test_context.py @@ -254,3 +254,137 @@ class TestUpdateSessionContext: # session_context should exist but have empty summary assert session_state.session_context is not None assert session_state.session_context.summary == "" + + +class TestSessionContextCache: + """Tests for server-side session context caching.""" + + def test_cache_and_retrieve_session_context(self): + """Test caching and retrieving a session context.""" + from haiku.rag.agents.chat.context import ( + _session_context_cache, + cache_session_context, + get_cached_session_context, + ) + + # Clear cache + _session_context_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_context_cache, + get_cached_session_context, + ) + + _session_context_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_context_cache, + cache_session_context, + get_cached_session_context, + ) + + _session_context_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_context_cache + + @pytest.mark.asyncio + async def test_update_session_context_caches_result(self): + """Test update_session_context stores result in cache.""" + from unittest.mock import AsyncMock, patch + + from haiku.rag.agents.chat.context import ( + _session_context_cache, + get_cached_session_context, + update_session_context, + ) + from haiku.rag.agents.chat.state import ChatSessionState + + _session_context_cache.clear() + + session_state = ChatSessionState(session_id="cache-test-session") + + qa_history = [ + QAResponse( + question="What is Python?", + answer="A programming language.", + confidence=0.95, + ) + ] + + # Mock summarize_session to avoid LLM call (we're testing caching, not summarization) + with patch( + "haiku.rag.agents.chat.context.summarize_session", + new=AsyncMock(return_value="Mocked summary"), + ): + await update_session_context( + qa_history=qa_history, + config=Config, + session_state=session_state, + ) + + # Should be cached + cached = get_cached_session_context("cache-test-session") + assert cached is not None + assert cached.summary == "Mocked summary" + assert cached.summary == session_state.session_context.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_context_cache, + get_cached_session_context, + update_session_context, + ) + from haiku.rag.agents.chat.state import ChatSessionState + + _session_context_cache.clear() + + # No session_id + session_state = ChatSessionState() + + await update_session_context( + qa_history=[], + config=Config, + session_state=session_state, + ) + + # Nothing should be cached (empty session_id) + cached = get_cached_session_context("") + assert cached is None diff --git a/tests/agents/chat/test_state.py b/tests/agents/chat/test_state.py index e54aa631..3fa0277b 100644 --- a/tests/agents/chat/test_state.py +++ b/tests/agents/chat/test_state.py @@ -297,45 +297,12 @@ def test_chat_deps_state_getter_includes_session_context(): ) -def test_chat_deps_state_setter_restores_session_context(): - """Test ChatDeps.state setter restores session_context from incoming state.""" - from unittest.mock import MagicMock +def test_chat_deps_state_setter_ignores_session_context(): + """Test ChatDeps.state setter ignores session_context from client. - from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState - - mock_client = MagicMock() - mock_config = MagicMock() - - session_state = ChatSessionState(session_id="initial") - deps = ChatDeps( - client=mock_client, - config=mock_config, - session_state=session_state, - state_key=AGUI_STATE_KEY, - ) - - incoming_state = { - AGUI_STATE_KEY: { - "session_id": "test-123", - "qa_history": [], - "citations": [], - "background_context": None, - "session_context": { - "summary": "Restored context summary.", - "last_updated": "2025-01-15T10:30:00", - }, - } - } - - deps.state = incoming_state - - assert deps.session_state is not None - assert deps.session_state.session_context is not None - assert deps.session_state.session_context.summary == "Restored context summary." - - -def test_chat_deps_state_setter_handles_null_session_context(): - """Test ChatDeps.state setter handles null session_context.""" + The agent owns session_context via server-side cache, so client-provided + session_context should be ignored to prevent stale state overwriting. + """ from unittest.mock import MagicMock from haiku.rag.agents.chat.state import ( @@ -348,10 +315,10 @@ def test_chat_deps_state_setter_handles_null_session_context(): mock_client = MagicMock() mock_config = MagicMock() - # Start with a session_context + # Start with a session_context (e.g., from cache) session_state = ChatSessionState( session_id="test", - session_context=SessionContext(summary="Initial summary"), + session_context=SessionContext(summary="Server-side context"), ) deps = ChatDeps( client=mock_client, @@ -360,17 +327,22 @@ def test_chat_deps_state_setter_handles_null_session_context(): state_key=AGUI_STATE_KEY, ) - # Send null to clear it + # Client sends different session_context (stale) incoming_state = { AGUI_STATE_KEY: { "session_id": "test", "qa_history": [], "citations": [], - "session_context": None, + "session_context": { + "summary": "Client-provided stale context", + "last_updated": "2025-01-15T10:30:00", + }, } } deps.state = incoming_state + # session_context should NOT be overwritten assert deps.session_state is not None - assert deps.session_state.session_context is None + assert deps.session_state.session_context is not None + assert deps.session_state.session_context.summary == "Server-side context"