From d67df09cce88d5c0dd644482ebc485c9d0177738 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 11 Feb 2026 14:18:26 +0200 Subject: [PATCH] Introduce ToolContextCache --- haiku_rag_slim/haiku/rag/agents/chat/agent.py | 6 +- haiku_rag_slim/haiku/rag/tools/__init__.py | 3 +- haiku_rag_slim/haiku/rag/tools/context.py | 47 ++++++++ haiku_rag_slim/haiku/rag/tools/qa.py | 4 +- haiku_rag_slim/haiku/rag/tools/search.py | 8 +- haiku_rag_slim/haiku/rag/tools/session.py | 5 +- tests/tools/test_context.py | 111 ++++++++++++++++++ 7 files changed, 175 insertions(+), 9 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index c380cee7..ef0aa486 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py @@ -180,9 +180,9 @@ def create_chat_agent( existing = context.get(SESSION_NAMESPACE, SessionState) if existing is None: - context.register(SESSION_NAMESPACE, SessionState(state_key=AGUI_STATE_KEY)) - elif existing.state_key is None: - existing.state_key = AGUI_STATE_KEY + context.register(SESSION_NAMESPACE, SessionState()) + if context.state_key is None: + context.state_key = AGUI_STATE_KEY if FEATURE_QA in features: if context.get(QA_SESSION_NAMESPACE, QASessionState) is None: diff --git a/haiku_rag_slim/haiku/rag/tools/__init__.py b/haiku_rag_slim/haiku/rag/tools/__init__.py index 633cd7b1..aa2a80cb 100644 --- a/haiku_rag_slim/haiku/rag/tools/__init__.py +++ b/haiku_rag_slim/haiku/rag/tools/__init__.py @@ -1,5 +1,5 @@ from haiku.rag.tools.analysis import create_analysis_toolset -from haiku.rag.tools.context import ToolContext +from haiku.rag.tools.context import ToolContext, ToolContextCache from haiku.rag.tools.document import ( DocumentInfo, DocumentListResponse, @@ -30,6 +30,7 @@ from haiku.rag.tools.session import ( __all__ = [ "ToolContext", + "ToolContextCache", "QAResult", "AnalysisResult", "build_document_filter", diff --git a/haiku_rag_slim/haiku/rag/tools/context.py b/haiku_rag_slim/haiku/rag/tools/context.py index 9e9dc2d8..782b6a1c 100644 --- a/haiku_rag_slim/haiku/rag/tools/context.py +++ b/haiku_rag_slim/haiku/rag/tools/context.py @@ -1,3 +1,4 @@ +from datetime import datetime, timedelta from typing import Any, TypeVar, overload from pydantic import BaseModel, PrivateAttr @@ -44,6 +45,7 @@ class ToolContext(BaseModel): ns_data = context.dump_namespaces() """ + state_key: str | None = None _namespaces: dict[str, BaseModel] = PrivateAttr(default_factory=dict) def register(self, namespace: str, state: BaseModel) -> None: @@ -126,3 +128,48 @@ class ToolContext(BaseModel): state = state_type.model_validate(data) self._namespaces[namespace] = state return state + + +class ToolContextCache: + """In-memory cache for ToolContext instances, keyed by external session/thread ID.""" + + def __init__(self, ttl: timedelta = timedelta(hours=1)) -> None: + self._cache: dict[str, ToolContext] = {} + self._timestamps: dict[str, datetime] = {} + self._ttl = ttl + + def get_or_create(self, key: str) -> tuple[ToolContext, bool]: + """Get an existing context or create a new one. + + Returns: + Tuple of (context, is_new) where is_new is True if a new context was created. + """ + self._cleanup() + if key in self._cache: + self._timestamps[key] = datetime.now() + return self._cache[key], False + + context = ToolContext() + self._cache[key] = context + self._timestamps[key] = datetime.now() + return context, True + + def remove(self, key: str) -> None: + """Remove a specific key from the cache.""" + self._cache.pop(key, None) + self._timestamps.pop(key, None) + + def clear(self) -> None: + """Clear all entries.""" + self._cache.clear() + self._timestamps.clear() + + def _cleanup(self) -> None: + """Remove entries older than TTL.""" + now = datetime.now() + expired = [ + key for key, ts in self._timestamps.items() if (now - ts) >= self._ttl + ] + for key in expired: + self._cache.pop(key, None) + self._timestamps.pop(key, None) diff --git a/haiku_rag_slim/haiku/rag/tools/qa.py b/haiku_rag_slim/haiku/rag/tools/qa.py index 72f62e6b..20a77cbd 100644 --- a/haiku_rag_slim/haiku/rag/tools/qa.py +++ b/haiku_rag_slim/haiku/rag/tools/qa.py @@ -281,10 +281,12 @@ def create_qa_toolset( session_state: SessionState | None = None qa_session_state: QASessionState | None = None old_state_snapshot: dict | None = None + state_key: str | None = None if context is not None: session_state = context.get(SESSION_NAMESPACE, SessionState) qa_session_state = context.get(QA_SESSION_NAMESPACE, QASessionState) + state_key = context.state_key # Use incoming values (what client sent) so delta shows server-side updates if session_state is not None: @@ -315,7 +317,7 @@ def create_qa_toolset( state_event = build_chat_state_delta( old_state_snapshot, new_state_snapshot, - state_key=session_state.state_key, + state_key=state_key, ) answer_text = qa_result.answer diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index cdca8e16..a019c35b 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/haiku_rag_slim/haiku/rag/tools/search.py @@ -68,8 +68,10 @@ def create_search_toolset( """ session_state: SessionState | None = None old_session_state: SessionState | None = None + state_key: str | None = None if context is not None: session_state = context.get(SESSION_NAMESPACE, SessionState) + state_key = context.state_key if session_state is not None: old_session_state = session_state.model_copy(deep=True) @@ -130,7 +132,11 @@ def create_search_toolset( formatted = f"Found {len(results)} results:\n\n" + "\n\n".join(result_lines) if old_session_state is not None: - state_event = compute_state_delta(old_session_state, session_state) + state_event = compute_state_delta( + old_session_state, + session_state, + state_key=state_key, + ) if state_event is not None: return ToolReturn( return_value=formatted, diff --git a/haiku_rag_slim/haiku/rag/tools/session.py b/haiku_rag_slim/haiku/rag/tools/session.py index 4777a59c..46ea9f01 100644 --- a/haiku_rag_slim/haiku/rag/tools/session.py +++ b/haiku_rag_slim/haiku/rag/tools/session.py @@ -24,7 +24,6 @@ class SessionState(BaseModel): document_filter: list[str] = [] citation_registry: dict[str, int] = {} citations: list[Citation] = [] - state_key: str | None = Field(default=None, exclude=True) def get_or_assign_index(self, chunk_id: str) -> int: """Get or assign a stable citation index for a chunk_id. @@ -44,16 +43,16 @@ class SessionState(BaseModel): def compute_state_delta( old_state: SessionState, new_state: SessionState, + state_key: str | None = None, ) -> StateDeltaEvent | None: """Compute state delta between old and new session state. Returns a StateDeltaEvent if there are changes, None otherwise. - The state_key from new_state is used for namespacing. """ return compute_combined_state_delta( old_state.model_dump(mode="json"), new_state.model_dump(mode="json"), - state_key=new_state.state_key, + state_key=state_key, ) diff --git a/tests/tools/test_context.py b/tests/tools/test_context.py index 9392e98a..de2b6851 100644 --- a/tests/tools/test_context.py +++ b/tests/tools/test_context.py @@ -201,3 +201,114 @@ def test_get_without_type(): result = ctx.get("ns") assert result is state + + +def test_tool_context_state_key_default_none(): + """Test ToolContext state_key defaults to None.""" + ctx = ToolContext() + assert ctx.state_key is None + + +def test_tool_context_state_key_set(): + """Test ToolContext state_key can be set.""" + ctx = ToolContext() + ctx.state_key = "haiku.rag.chat" + assert ctx.state_key == "haiku.rag.chat" + + +# ============================================================================= +# ToolContextCache Tests +# ============================================================================= + + +def test_tool_context_cache_get_or_create_new(): + """Test get_or_create returns a new context with is_new=True.""" + from haiku.rag.tools.context import ToolContextCache + + cache = ToolContextCache() + context, is_new = cache.get_or_create("thread-1") + + assert isinstance(context, ToolContext) + assert is_new is True + + +def test_tool_context_cache_get_or_create_existing(): + """Test get_or_create returns existing context with is_new=False.""" + from haiku.rag.tools.context import ToolContextCache + + cache = ToolContextCache() + ctx1, is_new1 = cache.get_or_create("thread-1") + ctx2, is_new2 = cache.get_or_create("thread-1") + + assert ctx2 is ctx1 + assert is_new1 is True + assert is_new2 is False + + +def test_tool_context_cache_different_keys(): + """Test get_or_create returns different contexts for different keys.""" + from haiku.rag.tools.context import ToolContextCache + + cache = ToolContextCache() + ctx1, _ = cache.get_or_create("thread-1") + ctx2, _ = cache.get_or_create("thread-2") + + assert ctx1 is not ctx2 + + +def test_tool_context_cache_ttl_expiry(): + """Test that contexts are evicted after TTL expires.""" + from datetime import timedelta + + from haiku.rag.tools.context import ToolContextCache + + cache = ToolContextCache(ttl=timedelta(seconds=0)) + ctx1, _ = cache.get_or_create("thread-1") + + # With zero TTL, next access should create a new context + ctx2, is_new = cache.get_or_create("thread-1") + + assert ctx2 is not ctx1 + assert is_new is True + + +def test_tool_context_cache_remove(): + """Test remove deletes a specific key.""" + from haiku.rag.tools.context import ToolContextCache + + cache = ToolContextCache() + cache.get_or_create("thread-1") + cache.get_or_create("thread-2") + + cache.remove("thread-1") + + ctx, is_new = cache.get_or_create("thread-1") + assert is_new is True + + # thread-2 should still exist + ctx2, is_new2 = cache.get_or_create("thread-2") + assert is_new2 is False + + +def test_tool_context_cache_remove_nonexistent(): + """Test remove handles nonexistent key gracefully.""" + from haiku.rag.tools.context import ToolContextCache + + cache = ToolContextCache() + cache.remove("nonexistent") # Should not raise + + +def test_tool_context_cache_clear(): + """Test clear removes all entries.""" + from haiku.rag.tools.context import ToolContextCache + + cache = ToolContextCache() + cache.get_or_create("thread-1") + cache.get_or_create("thread-2") + + cache.clear() + + ctx1, is_new1 = cache.get_or_create("thread-1") + ctx2, is_new2 = cache.get_or_create("thread-2") + assert is_new1 is True + assert is_new2 is True