agentengine-sdk-python 0.6.2__py3-none-any.whl → 0.6.9__py3-none-any.whl

ksadk/sessions/continuity.py

@@ -0,0 +1,313 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Optional
+
+from ksadk.sessions.base import BaseSessionService, Session
+
+
+class SessionContinuityLevel(str, Enum):
+    UI_ONLY = "ui_only"
+    SEMANTIC = "semantic"
+    RUNTIME = "runtime"
+    EXACT = "exact"
+
+
+@dataclass
+class SessionContinuityStatus:
+    level: SessionContinuityLevel
+    path: str
+    runner: str
+    details: dict[str, Any] = field(default_factory=dict)
+
+    def to_payload(self) -> dict[str, Any]:
+        payload = {
+            "Level": self.level.value,
+            "Path": self.path,
+            "Runner": self.runner,
+        }
+        if self.details:
+            payload["Details"] = self.details
+        return payload
+
+
+class ConversationSessionCore:
+    def __init__(self, service: BaseSessionService):
+        self.service = service
+
+    @staticmethod
+    def binding_scope(runner_key: str) -> str:
+        return f"runner_binding:{runner_key}"
+
+    @staticmethod
+    def runtime_scope(runner_key: str) -> str:
+        return f"runner_runtime:{runner_key}"
+
+    async def _load_session(self, session_id: str) -> Optional[Session]:
+        if not session_id:
+            return None
+        return await self.service.get_session(session_id)
+
+    async def get_binding_by_session_id(self, session_id: str, runner_key: str) -> dict[str, Any]:
+        session = await self._load_session(session_id)
+        if session is None:
+            return {}
+        state = await self.service.get_state(
+            agent_id=session.agent_id,
+            user_id=session.user_id,
+            session_id=session.id,
+            scope=self.binding_scope(runner_key),
+        )
+        return dict(state.state) if state else {}
+
+    async def set_binding_by_session_id(
+        self,
+        session_id: str,
+        runner_key: str,
+        state_delta: dict[str, Any],
+    ) -> dict[str, Any]:
+        session = await self._load_session(session_id)
+        if session is None:
+            return {}
+        state = await self.service.update_state(
+            agent_id=session.agent_id,
+            user_id=session.user_id,
+            session_id=session.id,
+            scope=self.binding_scope(runner_key),
+            state_delta=state_delta,
+        )
+        return dict(state.state)
+
+    async def get_runtime_state_by_session_id(self, session_id: str, runner_key: str) -> dict[str, Any]:
+        session = await self._load_session(session_id)
+        if session is None:
+            return {}
+        state = await self.service.get_state(
+            agent_id=session.agent_id,
+            user_id=session.user_id,
+            session_id=session.id,
+            scope=self.runtime_scope(runner_key),
+        )
+        return dict(state.state) if state else {}
+
+    async def set_runtime_state_by_session_id(
+        self,
+        session_id: str,
+        runner_key: str,
+        state_delta: dict[str, Any],
+    ) -> dict[str, Any]:
+        session = await self._load_session(session_id)
+        if session is None:
+            return {}
+        state = await self.service.update_state(
+            agent_id=session.agent_id,
+            user_id=session.user_id,
+            session_id=session.id,
+            scope=self.runtime_scope(runner_key),
+            state_delta=state_delta,
+        )
+        return dict(state.state)
+
+
+class RunnerSessionAdapter:
+    def runner_key(self, runner: Any) -> str:
+        runner_type = getattr(getattr(runner, "detection_result", None), "type", None)
+        runner_value = getattr(runner_type, "value", "") if runner_type is not None else ""
+        if isinstance(runner_value, str) and runner_value.strip():
+            return runner_value.strip()
+        runner_name = getattr(getattr(runner, "detection_result", None), "name", "")
+        return str(runner_name or runner.__class__.__name__).strip().lower()
+
+    def continuity_status(
+        self,
+        *,
+        runner: Any,
+        binding_state: dict[str, Any],
+        runtime_state: dict[str, Any],
+    ) -> SessionContinuityStatus:
+        path = str(runtime_state.get("path") or "replay")
+        level = runtime_state.get("level") or SessionContinuityLevel.SEMANTIC.value
+        return SessionContinuityStatus(
+            level=SessionContinuityLevel(str(level)),
+            path=path,
+            runner=self.runner_key(runner),
+            details={k: v for k, v in runtime_state.items() if k not in {"path", "level"}},
+        )
+
+    async def describe_continuity(
+        self,
+        *,
+        runner: Any,
+        session: Session,
+        core: ConversationSessionCore,
+    ) -> SessionContinuityStatus:
+        runner_key = self.runner_key(runner)
+        binding_state = await core.get_binding_by_session_id(session.id, runner_key)
+        runtime_state = await core.get_runtime_state_by_session_id(session.id, runner_key)
+        return self.continuity_status(
+            runner=runner,
+            binding_state=binding_state,
+            runtime_state=runtime_state,
+        )
+
+    async def bind_session(
+        self,
+        *,
+        runner: Any,
+        core: ConversationSessionCore,
+        session_id: str,
+        external_session_id: str,
+        internal_session_id: str,
+    ) -> dict[str, Any]:
+        if not session_id:
+            return {}
+        return await core.set_binding_by_session_id(
+            session_id,
+            self.runner_key(runner),
+            {
+                "external_session_id": external_session_id,
+                "internal_session_id": internal_session_id,
+            },
+        )
+
+    async def persist_turn(
+        self,
+        *,
+        runner: Any,
+        core: ConversationSessionCore,
+        session_id: str,
+        state_delta: dict[str, Any],
+    ) -> dict[str, Any]:
+        if not session_id or not state_delta:
+            return {}
+        return await core.set_runtime_state_by_session_id(
+            session_id,
+            self.runner_key(runner),
+            state_delta,
+        )
+
+    def prepare_request(self, *, runner: Any, input_data: dict[str, Any]) -> dict[str, Any]:
+        return dict(input_data)
+
+    async def clear_session(
+        self,
+        *,
+        runner: Any,
+        core: ConversationSessionCore,
+        session_id: str,
+    ) -> None:
+        await self.persist_turn(
+            runner=runner,
+            core=core,
+            session_id=session_id,
+            state_delta={"path": "replay", "level": SessionContinuityLevel.SEMANTIC.value},
+        )
+
+
+class TranscriptReplayAdapter(RunnerSessionAdapter):
+    pass
+
+
+class LangChainSessionAdapter(RunnerSessionAdapter):
+    def continuity_status(
+        self,
+        *,
+        runner: Any,
+        binding_state: dict[str, Any],
+        runtime_state: dict[str, Any],
+    ) -> SessionContinuityStatus:
+        if runtime_state:
+            return super().continuity_status(
+                runner=runner,
+                binding_state=binding_state,
+                runtime_state=runtime_state,
+            )
+
+        path = "replay"
+        module = getattr(runner, "_module", None)
+        if callable(getattr(module, "ksadk_prepare_input", None)):
+            path = "standard_hook"
+        else:
+            try:
+                from langchain_core.runnables.history import RunnableWithMessageHistory
+
+                if isinstance(getattr(runner, "_agent", None), RunnableWithMessageHistory):
+                    path = "runnable_with_message_history"
+            except Exception:
+                pass
+
+        return SessionContinuityStatus(
+            level=SessionContinuityLevel.SEMANTIC,
+            path=path,
+            runner=self.runner_key(runner),
+        )
+
+
+class LangGraphSessionAdapter(RunnerSessionAdapter):
+    def continuity_status(
+        self,
+        *,
+        runner: Any,
+        binding_state: dict[str, Any],
+        runtime_state: dict[str, Any],
+    ) -> SessionContinuityStatus:
+        if runtime_state:
+            return super().continuity_status(
+                runner=runner,
+                binding_state=binding_state,
+                runtime_state=runtime_state,
+            )
+        agent = getattr(runner, "_agent", None)
+        has_checkpointer = bool(
+            getattr(agent, "checkpointer", None) or getattr(agent, "_checkpointer", None)
+        )
+        module = getattr(runner, "_module", None)
+        has_hook = callable(getattr(module, "ksadk_prepare_state", None))
+        if has_hook:
+            path = "standard_hook"
+        elif has_checkpointer:
+            path = "checkpoint"
+        else:
+            path = "replay"
+        return SessionContinuityStatus(
+            level=SessionContinuityLevel.RUNTIME if has_checkpointer else SessionContinuityLevel.SEMANTIC,
+            path=path,
+            runner=self.runner_key(runner),
+        )
+
+
+class ADKSessionAdapter(RunnerSessionAdapter):
+    def continuity_status(
+        self,
+        *,
+        runner: Any,
+        binding_state: dict[str, Any],
+        runtime_state: dict[str, Any],
+    ) -> SessionContinuityStatus:
+        if runtime_state:
+            return super().continuity_status(
+                runner=runner,
+                binding_state=binding_state,
+                runtime_state=runtime_state,
+            )
+        has_native_session = bool(getattr(runner, "_short_term_memory", None)) or any(
+            str((__import__("os").environ.get(name) or "")).strip()
+            for name in (
+                "KSADK_ADK_SESSION_BACKEND",
+                "KSADK_ADK_SESSION_PATH",
+                "KSADK_ADK_SESSION_URL",
+                "KSADK_STM_BACKEND",
+                "KSADK_STM_PATH",
+                "KSADK_STM_URL",
+                "KSADK_STM_DB_PATH",
+                "KSADK_STM_DB_URL",
+                "KSADK_SESSION_BACKEND",
+                "KSADK_SESSION_DSN",
+            )
+        )
+        return SessionContinuityStatus(
+            level=SessionContinuityLevel.SEMANTIC,
+            path="native_session" if has_native_session else "replay",
+            runner=self.runner_key(runner),
+        )

