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

Clean up server vs client session setting priorities

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-11 16:54:00 +02:00
parent 0434cd068b
commit 58c2bd49ea
No known key found for this signature in database
7 changed files with 84 additions and 49 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
app/backend/main.py
View file

@ -79,13 +79,12 @@ async def stream_chat(request: Request) -> Response:
run_input = AGUIAdapter.build_run_input(body)
thread_id = getattr(run_input, "thread_id", None) or "default"
context, is_new = context_cache.get_or_create(thread_id)
context, _is_new = context_cache.get_or_create(thread_id)
agent = create_chat_agent(Config, get_client(), context)
deps = ChatDeps(
config=Config,
tool_context=context,
is_new=is_new,
state_key=AGUI_STATE_KEY,
)

22
docs/agents/index.md
View file

@ -127,13 +127,27 @@ The system prompt is automatically composed to match the selected features. See
### Session State
The `ChatSessionState` maintains:
Session state is managed through `ToolContext` — a namespace-based state container shared across all toolsets. The chat agent uses two namespaces:
**`SessionState`** (session management):
- `session_id` — Unique identifier for the session
- `qa_history` — List of previous Q/A pairs
- `session_context` — Automatically maintained session context summary
- `document_filter` — List of document titles/URIs to restrict searches
- `citation_registry` — Stable mapping of chunk IDs to citation indices
- `citations` — Citations from the current query
**`QASessionState`** (QA history and context):
- `qa_history` — List of previous Q/A pairs with embeddings
- `session_context` — Automatically maintained session context summary
For multi-session applications (e.g., web backends), use `ToolContextCache` to cache `ToolContext` instances by external session/thread ID:
```python
from haiku.rag.tools import ToolContext, ToolContextCache
cache = ToolContextCache() # TTL-based, defaults to 1 hour
context, _is_new = cache.get_or_create(thread_id)
```
**Citation Registry**: Citation indices persist across tool calls within a session. The same `chunk_id` always returns the same citation index (first-occurrence-wins). This ensures consistent citation numbering in multi-turn conversations — `[1]` always refers to the same source.

10
docs/tools.md
View file

@ -227,9 +227,12 @@ When using the chat agent with [AG-UI](https://docs.ag-ui.com) streaming, `ChatD
```python
from haiku.rag.agents.chat import AGUI_STATE_KEY, ChatDeps, create_chat_agent
from haiku.rag.tools import ToolContext
from haiku.rag.tools import ToolContext, ToolContextCache
# For multi-session apps, cache ToolContext per thread
cache = ToolContextCache()
context, _is_new = cache.get_or_create(thread_id)
context = ToolContext()
agent = create_chat_agent(config, client, context)
deps = ChatDeps(
config=config,
@ -243,7 +246,6 @@ The emitted state structure:
```json
{
"haiku.rag.chat": {
"session_id": "uuid",
"citations": [],
"qa_history": [],
"session_context": null,
@ -253,7 +255,7 @@ The emitted state structure:
}
```
State flows bidirectionally — the frontend sends its current state on each request, and the agent emits deltas (JSON Patch) reflecting server-side updates (new citations, QA history entries, session context). See the [Web Application](apps.md#web-application) for a complete implementation.
State flows bidirectionally — the frontend sends its current state on each request, and the agent emits deltas (JSON Patch) reflecting server-side updates (new citations, QA history entries, session context). The server always prefers its own `session_context` over the client's value, since background summarization may have updated it between requests. See the [Web Application](apps.md#web-application) for a complete implementation.
## Filter Helpers

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

@ -43,7 +43,6 @@ class ChatDeps:
config: AppConfig
tool_context: ToolContext
is_new: bool = True
state_key: str | None = None
@property
@ -74,57 +73,43 @@ class ChatDeps:
state_data = nested
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", [])
if "citation_registry" in state_data:
session_state.citation_registry = state_data["citation_registry"]
if "citations" in state_data:
from haiku.rag.agents.research.models import Citation
if self.is_new:
# First request for this context: fully populate from client state
if session_state is not None:
if "document_filter" in state_data:
session_state.document_filter = state_data.get(
"document_filter", []
)
if "citation_registry" in state_data:
session_state.citation_registry = state_data["citation_registry"]
if "citations" in state_data:
from haiku.rag.agents.research.models import Citation
session_state.citations = [
Citation(**c) if isinstance(c, dict) else c
for c in state_data.get("citations", [])
]
session_state.citations = [
Citation(**c) if isinstance(c, dict) else c
for c in state_data.get("citations", [])
]
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
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
qa_session_state.qa_history = [
QAHistoryEntry(**qa) if isinstance(qa, dict) else qa
for qa in state_data.get("qa_history", [])
]
qa_session_state.qa_history = [
QAHistoryEntry(**qa) if isinstance(qa, dict) else qa
for qa in state_data.get("qa_history", [])
]
# Restore session_context from client
# Prefer server's session_context (background summarizer may
# have updated it since the client's last snapshot).
if not qa_session_state.session_context:
session_context = state_data.get("session_context")
if isinstance(session_context, dict):
qa_session_state.session_context = SessionContext(
**session_context
).summary
elif session_context is None:
qa_session_state.session_context = None
# Handle initial_context -> session_context for first message
if "initial_context" in state_data:
initial = state_data.get("initial_context")
if initial and not qa_session_state.session_context:
qa_session_state.session_context = initial
else:
# Returning request: only merge client-controlled fields
if session_state is not None:
if "document_filter" in state_data:
session_state.document_filter = state_data.get(
"document_filter", []
)
def create_chat_agent(

1
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -251,7 +251,6 @@ class ChatApp(App):
deps = ChatDeps(
config=self.config,
tool_context=self.tool_context,
is_new=False,
)
async with self.agent.run_stream(

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

@ -48,6 +48,7 @@ __all__ = [
"QASessionState",
"QAHistoryEntry",
"create_qa_toolset",
"run_qa_core",
"create_analysis_toolset",
"SESSION_NAMESPACE",
"SessionState",

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

@ -64,7 +64,6 @@ def test_chat_deps_initialization(temp_db_path):
assert deps.config is Config
assert deps.tool_context is context
assert deps.is_new is True
assert deps.state_key is None
@ -153,6 +152,42 @@ def test_chat_deps_state_setter_parses_session_context_dict():
assert qa_session_state.session_context == "Previous conversation summary"
def test_chat_deps_state_setter_preserves_server_session_context():
"""Test that server's session_context is preferred over client's stale value."""
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
context = ToolContext()
qa_state = QASessionState()
qa_state.session_context = "Fresh summary from background summarizer"
context.register(QA_SESSION_NAMESPACE, qa_state)
context.register(SESSION_NAMESPACE, SessionState())
deps = ChatDeps(config=Config, tool_context=context, state_key=AGUI_STATE_KEY)
# Client sends stale session_context
incoming_state = {
AGUI_STATE_KEY: {
"session_context": {
"summary": "Stale summary from client",
"last_updated": "2025-01-27T12:00:00",
},
"qa_history": [],
"citations": [],
"document_filter": [],
"citation_registry": {},
}
}
deps.state = incoming_state
# Server's fresher session_context should be preserved
qa_session_state = context.get(QA_SESSION_NAMESPACE)
assert isinstance(qa_session_state, QASessionState)
assert (
qa_session_state.session_context == "Fresh summary from background summarizer"
)
def test_chat_session_state():
"""Test ChatSessionState model."""
state = ChatSessionState()