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

ksadk/sessions/__init__.py

@@ -0,0 +1,283 @@
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass
+from typing import Callable
+from urllib.parse import urlsplit, urlunsplit
+
+from ksadk.sessions.base import BaseSessionService, Session, SessionEvent, SessionState
+from ksadk.sessions.continuity import (
+    ADKSessionAdapter,
+    ConversationSessionCore,
+    LangChainSessionAdapter,
+    LangGraphSessionAdapter,
+    RunnerSessionAdapter,
+    SessionContinuityLevel,
+    SessionContinuityStatus,
+    TranscriptReplayAdapter,
+)
+from ksadk.sessions.in_memory import InMemorySessionService
+from ksadk.sessions.local_service import create_local_session_service
+
+_cached_session_service: BaseSessionService | None = None
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class SessionBackendConfig:
+    backend: str
+    dsn: str = ""
+    path: str = ""
+    namespace: str = "default"
+    tenant_id: str = "default"
+    workspace_id: str = "default"
+
+
+SessionBackendFactory = Callable[[SessionBackendConfig, str | None], BaseSessionService]
+_backend_factories: dict[str, SessionBackendFactory] = {}
+
+
+def register_session_backend(name: str, factory: SessionBackendFactory) -> None:
+    normalized = name.strip().lower()
+    if not normalized:
+        raise ValueError("Session backend name must not be empty")
+    _backend_factories[normalized] = factory
+
+
+def resolve_session_backend_config(*, backend: str | None = None) -> SessionBackendConfig:
+    _register_builtin_backends()
+    resolved_backend = (
+        backend
+        or os.getenv("KSADK_SESSION_BACKEND")
+        or os.getenv("AGENTENGINE_SESSION_BACKEND")
+        or os.getenv("KSADK_STM_BACKEND")
+        or ""
+    ).strip().lower()
+    if not resolved_backend:
+        resolved_backend = "local"
+    if resolved_backend == "sqlite":
+        resolved_backend = "local"
+    if resolved_backend not in _backend_factories:
+        supported = ", ".join(sorted({*list(_backend_factories), "sqlite"}))
+        raise ValueError(
+            "Unsupported KSADK session backend "
+            f"{resolved_backend!r}; supported backends are {supported}"
+        )
+
+    dsn = (
+        os.getenv("KSADK_SESSION_DSN")
+        or os.getenv("KSADK_STM_URL")
+        or os.getenv("KSADK_STM_DB_URL")
+        or ""
+    ).strip()
+    path = (
+        os.getenv("KSADK_SESSION_PATH")
+        or os.getenv("KSADK_STM_PATH")
+        or os.getenv("KSADK_STM_DB_PATH")
+        or ""
+    ).strip()
+    namespace = (
+        os.getenv("KSADK_SESSION_NAMESPACE")
+        or os.getenv("KSADK_WORKSPACE_ID")
+        or os.getenv("AGENTENGINE_WORKSPACE_ID")
+        or os.getenv("KSADK_TENANT_ID")
+        or os.getenv("AGENTENGINE_TENANT_ID")
+        or "default"
+    ).strip()
+    tenant_id = (
+        os.getenv("KSADK_TENANT_ID")
+        or os.getenv("AGENTENGINE_TENANT_ID")
+        or "default"
+    ).strip()
+    workspace_id = (
+        os.getenv("KSADK_WORKSPACE_ID")
+        or os.getenv("AGENTENGINE_WORKSPACE_ID")
+        or "default"
+    ).strip()
+    return SessionBackendConfig(
+        backend=resolved_backend,
+        dsn=dsn,
+        path=path,
+        namespace=namespace or "default",
+        tenant_id=tenant_id or "default",
+        workspace_id=workspace_id or "default",
+    )
+
+
+def _create_memory_backend(
+    config: SessionBackendConfig,
+    project_dir: str | None,
+) -> BaseSessionService:
+    del config, project_dir
+    return InMemorySessionService()
+
+
+def _create_local_backend(
+    config: SessionBackendConfig,
+    project_dir: str | None,
+) -> BaseSessionService:
+    if config.path:
+        from pathlib import Path
+
+        from ksadk.sessions.local_service import LocalSessionService
+
+        return LocalSessionService(db_path=Path(config.path))
+    return create_local_session_service(project_dir=project_dir)
+
+
+def _create_postgres_backend(
+    config: SessionBackendConfig,
+    project_dir: str | None,
+) -> BaseSessionService:
+    del project_dir
+    if not config.dsn:
+        raise ValueError("KSADK_SESSION_DSN is required when KSADK_SESSION_BACKEND=postgres")
+    from ksadk.sessions.postgres_service import PostgresSessionService
+
+    return PostgresSessionService(
+        dsn=config.dsn,
+        namespace=config.namespace,
+        tenant_id=config.tenant_id,
+        workspace_id=config.workspace_id,
+        connect_timeout=_postgres_connect_timeout_seconds(),
+    )
+
+
+def _postgres_connect_timeout_seconds() -> float:
+    raw = (
+        os.getenv("KSADK_SESSION_CONNECT_TIMEOUT")
+        or os.getenv("KSADK_SESSION_PG_CONNECT_TIMEOUT")
+        or ""
+    ).strip()
+    if not raw:
+        return 5.0
+    try:
+        value = float(raw)
+    except ValueError:
+        logger.warning("Invalid KSADK_SESSION_CONNECT_TIMEOUT=%r; using 5 seconds", raw)
+        return 5.0
+    return max(0.1, value)
+
+
+def _register_builtin_backends() -> None:
+    _backend_factories.setdefault("memory", _create_memory_backend)
+    _backend_factories.setdefault("local", _create_local_backend)
+    _backend_factories.setdefault("postgres", _create_postgres_backend)
+
+
+def describe_session_backend(*, backend: str | None = None) -> dict[str, object]:
+    config = resolve_session_backend_config(backend=backend)
+    return {
+        "Backend": config.backend,
+        "Shared": config.backend == "postgres",
+        "ProductionSafe": config.backend == "postgres",
+        "ContinuityDefault": "semantic/replay" if config.backend == "postgres" else "local_only",
+    }
+
+
+def log_session_backend_diagnostics(*, backend: str | None = None) -> None:
+    config = resolve_session_backend_config(backend=backend)
+    payload = describe_session_backend(backend=config.backend)
+    if config.backend == "postgres":
+        payload = {
+            **payload,
+            "Dsn": mask_session_dsn(config.dsn),
+            "Namespace": config.namespace,
+            "TenantId": config.tenant_id,
+            "WorkspaceId": config.workspace_id,
+        }
+    elif config.backend == "local" and config.path:
+        payload = {**payload, "Path": config.path}
+    logger.info("KSADK session backend: %s", payload)
+    if not bool(payload.get("ProductionSafe")):
+        logger.warning(
+            "KSADK session backend %s is not cross-pod recoverable; use postgres for K8s multi-replica deployments.",
+            payload.get("Backend"),
+        )
+
+
+def mask_session_dsn(dsn: str) -> str:
+    if not dsn:
+        return ""
+    try:
+        parts = urlsplit(dsn)
+    except ValueError:
+        return "***"
+    if not parts.password:
+        return dsn
+    username = parts.username or ""
+    host = parts.hostname or ""
+    port = f":{parts.port}" if parts.port else ""
+    auth = f"{username}:***@" if username else "***@"
+    netloc = f"{auth}{host}{port}"
+    return urlunsplit((parts.scheme, netloc, parts.path, parts.query, parts.fragment))
+
+
+def create_session_service(
+    endpoint: str | None = None,
+    *,
+    backend: str | None = None,
+    project_dir: str | None = None,
+) -> BaseSessionService:
+    del endpoint
+    _register_builtin_backends()
+    config = resolve_session_backend_config(backend=backend)
+    factory = _backend_factories[config.backend]
+    service = factory(config, project_dir)
+    log_session_backend_diagnostics(backend=config.backend)
+    return service
+
+
+def resolve_session_service() -> BaseSessionService:
+    global _cached_session_service
+    if _cached_session_service is not None:
+        return _cached_session_service
+    _cached_session_service = create_session_service()
+    return _cached_session_service
+
+
+async def reset_session_service() -> None:
+    global _cached_session_service
+    if _cached_session_service is None:
+        return
+
+    close = getattr(_cached_session_service, "aclose", None)
+    if close is not None:
+        await close()
+    _cached_session_service = None
+
+
+def get_session_service() -> BaseSessionService:
+    return resolve_session_service()
+
+
+async def close_session_service() -> None:
+    await reset_session_service()
+
+
+__all__ = [
+    "ADKSessionAdapter",
+    "BaseSessionService",
+    "ConversationSessionCore",
+    "InMemorySessionService",
+    "LangChainSessionAdapter",
+    "LangGraphSessionAdapter",
+    "RunnerSessionAdapter",
+    "Session",
+    "SessionContinuityLevel",
+    "SessionContinuityStatus",
+    "SessionEvent",
+    "SessionState",
+    "TranscriptReplayAdapter",
+    "close_session_service",
+    "create_session_service",
+    "describe_session_backend",
+    "get_session_service",
+    "log_session_backend_diagnostics",
+    "mask_session_dsn",
+    "register_session_backend",
+    "reset_session_service",
+    "resolve_session_backend_config",
+    "resolve_session_service",
+]

