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

Fix get_or_create type annotation, export compute_combined_state_delta

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-06 14:56:40 +01:00
parent bd3b26f78b
commit dd7162e72e
No known key found for this signature in database
2 changed files with 5 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
haiku_rag_slim/haiku/rag/tools/__init__.py
View file

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

7
haiku_rag_slim/haiku/rag/tools/context.py
View file

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