126 lines
4.2 KiB
Python
126 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib
|
|
import json
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
from typing import Protocol
|
|
|
|
|
|
@dataclass
|
|
class SessionState:
|
|
session_id: str
|
|
user_id: str
|
|
channel: str
|
|
status: str = "idle"
|
|
current_intent: str | None = None
|
|
pending_slots: list[str] = field(default_factory=list)
|
|
slots: dict[str, Any] = field(default_factory=dict)
|
|
workflow: dict[str, Any] | None = None
|
|
routing_debug: dict[str, Any] | None = None
|
|
last_user_text: str | None = None
|
|
last_agent_text: str | None = None
|
|
context_memory: dict[str, Any] = field(default_factory=dict)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"session_id": self.session_id,
|
|
"user_id": self.user_id,
|
|
"channel": self.channel,
|
|
"status": self.status,
|
|
"current_intent": self.current_intent,
|
|
"pending_slots": self.pending_slots,
|
|
"slots": self.slots,
|
|
"workflow": self.workflow,
|
|
"routing_debug": self.routing_debug,
|
|
"last_user_text": self.last_user_text,
|
|
"last_agent_text": self.last_agent_text,
|
|
"context_memory": self.context_memory,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: dict[str, Any]) -> "SessionState":
|
|
return cls(
|
|
session_id=data["session_id"],
|
|
user_id=data["user_id"],
|
|
channel=data.get("channel", "app"),
|
|
status=data.get("status", "idle"),
|
|
current_intent=data.get("current_intent"),
|
|
pending_slots=list(data.get("pending_slots", [])),
|
|
slots=dict(data.get("slots", {})),
|
|
workflow=data.get("workflow"),
|
|
routing_debug=data.get("routing_debug"),
|
|
last_user_text=data.get("last_user_text"),
|
|
last_agent_text=data.get("last_agent_text"),
|
|
context_memory=dict(data.get("context_memory", {})),
|
|
)
|
|
|
|
|
|
class SessionStore(Protocol):
|
|
def get_or_create(self, session_id: str, user_id: str, channel: str = "app") -> SessionState:
|
|
...
|
|
|
|
def get(self, session_id: str) -> SessionState | None:
|
|
...
|
|
|
|
def save(self, session: SessionState) -> SessionState:
|
|
...
|
|
|
|
|
|
class InMemorySessionStore:
|
|
def __init__(self) -> None:
|
|
self._sessions: dict[str, SessionState] = {}
|
|
|
|
def get_or_create(self, session_id: str, user_id: str, channel: str = "app") -> SessionState:
|
|
session = self._sessions.get(session_id)
|
|
if session is None:
|
|
session = SessionState(session_id=session_id, user_id=user_id, channel=channel)
|
|
self._sessions[session_id] = session
|
|
return session
|
|
|
|
def get(self, session_id: str) -> SessionState | None:
|
|
return self._sessions.get(session_id)
|
|
|
|
def save(self, session: SessionState) -> SessionState:
|
|
self._sessions[session.session_id] = session
|
|
return session
|
|
|
|
|
|
class RedisSessionStore:
|
|
def __init__(
|
|
self,
|
|
redis_url: str,
|
|
key_prefix: str = "agent:session",
|
|
ttl_seconds: int = 86400,
|
|
) -> None:
|
|
redis_module = importlib.import_module("redis")
|
|
self._client = redis_module.Redis.from_url(redis_url, decode_responses=True)
|
|
self._key_prefix = key_prefix
|
|
self._ttl_seconds = ttl_seconds
|
|
|
|
def get_or_create(self, session_id: str, user_id: str, channel: str = "app") -> SessionState:
|
|
session = self.get(session_id)
|
|
if session is not None:
|
|
return session
|
|
|
|
session = SessionState(session_id=session_id, user_id=user_id, channel=channel)
|
|
self.save(session)
|
|
return session
|
|
|
|
def get(self, session_id: str) -> SessionState | None:
|
|
payload = self._client.get(self._build_key(session_id))
|
|
if payload is None:
|
|
return None
|
|
return SessionState.from_dict(json.loads(payload))
|
|
|
|
def save(self, session: SessionState) -> SessionState:
|
|
self._client.set(
|
|
self._build_key(session.session_id),
|
|
json.dumps(session.to_dict(), ensure_ascii=False),
|
|
ex=self._ttl_seconds,
|
|
)
|
|
return session
|
|
|
|
def _build_key(self, session_id: str) -> str:
|
|
return f"{self._key_prefix}:{session_id}"
|