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

handle state_key duplication between deps and ToolContext

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-13 10:23:12 +02:00
parent 2496b02a3b
commit 5f6d9e9812
No known key found for this signature in database
7 changed files with 66 additions and 93 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
app/backend/main.py
View file

@ -13,7 +13,6 @@ from starlette.responses import JSONResponse, Response, StreamingResponse
from starlette.routing import Route
from haiku.rag.agents.chat import (
AGUI_STATE_KEY,
ChatDeps,
create_chat_agent,
prepare_chat_context,
@ -92,7 +91,6 @@ async def stream_chat(request: Request) -> Response:
config=Config,
client=get_client(),
tool_context=context,
state_key=AGUI_STATE_KEY,
)
adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)

15
docs/tools.md
View file

@ -235,10 +235,12 @@ async with HaikuRAG("path/to/db.lancedb") as client:
print(f"Total search results: {len(search_state.results)}")
```
`AgentDeps` satisfies the `RAGDeps` protocol and implements the AG-UI state protocol (`state` getter/setter). For AG-UI streaming, pass a `state_key`:
`AgentDeps` satisfies the `RAGDeps` protocol and implements the AG-UI state protocol (`state` getter/setter). For AG-UI streaming, set `state_key` on the `ToolContext` (via `prepare_context`):
```python
deps = AgentDeps(client=client, tool_context=context, state_key="my_app")
context = ToolContext()
prepare_context(context, features=["search", "qa"], state_key="my_app")
deps = AgentDeps(client=client, tool_context=context)
```
Tool functions access `client` and `tool_context` via pydantic-ai's `RunContext.deps`, so toolsets can be created once and reused across requests.
@ -247,7 +249,7 @@ All toolsets respect session-level document filters when a `SessionState` is reg
## AG-UI State Management
Both `AgentDeps` and `ChatDeps` implement the AG-UI `StateHandler` protocol. State is emitted under a namespaced key via `state_key`.
Both `AgentDeps` and `ChatDeps` implement the AG-UI `StateHandler` protocol. `ChatDeps` extends `AgentDeps` with chat-specific config and state handling. State is emitted under a namespaced key via `state_key` on the `ToolContext` — set it once via `prepare_context()`.
**Custom agents** use `AgentDeps` + `prepare_context`:
@ -256,14 +258,14 @@ from haiku.rag.tools import AgentDeps, ToolContext, ToolContextCache, prepare_co
context = ToolContext()
prepare_context(context, features=["search", "qa"], state_key="my_app")
deps = AgentDeps(client=client, tool_context=context, state_key="my_app")
deps = AgentDeps(client=client, tool_context=context)
```
**Chat agent** uses `ChatDeps` + `prepare_chat_context` (adds chat-specific overrides like background summarization and initial context handling):
```python
from haiku.rag.agents.chat import (
AGUI_STATE_KEY, ChatDeps, create_chat_agent, prepare_chat_context,
ChatDeps, create_chat_agent, prepare_chat_context,
)
from haiku.rag.tools import ToolContext, ToolContextCache
@ -272,13 +274,12 @@ agent = create_chat_agent(config)
# For multi-session apps, cache ToolContext per thread
cache = ToolContextCache()
context, _is_new = cache.get_or_create(thread_id)
prepare_chat_context(context) # idempotent namespace registration
prepare_chat_context(context) # idempotent; sets state_key="haiku.rag.chat"
deps = ChatDeps(
config=config,
client=client,
tool_context=context,
state_key=AGUI_STATE_KEY, # "haiku.rag.chat"
)
```

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

@ -1,4 +1,4 @@
from dataclasses import dataclass
from dataclasses import dataclass, field
from typing import Any
from pydantic_ai import Agent
@ -8,9 +8,9 @@ from haiku.rag.agents.chat.context import (
)
from haiku.rag.agents.chat.prompts import build_chat_prompt
from haiku.rag.agents.chat.state import AGUI_STATE_KEY
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.tools.context import ToolContext, prepare_context
from haiku.rag.tools.deps import AgentDeps
from haiku.rag.tools.document import create_document_toolset
from haiku.rag.tools.qa import (
QA_SESSION_NAMESPACE,
@ -34,41 +34,21 @@ def _on_qa_complete(qa_session_state: QASessionState, config: AppConfig) -> None
@dataclass
class ChatDeps:
class ChatDeps(AgentDeps):
"""Dependencies for chat agent.
Implements RAGDeps protocol and StateHandler protocol for AG-UI state management.
Extends AgentDeps with chat-specific config and state handling.
"""
config: AppConfig
client: HaikuRAG
tool_context: ToolContext
state_key: str | None = None
config: AppConfig = field(default_factory=AppConfig)
@property
def state(self) -> dict[str, Any]:
"""Get current state for AG-UI protocol.
Combines all registered namespace states into a single flat dict,
matching the ChatSessionState schema expected by AG-UI clients.
"""
snapshot = self.tool_context.build_state_snapshot()
if self.state_key:
return {self.state_key: snapshot}
return snapshot
@state.setter
@AgentDeps.state.setter
def state(self, value: dict[str, Any] | None) -> None:
"""Set state from AG-UI protocol."""
"""Set state from AG-UI protocol with chat-specific overrides."""
if value is None:
return
# Extract from namespaced key if present
state_data: dict[str, Any] = value
if self.state_key and self.state_key in value:
nested = value[self.state_key]
if isinstance(nested, dict):
state_data = nested
state_data = self._extract_state_data(value)
# Preserve server's session_context before restore overwrites it
qa_session_state = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState)

