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

Move session context to agent, make AGUI work

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-01-23 16:46:51 +02:00
parent 93737852a8
commit 9d71ccb213
No known key found for this signature in database
8 changed files with 244 additions and 98 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
app/backend/main.py
View file

@ -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,
)

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

@ -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,

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

@ -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

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

@ -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:

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

@ -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

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

@ -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:

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

@ -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

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

@ -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"