89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
from typing import Any
|
|
|
|
import jsonpatch
|
|
from ag_ui.core import EventType, StateDeltaEvent
|
|
from pydantic import BaseModel, Field
|
|
|
|
from haiku.rag.agents.research.models import Citation
|
|
|
|
SESSION_NAMESPACE = "haiku.rag.session"
|
|
|
|
|
|
class SessionState(BaseModel):
|
|
"""Session-level state for AG-UI integration.
|
|
|
|
This state is shared across toolsets and enables:
|
|
- Session identification
|
|
- Dynamic document filtering
|
|
- Stable citation indices across tool calls
|
|
- AG-UI state synchronization
|
|
"""
|
|
|
|
session_id: str = ""
|
|
incoming_session_id: str = Field(default="", exclude=True) # Track what client sent
|
|
document_filter: list[str] = []
|
|
citation_registry: dict[str, int] = {}
|
|
citations: list[Citation] = []
|
|
state_key: str | None = Field(default=None, exclude=True)
|
|
|
|
def get_or_assign_index(self, chunk_id: str) -> int:
|
|
"""Get or assign a stable citation index for a chunk_id.
|
|
|
|
Citation indices persist across tool calls within a session.
|
|
The first chunk gets index 1, subsequent new chunks get incrementing indices.
|
|
Same chunk_id always returns the same index.
|
|
"""
|
|
if chunk_id in self.citation_registry:
|
|
return self.citation_registry[chunk_id]
|
|
|
|
new_index = len(self.citation_registry) + 1
|
|
self.citation_registry[chunk_id] = new_index
|
|
return new_index
|
|
|
|
|
|
def compute_state_delta(
|
|
old_state: SessionState,
|
|
new_state: SessionState,
|
|
) -> StateDeltaEvent | None:
|
|
"""Compute state delta between old and new session state.
|
|
|
|
Returns a StateDeltaEvent if there are changes, None otherwise.
|
|
The state_key from new_state is used for namespacing.
|
|
"""
|
|
return compute_combined_state_delta(
|
|
old_state.model_dump(mode="json"),
|
|
new_state.model_dump(mode="json"),
|
|
state_key=new_state.state_key,
|
|
)
|
|
|
|
|
|
def compute_combined_state_delta(
|
|
old_snapshot: dict[str, Any],
|
|
new_snapshot: dict[str, Any],
|
|
state_key: str | None = None,
|
|
) -> StateDeltaEvent | None:
|
|
"""Compute state delta between old and new combined state snapshots.
|
|
|
|
This function computes delta for the combined chat state that includes
|
|
both SessionState and QASessionState fields.
|
|
|
|
Args:
|
|
old_snapshot: Previous state dict (e.g., from ChatDeps.state format).
|
|
new_snapshot: New state dict.
|
|
state_key: Optional namespace key for the state (e.g., "haiku.rag.chat").
|
|
|
|
Returns:
|
|
StateDeltaEvent if there are changes, None otherwise.
|
|
"""
|
|
wrapped_old = {state_key: old_snapshot} if state_key else old_snapshot
|
|
wrapped_new = {state_key: new_snapshot} if state_key else new_snapshot
|
|
|
|
patch = jsonpatch.make_patch(wrapped_old, wrapped_new)
|
|
|
|
if not patch.patch:
|
|
return None
|
|
|
|
return StateDeltaEvent(
|
|
type=EventType.STATE_DELTA,
|
|
delta=patch.patch,
|
|
)
|