22
haiku_rag_slim/haiku/rag/tools/deps.py
View file

@ -14,14 +14,14 @@ class AgentDeps:
client: HaikuRAG
tool_context: ToolContext
state_key: str | None = None
@property
def state(self) -> dict[str, Any]:
"""Get current state for AG-UI protocol."""
snapshot = self.tool_context.build_state_snapshot()
if self.state_key:
return {self.state_key: snapshot}
state_key = self.tool_context.state_key
if state_key:
return {state_key: snapshot}
return snapshot
@state.setter
@ -29,10 +29,14 @@ class AgentDeps:
"""Set state from AG-UI protocol."""
if value is None:
return
data: dict[str, Any] = value
if self.state_key and self.state_key in value:
nested = value[self.state_key]
if isinstance(nested, dict):
data = nested
data = self._extract_state_data(value)
self.tool_context.restore_state_snapshot(data)
def _extract_state_data(self, value: dict[str, Any]) -> dict[str, Any]:
"""Extract flat state dict, unwrapping state_key if present."""
state_key = self.tool_context.state_key
if state_key and state_key in value:
nested = value[state_key]
if isinstance(nested, dict):
return nested
return value

54
tests/agents/chat/test_chat_agent.py
View file

@ -64,7 +64,17 @@ def test_chat_deps_initialization(temp_db_path):
assert deps.config is Config
assert deps.client is client
assert deps.tool_context is context
assert deps.state_key is None
client.close()
def test_chat_deps_is_agent_deps(temp_db_path):
"""Test ChatDeps is a subclass of AgentDeps."""
from haiku.rag.tools.deps import AgentDeps
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
deps = ChatDeps(config=Config, client=client, tool_context=context)
assert isinstance(deps, AgentDeps)
client.close()
@ -73,42 +83,18 @@ def test_agui_state_key_constant():
assert AGUI_STATE_KEY == "haiku.rag.chat"
def test_chat_deps_with_state_key(temp_db_path):
"""Test ChatDeps can be initialized with state_key for keyed state emission."""
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
deps = ChatDeps(
config=Config, client=client, tool_context=context, state_key="my_state"
)
assert deps.config is Config
assert deps.state_key == "my_state"
client.close()
def test_chat_deps_state_key_default_none(temp_db_path):
"""Test ChatDeps state_key defaults to None."""
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
deps = ChatDeps(config=Config, client=client, tool_context=context)
assert deps.state_key is None
client.close()
def test_chat_deps_state_setter_handles_initial_context(temp_db_path):
"""Test ChatDeps.state setter transfers initial_context to qa_session_state."""
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
context.state_key = AGUI_STATE_KEY
# Register QASessionState (normally done by prepare_chat_context)
context.register(QA_SESSION_NAMESPACE, QASessionState())
context.register(SESSION_NAMESPACE, SessionState())
deps = ChatDeps(
config=Config, client=client, tool_context=context, state_key=AGUI_STATE_KEY
)
deps = ChatDeps(config=Config, client=client, tool_context=context)
# Client sends initial_context with no session_context
incoming_state = {
@ -140,12 +126,11 @@ def test_chat_deps_state_setter_parses_session_context_dict(temp_db_path):
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
context.state_key = AGUI_STATE_KEY
context.register(QA_SESSION_NAMESPACE, QASessionState())
context.register(SESSION_NAMESPACE, SessionState())
deps = ChatDeps(
config=Config, client=client, tool_context=context, state_key=AGUI_STATE_KEY
)
deps = ChatDeps(config=Config, client=client, tool_context=context)
# Client sends session_context as a dict (as it comes from JSON)
incoming_state = {
@ -178,6 +163,7 @@ def test_chat_deps_state_setter_preserves_server_session_context(temp_db_path):
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
context.state_key = AGUI_STATE_KEY
qa_state = QASessionState()
qa_state.session_context = SessionContext(
summary="Fresh summary from background summarizer"
@ -185,9 +171,7 @@ def test_chat_deps_state_setter_preserves_server_session_context(temp_db_path):
context.register(QA_SESSION_NAMESPACE, qa_state)
context.register(SESSION_NAMESPACE, SessionState())
deps = ChatDeps(
config=Config, client=client, tool_context=context, state_key=AGUI_STATE_KEY
)
deps = ChatDeps(config=Config, client=client, tool_context=context)
# Client sends stale session_context
incoming_state = {
@ -500,7 +484,6 @@ async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path)
config=Config,
client=client,
tool_context=context,
state_key=AGUI_STATE_KEY,
)
# Ask a question that should use the ask tool
@ -557,7 +540,6 @@ async def test_chat_agent_ask_triggers_background_summarization(
config=Config,
client=client,
tool_context=context,
state_key=AGUI_STATE_KEY,
)
# Patch internal trigger to avoid concurrent HTTP calls during VCR
@ -623,7 +605,6 @@ async def test_chat_agent_multi_turn_with_context(allow_model_requests, temp_db_
config=Config,
client=client,
tool_context=context,
state_key=AGUI_STATE_KEY,
)
# Set initial state with initial_context (mimicking AG-UI client)
@ -725,7 +706,6 @@ async def test_chat_agent_ask_with_prior_answer_retrieval(
config=Config,
client=client,
tool_context=context,
state_key=AGUI_STATE_KEY,
)
# First ask - establishes qa_history

16
tests/agents/chat/test_features.py
View file

@ -95,6 +95,8 @@ def test_no_qa_skips_qa_session_state(temp_db_path):
def test_chat_deps_state_without_qa(temp_db_path):
"""ChatDeps.state getter omits qa_history/session_context when QASessionState absent."""
from haiku.rag.agents.chat.state import AGUI_STATE_KEY
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
prepare_chat_context(context, features=[FEATURE_SEARCH])
@ -102,13 +104,17 @@ def test_chat_deps_state_without_qa(temp_db_path):
deps = ChatDeps(config=Config, client=client, tool_context=context)
state = deps.state
# State is wrapped under the AGUI state key
assert AGUI_STATE_KEY in state
inner = state[AGUI_STATE_KEY]
# SessionState fields should be present
assert "document_filter" in state
assert "citation_registry" in state
assert "citations" in state
assert "document_filter" in inner
assert "citation_registry" in inner
assert "citations" in inner
# QA fields should NOT be present
assert "qa_history" not in state
assert "session_context" not in state
assert "qa_history" not in inner
assert "session_context" not in inner
client.close()

14
tests/tools/test_deps.py
View file

@ -31,10 +31,11 @@ def test_agent_deps_state_getter_with_session(mock_client):
def test_agent_deps_state_getter_with_state_key(mock_client):
"""state wraps snapshot under state_key when set."""
"""state wraps snapshot under state_key when set on context."""
ctx = ToolContext()
ctx.state_key = "my_app"
ctx.register(SESSION_NAMESPACE, SessionState())
deps = AgentDeps(client=mock_client, tool_context=ctx, state_key="my_app")
deps = AgentDeps(client=mock_client, tool_context=ctx)
snapshot = deps.state
assert "my_app" in snapshot
assert "citations" in snapshot["my_app"]
@ -54,8 +55,9 @@ def test_agent_deps_state_setter_restores(mock_client):
def test_agent_deps_state_setter_with_state_key(mock_client):
"""state setter extracts data from namespaced key."""
ctx = ToolContext()
ctx.state_key = "my_app"
ctx.register(SESSION_NAMESPACE, SessionState())
deps = AgentDeps(client=mock_client, tool_context=ctx, state_key="my_app")
deps = AgentDeps(client=mock_client, tool_context=ctx)
deps.state = {"my_app": {"document_filter": ["doc1"]}}
session = ctx.get(SESSION_NAMESPACE, SessionState)
assert session is not None
@ -76,16 +78,18 @@ def test_agent_deps_state_setter_ignores_none(mock_client):
def test_agent_deps_state_roundtrip(mock_client):
"""Build snapshot then restore produces equivalent state."""
ctx = ToolContext()
ctx.state_key = "app"
ctx.register(SESSION_NAMESPACE, SessionState(document_filter=["doc1"]))
ctx.register(QA_SESSION_NAMESPACE, QASessionState())
deps = AgentDeps(client=mock_client, tool_context=ctx, state_key="app")
deps = AgentDeps(client=mock_client, tool_context=ctx)
snapshot = deps.state
ctx2 = ToolContext()
ctx2.state_key = "app"
ctx2.register(SESSION_NAMESPACE, SessionState())
ctx2.register(QA_SESSION_NAMESPACE, QASessionState())
deps2 = AgentDeps(client=mock_client, tool_context=ctx2, state_key="app")
deps2 = AgentDeps(client=mock_client, tool_context=ctx2)
deps2.state = snapshot
session = ctx2.get(SESSION_NAMESPACE, SessionState)