Source code for sqlspec.extensions.adk.memory.service

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

import inspect
from typing import TYPE_CHECKING, Any, Literal, cast

from google.adk.memory.base_memory_service import BaseMemoryService, SearchMemoryResponse

from sqlspec.extensions.adk.memory.converters import (
    memory_entry_to_record,
    records_to_memory_entries,
    session_to_memory_records,
)
from sqlspec.utils.logging import get_logger
from sqlspec.utils.sync_tools import async_

if TYPE_CHECKING:
    from collections.abc import Callable, Mapping, Sequence

    from google.adk.events.event import Event
    from google.adk.memory.memory_entry import MemoryEntry
    from google.adk.sessions import Session

    from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore

__all__ = ("SQLSpecMemoryService", "SQLSpecSyncMemoryService")

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


class SQLSpecMemoryService(BaseMemoryService):
    """SQLSpec-backed implementation of BaseMemoryService.

    Provides memory entry storage using SQLSpec database adapters.
    Delegates all database operations to a store implementation.

    ADK BaseMemoryService defines two core methods:
        - add_session_to_memory(session) - Ingests session into memory (returns void)
        - search_memory(app_name, user_id, query) - Searches stored memories

    Args:
        store: Database store implementation.
    """



[docs]

    def __init__(self, store: "BaseAsyncADKMemoryStore | BaseSyncADKMemoryStore") -> None:
        """Initialize the memory service.

        Args:
            store: Database store implementation.
        """
        self._store = store
@property def store(self) -> "BaseAsyncADKMemoryStore | BaseSyncADKMemoryStore": """Return the database store.""" return self._store async def add_session_to_memory(self, session: "Session", scope: str = "user") -> None: """Add a completed session to the memory store. Extracts all events with content from the session and stores them as searchable memory entries. Uses UPSERT to skip duplicates. The Session object contains app_name and user_id properties. Events are converted to memory records and bulk inserted via store. Returns void per ADK BaseMemoryService contract. Args: session: Completed ADK Session with events. scope: Visibility scope ('user' or 'app'). """ records = session_to_memory_records(session, scope=scope) if not records: logger.debug( "No content to store for session %s (app=%s, user=%s)", session.id, session.app_name, session.user_id ) return inserted_count = await self._call_store("insert_memory_entries", records) logger.debug( "Stored %d memory entries for session %s (total events: %d)", inserted_count, session.id, len(records) )