ksadk/sessions/base.py

@@ -0,0 +1,354 @@
+from __future__ import annotations
+
+import abc
+import time
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Optional
+
+
+def generate_id() -> str:
+    return uuid.uuid4().hex[:16]
+
+
+def normalize_timestamp(value: Any, default: Optional[float] = None) -> float:
+    if value is None:
+        return default if default is not None else time.time()
+    if isinstance(value, (int, float)):
+        numeric = float(value)
+        if numeric > 1_000_000_000_000:
+            return numeric / 1000.0
+        return numeric
+    if isinstance(value, datetime):
+        return value.timestamp()
+    if isinstance(value, str):
+        raw = value.strip()
+        if not raw:
+            return default if default is not None else time.time()
+        try:
+            numeric = float(raw)
+            return normalize_timestamp(numeric, default=default)
+        except ValueError:
+            dt = datetime.fromisoformat(raw.replace("Z", "+00:00"))
+            return dt.timestamp()
+    return default if default is not None else time.time()
+
+
+@dataclass
+class SessionEvent:
+    id: str = field(default_factory=generate_id)
+    session_id: str = ""
+    author: str = ""
+    event_type: str = ""
+    content: dict[str, Any] = field(default_factory=dict)
+    timestamp: float = field(default_factory=time.time)
+    state_delta: dict[str, Any] = field(default_factory=dict)
+    seq_id: int = 0
+    invocation_id: Optional[str] = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(
+        cls,
+        payload: dict[str, Any],
+        *,
+        session_id: str = "",
+        seq_id: int = 0,
+    ) -> "SessionEvent":
+        base_keys = {
+            "id",
+            "session_id",
+            "sessionId",
+            "author",
+            "event_type",
+            "eventType",
+            "content",
+            "timestamp",
+            "state_delta",
+            "stateDelta",
+            "seq_id",
+            "seqId",
+            "invocation_id",
+            "invocationId",
+            "metadata",
+        }
+        metadata = dict(payload.get("metadata") or {})
+        metadata.update({key: value for key, value in payload.items() if key not in base_keys})
+        return cls(
+            id=str(payload.get("id") or generate_id()),
+            session_id=str(
+                payload.get("session_id")
+                or payload.get("sessionId")
+                or session_id
+                or ""
+            ),
+            author=str(payload.get("author") or ""),
+            event_type=str(
+                payload.get("event_type")
+                or payload.get("eventType")
+                or _infer_event_type(payload)
+            ),
+            content=dict(payload.get("content") or {}),
+            timestamp=normalize_timestamp(payload.get("timestamp")),
+            state_delta=dict(payload.get("state_delta") or payload.get("stateDelta") or {}),
+            seq_id=int(payload.get("seq_id") or payload.get("seqId") or seq_id or 0),
+            invocation_id=payload.get("invocation_id") or payload.get("invocationId"),
+            metadata=metadata,
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        payload: dict[str, Any] = {
+            "id": self.id,
+            "session_id": self.session_id,
+            "author": self.author,
+            "event_type": self.event_type,
+            "content": self.content,
+            "timestamp": self.timestamp,
+            "state_delta": self.state_delta,
+            "seq_id": self.seq_id,
+            "metadata": self.metadata,
+        }
+        if self.invocation_id:
+            payload["invocation_id"] = self.invocation_id
+        return payload
+
+    def to_legacy_dict(self) -> dict[str, Any]:
+        payload = dict(self.metadata)
+        payload.update(
+            {
+                "id": self.id,
+                "author": self.author,
+                "invocationId": self.invocation_id,
+                "content": self.content,
+                "timestamp": int(self.timestamp * 1000),
+            }
+        )
+        if self.state_delta:
+            payload["stateDelta"] = self.state_delta
+        if self.event_type:
+            payload["eventType"] = self.event_type
+        return payload
+
+
+@dataclass
+class SessionState:
+    scope: str
+    agent_id: str
+    user_id: str = ""
+    session_id: str = ""
+    state: dict[str, Any] = field(default_factory=dict)
+    version: int = 0
+    updated_at: float = field(default_factory=time.time)
+
+    @classmethod
+    def from_dict(cls, payload: dict[str, Any]) -> "SessionState":
+        return cls(
+            scope=str(payload.get("scope") or "session"),
+            agent_id=str(payload.get("agent_id") or ""),
+            user_id=str(payload.get("user_id") or ""),
+            session_id=str(payload.get("session_id") or ""),
+            state=dict(payload.get("state") or {}),
+            version=int(payload.get("version") or 0),
+            updated_at=normalize_timestamp(payload.get("updated_at")),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "scope": self.scope,
+            "agent_id": self.agent_id,
+            "user_id": self.user_id,
+            "session_id": self.session_id,
+            "state": self.state,
+            "version": self.version,
+            "updated_at": self.updated_at,
+        }
+
+
+@dataclass
+class Session:
+    id: str = field(default_factory=generate_id)
+    agent_id: str = ""
+    user_id: str = ""
+    title: str = ""
+    title_source: str = ""
+    summary: str = ""
+    first_prompt: str = ""
+    last_prompt: str = ""
+    state: dict[str, Any] = field(default_factory=dict)
+    events: list[SessionEvent] = field(default_factory=list)
+    created_at: float = field(default_factory=time.time)
+    updated_at: float = field(default_factory=time.time)
+    version: int = 0
+
+    @classmethod
+    def from_dict(cls, payload: dict[str, Any]) -> "Session":
+        events = [
+            item if isinstance(item, SessionEvent) else SessionEvent.from_dict(item)
+            for item in payload.get("events", [])
+        ]
+        return cls(
+            id=str(payload.get("id") or generate_id()),
+            agent_id=str(
+                payload.get("agent_id")
+                or payload.get("app_name")
+                or payload.get("appName")
+                or ""
+            ),
+            user_id=str(payload.get("user_id") or payload.get("userId") or ""),
+            title=str(payload.get("title") or payload.get("Title") or ""),
+            title_source=str(payload.get("title_source") or payload.get("TitleSource") or ""),
+            summary=str(payload.get("summary") or payload.get("Summary") or ""),
+            first_prompt=str(payload.get("first_prompt") or payload.get("FirstPrompt") or ""),
+            last_prompt=str(payload.get("last_prompt") or payload.get("LastPrompt") or ""),
+            state=dict(payload.get("state") or {}),
+            events=events,
+            created_at=normalize_timestamp(payload.get("created_at") or payload.get("createdAt")),
+            updated_at=normalize_timestamp(payload.get("updated_at") or payload.get("updatedAt")),
+            version=int(payload.get("version") or 0),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "id": self.id,
+            "agent_id": self.agent_id,
+            "user_id": self.user_id,
+            "title": self.title,
+            "title_source": self.title_source,
+            "summary": self.summary,
+            "first_prompt": self.first_prompt,
+            "last_prompt": self.last_prompt,
+            "state": self.state,
+            "events": [event.to_dict() for event in self.events],
+            "created_at": self.created_at,
+            "updated_at": self.updated_at,
+            "version": self.version,
+        }
+
+    def to_legacy_dict(self) -> dict[str, Any]:
+        return {
+            "id": self.id,
+            "appName": self.agent_id,
+            "userId": self.user_id,
+            "title": self.title,
+            "summary": self.summary,
+            "events": [event.to_legacy_dict() for event in self.events],
+            "state": self.state,
+            "createdAt": self.created_at,
+            "updatedAt": self.updated_at,
+        }
+
+
+def _infer_event_type(payload: dict[str, Any]) -> str:
+    content = payload.get("content") or {}
+    role = str(content.get("role") or "")
+    author = str(payload.get("author") or "")
+    parts = content.get("parts") or []
+    for part in parts:
+        if isinstance(part, dict) and part.get("functionCall"):
+            return "tool_call"
+        if isinstance(part, dict) and part.get("functionResponse"):
+            return "tool_result"
+    try:
+        from ksadk.conversations.context import canonical_event_type
+    except Exception:
+        if role in {"assistant", "model"} or author in {"assistant", "model"}:
+            return "assistant_message"
+        return "user_message"
+    return canonical_event_type(None, author=author, role=role)
+
+
+class BaseSessionService(abc.ABC):
+    @abc.abstractmethod
+    async def create_session(
+        self,
+        agent_id: str,
+        user_id: str,
+        session_id: Optional[str] = None,
+    ) -> Session:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def get_session(self, session_id: str) -> Optional[Session]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def list_sessions(
+        self,
+        agent_id: str,
+        user_id: Optional[str] = None,
+        offset: Optional[int] = None,
+        limit: Optional[int] = None,
+    ) -> list[Session]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def count_sessions(
+        self,
+        agent_id: str,
+        user_id: Optional[str] = None,
+    ) -> int:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def delete_session(self, session_id: str) -> bool:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    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:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def append_event(self, session_id: str, event: SessionEvent) -> SessionEvent:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    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]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def count_events(
+        self,
+        session_id: str,
+        after_seq_id: Optional[int] = None,
+        before_seq_id: Optional[int] = None,
+    ) -> int:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def get_state(
+        self,
+        agent_id: str,
+        user_id: Optional[str],
+        session_id: Optional[str],
+        scope: str = "session",
+    ) -> Optional[SessionState]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def update_state(
+        self,
+        *,
+        agent_id: str,
+        user_id: Optional[str],
+        session_id: Optional[str],
+        scope: str,
+        state_delta: dict[str, Any],
+    ) -> SessionState:
+        raise NotImplementedError