ksadk/sessions/errors.py

@@ -0,0 +1,7 @@
+"""Session backend exception types."""
+
+from __future__ import annotations
+
+
+class SessionBackendUnavailable(RuntimeError):
+    """Raised when the configured session backend cannot be reached."""

ksadk/sessions/in_memory.py

@@ -0,0 +1,273 @@
+from __future__ import annotations
+
+import asyncio
+import copy
+import time
+from typing import Optional
+
+from ksadk.sessions.base import BaseSessionService, Session, SessionEvent, SessionState, generate_id
+
+
+class InMemorySessionService(BaseSessionService):
+    def __init__(self):
+        self._sessions: dict[str, Session] = {}
+        self._states: dict[tuple[str, str, str, str], SessionState] = {}
+        self._lock = asyncio.Lock()
+
+    async def create_session(
+        self,
+        agent_id: str,
+        user_id: str,
+        session_id: Optional[str] = None,
+    ) -> Session:
+        async with self._lock:
+            if session_id and session_id in self._sessions:
+                return copy.deepcopy(self._sessions[session_id])
+
+            session = Session(
+                id=session_id or generate_id(),
+                agent_id=agent_id,
+                user_id=user_id,
+            )
+            self._sessions[session.id] = session
+            self._states[self._state_key("session", agent_id, user_id, session.id)] = SessionState(
+                scope="session",
+                agent_id=agent_id,
+                user_id=user_id,
+                session_id=session.id,
+            )
+            return copy.deepcopy(session)
+
+    async def get_session(self, session_id: str) -> Optional[Session]:
+        async with self._lock:
+            session = self._sessions.get(session_id)
+            return copy.deepcopy(session) if session else None
+
+    async def list_sessions(
+        self,
+        agent_id: str,
+        user_id: Optional[str] = None,
+        offset: Optional[int] = None,
+        limit: Optional[int] = None,
+    ) -> list[Session]:
+        async with self._lock:
+            sessions = [
+                copy.deepcopy(session)
+                for session in self._sessions.values()
+                if session.agent_id == agent_id and (user_id is None or session.user_id == user_id)
+            ]
+            sessions.sort(key=lambda item: (item.updated_at, item.created_at), reverse=True)
+            start = offset or 0
+            end = None if limit is None else start + limit
+            return sessions[start:end]
+
+    async def count_sessions(
+        self,
+        agent_id: str,
+        user_id: Optional[str] = None,
+    ) -> int:
+        async with self._lock:
+            return sum(
+                1
+                for session in self._sessions.values()
+                if session.agent_id == agent_id and (user_id is None or session.user_id == user_id)
+            )
+
+    async def delete_session(self, session_id: str) -> bool:
+        async with self._lock:
+            session = self._sessions.pop(session_id, None)
+            if not session:
+                return False
+            self._states.pop(
+                self._state_key(
+                    "session",
+                    session.agent_id,
+                    session.user_id,
+                    session_id,
+                ),
+                None,
+            )
+            return True
+
+    async def update_session_metadata(
+        self,
+        session_id: str,
+        *,
+        title: Optional[str] = None,
+        title_source: Optional[str] = None,
+        summary: Optional[str] = None,
+        first_prompt: Optional[str] = None,
+        last_prompt: Optional[str] = None,
+    ) -> Session:
+        async with self._lock:
+            session = self._sessions.get(session_id)
+            if not session:
+                raise ValueError(f"Session {session_id} not found")
+            if title is not None:
+                session.title = title
+            if title_source is not None:
+                session.title_source = title_source
+            if summary is not None:
+                session.summary = summary
+            if first_prompt is not None:
+                session.first_prompt = first_prompt
+            if last_prompt is not None:
+                session.last_prompt = last_prompt
+            session.updated_at = time.time()
+            return copy.deepcopy(session)
+
+    async def append_event(self, session_id: str, event: SessionEvent) -> SessionEvent:
+        async with self._lock:
+            session = self._sessions.get(session_id)
+            if not session:
+                raise ValueError(f"Session {session_id} not found")
+
+            stored = copy.deepcopy(event)
+            stored.session_id = session_id
+            stored.seq_id = len(session.events) + 1
+            if not stored.id:
+                stored.id = generate_id()
+            session.events.append(stored)
+            session.updated_at = time.time()
+
+            if stored.state_delta:
+                session.state.update(stored.state_delta)
+                session.version += 1
+                self._states[
+                    self._state_key(
+                        "session",
+                        session.agent_id,
+                        session.user_id,
+                        session.id,
+                    )
+                ] = SessionState(
+                    scope="session",
+                    agent_id=session.agent_id,
+                    user_id=session.user_id,
+                    session_id=session.id,
+                    state=copy.deepcopy(session.state),
+                    version=session.version,
+                    updated_at=session.updated_at,
+                )
+
+            return copy.deepcopy(stored)
+
+    async def get_events(
+        self,
+        session_id: str,
+        offset: Optional[int] = None,
+        limit: Optional[int] = None,
+        after_seq_id: Optional[int] = None,
+        before_seq_id: Optional[int] = None,
+    ) -> list[SessionEvent]:
+        async with self._lock:
+            session = self._sessions.get(session_id)
+            if not session:
+                return []
+            events = list(session.events)
+            if after_seq_id is not None:
+                events = [event for event in events if event.seq_id > after_seq_id]
+            if before_seq_id is not None:
+                events = [event for event in events if event.seq_id < before_seq_id]
+            end = max(len(events) - (offset or 0), 0)
+            start = 0 if limit is None else max(end - limit, 0)
+            sliced = events[start:end]
+            return copy.deepcopy(sliced)
+
+    async def count_events(
+        self,
+        session_id: str,
+        after_seq_id: Optional[int] = None,
+        before_seq_id: Optional[int] = None,
+    ) -> int:
+        async with self._lock:
+            session = self._sessions.get(session_id)
+            if not session:
+                return 0
+            events = list(session.events)
+            if after_seq_id is not None:
+                events = [event for event in events if event.seq_id > after_seq_id]
+            if before_seq_id is not None:
+                events = [event for event in events if event.seq_id < before_seq_id]
+            return len(events)
+
+    async def get_state(
+        self,
+        agent_id: str,
+        user_id: Optional[str],
+        session_id: Optional[str],
+        scope: str = "session",
+    ) -> Optional[SessionState]:
+        async with self._lock:
+            if scope == "session" and session_id:
+                session = self._sessions.get(session_id)
+                if session:
+                    return SessionState(
+                        scope="session",
+                        agent_id=session.agent_id,
+                        user_id=session.user_id,
+                        session_id=session.id,
+                        state=copy.deepcopy(session.state),
+                        version=session.version,
+                        updated_at=session.updated_at,
+                    )
+            state = self._states.get(
+                self._state_key(scope, agent_id, user_id or "", session_id or "")
+            )
+            return copy.deepcopy(state) if state else None
+
+    async def update_state(
+        self,
+        *,
+        agent_id: str,
+        user_id: Optional[str],
+        session_id: Optional[str],
+        scope: str,
+        state_delta: dict,
+    ) -> SessionState:
+        async with self._lock:
+            user_key = user_id or ""
+            session_key = session_id or ""
+            state_key = self._state_key(scope, agent_id, user_key, session_key)
+            current = self._states.get(state_key)
+
+            if scope == "session":
+                session = self._sessions.get(session_key)
+                if not session:
+                    raise ValueError(f"Session {session_key} not found")
+                session.state.update(state_delta)
+                session.version += 1
+                session.updated_at = time.time()
+                updated = SessionState(
+                    scope="session",
+                    agent_id=session.agent_id,
+                    user_id=session.user_id,
+                    session_id=session.id,
+                    state=copy.deepcopy(session.state),
+                    version=session.version,
+                    updated_at=session.updated_at,
+                )
+            else:
+                next_state = copy.deepcopy(current.state) if current else {}
+                next_state.update(state_delta)
+                updated = SessionState(
+                    scope=scope,
+                    agent_id=agent_id,
+                    user_id=user_key,
+                    session_id=session_key,
+                    state=next_state,
+                    version=(current.version if current else 0) + 1,
+                    updated_at=time.time(),
+                )
+
+            self._states[state_key] = updated
+            return copy.deepcopy(updated)
+
+    @staticmethod
+    def _state_key(
+        scope: str,
+        agent_id: str,
+        user_id: str,
+        session_id: str,
+    ) -> tuple[str, str, str, str]:
+        return (scope, agent_id, user_id, session_id)