Source code for sqlspec.extensions.adk.service

"""SQLSpec-backed session service for Google ADK."""

import inspect
import logging
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, cast

from google.adk.sessions.base_session_service import BaseSessionService, GetSessionConfig, ListSessionsResponse

from sqlspec.extensions.adk.converters import (
    compute_update_marker,
    event_to_record,
    filter_temp_state,
    merge_scoped_state,
    record_to_session,
    split_scoped_state,
)
from sqlspec.extensions.adk.store import normalize_session_list_options
from sqlspec.utils.logging import get_logger, log_with_context
from sqlspec.utils.sync_tools import async_
from sqlspec.utils.uuids import uuid4

if TYPE_CHECKING:
    from collections.abc import Callable

    from google.adk.events.event import Event
    from google.adk.sessions import Session

    from sqlspec.extensions.adk._types import SessionOrderBy
    from sqlspec.extensions.adk.store import BaseAsyncADKStore, BaseSyncADKStore

    ADKStore = BaseAsyncADKStore | BaseSyncADKStore

__all__ = ("SQLSpecSessionService",)

logger = get_logger("sqlspec.extensions.adk.service")




class SQLSpecSessionService(BaseSessionService):
    """SQLSpec-backed implementation of BaseSessionService.

    Provides session and event storage using SQLSpec database adapters.
    Delegates all database operations to a store implementation.

    Args:
        store: Database store implementation.
    """



    def __init__(self, store: "ADKStore") -> None:
        """Initialize the session service.

        Args:
            store: Database store implementation.
        """
        self._store = store

@property def store(self) -> "ADKStore": """Return the database store.""" return self._store

async def create_session( self, *, app_name: str, user_id: str, state: "dict[str, Any] | None" = None, session_id: "str | None" = None ) -> "Session": """Create a new session. Args: app_name: Name of the application. user_id: ID of the user. state: Initial state of the session. session_id: Client-provided session ID. If None, generates a UUID. Returns: The newly created session. """ if session_id is None: session_id = str(uuid4()) if state is None: state = {} persisted_state = filter_temp_state(state) app_state_delta, user_state_delta, session_state = split_scoped_state(persisted_state) current_app_state = await self._call_store("get_app_state", app_name) current_user_state = await self._call_store("get_user_state", app_name, user_id) app_state = dict(current_app_state or {}) if app_state_delta: app_state.update(app_state_delta) user_state = dict(current_user_state or {}) if user_state_delta: user_state.update(user_state_delta) record = await self._call_store( "create_session", session_id=session_id, app_name=app_name, user_id=user_id, state=session_state ) if app_state_delta: await self._call_store("upsert_app_state", app_name, app_state) if user_state_delta: await self._call_store("upsert_user_state", app_name, user_id, user_state) record["state"] = merge_scoped_state(record["state"], app_state, user_state) log_with_context( logger, logging.DEBUG, "adk.session.create", app_name=app_name, session_id=session_id, has_state=bool(state) ) return record_to_session(record, events=[])