"""Base store classes for ADK memory backend (sync and async)."""
import logging
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Final, Generic, Literal, TypeVar, cast
from sqlspec.extensions.adk._config_utils import _adk_memory_store_config, _ADKMemoryStoreConfig
from sqlspec.extensions.adk._table_utils import owner_id_column_name, unique_statements
from sqlspec.extensions.adk.store import _reconcile_adk_schema_sync
from sqlspec.migrations.schema import SchemaTarget, ensure_schema_async
from sqlspec.observability import resolve_db_system
from sqlspec.utils.logging import get_logger, log_with_context
from sqlspec.utils.sync_tools import async_
if TYPE_CHECKING:
from collections.abc import Callable, Sequence
from sqlspec.config import DatabaseConfigProtocol
from sqlspec.extensions.adk.memory._types import StoredMemory
__all__ = ("BaseAsyncADKMemoryStore", "BaseSyncADKMemoryStore")
ConfigT = TypeVar("ConfigT", bound="DatabaseConfigProtocol[Any, Any, Any]")
logger = get_logger("sqlspec.extensions.adk.memory.store")
ADK_RESET_MEMORY_TABLES: Final = ("adk_memory", "adk_memory_entries")
class _ADKMemoryStoreCommon(Generic[ConfigT]):
"""Shared non-async ADK store state and helpers."""
if TYPE_CHECKING:
_drop_memory_table_sql: "Callable[[], list[str]]"
__slots__ = (
"_config",
"_enable_bm25",
"_enabled",
"_max_results",
"_memory_table",
"_owner_id_column_ddl",
"_owner_id_column_name",
"_scann_num_leaves",
"_scann_quantizer",
"_use_fts",
"_vector_dimensions",
"_vector_index_type",
)
def __init__(self, config: ConfigT) -> None:
"""Initialize the ADK memory store.
Args:
config: SQLSpec database configuration.
"""
self._config = config
store_config = self._store_config_from_extension()
self._enabled: bool = store_config.get("enable_memory", True)
self._memory_table: str = str(store_config["memory_table"])
self._use_fts: bool = bool(store_config.get("use_fts", False))
self._max_results: int = store_config.get("max_results", 20)
self._vector_index_type: str = store_config.get("vector_index_type", "hnsw")
self._vector_dimensions: int = store_config.get("vector_dimensions", 768)
self._enable_bm25: bool = bool(store_config.get("enable_bm25", False))
self._scann_num_leaves: int = store_config.get("scann_num_leaves", 100)
self._scann_quantizer: str = store_config.get("scann_quantizer", "SQ8")
self._owner_id_column_ddl: str | None = store_config.get("owner_id_column")
self._owner_id_column_name: str | None = (
owner_id_column_name(self._owner_id_column_ddl) if self._owner_id_column_ddl else None
)
def _store_config_from_extension(self) -> _ADKMemoryStoreConfig:
return _adk_memory_store_config(self._config)
@property
def config(self) -> ConfigT:
"""Return the database configuration."""
return self._config
@property
def is_enabled(self) -> bool:
"""Return whether memory storage is enabled."""
return self._enabled
@property
def memory_table(self) -> str:
"""Return the configured memory table name."""
return self._memory_table
@property
def vector_index_type(self) -> str:
"""Return the configured vector index type."""
return self._vector_index_type
@property
def vector_dimensions(self) -> int:
"""Return the vector dimensionality."""
return self._vector_dimensions
@property
def enable_bm25(self) -> bool:
"""Return whether BM25 text search is enabled."""
return self._enable_bm25
@property
def scann_num_leaves(self) -> int:
"""Return the number of leaves for ScaNN tree quantization."""
return self._scann_num_leaves
@property
def scann_quantizer(self) -> str:
"""Return the ScaNN quantizer."""
return self._scann_quantizer
@property
def use_fts(self) -> bool:
"""Return whether full-text search is enabled."""
return self._use_fts
@property
def max_results(self) -> int:
"""Return the default maximum results for search."""
return self._max_results
@property
def owner_id_column_ddl(self) -> str | None:
"""Return the configured owner column DDL snippet, if any."""
return self._owner_id_column_ddl
@property
def owner_id_column_name(self) -> str | None:
"""Return the extracted owner column name, if configured."""
return self._owner_id_column_name
def _schema_management_flags(self) -> tuple[bool, bool]:
extension_config = getattr(self._config, "extension_config", {})
adk_config = extension_config.get("adk", {}) if isinstance(extension_config, dict) else {}
manage_schema = adk_config.get("manage_schema", True) if isinstance(adk_config, dict) else True
create_schema = adk_config.get("create_schema", True) if isinstance(adk_config, dict) else True
return bool(manage_schema), bool(create_schema)
@property
def create_schema_enabled(self) -> bool:
"""Return whether adapter-level table creation should run."""
manage_schema, create_schema = self._schema_management_flags()
return manage_schema and create_schema
def _drop_sql_for_table(self, table_name: str) -> list[str]:
current_table = self._memory_table
self._memory_table = table_name
try:
return list(self._drop_memory_table_sql())
finally:
self._memory_table = current_table
def _reset_drop_memory_table_sql(self) -> list[str]:
configured = self._memory_table
candidates = (configured, *[name for name in ADK_RESET_MEMORY_TABLES if name != configured])
statements: list[str] = []
for cand in candidates:
statements.extend(self._drop_sql_for_table(cand))
return unique_statements(statements)
def _require_enabled(self) -> None:
if not self._enabled:
msg = "ADK memory store is disabled for this database configuration"
raise RuntimeError(msg)
def _effective_limit(self, limit: int | None) -> int:
return limit if limit is not None else self._max_results
def _log_operation(self, event: str, **kwargs: Any) -> None:
log_with_context(
logger,
logging.DEBUG,
event,
table_name=self._memory_table,
db_system=resolve_db_system(type(self).__name__),
**kwargs,
)
class BaseAsyncADKMemoryStore(_ADKMemoryStoreCommon[ConfigT], ABC):
"""Base class for async SQLSpec-backed ADK memory stores.
Implements storage operations for Google ADK memory entries using
SQLSpec database adapters with async/await.
"""
__slots__ = ()
@abstractmethod
async def create_tables(self) -> None:
"""Create the memory table and indexes if they don't exist.
Should check self._enabled and skip table creation if False.
"""
raise NotImplementedError