From dd7162e72e90c72c8a5e122aee57d941c3b9573a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 6 Feb 2026 14:56:40 +0100 Subject: [PATCH] Fix get_or_create type annotation, export compute_combined_state_delta --- haiku_rag_slim/haiku/rag/tools/__init__.py | 2 ++ haiku_rag_slim/haiku/rag/tools/context.py | 7 +++---- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/tools/__init__.py b/haiku_rag_slim/haiku/rag/tools/__init__.py index fb30f93e..795946b1 100644 --- a/haiku_rag_slim/haiku/rag/tools/__init__.py +++ b/haiku_rag_slim/haiku/rag/tools/__init__.py @@ -31,6 +31,7 @@ from haiku.rag.tools.search import SEARCH_NAMESPACE, SearchState, create_search_ from haiku.rag.tools.session import ( SESSION_NAMESPACE, SessionState, + compute_combined_state_delta, compute_state_delta, ) @@ -63,4 +64,5 @@ __all__ = [ "SESSION_NAMESPACE", "SessionState", "compute_state_delta", + "compute_combined_state_delta", ] diff --git a/haiku_rag_slim/haiku/rag/tools/context.py b/haiku_rag_slim/haiku/rag/tools/context.py index 1e4fdc0c..3b78ab3a 100644 --- a/haiku_rag_slim/haiku/rag/tools/context.py +++ b/haiku_rag_slim/haiku/rag/tools/context.py @@ -1,4 +1,3 @@ -from collections.abc import Callable from typing import Any, TypeVar from pydantic import BaseModel, PrivateAttr @@ -72,18 +71,18 @@ class ToolContext(BaseModel): return state return None - def get_or_create(self, namespace: str, factory: Callable[[], T]) -> T: + def get_or_create(self, namespace: str, state_type: type[T]) -> T: """Get state for a namespace, creating it if not registered. Args: namespace: The namespace to get or create state for. - factory: A callable that returns a new Pydantic model instance. + state_type: A Pydantic BaseModel subclass to instantiate if needed. Returns: The state for the namespace. """ if namespace not in self._namespaces: - self._namespaces[namespace] = factory() + self._namespaces[namespace] = state_type() return self._namespaces[namespace] # type: ignore[return-value] def clear_namespace(self, namespace: str) -> None: