diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index e06cb860..4f3df052 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py @@ -52,16 +52,14 @@ class ChatDeps: snapshot: dict[str, Any] = {"session_id": self.session_id} # Add SessionState fields - session_state = self.tool_context.get_typed(SESSION_NAMESPACE, SessionState) + session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) if session_state is not None: snapshot["document_filter"] = session_state.document_filter snapshot["citation_registry"] = session_state.citation_registry snapshot["citations"] = [c.model_dump() for c in session_state.citations] # Add QASessionState fields - qa_session_state = self.tool_context.get_typed( - QA_SESSION_NAMESPACE, QASessionState - ) + qa_session_state = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState) if qa_session_state is not None: snapshot["qa_history"] = [ qa.model_dump() for qa in qa_session_state.qa_history @@ -92,7 +90,7 @@ class ChatDeps: state_data = nested # Update SessionState from incoming state - session_state = self.tool_context.get_typed(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", []) @@ -122,9 +120,7 @@ class ChatDeps: session_state.incoming_session_id = incoming_session_id # Update QASessionState from incoming state - qa_session_state = self.tool_context.get_typed( - QA_SESSION_NAMESPACE, QASessionState - ) + qa_session_state = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState) if qa_session_state is not None: if "qa_history" in state_data: from haiku.rag.tools.qa import QAHistoryEntry @@ -187,12 +183,12 @@ def create_chat_agent( result = await agent.run("Search for X", deps=deps) """ # Ensure session states are registered with proper AG-UI state key - existing = context.get_typed(SESSION_NAMESPACE, SessionState) + 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 - if context.get_typed(QA_SESSION_NAMESPACE, QASessionState) is None: + if context.get(QA_SESSION_NAMESPACE, QASessionState) is None: context.register(QA_SESSION_NAMESPACE, QASessionState()) # Create toolsets - these capture client, config, and context in closures @@ -230,7 +226,7 @@ def trigger_background_summarization(deps: ChatDeps) -> None: Args: deps: Chat dependencies with tool_context containing QASessionState. """ - qa_session_state = deps.tool_context.get_typed(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: return if not deps.session_id: diff --git a/haiku_rag_slim/haiku/rag/agents/chat/state.py b/haiku_rag_slim/haiku/rag/agents/chat/state.py index 2a55bfff..be9e3795 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/state.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/state.py @@ -1,7 +1,7 @@ from datetime import datetime from typing import TYPE_CHECKING -from pydantic import BaseModel, Field +from pydantic import BaseModel from haiku.rag.agents.research.models import Citation diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 0ebbcd00..56b36d6b 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -163,7 +163,7 @@ class ChatApp(App): self.agent = create_chat_agent(self.config, self.client, self.tool_context) # Initialize session state in tool context - session_state = self.tool_context.get_typed(SESSION_NAMESPACE, SessionState) + session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) if session_state is not None: session_state.document_filter = self._document_filter @@ -320,7 +320,7 @@ class ChatApp(App): } # Sync session state to tool context before running - session_state = self.tool_context.get_typed(SESSION_NAMESPACE, SessionState) + session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState) if session_state is not None: session_state.session_id = self.session_state.session_id session_state.document_filter = self.session_state.document_filter @@ -330,7 +330,7 @@ class ChatApp(App): # Sync initial_context to QA session state from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState - qa_session_state = self.tool_context.get_typed( + qa_session_state = self.tool_context.get( QA_SESSION_NAMESPACE, QASessionState ) if qa_session_state is not None: @@ -379,7 +379,7 @@ class ChatApp(App): trigger_background_summarization(deps) # Sync session context from QASessionState to ChatSessionState for modal - qa_session_state = self.tool_context.get_typed( + qa_session_state = self.tool_context.get( QA_SESSION_NAMESPACE, QASessionState ) if qa_session_state is not None and qa_session_state.session_context: @@ -471,9 +471,7 @@ class ChatApp(App): from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState # Sync session context from QASessionState before showing modal - qa_session_state = self.tool_context.get_typed( - QA_SESSION_NAMESPACE, QASessionState - ) + qa_session_state = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState) if qa_session_state is not None and qa_session_state.session_context: from datetime import datetime diff --git a/haiku_rag_slim/haiku/rag/tools/context.py b/haiku_rag_slim/haiku/rag/tools/context.py index 3b78ab3a..9e9dc2d8 100644 --- a/haiku_rag_slim/haiku/rag/tools/context.py +++ b/haiku_rag_slim/haiku/rag/tools/context.py @@ -1,4 +1,4 @@ -from typing import Any, TypeVar +from typing import Any, TypeVar, overload from pydantic import BaseModel, PrivateAttr @@ -57,19 +57,24 @@ class ToolContext(BaseModel): """ self._namespaces[namespace] = state - def get(self, namespace: str) -> BaseModel | None: - """Get state for a namespace, or None if not registered.""" - return self._namespaces.get(namespace) + @overload + def get(self, namespace: str) -> BaseModel | None: ... - def get_typed(self, namespace: str, expected_type: type[T]) -> T | None: - """Get state for a namespace with type checking. + @overload + def get(self, namespace: str, state_type: type[T]) -> T | None: ... - Returns the state cast to expected_type if it matches, None otherwise. + def get( + self, namespace: str, state_type: type[T] | None = None + ) -> BaseModel | T | None: + """Get state for a namespace, or None if not registered. + + When state_type is provided, returns the state only if it matches + the expected type, otherwise returns None. """ state = self._namespaces.get(namespace) - if isinstance(state, expected_type): - return state - return None + if state_type is not None: + return state if isinstance(state, state_type) else None + return state def get_or_create(self, namespace: str, state_type: type[T]) -> T: """Get state for a namespace, creating it if not registered. diff --git a/haiku_rag_slim/haiku/rag/tools/document.py b/haiku_rag_slim/haiku/rag/tools/document.py index 14fc52af..e530a24f 100644 --- a/haiku_rag_slim/haiku/rag/tools/document.py +++ b/haiku_rag_slim/haiku/rag/tools/document.py @@ -65,16 +65,16 @@ async def find_document(client: HaikuRAG, query: str): limit=1, filter=f"LOWER(uri) LIKE LOWER('%{escaped_query}%') OR LOWER(uri) LIKE LOWER('%{no_spaces}%')", ) - if docs: - return docs[0] + if docs and docs[0].id: + return await client.get_document_by_id(docs[0].id) # Try partial title match (with and without spaces) docs = await client.list_documents( limit=1, filter=f"LOWER(title) LIKE LOWER('%{escaped_query}%') OR LOWER(title) LIKE LOWER('%{no_spaces}%')", ) - if docs: - return docs[0] + if docs and docs[0].id: + return await client.get_document_by_id(docs[0].id) return None diff --git a/haiku_rag_slim/haiku/rag/tools/filters.py b/haiku_rag_slim/haiku/rag/tools/filters.py index f7b1a729..3a74429b 100644 --- a/haiku_rag_slim/haiku/rag/tools/filters.py +++ b/haiku_rag_slim/haiku/rag/tools/filters.py @@ -52,7 +52,7 @@ def get_session_filter( from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState - session_state = context.get_typed(SESSION_NAMESPACE, SessionState) + session_state = context.get(SESSION_NAMESPACE, SessionState) if session_state is None or not session_state.document_filter: return base_filter diff --git a/haiku_rag_slim/haiku/rag/tools/qa.py b/haiku_rag_slim/haiku/rag/tools/qa.py index bc0bce1f..e99bd04b 100644 --- a/haiku_rag_slim/haiku/rag/tools/qa.py +++ b/haiku_rag_slim/haiku/rag/tools/qa.py @@ -154,8 +154,8 @@ def create_qa_toolset( old_state_snapshot: dict | None = None if context is not None: - session_state = context.get_typed(SESSION_NAMESPACE, SessionState) - qa_session_state = context.get_typed(QA_SESSION_NAMESPACE, QASessionState) + session_state = context.get(SESSION_NAMESPACE, SessionState) + qa_session_state = context.get(QA_SESSION_NAMESPACE, QASessionState) # Capture combined state snapshot before changes # Use incoming values (what client sent) so delta shows server-side updates diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index 17ff1f88..d5e4bb56 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/haiku_rag_slim/haiku/rag/tools/search.py @@ -71,7 +71,7 @@ def create_search_toolset( session_state: SessionState | None = None old_session_state: SessionState | None = None if context is not None: - session_state = context.get_typed(SESSION_NAMESPACE, SessionState) + session_state = context.get(SESSION_NAMESPACE, SessionState) if session_state is not None: old_session_state = session_state.model_copy(deep=True) diff --git a/tests/tools/test_context.py b/tests/tools/test_context.py index 2adf2b24..9392e98a 100644 --- a/tests/tools/test_context.py +++ b/tests/tools/test_context.py @@ -171,3 +171,33 @@ def test_serialization_roundtrip(): qa_state = restored.get("qa") assert isinstance(qa_state, TestState) assert qa_state.value == 99 + + +def test_get_with_type_match(): + """Test get with state_type returns typed state when type matches.""" + ctx = ToolContext() + state = TestState(value=42) + ctx.register("ns", state) + + result = ctx.get("ns", TestState) + assert result is state + assert result.value == 42 + + +def test_get_with_type_mismatch(): + """Test get with state_type returns None when type doesn't match.""" + ctx = ToolContext() + ctx.register("ns", TestState(value=42)) + + result = ctx.get("ns", TestStateWithList) + assert result is None + + +def test_get_without_type(): + """Test get without state_type returns BaseModel (unchanged behavior).""" + ctx = ToolContext() + state = TestState(value=42) + ctx.register("ns", state) + + result = ctx.get("ns") + assert result is state