From 075b5ea76085cc307fec1d006093788f24164cda Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sat, 15 Aug 2026 21:07:22 +0000 Subject: [PATCH] feat(adk): scoped memory recall and domain type modernization (#704) --- docs/changelog.rst | 4 +- docs/examples/extensions/adk/memory_store.py | 4 +- docs/extensions/adk/api.rst | 8 +- docs/extensions/adk/index.rst | 4 +- docs/extensions/adk/installation.rst | 2 +- docs/extensions/adk/quickstart.rst | 2 +- docs/extensions/adk/schema.rst | 6 +- docs/reference/extensions/adk.rst | 8 +- sqlspec/adapters/adbc/adk/store.py | 158 +++++--- sqlspec/adapters/aiomysql/adk/store.py | 102 +++-- sqlspec/adapters/aiosqlite/adk/store.py | 102 +++-- sqlspec/adapters/arrow_odbc/adk/store.py | 114 ++++-- sqlspec/adapters/asyncmy/adk/store.py | 102 +++-- sqlspec/adapters/asyncpg/adk/store.py | 113 ++++-- sqlspec/adapters/bigquery/adk/store.py | 34 +- .../adapters/cockroach_asyncpg/adk/store.py | 116 ++++-- .../adapters/cockroach_psycopg/adk/store.py | 178 ++++++--- sqlspec/adapters/duckdb/adk/store.py | 168 ++++---- sqlspec/adapters/mssql_python/adk/store.py | 26 +- sqlspec/adapters/mysqlconnector/adk/store.py | 178 ++++++--- sqlspec/adapters/oracledb/adk/store.py | 262 +++++++----- sqlspec/adapters/psqlpy/adk/store.py | 108 +++-- sqlspec/adapters/psycopg/adk/store.py | 260 +++++++----- sqlspec/adapters/pymssql/adk/store.py | 106 +++-- sqlspec/adapters/pymysql/adk/store.py | 130 +++--- sqlspec/adapters/spanner/adk/store.py | 184 +++++---- sqlspec/adapters/sqlite/adk/store.py | 162 ++++---- sqlspec/extensions/adk/__init__.py | 48 +-- sqlspec/extensions/adk/_types.py | 6 +- sqlspec/extensions/adk/artifact/__init__.py | 6 +- sqlspec/extensions/adk/artifact/_types.py | 4 +- sqlspec/extensions/adk/artifact/service.py | 6 +- sqlspec/extensions/adk/artifact/store.py | 18 +- sqlspec/extensions/adk/converters.py | 18 +- sqlspec/extensions/adk/memory/__init__.py | 39 +- sqlspec/extensions/adk/memory/_types.py | 11 +- sqlspec/extensions/adk/memory/converters.py | 57 ++- sqlspec/extensions/adk/memory/service.py | 40 +- sqlspec/extensions/adk/memory/store.py | 377 ++++++++---------- .../adk/migrations/0001_create_adk_tables.py | 129 +++++- .../adk/migrations/0002_reset_adk_tables.py | 9 +- sqlspec/extensions/adk/store.py | 40 +- .../adapters/_shared/adk_behaviors.py | 4 +- .../extensions/adk/test_memory_store.py | 82 +++- .../duckdb/extensions/adk/test_store.py | 12 +- .../aiomysql/extensions/adk/test_store.py | 4 +- .../asyncmy/extensions/adk/test_store.py | 4 +- .../extensions/adk/test_store.py | 4 +- .../extensions/adk/test_oracle_specific.py | 10 +- .../spanner/extensions/adk/test_adk_store.py | 6 +- .../adbc/extensions/adk/test_memory_store.py | 10 +- .../extensions/adk/test_memory_store.py | 99 ++++- .../extensions/adk/test_memory_store.py | 99 ++++- .../test_oracledb/test_oracle_adk_store.py | 6 + .../adapters/test_spanner/test_adk_store.py | 16 +- .../extensions/test_adk/test_converters.py | 12 +- .../unit/extensions/test_adk/test_service.py | 24 +- .../test_adk/test_store_common_mixins.py | 10 +- .../extensions/test_adk/test_store_config.py | 61 +-- .../test_adk/test_store_instantiation.py | 14 +- tests/unit/utils/test_mypyc_smoke.py | 6 +- tools/scripts/mypyc_smoke.py | 2 +- 62 files changed, 2458 insertions(+), 1476 deletions(-) diff --git a/docs/changelog.rst b/docs/changelog.rst index 42a2f77fc..4fcd5dd0d 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -633,7 +633,7 @@ v0.51.0 - ADK 2.0 clean-break store contract ``get_metadata``, and ``set_metadata``. * ``append_event_and_update_state()`` accepts optional ``app_state`` and ``user_state`` deltas and applies them atomically with the session and event - write, returning the updated ``SessionRecord``. + write, returning the updated ``StoredSession``. **Fixed:** @@ -750,7 +750,7 @@ v0.47.0 - Persistent listeners, schema builders, and performance polish Pydantic, dataclasses, and attrs. Pass ``wire_format=True`` to keep wire-aligned names. * Third-party ADK stores implementing ``append_event_and_update_state()`` must - return the updated ``SessionRecord``. + return the updated ``StoredSession``. * Data dictionary metadata/version helpers now live under ``sqlspec.data_dictionary``. ``ColumnMetadata``, ``ForeignKeyMetadata``, ``IndexMetadata``, ``TableMetadata``, ``VersionInfo``, and diff --git a/docs/examples/extensions/adk/memory_store.py b/docs/examples/extensions/adk/memory_store.py index e583b1a5b..3f37b2214 100644 --- a/docs/examples/extensions/adk/memory_store.py +++ b/docs/examples/extensions/adk/memory_store.py @@ -7,7 +7,7 @@ import pytest if TYPE_CHECKING: - from sqlspec.extensions.adk import EventRecord + from sqlspec.extensions.adk import StoredEvent __all__ = ("test_adk_memory_store",) @@ -16,7 +16,7 @@ def test_adk_memory_store() -> None: pytest.importorskip("aiosqlite") pytest.importorskip("google.adk") - async def _run() -> list[EventRecord]: + async def _run() -> list[StoredEvent]: # start-example from google.adk.events.event import Event from google.genai import types diff --git a/docs/extensions/adk/api.rst b/docs/extensions/adk/api.rst index e95a28a31..4b1bb768a 100644 --- a/docs/extensions/adk/api.rst +++ b/docs/extensions/adk/api.rst @@ -79,22 +79,22 @@ Artifact Stores Record Types ============ -.. autoclass:: SessionRecord +.. autoclass:: StoredSession :members: :show-inheritance: :no-index: -.. autoclass:: EventRecord +.. autoclass:: StoredEvent :members: :show-inheritance: :no-index: -.. autoclass:: MemoryRecord +.. autoclass:: StoredMemory :members: :show-inheritance: :no-index: -.. autoclass:: ArtifactRecord +.. autoclass:: StoredArtifact :members: :show-inheritance: :no-index: diff --git a/docs/extensions/adk/index.rst b/docs/extensions/adk/index.rst index 45d8f53b5..6fa19b259 100644 --- a/docs/extensions/adk/index.rst +++ b/docs/extensions/adk/index.rst @@ -12,7 +12,7 @@ Key capabilities: - **Session and event storage** with atomic ``append_event_and_update_state()`` ensuring events and state are always consistent. -- **Full-event JSON storage** (EventRecord) that captures the entire ADK Event +- **Full-event JSON storage** (StoredEvent) that captures the entire ADK Event in a single column, eliminating schema drift with upstream ADK releases. - **Scoped state semantics** (``app:``, ``user:``, ``temp:``) for controlling state visibility and persistence across sessions. @@ -56,7 +56,7 @@ Choose a guide :link: schema :link-type: doc - Table layouts, EventRecord, scoped state, and artifact metadata. + Table layouts, StoredEvent, scoped state, and artifact metadata. .. grid-item-card:: API Reference :link: api diff --git a/docs/extensions/adk/installation.rst b/docs/extensions/adk/installation.rst index 271ce142d..e9da02ae9 100644 --- a/docs/extensions/adk/installation.rst +++ b/docs/extensions/adk/installation.rst @@ -161,7 +161,7 @@ The ``adk`` extra includes the Google ADK SDK (``google-genai``). SQLSpec provid for context retrieval across conversations. - **Artifact Service contracts** -- Compose a concrete metadata store with pluggable object storage backends when your deployment needs ADK artifacts. -- **Event Storage** -- Full-event JSON storage (EventRecord) that captures the +- **Event Storage** -- Full-event JSON storage (StoredEvent) that captures the entire ADK Event without schema drift. Next Steps diff --git a/docs/extensions/adk/quickstart.rst b/docs/extensions/adk/quickstart.rst index 63ac95359..291d62435 100644 --- a/docs/extensions/adk/quickstart.rst +++ b/docs/extensions/adk/quickstart.rst @@ -117,5 +117,5 @@ Next Steps ========== - :doc:`backends` for the full support matrix and backend-specific details. -- :doc:`schema` for table layouts, EventRecord format, and scoped state semantics. +- :doc:`schema` for table layouts, StoredEvent format, and scoped state semantics. - :doc:`api` for the complete API reference. diff --git a/docs/extensions/adk/schema.rst b/docs/extensions/adk/schema.rst index f48564f51..04278a4e1 100644 --- a/docs/extensions/adk/schema.rst +++ b/docs/extensions/adk/schema.rst @@ -56,7 +56,7 @@ config for multi-tenant deployments. .. _event-record: -Events Table (EventRecord) +Events Table (StoredEvent) ========================== The events table uses **full-event JSON storage**: the entire ADK ``Event`` is @@ -107,7 +107,7 @@ via ``Event.model_validate()``. from sqlspec.extensions.adk.converters import event_to_record, record_to_event - # Serialize: Event -> EventRecord + # Serialize: Event -> StoredEvent record = event_to_record( event=adk_event, app_name="my_agent", @@ -115,7 +115,7 @@ via ``Event.model_validate()``. session_id="sess_123", ) - # Reconstruct: EventRecord -> Event + # Reconstruct: StoredEvent -> Event restored_event = record_to_event(record) .. _scoped-state: diff --git a/docs/reference/extensions/adk.rst b/docs/reference/extensions/adk.rst index 0f2bfaf09..5147a0a0c 100644 --- a/docs/reference/extensions/adk.rst +++ b/docs/reference/extensions/adk.rst @@ -67,19 +67,19 @@ Artifact Store Base Classes Record Types ============ -.. autoclass:: sqlspec.extensions.adk.SessionRecord +.. autoclass:: sqlspec.extensions.adk.StoredSession :members: :show-inheritance: -.. autoclass:: sqlspec.extensions.adk.EventRecord +.. autoclass:: sqlspec.extensions.adk.StoredEvent :members: :show-inheritance: -.. autoclass:: sqlspec.extensions.adk.MemoryRecord +.. autoclass:: sqlspec.extensions.adk.StoredMemory :members: :show-inheritance: -.. autoclass:: sqlspec.extensions.adk.ArtifactRecord +.. autoclass:: sqlspec.extensions.adk.StoredArtifact :members: :show-inheritance: diff --git a/sqlspec/adapters/adbc/adk/store.py b/sqlspec/adapters/adbc/adk/store.py index 7cbcfcf11..2d621c88a 100644 --- a/sqlspec/adapters/adbc/adk/store.py +++ b/sqlspec/adapters/adbc/adk/store.py @@ -3,16 +3,16 @@ import contextlib import re from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Literal -from sqlspec.extensions.adk import BaseSyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.utils.logging import get_logger from sqlspec.utils.serializers import from_json, to_json if TYPE_CHECKING: from sqlspec.adapters.adbc.config import AdbcConfig - from sqlspec.extensions.adk import MemoryRecord + from sqlspec.extensions.adk import StoredMemory __all__ = ("AdbcADKMemoryStore", "AdbcADKStore") @@ -80,13 +80,13 @@ def create_tables(self) -> None: def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new session.""" return self._create_session(session_id, app_name, user_id, state, owner_id) def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get session by ID.""" return self._get_session(app_name, user_id, session_id, renew_for=renew_for) @@ -94,7 +94,7 @@ def update_session_state(self, app_name: str, user_id: str, session_id: str, sta """Update session state.""" self._update_session_state(app_name, user_id, session_id, state) - def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": """List sessions for an app.""" return self._list_sessions(app_name, user_id) @@ -102,13 +102,13 @@ def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: """Delete session and associated events.""" self._delete_session(app_name, user_id, session_id) - def append_event(self, event_record: EventRecord) -> None: + def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session.""" self._append_event(event_record) def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -116,7 +116,7 @@ def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically append an event and update the session's durable state.""" return self._append_event_and_update_state( event_record, app_name, user_id, session_id, state, app_state=app_state, user_state=user_state @@ -129,7 +129,7 @@ def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Get events for a session.""" return self._get_events(app_name, user_id, session_id, after_timestamp, limit) @@ -599,7 +599,7 @@ def _json_placeholder(self) -> str: def _create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new session. Args: @@ -646,7 +646,7 @@ def _create_session( def _get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get session by ID. Args: @@ -696,7 +696,7 @@ def _get_session( if row is None: return None - return SessionRecord( + return StoredSession( id=row[0], app_name=row[1], user_id=row[2], @@ -759,7 +759,7 @@ def _delete_session(self, app_name: str, user_id: str, session_id: str) -> None: finally: cursor.close() - def _list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + def _list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": """List sessions for an app, optionally filtered by user. Args: @@ -794,7 +794,7 @@ def _list_sessions(self, app_name: str, user_id: str | None = None) -> "list[Ses rows = cursor.fetchall() return [ - SessionRecord( + StoredSession( id=row[0], app_name=row[1], user_id=row[2], @@ -812,7 +812,7 @@ def _list_sessions(self, app_name: str, user_id: str | None = None) -> "list[Ses return [] raise - def _insert_event(self, event_record: "EventRecord") -> None: + def _insert_event(self, event_record: "StoredEvent") -> None: """Insert an event record into the events table. Args: @@ -845,7 +845,7 @@ def _insert_event(self, event_record: "EventRecord") -> None: def _append_event_and_update_state( self, - event_record: "EventRecord", + event_record: "StoredEvent", app_name: str, user_id: str, session_id: str, @@ -853,7 +853,7 @@ def _append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically insert an event and update the session's durable state. The event insert, state update, and refresh-SELECT are executed within @@ -943,7 +943,7 @@ def _append_event_and_update_state( msg = f"Session {session_id} not found during append_event_and_update_state." raise ValueError(msg) - return SessionRecord( + return StoredSession( id=row[0], app_name=row[1], user_id=row[2], @@ -959,7 +959,7 @@ def _get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """List events for a session ordered by timestamp. Args: @@ -1000,7 +1000,7 @@ def _get_events( rows = cursor.fetchall() return [ - EventRecord( + StoredEvent( id=row[0], session_id=row[1], invocation_id=row[2] or "", @@ -1160,7 +1160,7 @@ def _set_metadata(self, key: str, value: str) -> None: finally: cursor.close() - def _append_event(self, event_record: EventRecord) -> None: + def _append_event(self, event_record: StoredEvent) -> None: """Synchronous implementation of append_event.""" self._insert_event(event_record) @@ -1186,23 +1186,28 @@ def create_tables(self) -> None: self._create_tables() - def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with deduplication.""" return self._insert_memory_entries(entries, owner_id) def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries by text query.""" - return self._search_entries(query, app_name, user_id, limit) + return self._search_entries(query, app_name, user_id, limit, scope_filter) def delete_entries_by_session(self, session_id: str) -> int: """Delete all memory entries for a specific session.""" return self._delete_entries_by_session(session_id) - def delete_entries_older_than(self, days: int) -> int: + def delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: """Delete memory entries older than specified days.""" - return self._delete_entries_older_than(days) + return self._delete_entries_older_than(days, app_name, scope) def _detect_dialect(self) -> str: driver_name = self._config.connection_config.get("driver_name", "").lower() @@ -1255,6 +1260,7 @@ def _memory_ddl_postgresql(self) -> str: session_id VARCHAR(128) NOT NULL, app_name VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, + scope VARCHAR(16) NOT NULL DEFAULT 'user', event_id VARCHAR(128) NOT NULL UNIQUE, author VARCHAR(256){owner_id_ddl}, timestamp TIMESTAMPTZ NOT NULL, @@ -1273,6 +1279,7 @@ def _memory_ddl_sqlite(self) -> str: session_id TEXT NOT NULL, app_name TEXT NOT NULL, user_id TEXT NOT NULL, + scope TEXT NOT NULL DEFAULT 'user', event_id TEXT NOT NULL UNIQUE, author TEXT{owner_id_ddl}, timestamp REAL NOT NULL, @@ -1291,6 +1298,7 @@ def _memory_ddl_duckdb(self) -> str: session_id VARCHAR(128) NOT NULL, app_name VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, + scope VARCHAR(16) NOT NULL DEFAULT 'user', event_id VARCHAR(128) NOT NULL UNIQUE, author VARCHAR(256){owner_id_ddl}, timestamp TIMESTAMP NOT NULL, @@ -1309,6 +1317,7 @@ def _memory_ddl_snowflake(self) -> str: session_id VARCHAR NOT NULL, app_name VARCHAR NOT NULL, user_id VARCHAR NOT NULL, + scope VARCHAR NOT NULL DEFAULT 'user', event_id VARCHAR NOT NULL UNIQUE, author VARCHAR{owner_id_ddl}, timestamp TIMESTAMP_TZ NOT NULL, @@ -1327,6 +1336,7 @@ def _memory_ddl_generic(self) -> str: session_id VARCHAR(128) NOT NULL, app_name VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, + scope VARCHAR(16) NOT NULL DEFAULT 'user', event_id VARCHAR(128) NOT NULL UNIQUE, author VARCHAR(256){owner_id_ddl}, timestamp TIMESTAMP NOT NULL, @@ -1350,8 +1360,12 @@ def _create_tables(self) -> None: cursor.execute(self._memory_table_ddl()) conn.commit() - idx_app_user = f"CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time ON {self._memory_table}(app_name, user_id, timestamp DESC)" - cursor.execute(idx_app_user) + idx_app_scope_user = f"CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_scope_user_time ON {self._memory_table}(app_name, scope, user_id, timestamp DESC)" + cursor.execute(idx_app_scope_user) + conn.commit() + + idx_scope = f"CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_scope ON {self._memory_table}(app_name, scope)" + cursor.execute(idx_scope) conn.commit() idx_session = ( @@ -1362,7 +1376,7 @@ def _create_tables(self) -> None: finally: cursor.close() - def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def _insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -1377,39 +1391,39 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec if use_returning: sql = f""" INSERT INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, content_text, metadata_json, inserted_at ) VALUES ( - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? + ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? ) ON CONFLICT(event_id) DO NOTHING RETURNING 1 """ else: sql = f""" INSERT INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, content_text, metadata_json, inserted_at ) VALUES ( - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? + ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? ) """ elif use_returning: sql = f""" INSERT INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at ) VALUES ( - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? + ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? ) ON CONFLICT(event_id) DO NOTHING RETURNING 1 """ else: sql = f""" INSERT INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at ) VALUES ( - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? + ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? ) """ @@ -1426,6 +1440,7 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], owner_id, @@ -1441,6 +1456,7 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], self._encode_timestamp(entry["timestamp"]), @@ -1469,8 +1485,13 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec return inserted_count def _search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -1480,13 +1501,13 @@ def _search_entries( effective_limit = limit if limit is not None else self._max_results pattern = f"%{query}%" + where_scope, scope_params = _build_adbc_scope_where(app_name, user_id, scope_filter) sql = f""" - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM {self._memory_table} - WHERE app_name = ? - AND user_id = ? + WHERE {where_scope} AND content_text LIKE ? ORDER BY timestamp DESC LIMIT ? @@ -1496,7 +1517,7 @@ def _search_entries( with self._config.provide_connection() as conn: cursor = conn.cursor() try: - cursor.execute(sql, (app_name, user_id, pattern, effective_limit)) + cursor.execute(sql, (*scope_params, pattern, effective_limit)) rows = cursor.fetchall() finally: cursor.close() @@ -1527,17 +1548,26 @@ def _delete_entries_by_session(self, session_id: str) -> int: finally: cursor.close() - def _delete_entries_older_than(self, days: int) -> int: + def _delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: cutoff = self._encode_timestamp(datetime.now(timezone.utc) - timedelta(days=days)) use_returning = self._dialect in {DIALECT_SQLITE, DIALECT_POSTGRESQL, DIALECT_DUCKDB} + clauses = ["inserted_at < ?"] + params: list[Any] = [cutoff] + if app_name is not None: + clauses.append("app_name = ?") + params.append(app_name) + if scope is not None: + clauses.append("scope = ?") + params.append(scope) + where_sql = " AND ".join(clauses) if use_returning: - sql = f"DELETE FROM {self._memory_table} WHERE inserted_at < ? RETURNING 1" + sql = f"DELETE FROM {self._memory_table} WHERE {where_sql} RETURNING 1" else: - sql = f"DELETE FROM {self._memory_table} WHERE inserted_at < ?" + sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" with self._config.provide_connection() as conn: cursor = conn.cursor() try: - cursor.execute(sql, (cutoff,)) + cursor.execute(sql, tuple(params)) if use_returning: deleted_rows = cursor.fetchall() conn.commit() @@ -1547,16 +1577,16 @@ def _delete_entries_older_than(self, days: int) -> int: finally: cursor.close() - def _rows_to_records(self, rows: "list[Any]") -> "list[MemoryRecord]": - records: list[MemoryRecord] = [] + def _rows_to_records(self, rows: "list[Any]") -> "list[StoredMemory]": + records: list[StoredMemory] = [] for row in rows: - content_json = row[7] + content_json = row[8] if isinstance(content_json, dict): content_value = content_json else: content_value = from_json(content_json if isinstance(content_json, (str, bytes)) else str(content_json)) - metadata_json = row[9] + metadata_json = row[10] if metadata_json is None: metadata_value = None elif isinstance(metadata_json, dict): @@ -1571,12 +1601,24 @@ def _rows_to_records(self, rows: "list[Any]") -> "list[MemoryRecord]": "session_id": row[1], "app_name": row[2], "user_id": row[3], - "event_id": row[4], - "author": row[5], - "timestamp": self._decode_timestamp(row[6]), + "scope": row[4], + "event_id": row[5], + "author": row[6], + "timestamp": self._decode_timestamp(row[7]), "content_json": content_value, - "content_text": row[8], + "content_text": row[9], "metadata_json": metadata_value, - "inserted_at": self._decode_timestamp(row[10]), + "inserted_at": self._decode_timestamp(row[11]), + "embedding": None, }) return records + + +def _build_adbc_scope_where( + app_name: str, user_id: str, scope_filter: Literal["all", "user", "app"] +) -> tuple[str, tuple[Any, ...]]: + if scope_filter == "all": + return "app_name = ? AND ((scope = 'user' AND user_id = ?) OR scope = 'app')", (app_name, user_id) + if scope_filter == "user": + return "app_name = ? AND scope = 'user' AND user_id = ?", (app_name, user_id) + return "app_name = ? AND scope = 'app'", (app_name,) diff --git a/sqlspec/adapters/aiomysql/adk/store.py b/sqlspec/adapters/aiomysql/adk/store.py index 518f449e2..c417bf8e6 100644 --- a/sqlspec/adapters/aiomysql/adk/store.py +++ b/sqlspec/adapters/aiomysql/adk/store.py @@ -2,14 +2,14 @@ import re from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast import pymysql.err from typing_extensions import NotRequired from sqlspec.adapters.aiomysql._typing import AiomysqlCursor, AiomysqlRawCursor from sqlspec.config import ADKConfig -from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore from sqlspec.utils.serializers import from_json, to_json @@ -17,7 +17,7 @@ from datetime import datetime, timedelta from sqlspec.adapters.aiomysql.config import AiomysqlConfig - from sqlspec.extensions.adk import MemoryRecord + from sqlspec.extensions.adk import StoredMemory __all__ = ("AiomysqlADKConfig", "AiomysqlADKMemoryStore", "AiomysqlADKStore") @@ -77,7 +77,7 @@ async def create_tables(self) -> None: async def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new session.""" state_json = to_json(state) @@ -110,7 +110,7 @@ async def create_session( async def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get session by scoped identifiers.""" if renew_for is not None and self._calculate_expires_at(renew_for) is not None: sql = f""" @@ -163,7 +163,7 @@ async def update_session_state(self, app_name: str, user_id: str, session_id: st await cursor.execute(sql, (to_json(state), app_name, user_id, session_id)) await conn.commit() - async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": """List sessions for an app, optionally filtered by user.""" if user_id is None: sql = f""" @@ -206,7 +206,7 @@ async def delete_session(self, app_name: str, user_id: str, session_id: str) -> await cursor.execute(sql, (app_name, user_id, session_id)) await conn.commit() - async def append_event(self, event_record: EventRecord) -> None: + async def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session.""" sql = f""" INSERT INTO {self._events_table} ( @@ -223,7 +223,7 @@ async def append_event(self, event_record: EventRecord) -> None: async def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -231,7 +231,7 @@ async def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically append an event and update session + scoped state.""" insert_sql = f""" INSERT INTO {self._events_table} ( @@ -293,7 +293,7 @@ async def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Get events for a session.""" if limit == 0: return [] @@ -512,7 +512,7 @@ async def create_tables(self) -> None: async with self._config.provide_session() as driver: await driver.execute_script(await self._memory_table_ddl()) - async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with deduplication.""" if not self._enabled: msg = "Memory store is disabled" @@ -525,17 +525,17 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " if self._owner_id_column_name: sql = f""" INSERT IGNORE INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) + ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """ else: sql = f""" INSERT IGNORE INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) + ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """ async with self._config.provide_connection() as conn: @@ -548,6 +548,7 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], owner_id, @@ -563,6 +564,7 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], entry["timestamp"], @@ -577,8 +579,13 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " return inserted_count async def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries by text query.""" if not self._enabled: msg = "Memory store is disabled" @@ -588,23 +595,24 @@ async def search_entries( return [] limit_value = limit or self._max_results + where_scope, scope_params = _build_mysql_scope_where(app_name, user_id, scope_filter) if self._use_fts: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = %s AND user_id = %s + WHERE {where_scope} AND MATCH(content_text) AGAINST (%s IN NATURAL LANGUAGE MODE) ORDER BY timestamp DESC LIMIT %s """ - params = (app_name, user_id, query, limit_value) + params = (*scope_params, query, limit_value) else: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = %s AND user_id = %s AND content_text LIKE %s + WHERE {where_scope} AND content_text LIKE %s ORDER BY timestamp DESC LIMIT %s """ - params = (app_name, user_id, f"%{query}%", limit_value) + params = (*scope_params, f"%{query}%", limit_value) async with ( self._config.provide_connection() as conn, @@ -614,7 +622,12 @@ async def search_entries( rows = await cursor.fetchall() columns = [col[0] for col in cursor.description or []] - return [cast("MemoryRecord", dict(zip(columns, row, strict=False))) for row in rows] + records: list[StoredMemory] = [] + for row in rows: + rec = cast("StoredMemory", dict(zip(columns, row, strict=False))) + rec["embedding"] = None + records.append(rec) + return records async def delete_entries_by_session(self, session_id: str) -> int: """Delete all memory entries for a specific session.""" @@ -631,21 +644,30 @@ async def delete_entries_by_session(self, session_id: str) -> int: await conn.commit() return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - async def delete_entries_older_than(self, days: int) -> int: + async def delete_entries_older_than( + self, days: int, app_name: "str | None" = None, scope: "str | None" = None + ) -> int: """Delete memory entries older than specified days.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) - sql = f""" - DELETE FROM {self._memory_table} - WHERE inserted_at < (UTC_TIMESTAMP(6) - INTERVAL %s DAY) - """ + clauses = ["inserted_at < (UTC_TIMESTAMP(6) - INTERVAL %s DAY)"] + params: list[Any] = [days] + if app_name is not None: + clauses.append("app_name = %s") + params.append(app_name) + if scope is not None: + clauses.append("scope = %s") + params.append(scope) + + where_sql = " AND ".join(clauses) + sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" async with ( self._config.provide_connection() as conn, AiomysqlCursor(conn, cursor_class=AiomysqlRawCursor) as cursor, ): - await cursor.execute(sql, (days,)) + await cursor.execute(sql, tuple(params)) await conn.commit() return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 @@ -671,6 +693,7 @@ async def _memory_table_ddl(self) -> str: session_id VARCHAR(128) NOT NULL, app_name VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, + scope VARCHAR(16) NOT NULL DEFAULT 'user', event_id VARCHAR(128) NOT NULL UNIQUE, author VARCHAR(256){owner_id_line}, timestamp TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), @@ -678,7 +701,8 @@ async def _memory_table_ddl(self) -> str: content_text TEXT NOT NULL, metadata_json JSON, inserted_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), - INDEX idx_{self._memory_table}_app_user_time (app_name, user_id, timestamp), + INDEX idx_{self._memory_table}_app_scope_user_time (app_name, scope, user_id, timestamp), + INDEX idx_{self._memory_table}_scope (app_name, scope), INDEX idx_{self._memory_table}_session (session_id){fts_index}{fk_constraint} ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{table_options} """ @@ -745,14 +769,14 @@ def _json_dict(value: Any) -> "dict[str, Any]": return cast("dict[str, Any]", value) -def _session_record_from_row(row: Any) -> SessionRecord: - return SessionRecord( +def _session_record_from_row(row: Any) -> StoredSession: + return StoredSession( id=row[0], app_name=row[1], user_id=row[2], state=_json_dict(row[3]), create_time=row[4], update_time=row[5] ) -def _event_record_from_row(row: Any) -> EventRecord: - return EventRecord( +def _event_record_from_row(row: Any) -> StoredEvent: + return StoredEvent( id=row[0], app_name=row[1], user_id=row[2], @@ -763,7 +787,7 @@ def _event_record_from_row(row: Any) -> EventRecord: ) -def _event_insert_params(event_record: EventRecord) -> "tuple[Any, ...]": +def _event_insert_params(event_record: StoredEvent) -> "tuple[Any, ...]": return ( event_record["id"], event_record["app_name"], @@ -887,3 +911,13 @@ def _mysql_upsert_metadata_sql(metadata_table: str) -> str: def _raise_session_not_found(session_id: str) -> None: msg = f"Session {session_id} not found during append_event_and_update_state." raise ValueError(msg) + + +def _build_mysql_scope_where( + app_name: str, user_id: str, scope_filter: Literal["all", "user", "app"] +) -> tuple[str, tuple[Any, ...]]: + if scope_filter == "all": + return "app_name = %s AND ((scope = 'user' AND user_id = %s) OR scope = 'app')", (app_name, user_id) + if scope_filter == "user": + return "app_name = %s AND scope = 'user' AND user_id = %s", (app_name, user_id) + return "app_name = %s AND scope = 'app'", (app_name,) diff --git a/sqlspec/adapters/aiosqlite/adk/store.py b/sqlspec/adapters/aiosqlite/adk/store.py index 158778c86..9adc5978f 100644 --- a/sqlspec/adapters/aiosqlite/adk/store.py +++ b/sqlspec/adapters/aiosqlite/adk/store.py @@ -11,13 +11,13 @@ from sqlspec.adapters.aiosqlite.config import _render_pragmas from sqlspec.config import ADKConfig from sqlspec.exceptions import ImproperConfigurationError -from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore from sqlspec.utils.serializers import from_json, to_json if TYPE_CHECKING: from sqlspec.adapters.aiosqlite.config import AiosqliteConfig - from sqlspec.extensions.adk import MemoryRecord + from sqlspec.extensions.adk import StoredMemory __all__ = ("AiosqliteADKConfig", "AiosqliteADKMemoryStore", "AiosqliteADKStore") @@ -91,7 +91,7 @@ async def create_tables(self) -> None: async def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new session. Args: @@ -128,13 +128,13 @@ async def create_session( await conn.execute(sql, params) await conn.commit() - return SessionRecord( + return StoredSession( id=session_id, app_name=app_name, user_id=user_id, state=state, create_time=now, update_time=now ) async def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get session by ID. Args: @@ -170,7 +170,7 @@ async def get_session( if row is None: return None - return SessionRecord( + return StoredSession( id=row[0], app_name=row[1], user_id=row[2], @@ -206,7 +206,7 @@ async def update_session_state(self, app_name: str, user_id: str, session_id: st await conn.execute(sql, (state_json, now_julian, app_name, user_id, session_id)) await conn.commit() - async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": """List sessions for an app, optionally filtered by user. Args: @@ -240,7 +240,7 @@ async def list_sessions(self, app_name: str, user_id: str | None = None) -> "lis rows = await cursor.fetchall() return [ - SessionRecord( + StoredSession( id=row[0], app_name=row[1], user_id=row[2], @@ -270,7 +270,7 @@ async def delete_session(self, app_name: str, user_id: str, session_id: str) -> await conn.execute(sql, (app_name, user_id, session_id)) await conn.commit() - async def append_event(self, event_record: EventRecord) -> None: + async def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session. Args: @@ -303,7 +303,7 @@ async def append_event(self, event_record: EventRecord) -> None: async def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -311,12 +311,12 @@ async def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically append an event and update the session's durable state. Inserts the event and updates the session state + update_time in a single transaction. Both operations succeed or fail together. Returns - the updated SessionRecord via SQLite RETURNING (3.35+). + the updated StoredSession via SQLite RETURNING (3.35+). Args: event_record: Event record to store. @@ -397,7 +397,7 @@ async def append_event_and_update_state( msg = f"Session {session_id} not found during append_event_and_update_state." raise ValueError(msg) - return SessionRecord( + return StoredSession( id=row[0], app_name=row[1], user_id=row[2], @@ -413,7 +413,7 @@ async def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Get events for a session. Args: @@ -453,7 +453,7 @@ async def get_events( rows = await cursor.fetchall() return [ - EventRecord( + StoredEvent( id=row[0], app_name=row[1], user_id=row[2], @@ -751,7 +751,7 @@ async def create_tables(self) -> None: await driver.execute_script(await self._memory_table_ddl()) await driver.commit() - async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with deduplication. Uses INSERT OR IGNORE to skip duplicates based on event_id unique constraint. @@ -767,19 +767,21 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " async with self._config.provide_connection() as conn: for entry in entries: params: tuple[Any, ...] + scope = entry.get("scope", "user") if self._owner_id_column_name: sql = f""" INSERT OR IGNORE INTO {self._memory_table} - (id, session_id, app_name, user_id, event_id, author, + (id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, content_text, metadata_json, inserted_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """ params = ( entry["id"], entry["session_id"], entry["app_name"], entry["user_id"], + scope, entry["event_id"], entry["author"], owner_id, @@ -792,15 +794,16 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " else: sql = f""" INSERT OR IGNORE INTO {self._memory_table} - (id, session_id, app_name, user_id, event_id, author, + (id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """ params = ( entry["id"], entry["session_id"], entry["app_name"], entry["user_id"], + scope, entry["event_id"], entry["author"], _datetime_to_julian(entry["timestamp"]), @@ -816,8 +819,13 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " return inserted_count async def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries by text query.""" if not self._enabled: msg = "Memory store is disabled" @@ -828,36 +836,39 @@ async def search_entries( limit_value = limit or self._max_results if self._use_fts: + where_scope, scope_params = _build_sqlite_scope_clause("m.", app_name, user_id, scope_filter) sql = f""" SELECT m.* FROM {self._memory_table} AS m JOIN {self._memory_table}_fts AS fts ON m.rowid = fts.rowid - WHERE m.app_name = ? AND m.user_id = ? AND fts.content_text MATCH ? + WHERE {where_scope} AND fts.content_text MATCH ? ORDER BY m.timestamp DESC LIMIT ? """ - params = (app_name, user_id, query, limit_value) + params = (*scope_params, query, limit_value) else: + where_scope, scope_params = _build_sqlite_scope_clause("", app_name, user_id, scope_filter) sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = ? AND user_id = ? AND content_text LIKE ? + WHERE {where_scope} AND content_text LIKE ? ORDER BY timestamp DESC LIMIT ? """ - params = (app_name, user_id, f"%{query}%", limit_value) + params = (*scope_params, f"%{query}%", limit_value) async with self._config.provide_connection() as conn: cursor = await conn.execute(sql, params) rows = await cursor.fetchall() columns = [col[0] for col in cursor.description or []] await cursor.close() - records: list[MemoryRecord] = [] + records: list[StoredMemory] = [] for row in rows: raw = dict(zip(columns, row, strict=False)) raw["timestamp"] = _julian_to_datetime(raw["timestamp"]) raw["inserted_at"] = _julian_to_datetime(raw["inserted_at"]) raw["content_json"] = from_json(raw["content_json"]) raw["metadata_json"] = from_json(raw["metadata_json"]) if raw["metadata_json"] else None - records.append(cast("MemoryRecord", raw)) + raw["embedding"] = None + records.append(cast("StoredMemory", raw)) return records async def delete_entries_by_session(self, session_id: str) -> int: @@ -872,7 +883,9 @@ async def delete_entries_by_session(self, session_id: str) -> int: await conn.commit() return cursor.rowcount - async def delete_entries_older_than(self, days: int) -> int: + async def delete_entries_older_than( + self, days: int, app_name: "str | None" = None, scope: "str | None" = None + ) -> int: """Delete memory entries older than specified days.""" if not self._enabled: msg = "Memory store is disabled" @@ -880,8 +893,16 @@ async def delete_entries_older_than(self, days: int) -> int: cutoff = _datetime_to_julian(datetime.now(timezone.utc)) - days sql = f"DELETE FROM {self._memory_table} WHERE inserted_at < ?" + params: list[Any] = [cutoff] + if app_name is not None: + sql += " AND app_name = ?" + params.append(app_name) + if scope is not None: + sql += " AND scope = ?" + params.append(scope) + async with self._config.provide_connection() as conn: - cursor = await conn.execute(sql, (cutoff,)) + cursor = await conn.execute(sql, tuple(params)) await conn.commit() return cursor.rowcount @@ -927,6 +948,7 @@ async def _memory_table_ddl(self) -> str: session_id TEXT NOT NULL, app_name TEXT NOT NULL, user_id TEXT NOT NULL, + scope TEXT NOT NULL DEFAULT 'user', event_id TEXT NOT NULL UNIQUE, author TEXT{owner_id_line}, timestamp REAL NOT NULL, @@ -936,8 +958,11 @@ async def _memory_table_ddl(self) -> str: inserted_at REAL NOT NULL ); - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time - ON {self._memory_table}(app_name, user_id, timestamp DESC); + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_scope_user_time + ON {self._memory_table}(app_name, scope, user_id, timestamp DESC); + + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_scope + ON {self._memory_table}(app_name, scope); CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session ON {self._memory_table}(session_id); @@ -1039,3 +1064,16 @@ def _julian_to_datetime(julian: float) -> datetime: days_since_epoch = julian - JULIAN_EPOCH timestamp = days_since_epoch * SECONDS_PER_DAY return datetime.fromtimestamp(timestamp, tz=timezone.utc) + + +def _build_sqlite_scope_clause( + prefix: str, app_name: str, user_id: str, scope_filter: "Literal['all', 'user', 'app']" +) -> "tuple[str, tuple[Any, ...]]": + if scope_filter == "all": + return ( + f"{prefix}app_name = ? AND (({prefix}scope = 'user' AND {prefix}user_id = ?) OR {prefix}scope = 'app')", + (app_name, user_id), + ) + if scope_filter == "user": + return f"{prefix}app_name = ? AND {prefix}scope = 'user' AND {prefix}user_id = ?", (app_name, user_id) + return f"{prefix}app_name = ? AND {prefix}scope = 'app'", (app_name,) diff --git a/sqlspec/adapters/arrow_odbc/adk/store.py b/sqlspec/adapters/arrow_odbc/adk/store.py index a8ef0e1a8..7d314bef3 100644 --- a/sqlspec/adapters/arrow_odbc/adk/store.py +++ b/sqlspec/adapters/arrow_odbc/adk/store.py @@ -1,14 +1,14 @@ """arrow-odbc ADK stores for Google Agent Development Kit session storage.""" from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, ClassVar, Final, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast from typing_extensions import NotRequired from sqlspec.config import ADKConfig from sqlspec.exceptions import SQLSpecError -from sqlspec.extensions.adk import BaseSyncADKStore, EventRecord, SessionRecord -from sqlspec.extensions.adk.memory import BaseSyncADKMemoryStore, MemoryRecord +from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession +from sqlspec.extensions.adk.memory import BaseSyncADKMemoryStore, StoredMemory from sqlspec.utils.serializers import from_json, to_json if TYPE_CHECKING: @@ -65,7 +65,7 @@ def create_tables(self) -> None: def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new ADK session.""" owner_column = f", {_quote_identifier(self._owner_id_column_name)}" if self._owner_id_column_name else "" owner_param = ", ?" if self._owner_id_column_name else "" @@ -93,7 +93,7 @@ def create_session( def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Return a scoped session or ``None`` if absent.""" try: with self._config.provide_session() as driver: @@ -129,7 +129,7 @@ def update_session_state(self, app_name: str, user_id: str, session_id: str, sta commit=True, ) - def list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[StoredSession]": """List ADK sessions for an application, optionally scoped to a user.""" if user_id is None: sql = f""" @@ -163,13 +163,13 @@ def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: commit=True, ) - def append_event(self, event_record: EventRecord) -> None: + def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session.""" self._execute(_insert_event_sql(self._events_table), _event_insert_params(event_record), commit=True) def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -177,7 +177,7 @@ def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically append an event and update durable session/scoped state.""" with self._config.provide_session() as driver: driver.execute( @@ -206,7 +206,7 @@ def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Return events for a scoped session ordered by event timestamp.""" if limit is not None and limit <= 0: return [] @@ -414,7 +414,7 @@ def create_tables(self) -> None: driver.execute(_create_index_sql(index_table, index_name, columns)) driver.commit() - def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Insert memory entries, skipping duplicates by event_id.""" if not self._enabled: msg = "Memory store is disabled" @@ -442,10 +442,10 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object driver.execute( f""" INSERT INTO {_table_ref(self._memory_table)} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at{owner_column} ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?{owner_param}) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?{owner_param}) """, params, ) @@ -454,8 +454,13 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object return inserted_count def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries with SQL Server LIKE matching.""" if not self._enabled: msg = "Memory store is disabled" @@ -463,18 +468,18 @@ def search_entries( effective_limit = max(0, int(limit if limit is not None else self._max_results)) if effective_limit == 0: return [] + where_scope, scope_params = _build_arrow_odbc_scope_where(app_name, user_id, scope_filter) rows = self._execute_fetchall( f""" - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM {_table_ref(self._memory_table)} - WHERE app_name = ? - AND user_id = ? + WHERE {where_scope} AND content_text LIKE ? ORDER BY timestamp DESC OFFSET 0 ROWS FETCH NEXT {effective_limit} ROWS ONLY """, - (app_name, user_id, f"%{query}%"), + (*scope_params, f"%{query}%"), ) return [_memory_record_from_row(row) for row in rows] @@ -486,19 +491,23 @@ def delete_entries_by_session(self, session_id: str) -> int: self._execute(f"DELETE FROM {_table_ref(self._memory_table)} WHERE session_id = ?", (session_id,), commit=True) return count - def delete_entries_older_than(self, days: int) -> int: + def delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: """Delete memory entries older than ``days`` days.""" cutoff = datetime.now(timezone.utc).timestamp() - (days * 86_400) cutoff_dt = datetime.fromtimestamp(cutoff, tz=timezone.utc) + clauses = ["inserted_at < ?"] + params: list[Any] = [_format_datetime(cutoff_dt)] + if app_name is not None: + clauses.append("app_name = ?") + params.append(app_name) + if scope is not None: + clauses.append("scope = ?") + params.append(scope) + where_sql = " AND ".join(clauses) count = self._select_count( - f"SELECT COUNT(*) AS row_count FROM {_table_ref(self._memory_table)} WHERE inserted_at < ?", - (_format_datetime(cutoff_dt),), - ) - self._execute( - f"DELETE FROM {_table_ref(self._memory_table)} WHERE inserted_at < ?", - (_format_datetime(cutoff_dt),), - commit=True, + f"SELECT COUNT(*) AS row_count FROM {_table_ref(self._memory_table)} WHERE {where_sql}", tuple(params) ) + self._execute(f"DELETE FROM {_table_ref(self._memory_table)} WHERE {where_sql}", tuple(params), commit=True) return count def _memory_table_ddl(self) -> str: @@ -509,6 +518,7 @@ def _memory_table_ddl(self) -> str: session_id NVARCHAR(128) NOT NULL, app_name NVARCHAR(128) NOT NULL, user_id NVARCHAR(128) NOT NULL, + scope NVARCHAR(16) NOT NULL DEFAULT 'user', event_id NVARCHAR(128) NOT NULL, author NVARCHAR(256) NULL, timestamp DATETIME2(6) NOT NULL, @@ -524,7 +534,12 @@ def _memory_table_ddl(self) -> str: def _memory_index_specs(self) -> "list[tuple[str, str, str]]": """Return ``(index_name, table, columns)`` specs for memory-table indexes.""" return [ - (f"idx_{self._memory_table}_app_user_time", self._memory_table, "app_name, user_id, timestamp DESC"), + ( + f"idx_{self._memory_table}_app_scope_user_time", + self._memory_table, + "app_name, scope, user_id, timestamp DESC", + ), + (f"idx_{self._memory_table}_scope", self._memory_table, "app_name, scope"), (f"idx_{self._memory_table}_session", self._memory_table, "session_id"), ] @@ -717,7 +732,7 @@ def _events_query( return sql, tuple(params) -def _event_insert_params(event_record: EventRecord) -> "tuple[Any, ...]": +def _event_insert_params(event_record: StoredEvent) -> "tuple[Any, ...]": return ( event_record["id"], event_record["app_name"], @@ -729,8 +744,8 @@ def _event_insert_params(event_record: EventRecord) -> "tuple[Any, ...]": ) -def _session_record_from_row(row: Any) -> SessionRecord: - return SessionRecord( +def _session_record_from_row(row: Any) -> StoredSession: + return StoredSession( id=str(_row_value(row, "id", 0)), app_name=str(_row_value(row, "app_name", 1)), user_id=str(_row_value(row, "user_id", 2)), @@ -740,8 +755,8 @@ def _session_record_from_row(row: Any) -> SessionRecord: ) -def _event_record_from_row(row: Any) -> EventRecord: - return EventRecord( +def _event_record_from_row(row: Any) -> StoredEvent: + return StoredEvent( id=str(_row_value(row, "id", 0)), app_name=str(_row_value(row, "app_name", 1)), user_id=str(_row_value(row, "user_id", 2)), @@ -752,12 +767,13 @@ def _event_record_from_row(row: Any) -> EventRecord: ) -def _memory_insert_params(entry: MemoryRecord) -> "tuple[Any, ...]": +def _memory_insert_params(entry: StoredMemory) -> "tuple[Any, ...]": return ( entry["id"], entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], _format_datetime(entry["timestamp"]), @@ -768,19 +784,21 @@ def _memory_insert_params(entry: MemoryRecord) -> "tuple[Any, ...]": ) -def _memory_record_from_row(row: Any) -> MemoryRecord: - return MemoryRecord( +def _memory_record_from_row(row: Any) -> StoredMemory: + return StoredMemory( id=str(_row_value(row, "id", 0)), session_id=str(_row_value(row, "session_id", 1)), app_name=str(_row_value(row, "app_name", 2)), user_id=str(_row_value(row, "user_id", 3)), - event_id=str(_row_value(row, "event_id", 4)), - author=cast("str | None", _row_value(row, "author", 5)), - timestamp=_datetime_value(_row_value(row, "timestamp", 6)), - content_json=_json_dict(_row_value(row, "content_json", 7)), - content_text=str(_row_value(row, "content_text", 8) or ""), - metadata_json=_optional_json_dict(_row_value(row, "metadata_json", 9)), - inserted_at=_datetime_value(_row_value(row, "inserted_at", 10)), + scope=str(_row_value(row, "scope", 4) or "user"), + event_id=str(_row_value(row, "event_id", 5)), + author=cast("str | None", _row_value(row, "author", 6)), + timestamp=_datetime_value(_row_value(row, "timestamp", 7)), + content_json=_json_dict(_row_value(row, "content_json", 8)), + content_text=str(_row_value(row, "content_text", 9) or ""), + metadata_json=_optional_json_dict(_row_value(row, "metadata_json", 10)), + inserted_at=_datetime_value(_row_value(row, "inserted_at", 11)), + embedding=None, ) @@ -862,3 +880,13 @@ def _constraint_ref(prefix: str, table: str, suffix: str) -> str: def _raise_session_not_found(session_id: str) -> None: msg = f"Session {session_id} not found during append_event_and_update_state." raise ValueError(msg) + + +def _build_arrow_odbc_scope_where( + app_name: str, user_id: str, scope_filter: Literal["all", "user", "app"] +) -> tuple[str, tuple[Any, ...]]: + if scope_filter == "all": + return "app_name = ? AND ((scope = 'user' AND user_id = ?) OR scope = 'app')", (app_name, user_id) + if scope_filter == "user": + return "app_name = ? AND scope = 'user' AND user_id = ?", (app_name, user_id) + return "app_name = ? AND scope = 'app'", (app_name,) diff --git a/sqlspec/adapters/asyncmy/adk/store.py b/sqlspec/adapters/asyncmy/adk/store.py index 2416e6a4d..612fb83f7 100644 --- a/sqlspec/adapters/asyncmy/adk/store.py +++ b/sqlspec/adapters/asyncmy/adk/store.py @@ -2,13 +2,13 @@ import re from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast import asyncmy from typing_extensions import NotRequired from sqlspec.config import ADKConfig -from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore from sqlspec.utils.serializers import from_json, to_json @@ -16,7 +16,7 @@ from datetime import datetime, timedelta from sqlspec.adapters.asyncmy.config import AsyncmyConfig - from sqlspec.extensions.adk import MemoryRecord + from sqlspec.extensions.adk import StoredMemory __all__ = ("AsyncmyADKConfig", "AsyncmyADKMemoryStore", "AsyncmyADKStore") @@ -76,7 +76,7 @@ async def create_tables(self) -> None: async def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new session.""" params: tuple[Any, ...] if self._owner_id_column_name: @@ -104,7 +104,7 @@ async def create_session( async def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get session by scoped identifiers.""" try: async with self._config.provide_connection() as conn, conn.cursor() as cursor: @@ -146,7 +146,7 @@ async def update_session_state(self, app_name: str, user_id: str, session_id: st await cursor.execute(sql, (to_json(state), app_name, user_id, session_id)) await conn.commit() - async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": """List sessions for an app, optionally filtered by user.""" if user_id is None: sql = f""" @@ -183,7 +183,7 @@ async def delete_session(self, app_name: str, user_id: str, session_id: str) -> await cursor.execute(sql, (app_name, user_id, session_id)) await conn.commit() - async def append_event(self, event_record: EventRecord) -> None: + async def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session.""" sql = f""" INSERT INTO {self._events_table} ( @@ -197,7 +197,7 @@ async def append_event(self, event_record: EventRecord) -> None: async def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -205,7 +205,7 @@ async def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically append an event and update session + scoped state.""" insert_sql = f""" INSERT INTO {self._events_table} ( @@ -264,7 +264,7 @@ async def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Get events for a session.""" if limit == 0: return [] @@ -456,7 +456,7 @@ async def create_tables(self) -> None: async with self._config.provide_session() as driver: await driver.execute_script(await self._memory_table_ddl()) - async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with deduplication.""" if not self._enabled: msg = "Memory store is disabled" @@ -469,17 +469,17 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " if self._owner_id_column_name: sql = f""" INSERT IGNORE INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) + ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """ else: sql = f""" INSERT IGNORE INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) + ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """ async with self._config.provide_connection() as conn: @@ -492,6 +492,7 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], owner_id, @@ -507,6 +508,7 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], entry["timestamp"], @@ -521,8 +523,13 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " return inserted_count async def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries by text query.""" if not self._enabled: msg = "Memory store is disabled" @@ -532,30 +539,36 @@ async def search_entries( return [] limit_value = limit or self._max_results + where_scope, scope_params = _build_mysql_scope_where(app_name, user_id, scope_filter) if self._use_fts: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = %s AND user_id = %s + WHERE {where_scope} AND MATCH(content_text) AGAINST (%s IN NATURAL LANGUAGE MODE) ORDER BY timestamp DESC LIMIT %s """ - params = (app_name, user_id, query, limit_value) + params = (*scope_params, query, limit_value) else: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = %s AND user_id = %s AND content_text LIKE %s + WHERE {where_scope} AND content_text LIKE %s ORDER BY timestamp DESC LIMIT %s """ - params = (app_name, user_id, f"%{query}%", limit_value) + params = (*scope_params, f"%{query}%", limit_value) async with self._config.provide_connection() as conn, conn.cursor() as cursor: await cursor.execute(sql, params) rows = await cursor.fetchall() columns = [col[0] for col in cursor.description or []] - return [cast("MemoryRecord", dict(zip(columns, row, strict=False))) for row in rows] + records: list[StoredMemory] = [] + for row in rows: + rec = cast("StoredMemory", dict(zip(columns, row, strict=False))) + rec["embedding"] = None + records.append(rec) + return records async def delete_entries_by_session(self, session_id: str) -> int: """Delete all memory entries for a specific session.""" @@ -569,18 +582,27 @@ async def delete_entries_by_session(self, session_id: str) -> int: await conn.commit() return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - async def delete_entries_older_than(self, days: int) -> int: + async def delete_entries_older_than( + self, days: int, app_name: "str | None" = None, scope: "str | None" = None + ) -> int: """Delete memory entries older than specified days.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) - sql = f""" - DELETE FROM {self._memory_table} - WHERE inserted_at < (UTC_TIMESTAMP(6) - INTERVAL %s DAY) - """ + clauses = ["inserted_at < (UTC_TIMESTAMP(6) - INTERVAL %s DAY)"] + params: list[Any] = [days] + if app_name is not None: + clauses.append("app_name = %s") + params.append(app_name) + if scope is not None: + clauses.append("scope = %s") + params.append(scope) + + where_sql = " AND ".join(clauses) + sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" async with self._config.provide_connection() as conn, conn.cursor() as cursor: - await cursor.execute(sql, (days,)) + await cursor.execute(sql, tuple(params)) await conn.commit() return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 @@ -606,6 +628,7 @@ async def _memory_table_ddl(self) -> str: session_id VARCHAR(128) NOT NULL, app_name VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, + scope VARCHAR(16) NOT NULL DEFAULT 'user', event_id VARCHAR(128) NOT NULL UNIQUE, author VARCHAR(256){owner_id_line}, timestamp TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), @@ -613,7 +636,8 @@ async def _memory_table_ddl(self) -> str: content_text TEXT NOT NULL, metadata_json JSON, inserted_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), - INDEX idx_{self._memory_table}_app_user_time (app_name, user_id, timestamp), + INDEX idx_{self._memory_table}_app_scope_user_time (app_name, scope, user_id, timestamp), + INDEX idx_{self._memory_table}_scope (app_name, scope), INDEX idx_{self._memory_table}_session (session_id){fts_index}{fk_constraint} ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{table_options} """ @@ -685,14 +709,14 @@ def _json_dict(value: Any) -> "dict[str, Any]": return cast("dict[str, Any]", value) -def _session_record_from_row(row: Any) -> SessionRecord: - return SessionRecord( +def _session_record_from_row(row: Any) -> StoredSession: + return StoredSession( id=row[0], app_name=row[1], user_id=row[2], state=_json_dict(row[3]), create_time=row[4], update_time=row[5] ) -def _event_record_from_row(row: Any) -> EventRecord: - return EventRecord( +def _event_record_from_row(row: Any) -> StoredEvent: + return StoredEvent( id=row[0], app_name=row[1], user_id=row[2], @@ -703,7 +727,7 @@ def _event_record_from_row(row: Any) -> EventRecord: ) -def _event_insert_params(event_record: EventRecord) -> "tuple[Any, ...]": +def _event_insert_params(event_record: StoredEvent) -> "tuple[Any, ...]": return ( event_record["id"], event_record["app_name"], @@ -827,3 +851,13 @@ def _mysql_upsert_metadata_sql(metadata_table: str) -> str: def _raise_session_not_found(session_id: str) -> None: msg = f"Session {session_id} not found during append_event_and_update_state." raise ValueError(msg) + + +def _build_mysql_scope_where( + app_name: str, user_id: str, scope_filter: Literal["all", "user", "app"] +) -> tuple[str, tuple[Any, ...]]: + if scope_filter == "all": + return "app_name = %s AND ((scope = 'user' AND user_id = %s) OR scope = 'app')", (app_name, user_id) + if scope_filter == "user": + return "app_name = %s AND scope = 'user' AND user_id = %s", (app_name, user_id) + return "app_name = %s AND scope = 'app'", (app_name,) diff --git a/sqlspec/adapters/asyncpg/adk/store.py b/sqlspec/adapters/asyncpg/adk/store.py index f163970da..646340f46 100644 --- a/sqlspec/adapters/asyncpg/adk/store.py +++ b/sqlspec/adapters/asyncpg/adk/store.py @@ -1,19 +1,19 @@ """AsyncPG ADK store for Google Agent Development Kit session/event storage.""" -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast import asyncpg from typing_extensions import NotRequired from sqlspec.config import ADKConfig, AsyncConfigT -from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore if TYPE_CHECKING: from datetime import datetime, timedelta from sqlspec.adapters.asyncpg.config import AsyncpgConfig - from sqlspec.extensions.adk import MemoryRecord + from sqlspec.extensions.adk import StoredMemory __all__ = ("AsyncpgADKConfig", "AsyncpgADKMemoryStore", "AsyncpgADKStore") @@ -83,7 +83,7 @@ async def create_tables(self) -> None: async def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: async with self._config.provide_connection() as conn: if self._owner_id_column_name: sql = f""" @@ -107,7 +107,7 @@ async def create_session( async def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": if renew_for is not None and self._calculate_expires_at(renew_for) is not None: sql = f""" UPDATE {self._session_table} @@ -131,7 +131,7 @@ async def get_session( if row is None: return None - return SessionRecord( + return StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -158,7 +158,7 @@ async def delete_session(self, app_name: str, user_id: str, session_id: str) -> async with self._config.provide_connection() as conn: await conn.execute(sql, app_name, user_id, session_id) - async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": if user_id is None: sql = f""" SELECT id, app_name, user_id, state, create_time, update_time @@ -181,7 +181,7 @@ async def list_sessions(self, app_name: str, user_id: str | None = None) -> "lis rows = await conn.fetch(sql, *params) return [ - SessionRecord( + StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -194,7 +194,7 @@ async def list_sessions(self, app_name: str, user_id: str | None = None) -> "lis except asyncpg.exceptions.UndefinedTableError: return [] - async def append_event(self, event_record: EventRecord) -> None: + async def append_event(self, event_record: StoredEvent) -> None: sql = f""" INSERT INTO {self._events_table} ( id, session_id, invocation_id, timestamp, event_data @@ -213,7 +213,7 @@ async def append_event(self, event_record: EventRecord) -> None: async def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -221,7 +221,7 @@ async def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: insert_sql = f""" INSERT INTO {self._events_table} ( id, session_id, invocation_id, timestamp, event_data @@ -266,7 +266,7 @@ async def append_event_and_update_state( if user_state is not None: await conn.execute(user_upsert_sql, app_name, user_id, user_state) - return SessionRecord( + return StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -282,7 +282,7 @@ async def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": if limit == 0: return [] @@ -311,7 +311,7 @@ async def get_events( rows = await conn.fetch(sql, *params) return [ - EventRecord( + StoredEvent( id=row["id"], session_id=row["session_id"], invocation_id=row["invocation_id"], @@ -554,7 +554,7 @@ async def create_tables(self) -> None: async with self._config.provide_session() as driver: await driver.execute_script(await self._memory_table_ddl()) - async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -568,10 +568,10 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " if self._owner_id_column_name: sql = f""" INSERT INTO {self._memory_table} - (id, session_id, app_name, user_id, event_id, author, + (id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, content_text, metadata_json, inserted_at) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) ON CONFLICT (event_id) DO NOTHING """ result = await conn.execute( @@ -580,6 +580,7 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], owner_id, @@ -592,9 +593,9 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " else: sql = f""" INSERT INTO {self._memory_table} - (id, session_id, app_name, user_id, event_id, author, + (id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) ON CONFLICT (event_id) DO NOTHING """ result = await conn.execute( @@ -603,6 +604,7 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], entry["timestamp"], @@ -619,8 +621,13 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " return inserted_count async def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -631,27 +638,43 @@ async def search_entries( from typing import cast limit_value = limit or self._max_results + if scope_filter == "all": + where_scope = "app_name = $1 AND ((scope = 'user' AND user_id = $2) OR scope = 'app')" + scope_params: tuple[Any, ...] = (app_name, user_id) + p_q = "$3" + p_lim = "$4" + elif scope_filter == "user": + where_scope = "app_name = $1 AND scope = 'user' AND user_id = $2" + scope_params = (app_name, user_id) + p_q = "$3" + p_lim = "$4" + else: + where_scope = "app_name = $1 AND scope = 'app'" + scope_params = (app_name,) + p_q = "$2" + p_lim = "$3" + if self._use_fts: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = $1 AND user_id = $2 - AND to_tsvector('english', content_text) @@ plainto_tsquery('english', $3) + WHERE {where_scope} + AND to_tsvector('english', content_text) @@ plainto_tsquery('english', {p_q}) ORDER BY timestamp DESC - LIMIT $4 + LIMIT {p_lim} """ - params = (app_name, user_id, query, limit_value) + params = (*scope_params, query, limit_value) else: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = $1 AND user_id = $2 AND content_text ILIKE $3 + WHERE {where_scope} AND content_text ILIKE {p_q} ORDER BY timestamp DESC - LIMIT $4 + LIMIT {p_lim} """ - params = (app_name, user_id, f"%{query}%", limit_value) + params = (*scope_params, f"%{query}%", limit_value) async with self._config.provide_connection() as conn: rows = await conn.fetch(sql, *params) - return [cast("MemoryRecord", dict(row)) for row in rows] + return [cast("StoredMemory", dict(row)) for row in rows] async def delete_entries_by_session(self, session_id: str) -> int: if not self._enabled: @@ -666,17 +689,29 @@ async def delete_entries_by_session(self, session_id: str) -> int: except (IndexError, ValueError): return 0 - async def delete_entries_older_than(self, days: int) -> int: + async def delete_entries_older_than( + self, days: int, app_name: "str | None" = None, scope: "str | None" = None + ) -> int: if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) - sql = f""" - DELETE FROM {self._memory_table} - WHERE inserted_at < (CURRENT_TIMESTAMP - ($1::int * INTERVAL '1 day')) - """ + clauses = ["inserted_at < (CURRENT_TIMESTAMP - ($1::int * INTERVAL '1 day'))"] + params: list[Any] = [days] + idx = 2 + if app_name is not None: + clauses.append(f"app_name = ${idx}") + params.append(app_name) + idx += 1 + if scope is not None: + clauses.append(f"scope = ${idx}") + params.append(scope) + idx += 1 + + where_sql = " AND ".join(clauses) + sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" async with self._config.provide_connection() as conn: - result = await conn.execute(sql, days) + result = await conn.execute(sql, *params) try: return int(result.split(" ")[1]) except (IndexError, ValueError): @@ -700,6 +735,7 @@ async def _memory_table_ddl(self) -> str: session_id VARCHAR(128) NOT NULL, app_name VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, + scope VARCHAR(16) NOT NULL DEFAULT 'user', event_id VARCHAR(128) NOT NULL UNIQUE, author VARCHAR(256){owner_id_line}, timestamp TIMESTAMPTZ NOT NULL, @@ -709,8 +745,11 @@ async def _memory_table_ddl(self) -> str: inserted_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP ); - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time - ON {self._memory_table}(app_name, user_id, timestamp DESC); + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_scope_user_time + ON {self._memory_table}(app_name, scope, user_id, timestamp DESC); + + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_scope + ON {self._memory_table}(app_name, scope); CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session ON {self._memory_table}(session_id); diff --git a/sqlspec/adapters/bigquery/adk/store.py b/sqlspec/adapters/bigquery/adk/store.py index 749ce1f2c..80c5c469c 100644 --- a/sqlspec/adapters/bigquery/adk/store.py +++ b/sqlspec/adapters/bigquery/adk/store.py @@ -17,7 +17,7 @@ from sqlspec.adapters.bigquery.config import BigQueryConfig from sqlspec.config import ADKConfig -from sqlspec.extensions.adk import BaseSyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk._config_utils import _adk_config_from_extension from sqlspec.utils.serializers import from_json, to_json from sqlspec.utils.uuids import uuid4 @@ -95,13 +95,13 @@ def create_tables(self) -> None: def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create or refresh a session row for the analytics replica.""" return self._create_session(session_id, app_name, user_id, state, owner_id) def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get a session by app, user, and session identifier.""" return self._get_session(app_name, user_id, session_id, renew_for=renew_for) @@ -109,7 +109,7 @@ def update_session_state(self, app_name: str, user_id: str, session_id: str, sta """Replace the durable session state snapshot.""" self._update_session_state(app_name, user_id, session_id, state) - def list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[StoredSession]": """List sessions for an app, optionally filtered by user.""" return self._list_sessions(app_name, user_id) @@ -117,13 +117,13 @@ def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: """Delete a session and its replicated events.""" self._delete_session(app_name, user_id, session_id) - def append_event(self, event_record: EventRecord) -> None: + def append_event(self, event_record: StoredEvent) -> None: """Append an ADK event blob.""" self._append_event(event_record) def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -131,7 +131,7 @@ def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Append an event and then update analytics-replica state. BigQuery has no cross-statement transaction for this path. The method @@ -149,7 +149,7 @@ def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Get events for a session.""" return self._get_events(app_name, user_id, session_id, after_timestamp, limit) @@ -216,7 +216,7 @@ def _json_param(self, name: str, value: "dict[str, Any] | None") -> Any: def _create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: now = datetime.now(timezone.utc) owner_column = f", {self._owner_id_column_name}" if self._owner_id_column_name else "" owner_select = ", @owner_id AS owner_id" if self._owner_id_column_name else "" @@ -257,7 +257,7 @@ def _create_session( def _get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": if renew_for is not None and self._calculate_expires_at(renew_for) is not None: self._update_session_touch(app_name, user_id, session_id) @@ -313,7 +313,7 @@ def _update_session_state(self, app_name: str, user_id: str, session_id: str, st ], ) - def _list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[SessionRecord]": + def _list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[StoredSession]": window_start = datetime.now(timezone.utc) - timedelta(days=self._lookup_window_days) sql = f""" SELECT id, app_name, user_id, state, create_time, update_time @@ -343,7 +343,7 @@ def _delete_session(self, app_name: str, user_id: str, session_id: str) -> None: self._run_query(events_sql, params) self._run_query(sessions_sql, params) - def _append_event(self, event_record: EventRecord) -> None: + def _append_event(self, event_record: StoredEvent) -> None: sql = f""" INSERT INTO {self._qualified(self._events_table)} (id, app_name, user_id, session_id, invocation_id, timestamp, event_data) @@ -362,7 +362,7 @@ def _append_event(self, event_record: EventRecord) -> None: def _append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -370,7 +370,7 @@ def _append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: self._append_event(event_record) self._update_session_state(app_name, user_id, session_id, state) if app_state is not None: @@ -391,7 +391,7 @@ def _get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": sql = f""" SELECT e.id, e.app_name, e.user_id, e.session_id, e.invocation_id, e.timestamp, e.event_data FROM {self._qualified(self._events_table)} e @@ -601,7 +601,7 @@ def _new_id() -> str: return str(uuid4()) -def _session_record_from_row(row: "dict[str, Any]") -> SessionRecord: +def _session_record_from_row(row: "dict[str, Any]") -> StoredSession: return { "id": row["id"], "app_name": row["app_name"], @@ -612,7 +612,7 @@ def _session_record_from_row(row: "dict[str, Any]") -> SessionRecord: } -def _event_record_from_row(row: "dict[str, Any]") -> EventRecord: +def _event_record_from_row(row: "dict[str, Any]") -> StoredEvent: return { "id": row["id"], "app_name": row["app_name"], diff --git a/sqlspec/adapters/cockroach_asyncpg/adk/store.py b/sqlspec/adapters/cockroach_asyncpg/adk/store.py index 05fec237f..80831fa15 100644 --- a/sqlspec/adapters/cockroach_asyncpg/adk/store.py +++ b/sqlspec/adapters/cockroach_asyncpg/adk/store.py @@ -1,12 +1,12 @@ """CockroachDB ADK store for Google Agent Development Kit session/event storage (asyncpg).""" -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Literal, cast import asyncpg from typing_extensions import NotRequired from sqlspec.config import ADKConfig -from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore from sqlspec.utils.logging import get_logger @@ -14,7 +14,7 @@ from datetime import datetime, timedelta from sqlspec.adapters.cockroach_asyncpg.config import CockroachAsyncpgConfig - from sqlspec.extensions.adk import MemoryRecord + from sqlspec.extensions.adk import StoredMemory __all__ = ("CockroachAsyncpgADKConfig", "CockroachAsyncpgADKMemoryStore", "CockroachAsyncpgADKStore") @@ -84,7 +84,7 @@ async def create_tables(self) -> None: async def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: params: tuple[Any, ...] if self._owner_id_column_name: sql = f""" @@ -110,7 +110,7 @@ async def create_session( async def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": if renew_for is not None and self._calculate_expires_at(renew_for) is not None: sql = f""" UPDATE {self._session_table} @@ -133,7 +133,7 @@ async def get_session( if row is None: return None - return SessionRecord( + return StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -154,7 +154,7 @@ async def update_session_state(self, app_name: str, user_id: str, session_id: st async with self._config.provide_connection() as conn: await conn.execute(sql, state, app_name, user_id, session_id) - async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": if user_id is None: sql = f""" SELECT id, app_name, user_id, state, create_time, update_time @@ -179,7 +179,7 @@ async def list_sessions(self, app_name: str, user_id: str | None = None) -> "lis return [] return [ - SessionRecord( + StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -196,7 +196,7 @@ async def delete_session(self, app_name: str, user_id: str, session_id: str) -> async with self._config.provide_connection() as conn: await conn.execute(sql, app_name, user_id, session_id) - async def append_event(self, event_record: EventRecord) -> None: + async def append_event(self, event_record: StoredEvent) -> None: sql = f""" INSERT INTO {self._events_table} ( id, session_id, invocation_id, timestamp, event_data @@ -215,7 +215,7 @@ async def append_event(self, event_record: EventRecord) -> None: async def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -223,7 +223,7 @@ async def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: insert_sql = f""" INSERT INTO {self._events_table} ( id, session_id, invocation_id, timestamp, event_data @@ -270,7 +270,7 @@ async def append_event_and_update_state( msg = f"Session {session_id} not found during append_event_and_update_state." raise ValueError(msg) - return SessionRecord( + return StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -286,7 +286,7 @@ async def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": if limit == 0: return [] @@ -317,7 +317,7 @@ async def get_events( return [] return [ - EventRecord( + StoredEvent( id=row["id"], session_id=row["session_id"], invocation_id=row["invocation_id"], @@ -534,7 +534,7 @@ async def create_tables(self) -> None: async with self._config.provide_session() as driver: await driver.execute_script(await self._memory_table_ddl()) - async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -548,7 +548,7 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " if self._owner_id_column_name: sql = f""" INSERT INTO {self._memory_table} - (id, session_id, app_name, user_id, event_id, author, + (id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, content_text, metadata_json, inserted_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) @@ -560,6 +560,7 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], owner_id, @@ -572,7 +573,7 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " else: sql = f""" INSERT INTO {self._memory_table} - (id, session_id, app_name, user_id, event_id, author, + (id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) ON CONFLICT (event_id) DO NOTHING @@ -583,6 +584,7 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], entry["timestamp"], @@ -598,8 +600,13 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " return inserted_count async def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -609,28 +616,51 @@ async def search_entries( effective_limit = limit if limit is not None else self._max_results + if scope_filter == "all": + where_scope = "app_name = $1 AND ((scope = 'user' AND user_id = $2) OR scope = 'app')" + scope_params: tuple[Any, ...] = (app_name, user_id) + p_lim = "$4" + p_pat = "$3" + elif scope_filter == "user": + where_scope = "app_name = $1 AND scope = 'user' AND user_id = $2" + scope_params = (app_name, user_id) + p_lim = "$4" + p_pat = "$3" + else: + where_scope = "app_name = $1 AND scope = 'app'" + scope_params = (app_name,) + p_lim = "$3" + p_pat = "$2" + if self._use_fts: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = $1 AND user_id = $2 - AND to_tsvector('english', content_text) @@ plainto_tsquery('english', $3) + WHERE {where_scope} + AND to_tsvector('english', content_text) @@ plainto_tsquery('english', {p_pat}) ORDER BY timestamp DESC - LIMIT $4 + LIMIT {p_lim} """ - params: tuple[Any, ...] = (app_name, user_id, query, effective_limit) + search_param = query else: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = $1 AND user_id = $2 AND content_text ILIKE $3 + WHERE {where_scope} AND content_text ILIKE {p_pat} ORDER BY timestamp DESC - LIMIT $4 + LIMIT {p_lim} """ - params = (app_name, user_id, f"%{query}%", effective_limit) + search_param = f"%{query}%" + + params = (*scope_params, search_param, effective_limit) async with self._config.provide_connection() as conn: rows = await conn.fetch(sql, *params) - return [cast("MemoryRecord", dict(row)) for row in rows] + records: list[StoredMemory] = [] + for row in rows: + rec = cast("StoredMemory", dict(row)) + rec["embedding"] = None + records.append(rec) + return records async def delete_entries_by_session(self, session_id: str) -> int: if not self._enabled: @@ -642,17 +672,29 @@ async def delete_entries_by_session(self, session_id: str) -> int: result = await conn.execute(sql, session_id) return int(result.split()[-1]) if result else 0 - async def delete_entries_older_than(self, days: int) -> int: + async def delete_entries_older_than( + self, days: int, app_name: "str | None" = None, scope: "str | None" = None + ) -> int: if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) - sql = f""" - DELETE FROM {self._memory_table} - WHERE inserted_at < (CURRENT_TIMESTAMP - INTERVAL '{days} days') - """ + clauses = [f"inserted_at < (CURRENT_TIMESTAMP - INTERVAL '{days} days')"] + params: list[Any] = [] + idx = 1 + if app_name is not None: + clauses.append(f"app_name = ${idx}") + params.append(app_name) + idx += 1 + if scope is not None: + clauses.append(f"scope = ${idx}") + params.append(scope) + idx += 1 + + where_sql = " AND ".join(clauses) + sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" async with self._config.provide_connection() as conn: - result = await conn.execute(sql) + result = await conn.execute(sql, *params) return int(result.split()[-1]) if result else 0 async def _memory_table_ddl(self) -> str: @@ -682,6 +724,7 @@ async def _memory_table_ddl(self) -> str: session_id VARCHAR(128) NOT NULL, app_name VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, + scope VARCHAR(16) NOT NULL DEFAULT 'user', event_id VARCHAR(128) NOT NULL UNIQUE, author VARCHAR(256){owner_id_line}, timestamp TIMESTAMPTZ NOT NULL, @@ -691,8 +734,11 @@ async def _memory_table_ddl(self) -> str: inserted_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP ){memory_locality}; - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time - ON {self._memory_table}(app_name, user_id, timestamp DESC){hash_shard_clause}; + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_scope_user_time + ON {self._memory_table}(app_name, scope, user_id, timestamp DESC){hash_shard_clause}; + + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_scope + ON {self._memory_table}(app_name, scope); CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session ON {self._memory_table}(session_id); diff --git a/sqlspec/adapters/cockroach_psycopg/adk/store.py b/sqlspec/adapters/cockroach_psycopg/adk/store.py index 5721cabff..59e2bfc1b 100644 --- a/sqlspec/adapters/cockroach_psycopg/adk/store.py +++ b/sqlspec/adapters/cockroach_psycopg/adk/store.py @@ -1,6 +1,6 @@ """CockroachDB ADK store for Google Agent Development Kit session/event storage (psycopg).""" -from typing import TYPE_CHECKING, Any, NoReturn, cast +from typing import TYPE_CHECKING, Any, Literal, NoReturn, cast from psycopg import errors from psycopg import sql as pg_sql @@ -9,7 +9,7 @@ from typing_extensions import NotRequired from sqlspec.config import ADKConfig -from sqlspec.extensions.adk import BaseAsyncADKStore, BaseSyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseAsyncADKStore, BaseSyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore from sqlspec.utils.logging import get_logger @@ -17,7 +17,7 @@ from datetime import datetime, timedelta from sqlspec.adapters.cockroach_psycopg.config import CockroachPsycopgAsyncConfig, CockroachPsycopgSyncConfig - from sqlspec.extensions.adk import MemoryRecord + from sqlspec.extensions.adk import StoredMemory __all__ = ( @@ -210,7 +210,7 @@ async def create_tables(self) -> None: async def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: state_json = Jsonb(state) params: tuple[Any, ...] if self._owner_id_column_name: @@ -238,7 +238,7 @@ async def create_session( async def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": if renew_for is not None and self._calculate_expires_at(renew_for) is not None: sql = f""" UPDATE {self._session_table} @@ -261,7 +261,7 @@ async def get_session( if row is None: return None - return SessionRecord( + return StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -283,7 +283,7 @@ async def update_session_state(self, app_name: str, user_id: str, session_id: st await cur.execute(sql.encode(), (Jsonb(state), app_name, user_id, session_id)) await conn.commit() - async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": if user_id is None: sql = f""" SELECT id, app_name, user_id, state, create_time, update_time @@ -309,7 +309,7 @@ async def list_sessions(self, app_name: str, user_id: str | None = None) -> "lis return [] return [ - SessionRecord( + StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -327,7 +327,7 @@ async def delete_session(self, app_name: str, user_id: str, session_id: str) -> await cur.execute(sql.encode(), (app_name, user_id, session_id)) await conn.commit() - async def append_event(self, event_record: EventRecord) -> None: + async def append_event(self, event_record: StoredEvent) -> None: sql = f""" INSERT INTO {self._events_table} ( id, session_id, invocation_id, timestamp, event_data @@ -351,7 +351,7 @@ async def append_event(self, event_record: EventRecord) -> None: async def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -359,7 +359,7 @@ async def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: insert_sql = f""" INSERT INTO {self._events_table} ( id, session_id, invocation_id, timestamp, event_data @@ -407,7 +407,7 @@ async def append_event_and_update_state( raise await conn.commit() - return SessionRecord( + return StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -423,7 +423,7 @@ async def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": if limit == 0: return [] @@ -455,7 +455,7 @@ async def get_events( return [] return [ - EventRecord( + StoredEvent( id=row["id"], session_id=row["session_id"], invocation_id=row["invocation_id"], @@ -653,7 +653,7 @@ def create_tables(self) -> None: def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new session.""" state_json = Jsonb(state) params: tuple[Any, ...] @@ -682,7 +682,7 @@ def create_session( def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get session by ID.""" if renew_for is not None and self._calculate_expires_at(renew_for) is not None: sql = f""" @@ -706,7 +706,7 @@ def get_session( if row is None: return None - return SessionRecord( + return StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -729,7 +729,7 @@ def update_session_state(self, app_name: str, user_id: str, session_id: str, sta cur.execute(sql.encode(), (Jsonb(state), app_name, user_id, session_id)) conn.commit() - def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": """List sessions for an app.""" if user_id is None: sql = f""" @@ -754,7 +754,7 @@ def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[Sess rows = cur.fetchall() return [ - SessionRecord( + StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -775,14 +775,14 @@ def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: cur.execute(sql.encode(), (app_name, user_id, session_id)) conn.commit() - def append_event(self, event_record: EventRecord) -> None: + def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session.""" """Synchronous implementation of append_event.""" self._insert_event(event_record) def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -790,7 +790,7 @@ def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically append an event and update session + scoped state.""" insert_sql = f""" INSERT INTO {self._events_table} ( @@ -839,7 +839,7 @@ def append_event_and_update_state( raise conn.commit() - return SessionRecord( + return StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -855,7 +855,7 @@ def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Get events for a session.""" if limit == 0: return [] @@ -886,7 +886,7 @@ def get_events( rows = cur.fetchall() return [ - EventRecord( + StoredEvent( id=row["id"], session_id=row["session_id"], invocation_id=row["invocation_id"], @@ -1070,7 +1070,7 @@ def _drop_tables_sql(self) -> "list[str]": f"DROP TABLE IF EXISTS {self._session_table}", ] - def _insert_event(self, event_record: EventRecord) -> None: + def _insert_event(self, event_record: StoredEvent) -> None: sql = f""" INSERT INTO {self._events_table} ( id, session_id, invocation_id, timestamp, event_data @@ -1112,7 +1112,7 @@ async def create_tables(self) -> None: async with self._config.provide_session() as driver: await driver.execute_script(await self._memory_table_ddl()) - async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -1124,11 +1124,11 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " if self._owner_id_column_name: query = pg_sql.SQL(""" INSERT INTO {table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, {owner_id_col}, timestamp, content_json, content_text, metadata_json, inserted_at ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s + %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s ) ON CONFLICT (event_id) DO NOTHING """).format( @@ -1137,10 +1137,10 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " else: query = pg_sql.SQL(""" INSERT INTO {table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s + %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s ) ON CONFLICT (event_id) DO NOTHING """).format(table=pg_sql.Identifier(self._memory_table)) @@ -1158,8 +1158,13 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " return inserted_count async def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -1169,24 +1174,26 @@ async def search_entries( effective_limit = limit if limit is not None else self._max_results + where_scope, scope_params = _build_cockroach_scope_where(app_name, user_id, scope_filter) if self._use_fts: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = %s AND user_id = %s + WHERE {where_scope} AND to_tsvector('english', content_text) @@ plainto_tsquery('english', %s) ORDER BY timestamp DESC LIMIT %s """ + search_param = query else: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = %s AND user_id = %s AND content_text ILIKE %s + WHERE {where_scope} AND content_text ILIKE %s ORDER BY timestamp DESC LIMIT %s """ + search_param = f"%{query}%" - search_param = query if self._use_fts else f"%{query}%" - params = (app_name, user_id, search_param, effective_limit) + params = (*scope_params, search_param, effective_limit) try: async with self._config.provide_connection() as conn, conn.cursor() as cur: @@ -1196,7 +1203,12 @@ async def search_entries( except errors.UndefinedTable: return [] - return [cast("MemoryRecord", dict(zip(columns, row, strict=False))) for row in rows] + records: list[StoredMemory] = [] + for row in rows: + rec = cast("StoredMemory", dict(zip(columns, row, strict=False))) + rec["embedding"] = None + records.append(rec) + return records async def delete_entries_by_session(self, session_id: str) -> int: if not self._enabled: @@ -1209,17 +1221,25 @@ async def delete_entries_by_session(self, session_id: str) -> int: await conn.commit() return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 - async def delete_entries_older_than(self, days: int) -> int: + async def delete_entries_older_than( + self, days: int, app_name: "str | None" = None, scope: "str | None" = None + ) -> int: if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) - sql = f""" - DELETE FROM {self._memory_table} - WHERE inserted_at < CURRENT_TIMESTAMP - INTERVAL '{days} days' - """ + clauses = [f"inserted_at < CURRENT_TIMESTAMP - INTERVAL '{days} days'"] + params: list[Any] = [] + if app_name is not None: + clauses.append("app_name = %s") + params.append(app_name) + if scope is not None: + clauses.append("scope = %s") + params.append(scope) + where_sql = " AND ".join(clauses) + sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" async with self._config.provide_connection() as conn, conn.cursor() as cur: - await cur.execute(sql.encode()) + await cur.execute(sql.encode(), tuple(params) if params else None) await conn.commit() return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 @@ -1275,7 +1295,7 @@ def create_tables(self) -> None: with self._config.provide_session() as driver: driver.execute_script(self._memory_table_ddl()) - def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with deduplication.""" if not self._enabled: msg = "Memory store is disabled" @@ -1288,11 +1308,11 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object if self._owner_id_column_name: query = pg_sql.SQL(""" INSERT INTO {table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, {owner_id_col}, timestamp, content_json, content_text, metadata_json, inserted_at ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s + %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s ) ON CONFLICT (event_id) DO NOTHING """).format( @@ -1301,10 +1321,10 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object else: query = pg_sql.SQL(""" INSERT INTO {table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s + %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s ) ON CONFLICT (event_id) DO NOTHING """).format(table=pg_sql.Identifier(self._memory_table)) @@ -1320,8 +1340,13 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object return inserted_count def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries by text query.""" if not self._enabled: msg = "Memory store is disabled" @@ -1332,24 +1357,26 @@ def search_entries( effective_limit = limit if limit is not None else self._max_results + where_scope, scope_params = _build_cockroach_scope_where(app_name, user_id, scope_filter) if self._use_fts: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = %s AND user_id = %s + WHERE {where_scope} AND to_tsvector('english', content_text) @@ plainto_tsquery('english', %s) ORDER BY timestamp DESC LIMIT %s """ + search_param = query else: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = %s AND user_id = %s AND content_text ILIKE %s + WHERE {where_scope} AND content_text ILIKE %s ORDER BY timestamp DESC LIMIT %s """ + search_param = f"%{query}%" - search_param = query if self._use_fts else f"%{query}%" - params = (app_name, user_id, search_param, effective_limit) + params = (*scope_params, search_param, effective_limit) try: with self._config.provide_connection() as conn, conn.cursor() as cur: @@ -1359,7 +1386,12 @@ def search_entries( except errors.UndefinedTable: return [] - return [cast("MemoryRecord", dict(zip(columns, row, strict=False))) for row in rows] + records: list[StoredMemory] = [] + for row in rows: + rec = cast("StoredMemory", dict(zip(columns, row, strict=False))) + rec["embedding"] = None + records.append(rec) + return records def delete_entries_by_session(self, session_id: str) -> int: """Delete all memory entries for a specific session.""" @@ -1373,18 +1405,24 @@ def delete_entries_by_session(self, session_id: str) -> int: conn.commit() return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 - def delete_entries_older_than(self, days: int) -> int: + def delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: """Delete memory entries older than specified days.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) - sql = f""" - DELETE FROM {self._memory_table} - WHERE inserted_at < CURRENT_TIMESTAMP - INTERVAL '{days} days' - """ + clauses = [f"inserted_at < CURRENT_TIMESTAMP - INTERVAL '{days} days'"] + params: list[Any] = [] + if app_name is not None: + clauses.append("app_name = %s") + params.append(app_name) + if scope is not None: + clauses.append("scope = %s") + params.append(scope) + where_sql = " AND ".join(clauses) + sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" with self._config.provide_connection() as conn, conn.cursor() as cur: - cur.execute(sql.encode()) + cur.execute(sql.encode(), tuple(params) if params else None) conn.commit() return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 @@ -1420,12 +1458,13 @@ def _drop_memory_table_sql(self) -> "list[str]": return [f"DROP TABLE IF EXISTS {self._memory_table}"] -def _build_insert_params(entry: "MemoryRecord") -> "tuple[object, ...]": +def _build_insert_params(entry: "StoredMemory") -> "tuple[object, ...]": return ( entry["id"], entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], entry["timestamp"], @@ -1436,12 +1475,13 @@ def _build_insert_params(entry: "MemoryRecord") -> "tuple[object, ...]": ) -def _build_insert_params_with_owner(entry: "MemoryRecord", owner_id: "object | None") -> "tuple[object, ...]": +def _build_insert_params_with_owner(entry: "StoredMemory", owner_id: "object | None") -> "tuple[object, ...]": return ( entry["id"], entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], owner_id, @@ -1495,3 +1535,13 @@ def _cockroach_storing_clause(adk_config: CockroachPsycopgADKConfig, columns: tu if not adk_config.get("enable_storing_indexes", False): return "" return f" STORING ({', '.join(columns)})" + + +def _build_cockroach_scope_where( + app_name: str, user_id: str, scope_filter: Literal["all", "user", "app"] +) -> tuple[str, tuple[Any, ...]]: + if scope_filter == "all": + return "app_name = %s AND ((scope = 'user' AND user_id = %s) OR scope = 'app')", (app_name, user_id) + if scope_filter == "user": + return "app_name = %s AND scope = 'user' AND user_id = %s", (app_name, user_id) + return "app_name = %s AND scope = 'app'", (app_name,) diff --git a/sqlspec/adapters/duckdb/adk/store.py b/sqlspec/adapters/duckdb/adk/store.py index 0dea7bdc9..057a41aa6 100644 --- a/sqlspec/adapters/duckdb/adk/store.py +++ b/sqlspec/adapters/duckdb/adk/store.py @@ -10,19 +10,19 @@ import contextlib from collections.abc import Mapping from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast from typing_extensions import NotRequired, TypedDict from sqlspec.config import ADKConfig -from sqlspec.extensions.adk import BaseSyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.utils.logging import get_logger from sqlspec.utils.serializers import from_json, to_json if TYPE_CHECKING: from sqlspec.adapters.duckdb.config import DuckDBConfig - from sqlspec.extensions.adk import MemoryRecord + from sqlspec.extensions.adk import StoredMemory __all__ = ("DuckdbADKConfig", "DuckdbADKFTSOptions", "DuckdbADKMemoryStore", "DuckdbADKStore") @@ -117,7 +117,7 @@ def create_tables(self) -> None: def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new session. Args: @@ -134,7 +134,7 @@ def create_session( def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get session by ID. Args: @@ -159,7 +159,7 @@ def update_session_state(self, app_name: str, user_id: str, session_id: str, sta """ self._update_session_state(app_name, user_id, session_id, state) - def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": """List sessions for an app, optionally filtered by user. Args: @@ -181,7 +181,7 @@ def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: """ self._delete_session(app_name, user_id, session_id) - def append_event(self, event_record: EventRecord) -> None: + def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session. Args: @@ -191,7 +191,7 @@ def append_event(self, event_record: EventRecord) -> None: def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -199,11 +199,11 @@ def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically append an event and update the session's durable state. The event insert and state update succeed together or fail together - within a single DuckDB transaction; the updated SessionRecord is + within a single DuckDB transaction; the updated StoredSession is returned via UPDATE...RETURNING. Args: @@ -227,7 +227,7 @@ def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Get events for a session. Args: @@ -450,7 +450,7 @@ def _sync_events_table_ddl(self) -> str: def _create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Synchronous implementation of create_session.""" now = datetime.now(timezone.utc) state_json = to_json(state) @@ -474,13 +474,13 @@ def _create_session( conn.execute(sql, params) conn.commit() - return SessionRecord( + return StoredSession( id=session_id, app_name=app_name, user_id=user_id, state=state, create_time=now, update_time=now ) def _get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Synchronous implementation of get_session.""" if renew_for is not None and self._calculate_expires_at(renew_for) is not None: sql = f""" @@ -510,7 +510,7 @@ def _get_session( state = from_json(state_data) if state_data else {} - return SessionRecord( + return StoredSession( id=session_id_val, app_name=app_name, user_id=user_id, @@ -548,7 +548,7 @@ def _delete_session(self, app_name: str, user_id: str, session_id: str) -> None: conn.execute(delete_session_sql, (app_name, user_id, session_id)) conn.commit() - def _list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[SessionRecord]": + def _list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[StoredSession]": """Synchronous implementation of list_sessions.""" if user_id is None: sql = f""" @@ -573,7 +573,7 @@ def _list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[S rows = cursor.fetchall() return [ - SessionRecord( + StoredSession( id=row[0], app_name=row[1], user_id=row[2], @@ -588,7 +588,7 @@ def _list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[S return [] raise - def _append_event(self, event_record: EventRecord) -> None: + def _append_event(self, event_record: StoredEvent) -> None: """Synchronous implementation of append_event.""" event_data_str = to_json(event_record["event_data"]) @@ -613,7 +613,7 @@ def _append_event(self, event_record: EventRecord) -> None: def _append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -621,7 +621,7 @@ def _append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Synchronous implementation of append_event_and_update_state.""" now = datetime.now(timezone.utc) state_json = to_json(state) @@ -683,7 +683,7 @@ def _append_event_and_update_state( assert row is not None session_id_val, app_name, user_id, state_data, create_time, update_time = row - return SessionRecord( + return StoredSession( id=session_id_val, app_name=app_name, user_id=user_id, @@ -699,7 +699,7 @@ def _get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Synchronous implementation of get_events.""" if limit == 0: return [] @@ -728,7 +728,7 @@ def _get_events( rows = cursor.fetchall() return [ - EventRecord( + StoredEvent( id=row[0], session_id=row[1], invocation_id=row[2], @@ -890,7 +890,7 @@ def create_tables(self) -> None: self._create_tables() - def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with deduplication. After successful inserts, refreshes the FTS index if FTS is enabled. @@ -898,22 +898,27 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object return self._insert_memory_entries(entries, owner_id) def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries by text query. When FTS is enabled, uses ``match_bm25()`` for BM25-ranked results. Falls back to ILIKE for simple substring matching. """ - return self._search_entries(query, app_name, user_id, limit) + return self._search_entries(query, app_name, user_id, limit, scope_filter) def delete_entries_by_session(self, session_id: str) -> int: """Delete all memory entries for a specific session.""" return self._delete_entries_by_session(session_id) - def delete_entries_older_than(self, days: int) -> int: + def delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: """Delete memory entries older than specified days.""" - return self._delete_entries_older_than(days) + return self._delete_entries_older_than(days, app_name, scope) def _ensure_fts_extension(self, conn: Any) -> bool: """Ensure the DuckDB FTS extension is available for this connection.""" @@ -992,6 +997,7 @@ def _memory_table_ddl(self) -> str: session_id VARCHAR(128) NOT NULL, app_name VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, + scope VARCHAR(16) NOT NULL DEFAULT 'user', event_id VARCHAR(128) NOT NULL UNIQUE, author VARCHAR(256){owner_id_line}, timestamp TIMESTAMP NOT NULL, @@ -1001,8 +1007,11 @@ def _memory_table_ddl(self) -> str: inserted_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ); - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time - ON {self._memory_table}(app_name, user_id, timestamp DESC); + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_scope_user_time + ON {self._memory_table}(app_name, scope, user_id, timestamp DESC); + + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_scope + ON {self._memory_table}(app_name, scope); CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session ON {self._memory_table}(session_id); @@ -1025,33 +1034,9 @@ def _create_tables(self) -> None: def _sync_memory_table_ddl(self) -> str: """Synchronous version of DDL generation for use in _create_tables.""" - owner_id_line = "" - if self._owner_id_column_ddl: - owner_id_line = f",\n {self._owner_id_column_ddl}" + return self._memory_table_ddl() - return f""" - CREATE TABLE IF NOT EXISTS {self._memory_table} ( - id VARCHAR(128) PRIMARY KEY, - session_id VARCHAR(128) NOT NULL, - app_name VARCHAR(128) NOT NULL, - user_id VARCHAR(128) NOT NULL, - event_id VARCHAR(128) NOT NULL UNIQUE, - author VARCHAR(256){owner_id_line}, - timestamp TIMESTAMP NOT NULL, - content_json JSON NOT NULL, - content_text TEXT NOT NULL, - metadata_json JSON, - inserted_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP - ); - - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time - ON {self._memory_table}(app_name, user_id, timestamp DESC); - - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session - ON {self._memory_table}(session_id); - """ - - def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def _insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Synchronous implementation of insert_memory_entries.""" if not self._enabled: msg = "Memory store is disabled" @@ -1064,30 +1049,32 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec if self._owner_id_column_name: sql = f""" INSERT INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(event_id) DO NOTHING RETURNING 1 """ else: sql = f""" INSERT INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(event_id) DO NOTHING RETURNING 1 """ with self._config.provide_connection() as conn: for entry in entries: params: tuple[Any, ...] + scope_value = entry.get("scope", "user") if self._owner_id_column_name: params = ( entry["id"], entry["session_id"], entry["app_name"], entry["user_id"], + scope_value, entry["event_id"], entry["author"], owner_id, @@ -1103,6 +1090,7 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec entry["session_id"], entry["app_name"], entry["user_id"], + scope_value, entry["event_id"], entry["author"], entry["timestamp"], @@ -1122,8 +1110,13 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec return inserted_count def _search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Synchronous implementation of search_entries.""" if not self._enabled: msg = "Memory store is disabled" @@ -1141,6 +1134,7 @@ def _search_entries( if use_fts: # Use match_bm25() -- the correct DuckDB FTS syntax + where_scope, scope_params = _build_duckdb_scope_where(app_name, user_id, scope_filter, prefix="m") sql = f""" SELECT m.* FROM {self._memory_table} m @@ -1148,31 +1142,33 @@ def _search_entries( SELECT id, fts_main_{self._memory_table}.match_bm25(id, ?, fields := 'content_text') AS score FROM {self._memory_table} ) fts ON m.id = fts.id - WHERE m.app_name = ? AND m.user_id = ? AND fts.score IS NOT NULL + WHERE {where_scope} AND fts.score IS NOT NULL ORDER BY fts.score DESC LIMIT ? """ - params = (query, app_name, user_id, limit_value) + params = (query, *scope_params, limit_value) else: + where_scope, scope_params = _build_duckdb_scope_where(app_name, user_id, scope_filter) sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = ? AND user_id = ? AND content_text ILIKE ? + WHERE {where_scope} AND content_text ILIKE ? ORDER BY timestamp DESC LIMIT ? """ - params = (app_name, user_id, f"%{query}%", limit_value) + params = (*scope_params, f"%{query}%", limit_value) rows = conn.execute(sql, params).fetchall() columns = [col[0] for col in conn.description or []] - records: list[MemoryRecord] = [] + records: list[StoredMemory] = [] for row in rows: - record = cast("MemoryRecord", dict(zip(columns, row, strict=False))) + record = cast("StoredMemory", dict(zip(columns, row, strict=False))) content_value = record["content_json"] if isinstance(content_value, (str, bytes)): record["content_json"] = from_json(content_value) metadata_value = record.get("metadata_json") if isinstance(metadata_value, (str, bytes)): record["metadata_json"] = from_json(metadata_value) + record["embedding"] = None records.append(record) return records @@ -1191,19 +1187,25 @@ def _delete_entries_by_session(self, session_id: str) -> int: self._refresh_fts_index(conn) return deleted_count - def _delete_entries_older_than(self, days: int) -> int: + def _delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: """Synchronous implementation of delete_entries_older_than.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) - sql = f""" - DELETE FROM {self._memory_table} - WHERE inserted_at < (CURRENT_TIMESTAMP - INTERVAL '{days} days') - RETURNING 1 - """ + clauses = [f"inserted_at < (CURRENT_TIMESTAMP - INTERVAL '{days} days')"] + params: list[Any] = [] + if app_name is not None: + clauses.append("app_name = ?") + params.append(app_name) + if scope is not None: + clauses.append("scope = ?") + params.append(scope) + + where_sql = " AND ".join(clauses) + sql = f"DELETE FROM {self._memory_table} WHERE {where_sql} RETURNING 1" with self._config.provide_connection() as conn: - result = conn.execute(sql) + result = conn.execute(sql, tuple(params)) deleted_count = len(result.fetchall()) conn.commit() if self._use_fts and deleted_count > 0: @@ -1246,3 +1248,17 @@ def _format_duckdb_fts_option(value: object) -> str: return f"'{escaped_value}'" msg = f"DuckDB ADK memory_fts_options values must be str, int, or bool; got {type(value).__name__}" raise TypeError(msg) + + +def _build_duckdb_scope_where( + app_name: str, user_id: str, scope_filter: Literal["all", "user", "app"], *, prefix: str = "" +) -> tuple[str, tuple[Any, ...]]: + pfx = f"{prefix}." if prefix else "" + if scope_filter == "all": + return f"{pfx}app_name = ? AND (({pfx}scope = 'user' AND {pfx}user_id = ?) OR {pfx}scope = 'app')", ( + app_name, + user_id, + ) + if scope_filter == "user": + return f"{pfx}app_name = ? AND {pfx}scope = 'user' AND {pfx}user_id = ?", (app_name, user_id) + return f"{pfx}app_name = ? AND {pfx}scope = 'app'", (app_name,) diff --git a/sqlspec/adapters/mssql_python/adk/store.py b/sqlspec/adapters/mssql_python/adk/store.py index e558d74bf..ef151b1a7 100644 --- a/sqlspec/adapters/mssql_python/adk/store.py +++ b/sqlspec/adapters/mssql_python/adk/store.py @@ -9,7 +9,7 @@ from sqlspec.adapters.mssql_python._typing import MSSQL_PYTHON_MODULE, MssqlPythonCursor from sqlspec.adapters.mssql_python.data_dictionary import MssqlVersionInfo from sqlspec.config import ADKConfig -from sqlspec.extensions.adk import BaseSyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession from sqlspec.utils.serializers import from_json, to_json if TYPE_CHECKING: @@ -72,7 +72,7 @@ def create_tables(self) -> None: def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new ADK session.""" owner_column = f", {_quote_identifier(self._owner_id_column_name)}" if self._owner_id_column_name else "" owner_param = ", ?" if self._owner_id_column_name else "" @@ -96,7 +96,7 @@ def create_session( def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Return a scoped session or ``None`` if absent.""" try: if renew_for is not None and self._calculate_expires_at(renew_for) is not None: @@ -135,7 +135,7 @@ def update_session_state(self, app_name: str, user_id: str, session_id: str, sta commit=True, ) - def list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[StoredSession]": """List ADK sessions for an application, optionally scoped to a user.""" if user_id is None: sql = f""" @@ -169,13 +169,13 @@ def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: commit=True, ) - def append_event(self, event_record: EventRecord) -> None: + def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session.""" self._execute(_insert_event_sql(self._events_table), _event_insert_params(event_record), commit=True) def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -183,7 +183,7 @@ def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically append an event and update durable session/scoped state.""" update_sql = f""" UPDATE {_table_ref(self._session_table)} @@ -215,7 +215,7 @@ def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Return events for a scoped session ordered by event timestamp.""" if limit == 0: return [] @@ -596,7 +596,7 @@ def _events_query( return sql, tuple(params) -def _event_insert_params(event_record: EventRecord) -> "tuple[Any, ...]": +def _event_insert_params(event_record: StoredEvent) -> "tuple[Any, ...]": return ( event_record["id"], event_record["app_name"], @@ -608,14 +608,14 @@ def _event_insert_params(event_record: EventRecord) -> "tuple[Any, ...]": ) -def _session_record_from_row(row: Any) -> SessionRecord: - return SessionRecord( +def _session_record_from_row(row: Any) -> StoredSession: + return StoredSession( id=row[0], app_name=row[1], user_id=row[2], state=_json_dict(row[3]), create_time=row[4], update_time=row[5] ) -def _event_record_from_row(row: Any) -> EventRecord: - return EventRecord( +def _event_record_from_row(row: Any) -> StoredEvent: + return StoredEvent( id=row[0], app_name=row[1], user_id=row[2], diff --git a/sqlspec/adapters/mysqlconnector/adk/store.py b/sqlspec/adapters/mysqlconnector/adk/store.py index 2182d68f6..8ba50060e 100644 --- a/sqlspec/adapters/mysqlconnector/adk/store.py +++ b/sqlspec/adapters/mysqlconnector/adk/store.py @@ -2,12 +2,12 @@ import re from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast from typing_extensions import NotRequired from sqlspec.config import ADKConfig -from sqlspec.extensions.adk import BaseAsyncADKStore, BaseSyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseAsyncADKStore, BaseSyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore from sqlspec.protocols import HasErrnoProtocol from sqlspec.utils.serializers import from_json, to_json @@ -16,7 +16,7 @@ from datetime import datetime, timedelta from sqlspec.adapters.mysqlconnector.config import MysqlConnectorAsyncConfig, MysqlConnectorSyncConfig - from sqlspec.extensions.adk import MemoryRecord + from sqlspec.extensions.adk import StoredMemory __all__ = ( @@ -41,6 +41,7 @@ " session_id VARCHAR(128) NOT NULL,\n" " app_name VARCHAR(128) NOT NULL,\n" " user_id VARCHAR(128) NOT NULL,\n" + " scope VARCHAR(16) NOT NULL DEFAULT 'user',\n" " event_id VARCHAR(128) NOT NULL UNIQUE,\n" " author VARCHAR(256){1},\n" " timestamp TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),\n" @@ -48,9 +49,10 @@ " content_text TEXT NOT NULL,\n" " metadata_json JSON,\n" " inserted_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),\n" - " INDEX idx_{2}_app_user_time (app_name, user_id, timestamp),\n" - " INDEX idx_{3}_session (session_id){4}{5}\n" - " ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{6}\n" + " INDEX idx_{2}_app_scope_user_time (app_name, scope, user_id, timestamp),\n" + " INDEX idx_{3}_scope (app_name, scope),\n" + " INDEX idx_{4}_session (session_id){5}{6}\n" + " ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{7}\n" " " ) @@ -176,7 +178,7 @@ async def create_tables(self) -> None: async def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: params: tuple[Any, ...] if self._owner_id_column_name: sql = f""" @@ -207,7 +209,7 @@ async def create_session( async def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": import mysql.connector try: @@ -257,7 +259,7 @@ async def update_session_state(self, app_name: str, user_id: str, session_id: st await cursor.close() await conn.commit() - async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": import mysql.connector if user_id is None: @@ -301,7 +303,7 @@ async def delete_session(self, app_name: str, user_id: str, session_id: str) -> await cursor.close() await conn.commit() - async def append_event(self, event_record: EventRecord) -> None: + async def append_event(self, event_record: StoredEvent) -> None: sql = f""" INSERT INTO {self._events_table} ( id, app_name, user_id, session_id, invocation_id, timestamp, event_data @@ -317,7 +319,7 @@ async def append_event(self, event_record: EventRecord) -> None: async def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -325,7 +327,7 @@ async def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: insert_sql = f""" INSERT INTO {self._events_table} ( id, app_name, user_id, session_id, invocation_id, timestamp, event_data @@ -384,7 +386,7 @@ async def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": import mysql.connector if limit == 0: @@ -527,7 +529,7 @@ def create_tables(self) -> None: def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new session.""" params: tuple[Any, ...] if self._owner_id_column_name: @@ -559,7 +561,7 @@ def create_session( def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get session by ID.""" import mysql.connector @@ -610,7 +612,7 @@ def update_session_state(self, app_name: str, user_id: str, session_id: str, sta cursor.close() conn.commit() - def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": """List sessions for an app.""" import mysql.connector @@ -656,7 +658,7 @@ def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: cursor.close() conn.commit() - def append_event(self, event_record: EventRecord) -> None: + def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session.""" sql = f""" INSERT INTO {self._events_table} ( @@ -673,7 +675,7 @@ def append_event(self, event_record: EventRecord) -> None: def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -681,7 +683,7 @@ def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically append an event and update the session's durable state.""" insert_sql = f""" INSERT INTO {self._events_table} ( @@ -741,7 +743,7 @@ def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Get events for a session.""" import mysql.connector @@ -887,7 +889,7 @@ async def create_tables(self) -> None: async with self._config.provide_session() as driver: await driver.execute_script(await self._memory_table_ddl()) - async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -899,17 +901,17 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " if self._owner_id_column_name: sql = f""" INSERT IGNORE INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) + ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """ else: sql = f""" INSERT IGNORE INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) + ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """ async with self._config.provide_connection() as conn: @@ -923,6 +925,7 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], owner_id, @@ -938,6 +941,7 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], entry["timestamp"], @@ -954,8 +958,13 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " return inserted_count async def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -964,23 +973,24 @@ async def search_entries( return [] limit_value = limit or self._max_results + where_scope, scope_params = _build_mysql_scope_where(app_name, user_id, scope_filter) if self._use_fts: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = %s AND user_id = %s + WHERE {where_scope} AND MATCH(content_text) AGAINST (%s IN NATURAL LANGUAGE MODE) ORDER BY timestamp DESC LIMIT %s """ - params = (app_name, user_id, query, limit_value) + params = (*scope_params, query, limit_value) else: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = %s AND user_id = %s AND content_text LIKE %s + WHERE {where_scope} AND content_text LIKE %s ORDER BY timestamp DESC LIMIT %s """ - params = (app_name, user_id, f"%{query}%", limit_value) + params = (*scope_params, f"%{query}%", limit_value) async with self._config.provide_connection() as conn: cursor = await conn.cursor() @@ -991,7 +1001,12 @@ async def search_entries( finally: await cursor.close() - return [cast("MemoryRecord", dict(zip(columns, row, strict=False))) for row in rows] + records: list[StoredMemory] = [] + for row in rows: + rec = cast("StoredMemory", dict(zip(columns, row, strict=False))) + rec["embedding"] = None + records.append(rec) + return records async def delete_entries_by_session(self, session_id: str) -> int: if not self._enabled: @@ -1008,19 +1023,28 @@ async def delete_entries_by_session(self, session_id: str) -> int: finally: await cursor.close() - async def delete_entries_older_than(self, days: int) -> int: + async def delete_entries_older_than( + self, days: int, app_name: "str | None" = None, scope: "str | None" = None + ) -> int: if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) - sql = f""" - DELETE FROM {self._memory_table} - WHERE inserted_at < (UTC_TIMESTAMP(6) - INTERVAL %s DAY) - """ + clauses = ["inserted_at < (UTC_TIMESTAMP(6) - INTERVAL %s DAY)"] + params: list[Any] = [days] + if app_name is not None: + clauses.append("app_name = %s") + params.append(app_name) + if scope is not None: + clauses.append("scope = %s") + params.append(scope) + + where_sql = " AND ".join(clauses) + sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" async with self._config.provide_connection() as conn: cursor = await conn.cursor() try: - await cursor.execute(sql, (days,)) + await cursor.execute(sql, tuple(params)) await conn.commit() return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 finally: @@ -1046,6 +1070,7 @@ async def _memory_table_ddl(self) -> str: owner_id_line, self._memory_table, self._memory_table, + self._memory_table, fts_index, fk_constraint, table_options, @@ -1075,7 +1100,7 @@ def create_tables(self) -> None: with self._config.provide_session() as driver: driver.execute_script(self._memory_table_ddl()) - def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with deduplication.""" if not self._enabled: msg = "Memory store is disabled" @@ -1088,17 +1113,17 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object if self._owner_id_column_name: sql = f""" INSERT IGNORE INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) + ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """ else: sql = f""" INSERT IGNORE INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) + ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """ with self._config.provide_connection() as conn: @@ -1112,6 +1137,7 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], owner_id, @@ -1127,6 +1153,7 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], entry["timestamp"], @@ -1143,8 +1170,13 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object return inserted_count def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries by text query.""" if not self._enabled: msg = "Memory store is disabled" @@ -1154,23 +1186,24 @@ def search_entries( return [] limit_value = limit or self._max_results + where_scope, scope_params = _build_mysql_scope_where(app_name, user_id, scope_filter) if self._use_fts: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = %s AND user_id = %s + WHERE {where_scope} AND MATCH(content_text) AGAINST (%s IN NATURAL LANGUAGE MODE) ORDER BY timestamp DESC LIMIT %s """ - params = (app_name, user_id, query, limit_value) + params = (*scope_params, query, limit_value) else: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = %s AND user_id = %s AND content_text LIKE %s + WHERE {where_scope} AND content_text LIKE %s ORDER BY timestamp DESC LIMIT %s """ - params = (app_name, user_id, f"%{query}%", limit_value) + params = (*scope_params, f"%{query}%", limit_value) with self._config.provide_connection() as conn: cursor = conn.cursor() @@ -1181,7 +1214,12 @@ def search_entries( finally: cursor.close() - return [cast("MemoryRecord", dict(zip(columns, row, strict=False))) for row in rows] + records: list[StoredMemory] = [] + for row in rows: + rec = cast("StoredMemory", dict(zip(columns, row, strict=False))) + rec["embedding"] = None + records.append(rec) + return records def delete_entries_by_session(self, session_id: str) -> int: """Delete all memory entries for a specific session.""" @@ -1199,20 +1237,27 @@ def delete_entries_by_session(self, session_id: str) -> int: finally: cursor.close() - def delete_entries_older_than(self, days: int) -> int: + def delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: """Delete memory entries older than specified days.""" if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) - sql = f""" - DELETE FROM {self._memory_table} - WHERE inserted_at < (UTC_TIMESTAMP(6) - INTERVAL %s DAY) - """ + clauses = ["inserted_at < (UTC_TIMESTAMP(6) - INTERVAL %s DAY)"] + params: list[Any] = [days] + if app_name is not None: + clauses.append("app_name = %s") + params.append(app_name) + if scope is not None: + clauses.append("scope = %s") + params.append(scope) + + where_sql = " AND ".join(clauses) + sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" with self._config.provide_connection() as conn: cursor = conn.cursor() try: - cursor.execute(sql, (days,)) + cursor.execute(sql, tuple(params)) conn.commit() return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 finally: @@ -1238,6 +1283,7 @@ def _memory_table_ddl(self) -> str: owner_id_line, self._memory_table, self._memory_table, + self._memory_table, fts_index, fk_constraint, table_options, @@ -1301,14 +1347,14 @@ def _json_dict(value: Any) -> "dict[str, Any]": return cast("dict[str, Any]", value) -def _session_record_from_row(row: Any) -> SessionRecord: - return SessionRecord( +def _session_record_from_row(row: Any) -> StoredSession: + return StoredSession( id=row[0], app_name=row[1], user_id=row[2], state=_json_dict(row[3]), create_time=row[4], update_time=row[5] ) -def _event_record_from_row(row: Any) -> EventRecord: - return EventRecord( +def _event_record_from_row(row: Any) -> StoredEvent: + return StoredEvent( id=row[0], app_name=row[1], user_id=row[2], @@ -1319,7 +1365,7 @@ def _event_record_from_row(row: Any) -> EventRecord: ) -def _event_insert_params(event_record: EventRecord) -> "tuple[Any, ...]": +def _event_insert_params(event_record: StoredEvent) -> "tuple[Any, ...]": return ( event_record["id"], event_record["app_name"], @@ -1518,3 +1564,13 @@ def _mysql_upsert_metadata_sql(metadata_table: str) -> str: VALUES (%s, %s) ON DUPLICATE KEY UPDATE value = VALUES(value) """ + + +def _build_mysql_scope_where( + app_name: str, user_id: str, scope_filter: Literal["all", "user", "app"] +) -> tuple[str, tuple[Any, ...]]: + if scope_filter == "all": + return "app_name = %s AND ((scope = 'user' AND user_id = %s) OR scope = 'app')", (app_name, user_id) + if scope_filter == "user": + return "app_name = %s AND scope = 'user' AND user_id = %s", (app_name, user_id) + return "app_name = %s AND scope = 'app'", (app_name,) diff --git a/sqlspec/adapters/oracledb/adk/store.py b/sqlspec/adapters/oracledb/adk/store.py index 50933c576..cd3682358 100644 --- a/sqlspec/adapters/oracledb/adk/store.py +++ b/sqlspec/adapters/oracledb/adk/store.py @@ -1,7 +1,7 @@ """Oracle ADK store for Google Agent Development Kit session/event storage.""" from decimal import Decimal -from typing import TYPE_CHECKING, Any, Final, NoReturn, cast +from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, cast import oracledb from typing_extensions import NotRequired, TypedDict @@ -20,7 +20,7 @@ storage_type_from_version, ) from sqlspec.config import ADKConfig -from sqlspec.extensions.adk import BaseAsyncADKStore, BaseSyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseAsyncADKStore, BaseSyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore from sqlspec.utils.logging import get_logger from sqlspec.utils.serializers import from_json, to_json @@ -30,7 +30,7 @@ from datetime import datetime, timedelta from sqlspec.adapters.oracledb.config import OracleAsyncConfig, OracleSyncConfig - from sqlspec.extensions.adk import MemoryRecord + from sqlspec.extensions.adk import StoredMemory __all__ = ( "JSONStorageType", @@ -172,6 +172,7 @@ " session_id VARCHAR2(128) NOT NULL,\n" " app_name VARCHAR2(128) NOT NULL,\n" " user_id VARCHAR2(128) NOT NULL,\n" + " scope VARCHAR2(16) DEFAULT ''user'' NOT NULL,\n" " event_id VARCHAR2(128) NOT NULL UNIQUE,\n" " author VARCHAR2(256){1},\n" " timestamp TIMESTAMP WITH TIME ZONE NOT NULL,\n" @@ -182,15 +183,20 @@ " END;\n" "\n" " BEGIN\n" - " EXECUTE IMMEDIATE 'CREATE INDEX idx_{4}_app_user_time\n" - " ON {5}(app_name, user_id, timestamp DESC)';\n" + " EXECUTE IMMEDIATE 'CREATE INDEX idx_{4}_app_scope_user_time\n" + " ON {5}(app_name, scope, user_id, timestamp DESC)';\n" " END;\n" "\n" " BEGIN\n" - " EXECUTE IMMEDIATE 'CREATE INDEX idx_{6}_session\n" - " ON {7}(session_id)';\n" + " EXECUTE IMMEDIATE 'CREATE INDEX idx_{6}_scope\n" + " ON {7}(app_name, scope)';\n" " END;\n" - " {8}\n" + "\n" + " BEGIN\n" + " EXECUTE IMMEDIATE 'CREATE INDEX idx_{8}_session\n" + " ON {9}(session_id)';\n" + " END;\n" + " {10}\n" " " ) @@ -362,7 +368,7 @@ async def prepare_schema_async(self, driver: Any) -> None: async def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new session. Args: @@ -410,7 +416,7 @@ async def create_session( async def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get session by ID. Args: @@ -454,7 +460,7 @@ async def get_session( state = await self._deserialize_state(state_data) - return SessionRecord( + return StoredSession( id=session_id_val, app_name=app_name, user_id=user_id, @@ -495,7 +501,7 @@ async def update_session_state(self, app_name: str, user_id: str, session_id: st await cursor.execute(sql, {"state": state_data, "app_name": app_name, "user_id": user_id, "id": session_id}) await conn.commit() - async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": """List sessions for an app, optionally filtered by user. Args: @@ -538,7 +544,7 @@ async def list_sessions(self, app_name: str, user_id: str | None = None) -> "lis state = await self._deserialize_state(row[3]) results.append( - SessionRecord( + StoredSession( id=row[0], app_name=row[1], user_id=row[2], @@ -572,7 +578,7 @@ async def delete_session(self, app_name: str, user_id: str, session_id: str) -> await cursor.execute(sql, {"app_name": app_name, "user_id": user_id, "id": session_id}) await conn.commit() - async def append_event(self, event_record: EventRecord) -> None: + async def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session. Args: @@ -602,7 +608,7 @@ async def append_event(self, event_record: EventRecord) -> None: async def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -610,7 +616,7 @@ async def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically append an event and update session + scoped state. All writes are executed within a single transaction so they succeed or @@ -694,7 +700,7 @@ async def append_event_and_update_state( raise session_id_val, row_app_name, row_user_id, state_data_row, create_time, update_time = row - return SessionRecord( + return StoredSession( id=session_id_val, app_name=row_app_name, user_id=row_user_id, @@ -710,7 +716,7 @@ async def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Get events for a session. Args: @@ -754,7 +760,7 @@ async def get_events( rows = await cursor.fetchall() return [ - EventRecord( + StoredEvent( id=row[0], session_id=row[1], invocation_id=_oracle_text_value(row[2]), @@ -1331,7 +1337,7 @@ def prepare_schema_sync(self, driver: Any) -> None: def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new session.""" """Create a new session. @@ -1384,7 +1390,7 @@ def create_session( def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get session by ID.""" """Get session by ID. @@ -1428,7 +1434,7 @@ def get_session( state = self._deserialize_state(state_data) - return SessionRecord( + return StoredSession( id=session_id_val, app_name=app_name, user_id=user_id, @@ -1470,7 +1476,7 @@ def update_session_state(self, app_name: str, user_id: str, session_id: str, sta cursor.execute(sql, {"state": state_data, "app_name": app_name, "user_id": user_id, "id": session_id}) conn.commit() - def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": """List sessions for an app.""" """List sessions for an app, optionally filtered by user. @@ -1514,7 +1520,7 @@ def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[Sess state = self._deserialize_state(row[3]) results.append( - SessionRecord( + StoredSession( id=row[0], app_name=row[1], user_id=row[2], @@ -1549,7 +1555,7 @@ def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: cursor.execute(sql, {"app_name": app_name, "user_id": user_id, "id": session_id}) conn.commit() - def append_event(self, event_record: EventRecord) -> None: + def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session.""" """Synchronous implementation of append_event.""" sql = f""" @@ -1576,7 +1582,7 @@ def append_event(self, event_record: EventRecord) -> None: def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -1584,7 +1590,7 @@ def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically append an event and update session + scoped state.""" """Atomically create an event and update session + scoped state.""" insert_sql = f""" @@ -1663,7 +1669,7 @@ def append_event_and_update_state( raise session_id_val, row_app_name, row_user_id, state_data_row, create_time, update_time = row - return SessionRecord( + return StoredSession( id=session_id_val, app_name=row_app_name, user_id=row_user_id, @@ -1679,7 +1685,7 @@ def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Get events for a session.""" """List events for a session ordered by timestamp. @@ -1721,7 +1727,7 @@ def get_events( rows = cursor.fetchall() return [ - EventRecord( + StoredEvent( id=row[0], session_id=row[1], invocation_id=_oracle_text_value(row[2]), @@ -2246,7 +2252,7 @@ async def prepare_schema_async(self, driver: Any) -> None: """Resolve pool-scoped Oracle storage capabilities before DDL generation.""" await _resolve_oracle_storage_capabilities_async(driver) - async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -2258,10 +2264,10 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " owner_param = ", :owner_id" if self._owner_id_column_name else "" sql = f""" INSERT INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author{owner_column}, + id, session_id, app_name, user_id, scope, event_id, author{owner_column}, timestamp, content_json, content_text, metadata_json, inserted_at ) VALUES ( - :id, :session_id, :app_name, :user_id, :event_id, :author{owner_param}, + :id, :session_id, :app_name, :user_id, :scope, :event_id, :author{owner_param}, :timestamp, :content_json, :content_text, :metadata_json, :inserted_at ) """ @@ -2277,6 +2283,7 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " "session_id": entry["session_id"], "app_name": entry["app_name"], "user_id": entry["user_id"], + "scope": entry.get("scope", "user"), "event_id": entry["event_id"], "author": entry["author"], "timestamp": entry["timestamp"], @@ -2294,8 +2301,13 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " return inserted_count async def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -2304,8 +2316,8 @@ async def search_entries( try: if self._use_fts: - return await self._search_entries_fts(query, app_name, user_id, effective_limit) - return await self._search_entries_simple(query, app_name, user_id, effective_limit) + return await self._search_entries_fts(query, app_name, user_id, effective_limit, scope_filter) + return await self._search_entries_simple(query, app_name, user_id, effective_limit, scope_filter) except OracleDatabaseError as exc: error_obj = exc.args[0] if exc.args else None if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: @@ -2320,14 +2332,23 @@ async def delete_entries_by_session(self, session_id: str) -> int: await conn.commit() return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - async def delete_entries_older_than(self, days: int) -> int: - sql = f""" - DELETE FROM {self._memory_table} - WHERE inserted_at < SYSTIMESTAMP - NUMTODSINTERVAL(:days, 'DAY') - """ + async def delete_entries_older_than( + self, days: int, app_name: "str | None" = None, scope: "str | None" = None + ) -> int: + clauses = ["inserted_at < SYSTIMESTAMP - NUMTODSINTERVAL(:days, 'DAY')"] + params: dict[str, Any] = {"days": days} + if app_name is not None: + clauses.append("app_name = :app_name") + params["app_name"] = app_name + if scope is not None: + clauses.append("scope = :scope") + params["scope"] = scope + + where_sql = " AND ".join(clauses) + sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" async with self._config.provide_connection() as conn: cursor = conn.cursor() - await cursor.execute(sql, {"days": days}) + await cursor.execute(sql, params) await conn.commit() return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 @@ -2412,6 +2433,8 @@ def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType") -> str: self._memory_table, self._memory_table, self._memory_table, + self._memory_table, + self._memory_table, fts_index, ) @@ -2460,58 +2483,62 @@ async def _execute_insert_entry(self, cursor: Any, sql: str, params: "dict[str, raise return True - async def _search_entries_fts(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": + async def _search_entries_fts( + self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "list[StoredMemory]": + where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) sql = f""" - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM ( - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at, SCORE(1) AS score FROM {self._memory_table} - WHERE app_name = :app_name - AND user_id = :user_id + WHERE {where_scope} AND CONTAINS(content_text, :query, 1) > 0 ORDER BY score DESC, timestamp DESC ) WHERE ROWNUM <= :limit """ - params = {"app_name": app_name, "user_id": user_id, "query": query, "limit": limit} + params = {**scope_params, "query": query, "limit": limit} async with self._config.provide_connection() as conn: cursor = conn.cursor() await cursor.execute(sql, params) rows = await cursor.fetchall() return await self._rows_to_records(rows) - async def _search_entries_simple(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": + async def _search_entries_simple( + self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "list[StoredMemory]": + where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) sql = f""" - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM ( - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM {self._memory_table} - WHERE app_name = :app_name - AND user_id = :user_id + WHERE {where_scope} AND LOWER(content_text) LIKE :pattern ORDER BY timestamp DESC ) WHERE ROWNUM <= :limit """ pattern = f"%{query.lower()}%" - params = {"app_name": app_name, "user_id": user_id, "pattern": pattern, "limit": limit} + params = {**scope_params, "pattern": pattern, "limit": limit} async with self._config.provide_connection() as conn: cursor = conn.cursor() await cursor.execute(sql, params) rows = await cursor.fetchall() return await self._rows_to_records(rows) - async def _rows_to_records(self, rows: "list[Any]") -> "list[MemoryRecord]": - records: list[MemoryRecord] = [] + async def _rows_to_records(self, rows: "list[Any]") -> "list[StoredMemory]": + records: list[StoredMemory] = [] for row in rows: - content_json = await self._deserialize_json_field(row[7]) if row[7] is not None else {} - metadata_json = await self._deserialize_json_field(row[9]) - content_text = row[8] + content_json = await self._deserialize_json_field(row[8]) if row[8] is not None else {} + metadata_json = await self._deserialize_json_field(row[10]) + content_text = row[9] if is_async_readable(content_text) or is_readable(content_text): content_text = await _read_lob_async(content_text) records.append({ @@ -2519,13 +2546,15 @@ async def _rows_to_records(self, rows: "list[Any]") -> "list[MemoryRecord]": "session_id": row[1], "app_name": row[2], "user_id": row[3], - "event_id": row[4], - "author": row[5], - "timestamp": row[6], + "scope": row[4], + "event_id": row[5], + "author": row[6], + "timestamp": row[7], "content_json": cast("dict[str, Any]", content_json), "content_text": str(content_text), "metadata_json": metadata_json, - "inserted_at": row[10], + "inserted_at": row[11], + "embedding": None, }) return records @@ -2559,7 +2588,7 @@ def prepare_schema_sync(self, driver: Any) -> None: """Resolve pool-scoped Oracle storage capabilities before DDL generation.""" _resolve_oracle_storage_capabilities_sync(driver) - def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with deduplication.""" if not self._enabled: msg = "Memory store is disabled" @@ -2572,10 +2601,10 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object owner_param = ", :owner_id" if self._owner_id_column_name else "" sql = f""" INSERT INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author{owner_column}, + id, session_id, app_name, user_id, scope, event_id, author{owner_column}, timestamp, content_json, content_text, metadata_json, inserted_at ) VALUES ( - :id, :session_id, :app_name, :user_id, :event_id, :author{owner_param}, + :id, :session_id, :app_name, :user_id, :scope, :event_id, :author{owner_param}, :timestamp, :content_json, :content_text, :metadata_json, :inserted_at ) """ @@ -2591,6 +2620,7 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object "session_id": entry["session_id"], "app_name": entry["app_name"], "user_id": entry["user_id"], + "scope": entry.get("scope", "user"), "event_id": entry["event_id"], "author": entry["author"], "timestamp": entry["timestamp"], @@ -2608,8 +2638,13 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object return inserted_count def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries by text query.""" if not self._enabled: msg = "Memory store is disabled" @@ -2619,8 +2654,8 @@ def search_entries( try: if self._use_fts: - return self._search_entries_fts(query, app_name, user_id, effective_limit) - return self._search_entries_simple(query, app_name, user_id, effective_limit) + return self._search_entries_fts(query, app_name, user_id, effective_limit, scope_filter) + return self._search_entries_simple(query, app_name, user_id, effective_limit, scope_filter) except OracleDatabaseError as exc: error_obj = exc.args[0] if exc.args else None if error_obj and error_obj.code == ORACLE_TABLE_NOT_FOUND_ERROR: @@ -2636,15 +2671,22 @@ def delete_entries_by_session(self, session_id: str) -> int: conn.commit() return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 - def delete_entries_older_than(self, days: int) -> int: + def delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: """Delete memory entries older than specified days.""" - sql = f""" - DELETE FROM {self._memory_table} - WHERE inserted_at < SYSTIMESTAMP - NUMTODSINTERVAL(:days, 'DAY') - """ + clauses = ["inserted_at < SYSTIMESTAMP - NUMTODSINTERVAL(:days, 'DAY')"] + params: dict[str, Any] = {"days": days} + if app_name is not None: + clauses.append("app_name = :app_name") + params["app_name"] = app_name + if scope is not None: + clauses.append("scope = :scope") + params["scope"] = scope + + where_sql = " AND ".join(clauses) + sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" with self._config.provide_connection() as conn: cursor = conn.cursor() - cursor.execute(sql, {"days": days}) + cursor.execute(sql, params) conn.commit() return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 @@ -2729,6 +2771,8 @@ def _memory_table_ddl_for_type(self, storage_type: "JSONStorageType") -> str: self._memory_table, self._memory_table, self._memory_table, + self._memory_table, + self._memory_table, fts_index, ) @@ -2777,58 +2821,62 @@ def _execute_insert_entry(self, cursor: Any, sql: str, params: "dict[str, Any]") raise return True - def _search_entries_fts(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": + def _search_entries_fts( + self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "list[StoredMemory]": + where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) sql = f""" - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM ( - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at, SCORE(1) AS score FROM {self._memory_table} - WHERE app_name = :app_name - AND user_id = :user_id + WHERE {where_scope} AND CONTAINS(content_text, :query, 1) > 0 ORDER BY score DESC, timestamp DESC ) WHERE ROWNUM <= :limit """ - params = {"app_name": app_name, "user_id": user_id, "query": query, "limit": limit} + params = {**scope_params, "query": query, "limit": limit} with self._config.provide_connection() as conn: cursor = conn.cursor() cursor.execute(sql, params) rows = cursor.fetchall() return self._rows_to_records(rows) - def _search_entries_simple(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": + def _search_entries_simple( + self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "list[StoredMemory]": + where_scope, scope_params = _build_oracle_scope_where(app_name, user_id, scope_filter) sql = f""" - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM ( - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM {self._memory_table} - WHERE app_name = :app_name - AND user_id = :user_id + WHERE {where_scope} AND LOWER(content_text) LIKE :pattern ORDER BY timestamp DESC ) WHERE ROWNUM <= :limit """ pattern = f"%{query.lower()}%" - params = {"app_name": app_name, "user_id": user_id, "pattern": pattern, "limit": limit} + params = {**scope_params, "pattern": pattern, "limit": limit} with self._config.provide_connection() as conn: cursor = conn.cursor() cursor.execute(sql, params) rows = cursor.fetchall() return self._rows_to_records(rows) - def _rows_to_records(self, rows: "list[Any]") -> "list[MemoryRecord]": - records: list[MemoryRecord] = [] + def _rows_to_records(self, rows: "list[Any]") -> "list[StoredMemory]": + records: list[StoredMemory] = [] for row in rows: - content_json = self._deserialize_json_field(row[7]) if row[7] is not None else {} - metadata_json = self._deserialize_json_field(row[9]) - content_text = row[8] + content_json = self._deserialize_json_field(row[8]) if row[8] is not None else {} + metadata_json = self._deserialize_json_field(row[10]) + content_text = row[9] if is_readable(content_text): content_text = _read_lob_sync(content_text) records.append({ @@ -2836,13 +2884,15 @@ def _rows_to_records(self, rows: "list[Any]") -> "list[MemoryRecord]": "session_id": row[1], "app_name": row[2], "user_id": row[3], - "event_id": row[4], - "author": row[5], - "timestamp": row[6], + "scope": row[4], + "event_id": row[5], + "author": row[6], + "timestamp": row[7], "content_json": cast("dict[str, Any]", content_json), "content_text": str(content_text), "metadata_json": metadata_json, - "inserted_at": row[10], + "inserted_at": row[11], + "embedding": None, }) return records @@ -2979,3 +3029,19 @@ def _read_lob_sync(data: Any) -> Any: if is_readable(data): return data.read() return data + + +def _build_oracle_scope_where( + app_name: str, user_id: str, scope_filter: Literal["all", "user", "app"] +) -> tuple[str, dict[str, Any]]: + if scope_filter == "all": + return "app_name = :app_name AND ((scope = 'user' AND user_id = :user_id) OR scope = 'app')", { + "app_name": app_name, + "user_id": user_id, + } + if scope_filter == "user": + return "app_name = :app_name AND scope = 'user' AND user_id = :user_id", { + "app_name": app_name, + "user_id": user_id, + } + return "app_name = :app_name AND scope = 'app'", {"app_name": app_name} diff --git a/sqlspec/adapters/psqlpy/adk/store.py b/sqlspec/adapters/psqlpy/adk/store.py index 8b6c6abea..02ee2d2fd 100644 --- a/sqlspec/adapters/psqlpy/adk/store.py +++ b/sqlspec/adapters/psqlpy/adk/store.py @@ -1,12 +1,12 @@ """Psqlpy ADK store for Google Agent Development Kit session/event storage.""" import re -from typing import TYPE_CHECKING, Any, Final, NoReturn, cast +from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, cast from typing_extensions import NotRequired from sqlspec.config import ADKConfig -from sqlspec.extensions.adk import BaseAsyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore from sqlspec.utils.logging import get_logger from sqlspec.utils.type_guards import has_query_result_metadata @@ -15,7 +15,7 @@ from datetime import datetime, timedelta from sqlspec.adapters.psqlpy.config import PsqlpyConfig - from sqlspec.extensions.adk import MemoryRecord + from sqlspec.extensions.adk import StoredMemory __all__ = ("PsqlpyADKConfig", "PsqlpyADKMemoryStore", "PsqlpyADKStore") @@ -88,7 +88,7 @@ async def create_tables(self) -> None: async def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: async with self._config.provide_connection() as conn: # pyright: ignore[reportAttributeAccessIssue] if self._owner_id_column_name: sql = f""" @@ -112,7 +112,7 @@ async def create_session( async def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": if renew_for is not None and self._calculate_expires_at(renew_for) is not None: sql = f""" UPDATE {self._session_table} @@ -136,7 +136,7 @@ async def get_session( return None row = rows[0] - return SessionRecord( + return StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -159,7 +159,7 @@ async def update_session_state(self, app_name: str, user_id: str, session_id: st async with self._config.provide_connection() as conn: # pyright: ignore[reportAttributeAccessIssue] await conn.execute(sql, [state, app_name, user_id, session_id]) - async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": if user_id is None: sql = f""" SELECT id, app_name, user_id, state, create_time, update_time @@ -183,7 +183,7 @@ async def list_sessions(self, app_name: str, user_id: str | None = None) -> "lis rows: list[dict[str, Any]] = result.result() if result else [] return [ - SessionRecord( + StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -204,7 +204,7 @@ async def delete_session(self, app_name: str, user_id: str, session_id: str) -> async with self._config.provide_connection() as conn: # pyright: ignore[reportAttributeAccessIssue] await conn.execute(sql, [app_name, user_id, session_id]) - async def append_event(self, event_record: EventRecord) -> None: + async def append_event(self, event_record: StoredEvent) -> None: sql = f""" INSERT INTO {self._events_table} ( id, session_id, invocation_id, timestamp, event_data @@ -225,7 +225,7 @@ async def append_event(self, event_record: EventRecord) -> None: async def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -233,7 +233,7 @@ async def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: insert_sql = f""" INSERT INTO {self._events_table} ( id, session_id, invocation_id, timestamp, event_data @@ -287,7 +287,7 @@ async def append_event_and_update_state( await conn.execute("COMMIT") row = rows[0] - return SessionRecord( + return StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -303,7 +303,7 @@ async def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": if limit == 0: return [] @@ -333,7 +333,7 @@ async def get_events( rows: list[dict[str, Any]] = result.result() if result else [] return [ - EventRecord( + StoredEvent( id=row["id"], session_id=row["session_id"], invocation_id=row["invocation_id"], @@ -573,7 +573,7 @@ async def create_tables(self) -> None: async with self._config.provide_session() as driver: await driver.execute_script(await self._memory_table_ddl()) - async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with deduplication.""" if not self._enabled: msg = "Memory store is disabled" @@ -645,8 +645,13 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " return inserted_count async def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries by text query.""" if not self._enabled: msg = "Memory store is disabled" @@ -687,7 +692,9 @@ async def delete_entries_by_session(self, session_id: str) -> int: return 0 raise - async def delete_entries_older_than(self, days: int) -> int: + async def delete_entries_older_than( + self, days: int, app_name: "str | None" = None, scope: "str | None" = None + ) -> int: """Delete memory entries older than specified days.""" count_sql = f""" SELECT COUNT(*) AS count FROM {self._memory_table} @@ -752,39 +759,68 @@ def _drop_memory_table_sql(self) -> "list[str]": """Get PostgreSQL DROP TABLE SQL statements.""" return [f"DROP TABLE IF EXISTS {self._memory_table}"] - async def _search_entries_fts(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": + async def _search_entries_fts( + self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "list[StoredMemory]": + if scope_filter == "all": + where_scope = "app_name = $2 AND ((scope = 'user' AND user_id = $3) OR scope = 'app')" + scope_params: list[Any] = [query, app_name, user_id] + p_lim = "$4" + elif scope_filter == "user": + where_scope = "app_name = $2 AND scope = 'user' AND user_id = $3" + scope_params = [query, app_name, user_id] + p_lim = "$4" + else: + where_scope = "app_name = $2 AND scope = 'app'" + scope_params = [query, app_name] + p_lim = "$3" + sql = f""" - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at, ts_rank(to_tsvector('english', content_text), plainto_tsquery('english', $1)) as rank FROM {self._memory_table} - WHERE app_name = $2 - AND user_id = $3 + WHERE {where_scope} AND to_tsvector('english', content_text) @@ plainto_tsquery('english', $1) ORDER BY rank DESC, timestamp DESC - LIMIT $4 + LIMIT {p_lim} """ - params = [query, app_name, user_id, limit] async with self._config.provide_connection() as conn: # pyright: ignore[reportAttributeAccessIssue] - result = await conn.fetch(sql, params) + result = await conn.fetch(sql, [*scope_params, limit]) rows: list[dict[str, Any]] = result.result() if result else [] return _rows_to_records(rows) - async def _search_entries_simple(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": + async def _search_entries_simple( + self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "list[StoredMemory]": + pattern = f"%{query}%" + if scope_filter == "all": + where_scope = "app_name = $1 AND ((scope = 'user' AND user_id = $2) OR scope = 'app')" + scope_params = [app_name, user_id, pattern] + p_lim = "$4" + p_pat = "$3" + elif scope_filter == "user": + where_scope = "app_name = $1 AND scope = 'user' AND user_id = $2" + scope_params = [app_name, user_id, pattern] + p_lim = "$4" + p_pat = "$3" + else: + where_scope = "app_name = $1 AND scope = 'app'" + scope_params = [app_name, pattern] + p_lim = "$3" + p_pat = "$2" + sql = f""" - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM {self._memory_table} - WHERE app_name = $1 - AND user_id = $2 - AND content_text ILIKE $3 + WHERE {where_scope} + AND content_text ILIKE {p_pat} ORDER BY timestamp DESC - LIMIT $4 + LIMIT {p_lim} """ - pattern = f"%{query}%" - params = [app_name, user_id, pattern, limit] async with self._config.provide_connection() as conn: # pyright: ignore[reportAttributeAccessIssue] - result = await conn.fetch(sql, params) + result = await conn.fetch(sql, [*scope_params, limit]) rows: list[dict[str, Any]] = result.result() if result else [] return _rows_to_records(rows) @@ -817,13 +853,14 @@ def _parse_command_tag(self, tag: str) -> int: return -1 -def _rows_to_records(rows: "list[dict[str, Any]]") -> "list[MemoryRecord]": +def _rows_to_records(rows: "list[dict[str, Any]]") -> "list[StoredMemory]": return [ { "id": row["id"], "session_id": row["session_id"], "app_name": row["app_name"], "user_id": row["user_id"], + "scope": row.get("scope", "user"), "event_id": row["event_id"], "author": row["author"], "timestamp": row["timestamp"], @@ -831,6 +868,7 @@ def _rows_to_records(rows: "list[dict[str, Any]]") -> "list[MemoryRecord]": "content_text": row["content_text"], "metadata_json": row["metadata_json"], "inserted_at": row["inserted_at"], + "embedding": None, } for row in rows ] diff --git a/sqlspec/adapters/psycopg/adk/store.py b/sqlspec/adapters/psycopg/adk/store.py index 878976068..01d99acc6 100644 --- a/sqlspec/adapters/psycopg/adk/store.py +++ b/sqlspec/adapters/psycopg/adk/store.py @@ -1,6 +1,6 @@ """Psycopg ADK store for Google Agent Development Kit session/event storage.""" -from typing import TYPE_CHECKING, Any, NoReturn, cast +from typing import TYPE_CHECKING, Any, Literal, NoReturn, cast from psycopg import errors from psycopg import sql as pg_sql @@ -9,7 +9,7 @@ from typing_extensions import NotRequired from sqlspec.config import ADKConfig -from sqlspec.extensions.adk import BaseAsyncADKStore, BaseSyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseAsyncADKStore, BaseSyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore from sqlspec.utils.logging import get_logger @@ -17,7 +17,7 @@ from datetime import datetime, timedelta from sqlspec.adapters.psycopg.config import PsycopgAsyncConfig, PsycopgSyncConfig - from sqlspec.extensions.adk import MemoryRecord + from sqlspec.extensions.adk import StoredMemory __all__ = ( @@ -119,6 +119,7 @@ " session_id VARCHAR(128) NOT NULL,\n" " app_name VARCHAR(128) NOT NULL,\n" " user_id VARCHAR(128) NOT NULL,\n" + " scope VARCHAR(16) NOT NULL DEFAULT 'user',\n" " event_id VARCHAR(128) NOT NULL UNIQUE,\n" " author VARCHAR(256){1},\n" " timestamp TIMESTAMPTZ NOT NULL,\n" @@ -128,12 +129,15 @@ " inserted_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP\n" " );\n" "\n" - " CREATE INDEX IF NOT EXISTS idx_{2}_app_user_time\n" - " ON {3}(app_name, user_id, timestamp DESC);\n" + " CREATE INDEX IF NOT EXISTS idx_{2}_app_scope_user_time\n" + " ON {3}(app_name, scope, user_id, timestamp DESC);\n" "\n" - " CREATE INDEX IF NOT EXISTS idx_{4}_session\n" - " ON {5}(session_id);\n" - " {6}\n" + " CREATE INDEX IF NOT EXISTS idx_{4}_scope\n" + " ON {5}(app_name, scope);\n" + "\n" + " CREATE INDEX IF NOT EXISTS idx_{6}_session\n" + " ON {7}(session_id);\n" + " {8}\n" " " ) @@ -211,7 +215,7 @@ async def create_tables(self) -> None: async def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: params: tuple[Any, ...] if self._owner_id_column_name: query = pg_sql.SQL(""" @@ -239,7 +243,7 @@ async def create_session( async def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": if renew_for is not None and self._calculate_expires_at(renew_for) is not None: query = pg_sql.SQL(""" UPDATE {table} @@ -264,7 +268,7 @@ async def get_session( if row is None: return None - return SessionRecord( + return StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -285,7 +289,7 @@ async def update_session_state(self, app_name: str, user_id: str, session_id: st async with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur: await cur.execute(query, (Jsonb(state), app_name, user_id, session_id)) - async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + async def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": if user_id is None: query = pg_sql.SQL(""" SELECT id, app_name, user_id, state, create_time, update_time @@ -309,7 +313,7 @@ async def list_sessions(self, app_name: str, user_id: str | None = None) -> "lis rows = await cur.fetchall() return [ - SessionRecord( + StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -330,7 +334,7 @@ async def delete_session(self, app_name: str, user_id: str, session_id: str) -> async with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur: await cur.execute(query, (app_name, user_id, session_id)) - async def append_event(self, event_record: EventRecord) -> None: + async def append_event(self, event_record: StoredEvent) -> None: query = pg_sql.SQL(""" INSERT INTO {table} ( id, session_id, invocation_id, timestamp, event_data @@ -354,7 +358,7 @@ async def append_event(self, event_record: EventRecord) -> None: async def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -362,7 +366,7 @@ async def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: insert_query = pg_sql.SQL(""" INSERT INTO {table} ( id, session_id, invocation_id, timestamp, event_data @@ -420,7 +424,7 @@ async def append_event_and_update_state( raise await conn.commit() - return SessionRecord( + return StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -436,7 +440,7 @@ async def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": if limit == 0: return [] @@ -472,7 +476,7 @@ async def get_events( rows = await cur.fetchall() return [ - EventRecord( + StoredEvent( id=row["id"], session_id=row["session_id"], invocation_id=row["invocation_id"], @@ -674,7 +678,7 @@ def create_tables(self) -> None: def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new session.""" params: tuple[Any, ...] if self._owner_id_column_name: @@ -703,7 +707,7 @@ def create_session( def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get session by ID.""" if renew_for is not None and self._calculate_expires_at(renew_for) is not None: query = pg_sql.SQL(""" @@ -729,7 +733,7 @@ def get_session( if row is None: return None - return SessionRecord( + return StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -751,7 +755,7 @@ def update_session_state(self, app_name: str, user_id: str, session_id: str, sta with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur: cur.execute(query, (Jsonb(state), app_name, user_id, session_id)) - def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": """List sessions for an app.""" if user_id is None: query = pg_sql.SQL(""" @@ -776,7 +780,7 @@ def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[Sess rows = cur.fetchall() return [ - SessionRecord( + StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -798,14 +802,14 @@ def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur: cur.execute(query, (app_name, user_id, session_id)) - def append_event(self, event_record: EventRecord) -> None: + def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session.""" """Synchronous implementation of append_event.""" self._insert_event(event_record) def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -813,7 +817,7 @@ def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically append an event and update session + scoped state.""" insert_query = pg_sql.SQL(""" INSERT INTO {table} ( @@ -872,7 +876,7 @@ def append_event_and_update_state( raise conn.commit() - return SessionRecord( + return StoredSession( id=row["id"], app_name=row["app_name"], user_id=row["user_id"], @@ -888,7 +892,7 @@ def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Get events for a session.""" if limit == 0: return [] @@ -925,7 +929,7 @@ def get_events( rows = cur.fetchall() return [ - EventRecord( + StoredEvent( id=row["id"], session_id=row["session_id"], invocation_id=row["invocation_id"], @@ -1114,7 +1118,7 @@ def _drop_tables_sql(self) -> "list[str]": f"DROP TABLE IF EXISTS {self._session_table}", ] - def _insert_event(self, event_record: EventRecord) -> None: + def _insert_event(self, event_record: StoredEvent) -> None: insert_query = pg_sql.SQL(""" INSERT INTO {table} ( id, session_id, invocation_id, timestamp, event_data @@ -1159,7 +1163,7 @@ async def create_tables(self) -> None: async with self._config.provide_session() as driver: await driver.execute_script(await self._memory_table_ddl()) - async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with deduplication.""" if not self._enabled: msg = "Memory store is disabled" @@ -1172,11 +1176,11 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " if self._owner_id_column_name: query = pg_sql.SQL(""" INSERT INTO {table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, {owner_id_col}, timestamp, content_json, content_text, metadata_json, inserted_at ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s + %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s ) ON CONFLICT (event_id) DO NOTHING """).format( @@ -1185,10 +1189,10 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " else: query = pg_sql.SQL(""" INSERT INTO {table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s + %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s ) ON CONFLICT (event_id) DO NOTHING """).format(table=pg_sql.Identifier(self._memory_table)) @@ -1205,8 +1209,13 @@ async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: " return inserted_count async def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries by text query.""" if not self._enabled: msg = "Memory store is disabled" @@ -1217,10 +1226,10 @@ async def search_entries( try: if self._use_fts: try: - return await self._search_entries_fts(query, app_name, user_id, effective_limit) + return await self._search_entries_fts(query, app_name, user_id, effective_limit, scope_filter) except Exception as exc: # pragma: no cover logger.warning("FTS search failed; falling back to simple search: %s", exc) - return await self._search_entries_simple(query, app_name, user_id, effective_limit) + return await self._search_entries_simple(query, app_name, user_id, effective_limit, scope_filter) except errors.UndefinedTable: return [] @@ -1234,17 +1243,30 @@ async def delete_entries_by_session(self, session_id: str) -> int: await cur.execute(sql, (session_id,)) return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 - async def delete_entries_older_than(self, days: int) -> int: + async def delete_entries_older_than( + self, days: int, app_name: "str | None" = None, scope: "str | None" = None + ) -> int: """Delete memory entries older than specified days.""" - sql = pg_sql.SQL( - """ - DELETE FROM {table} - WHERE inserted_at < CURRENT_TIMESTAMP - {interval}::interval - """ - ).format(table=pg_sql.Identifier(self._memory_table), interval=pg_sql.Literal(f"{days} days")) + clauses: list[pg_sql.Composable] = [ + pg_sql.SQL("inserted_at < CURRENT_TIMESTAMP - {interval}::interval").format( + interval=pg_sql.Literal(f"{days} days") + ) + ] + params: list[Any] = [] + if app_name is not None: + clauses.append(pg_sql.SQL("app_name = %s")) + params.append(app_name) + if scope is not None: + clauses.append(pg_sql.SQL("scope = %s")) + params.append(scope) + + where_sql = pg_sql.SQL(" AND ").join(clauses) + sql = pg_sql.SQL("DELETE FROM {table} WHERE {where}").format( + table=pg_sql.Identifier(self._memory_table), where=where_sql + ) async with self._config.provide_connection() as conn, conn.cursor() as cur: - await cur.execute(sql) + await cur.execute(sql, tuple(params) if params else None) return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 async def _memory_table_ddl(self) -> str: @@ -1264,6 +1286,8 @@ async def _memory_table_ddl(self) -> str: self._memory_table, self._memory_table, self._memory_table, + self._memory_table, + self._memory_table, fts_index, ) @@ -1271,42 +1295,46 @@ def _drop_memory_table_sql(self) -> "list[str]": """Get PostgreSQL DROP TABLE SQL statements.""" return [f"DROP TABLE IF EXISTS {self._memory_table}"] - async def _search_entries_fts(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": + async def _search_entries_fts( + self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "list[StoredMemory]": + where_scope, scope_params = _build_psycopg_scope_where(app_name, user_id, scope_filter) sql = pg_sql.SQL( """ - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at, ts_rank(to_tsvector('english', content_text), plainto_tsquery('english', %s)) as rank FROM {table} - WHERE app_name = %s - AND user_id = %s + WHERE {where_scope} AND to_tsvector('english', content_text) @@ plainto_tsquery('english', %s) ORDER BY rank DESC, timestamp DESC LIMIT %s """ - ).format(table=pg_sql.Identifier(self._memory_table)) - params: tuple[str, str, str, str, int] = (query, app_name, user_id, query, limit) - async with self._config.provide_connection() as conn, conn.cursor() as cur: + ).format(table=pg_sql.Identifier(self._memory_table), where_scope=where_scope) + params = (query, *scope_params, query, limit) + async with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur: await cur.execute(sql, params) rows = await cur.fetchall() return _rows_to_records(rows) - async def _search_entries_simple(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": + async def _search_entries_simple( + self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "list[StoredMemory]": + where_scope, scope_params = _build_psycopg_scope_where(app_name, user_id, scope_filter) sql = pg_sql.SQL( """ - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM {table} - WHERE app_name = %s - AND user_id = %s + WHERE {where_scope} AND content_text ILIKE %s ORDER BY timestamp DESC LIMIT %s """ - ).format(table=pg_sql.Identifier(self._memory_table)) + ).format(table=pg_sql.Identifier(self._memory_table), where_scope=where_scope) pattern = f"%{query}%" - params: tuple[str, str, str, int] = (app_name, user_id, pattern, limit) - async with self._config.provide_connection() as conn, conn.cursor() as cur: + params = (*scope_params, pattern, limit) + async with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur: await cur.execute(sql, params) rows = await cur.fetchall() return _rows_to_records(rows) @@ -1334,7 +1362,7 @@ def create_tables(self) -> None: with self._config.provide_session() as driver: driver.execute_script(self._memory_table_ddl()) - def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with deduplication.""" """Bulk insert memory entries with deduplication.""" if not self._enabled: @@ -1348,11 +1376,11 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object if self._owner_id_column_name: query = pg_sql.SQL(""" INSERT INTO {table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, {owner_id_col}, timestamp, content_json, content_text, metadata_json, inserted_at ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s + %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s ) ON CONFLICT (event_id) DO NOTHING """).format( @@ -1361,10 +1389,10 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object else: query = pg_sql.SQL(""" INSERT INTO {table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s + %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s ) ON CONFLICT (event_id) DO NOTHING """).format(table=pg_sql.Identifier(self._memory_table)) @@ -1381,8 +1409,13 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object return inserted_count def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries by text query.""" """Search memory entries by text query.""" if not self._enabled: @@ -1394,10 +1427,10 @@ def search_entries( try: if self._use_fts: try: - return self._search_entries_fts(query, app_name, user_id, effective_limit) + return self._search_entries_fts(query, app_name, user_id, effective_limit, scope_filter) except Exception as exc: # pragma: no cover logger.warning("FTS search failed; falling back to simple search: %s", exc) - return self._search_entries_simple(query, app_name, user_id, effective_limit) + return self._search_entries_simple(query, app_name, user_id, effective_limit, scope_filter) except errors.UndefinedTable: return [] @@ -1412,18 +1445,28 @@ def delete_entries_by_session(self, session_id: str) -> int: cur.execute(sql, (session_id,)) return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 - def delete_entries_older_than(self, days: int) -> int: - """Delete memory entries older than specified days.""" + def delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: """Delete memory entries older than specified days.""" - sql = pg_sql.SQL( - """ - DELETE FROM {table} - WHERE inserted_at < CURRENT_TIMESTAMP - {interval}::interval - """ - ).format(table=pg_sql.Identifier(self._memory_table), interval=pg_sql.Literal(f"{days} days")) + clauses: list[pg_sql.Composable] = [ + pg_sql.SQL("inserted_at < CURRENT_TIMESTAMP - {interval}::interval").format( + interval=pg_sql.Literal(f"{days} days") + ) + ] + params: list[Any] = [] + if app_name is not None: + clauses.append(pg_sql.SQL("app_name = %s")) + params.append(app_name) + if scope is not None: + clauses.append(pg_sql.SQL("scope = %s")) + params.append(scope) + + where_sql = pg_sql.SQL(" AND ").join(clauses) + sql = pg_sql.SQL("DELETE FROM {table} WHERE {where}").format( + table=pg_sql.Identifier(self._memory_table), where=where_sql + ) with self._config.provide_connection() as conn, conn.cursor() as cur: - cur.execute(sql) + cur.execute(sql, tuple(params) if params else None) return cur.rowcount if cur.rowcount and cur.rowcount > 0 else 0 def _memory_table_ddl(self) -> str: @@ -1443,6 +1486,8 @@ def _memory_table_ddl(self) -> str: self._memory_table, self._memory_table, self._memory_table, + self._memory_table, + self._memory_table, fts_index, ) @@ -1450,53 +1495,58 @@ def _drop_memory_table_sql(self) -> "list[str]": """Get PostgreSQL DROP TABLE SQL statements.""" return [f"DROP TABLE IF EXISTS {self._memory_table}"] - def _search_entries_fts(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": + def _search_entries_fts( + self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "list[StoredMemory]": + where_scope, scope_params = _build_psycopg_scope_where(app_name, user_id, scope_filter) sql = pg_sql.SQL( """ - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at, ts_rank(to_tsvector('english', content_text), plainto_tsquery('english', %s)) as rank FROM {table} - WHERE app_name = %s - AND user_id = %s + WHERE {where_scope} AND to_tsvector('english', content_text) @@ plainto_tsquery('english', %s) ORDER BY rank DESC, timestamp DESC LIMIT %s """ - ).format(table=pg_sql.Identifier(self._memory_table)) - params: tuple[str, str, str, str, int] = (query, app_name, user_id, query, limit) - with self._config.provide_connection() as conn, conn.cursor() as cur: + ).format(table=pg_sql.Identifier(self._memory_table), where_scope=where_scope) + params = (query, *scope_params, query, limit) + with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur: cur.execute(sql, params) rows = cur.fetchall() return _rows_to_records(rows) - def _search_entries_simple(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": + def _search_entries_simple( + self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "list[StoredMemory]": + where_scope, scope_params = _build_psycopg_scope_where(app_name, user_id, scope_filter) sql = pg_sql.SQL( """ - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM {table} - WHERE app_name = %s - AND user_id = %s + WHERE {where_scope} AND content_text ILIKE %s ORDER BY timestamp DESC LIMIT %s """ - ).format(table=pg_sql.Identifier(self._memory_table)) + ).format(table=pg_sql.Identifier(self._memory_table), where_scope=where_scope) pattern = f"%{query}%" - params: tuple[str, str, str, int] = (app_name, user_id, pattern, limit) - with self._config.provide_connection() as conn, conn.cursor() as cur: + params = (*scope_params, pattern, limit) + with self._config.provide_connection() as conn, conn.cursor(row_factory=dict_row) as cur: cur.execute(sql, params) rows = cur.fetchall() return _rows_to_records(rows) -def _build_insert_params(entry: "MemoryRecord") -> "tuple[object, ...]": +def _build_insert_params(entry: "StoredMemory") -> "tuple[object, ...]": return ( entry["id"], entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], entry["timestamp"], @@ -1507,12 +1557,13 @@ def _build_insert_params(entry: "MemoryRecord") -> "tuple[object, ...]": ) -def _build_insert_params_with_owner(entry: "MemoryRecord", owner_id: "object | None") -> "tuple[object, ...]": +def _build_insert_params_with_owner(entry: "StoredMemory", owner_id: "object | None") -> "tuple[object, ...]": return ( entry["id"], entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], owner_id, @@ -1524,13 +1575,14 @@ def _build_insert_params_with_owner(entry: "MemoryRecord", owner_id: "object | N ) -def _rows_to_records(rows: "list[Any]") -> "list[MemoryRecord]": +def _rows_to_records(rows: "list[Any]") -> "list[StoredMemory]": return [ { "id": row["id"], "session_id": row["session_id"], "app_name": row["app_name"], "user_id": row["user_id"], + "scope": row.get("scope", "user"), "event_id": row["event_id"], "author": row["author"], "timestamp": row["timestamp"], @@ -1538,6 +1590,7 @@ def _rows_to_records(rows: "list[Any]") -> "list[MemoryRecord]": "content_text": row["content_text"], "metadata_json": row["metadata_json"], "inserted_at": row["inserted_at"], + "embedding": None, } for row in rows ] @@ -1603,3 +1656,16 @@ def _postgres_event_ddl_options(adk_config: PsycopgADKConfig, events_table: str) def _raise_missing_session(session_id: str) -> NoReturn: msg = f"Session {session_id} not found during append_event_and_update_state." raise ValueError(msg) + + +def _build_psycopg_scope_where( + app_name: str, user_id: str, scope_filter: Literal["all", "user", "app"] +) -> tuple[pg_sql.Composable, tuple[Any, ...]]: + if scope_filter == "all": + where = pg_sql.SQL("app_name = %s AND ((scope = 'user' AND user_id = %s) OR scope = 'app')") + return where, (app_name, user_id) + if scope_filter == "user": + where = pg_sql.SQL("app_name = %s AND scope = 'user' AND user_id = %s") + return where, (app_name, user_id) + where = pg_sql.SQL("app_name = %s AND scope = 'app'") + return where, (app_name,) diff --git a/sqlspec/adapters/pymssql/adk/store.py b/sqlspec/adapters/pymssql/adk/store.py index f5f15ec2e..9403ab9d4 100644 --- a/sqlspec/adapters/pymssql/adk/store.py +++ b/sqlspec/adapters/pymssql/adk/store.py @@ -2,14 +2,14 @@ import re from datetime import datetime -from typing import TYPE_CHECKING, Any, ClassVar, Final, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast from typing_extensions import NotRequired from sqlspec.adapters.pymssql._typing import PYMSSQL_MODULE, PymssqlCursor from sqlspec.adapters.pymssql.data_dictionary import MssqlVersionInfo from sqlspec.config import ADKConfig -from sqlspec.extensions.adk import BaseSyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.utils.serializers import from_json, to_json @@ -18,7 +18,7 @@ from sqlspec.adapters.pymssql.config import PymssqlConfig from sqlspec.adapters.pymssql.driver import PymssqlDriver - from sqlspec.extensions.adk.memory._types import MemoryRecord + from sqlspec.extensions.adk.memory._types import StoredMemory __all__ = ("PymssqlADKConfig", "PymssqlADKMemoryStore", "PymssqlADKStore") @@ -74,7 +74,7 @@ def create_tables(self) -> None: def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new ADK session.""" owner_column = f", {_quote_identifier(self._owner_id_column_name)}" if self._owner_id_column_name else "" owner_param = ", %s" if self._owner_id_column_name else "" @@ -98,7 +98,7 @@ def create_session( def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Return a scoped session or ``None`` if absent.""" try: if renew_for is not None and self._calculate_expires_at(renew_for) is not None: @@ -137,7 +137,7 @@ def update_session_state(self, app_name: str, user_id: str, session_id: str, sta commit=True, ) - def list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[StoredSession]": """List ADK sessions for an application, optionally scoped to a user.""" if user_id is None: sql = f""" @@ -171,13 +171,13 @@ def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: commit=True, ) - def append_event(self, event_record: EventRecord) -> None: + def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session.""" self._execute(_insert_event_sql(self._events_table), _event_insert_params(event_record), commit=True) def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -185,7 +185,7 @@ def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically append an event and update durable session/scoped state.""" update_sql = f""" UPDATE {_table_ref(self._session_table)} @@ -217,7 +217,7 @@ def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Return events for a scoped session ordered by event timestamp.""" if limit == 0: return [] @@ -422,7 +422,7 @@ def create_tables(self) -> None: driver.execute(_create_index_sql(index_table, index_name, columns)) driver.commit() - def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with event-id deduplication.""" if not self._enabled: msg = "ADK memory store is disabled" @@ -436,10 +436,10 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object IF NOT EXISTS (SELECT 1 FROM {_table_ref(self._memory_table)} WHERE event_id = %s) BEGIN INSERT INTO {_table_ref(self._memory_table)} ( - id, session_id, app_name, user_id, event_id, author, timestamp, + id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json{owner_column} ) - VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s{owner_value}); + VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s{owner_value}); END; """ inserted = 0 @@ -451,6 +451,7 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry.get("author"), entry["timestamp"], @@ -466,22 +467,28 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object return inserted def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries by text query.""" if not self._enabled: msg = "ADK memory store is disabled" raise RuntimeError(msg) limit_value = limit or self._max_results + where_scope, scope_params = _build_mssql_scope_where(app_name, user_id, scope_filter) sql = f""" SELECT TOP (%s) - id, session_id, app_name, user_id, event_id, author, timestamp, + id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM {_table_ref(self._memory_table)} - WHERE app_name = %s AND user_id = %s AND content_text LIKE %s + WHERE {where_scope} AND content_text LIKE %s ORDER BY timestamp DESC """ - rows = self._execute_fetchall(sql, (limit_value, app_name, user_id, f"%{query}%")) + rows = self._execute_fetchall(sql, (limit_value, *scope_params, f"%{query}%")) return [_memory_record_from_row(row) for row in rows] def delete_entries_by_session(self, session_id: str) -> int: @@ -490,12 +497,19 @@ def delete_entries_by_session(self, session_id: str) -> int: f"DELETE FROM {_table_ref(self._memory_table)} WHERE session_id = %s", (session_id,), commit=True ) - def delete_entries_older_than(self, days: int) -> int: + def delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: """Delete memory entries older than the retention window.""" + clauses = ["inserted_at < DATEADD(day, -%s, SYSUTCDATETIME())"] + params: list[Any] = [days] + if app_name is not None: + clauses.append("app_name = %s") + params.append(app_name) + if scope is not None: + clauses.append("scope = %s") + params.append(scope) + where_sql = " AND ".join(clauses) return self._execute( - f"DELETE FROM {_table_ref(self._memory_table)} WHERE inserted_at < DATEADD(day, -%s, SYSUTCDATETIME())", - (days,), - commit=True, + f"DELETE FROM {_table_ref(self._memory_table)} WHERE {where_sql}", tuple(params), commit=True ) def _memory_table_ddl(self) -> str: @@ -509,6 +523,7 @@ def _memory_table_ddl(self) -> str: session_id NVARCHAR(128) NOT NULL, app_name NVARCHAR(128) NOT NULL, user_id NVARCHAR(128) NOT NULL, + scope NVARCHAR(16) NOT NULL CONSTRAINT {_constraint_ref("df", self._memory_table, "scope")} DEFAULT N'user', event_id NVARCHAR(128) NOT NULL, author NVARCHAR(256) NULL, timestamp DATETIME2(6) NOT NULL, @@ -526,7 +541,12 @@ def _memory_table_ddl(self) -> str: def _memory_index_specs(self) -> "list[tuple[str, str, str]]": """Return ``(index_name, table, columns)`` specs for memory-table indexes.""" return [ - (f"idx_{self._memory_table}_scope", self._memory_table, "app_name, user_id"), + ( + f"idx_{self._memory_table}_app_scope_user_time", + self._memory_table, + "app_name, scope, user_id, timestamp DESC", + ), + (f"idx_{self._memory_table}_scope", self._memory_table, "app_name, scope"), (f"idx_{self._memory_table}_session", self._memory_table, "session_id"), (f"idx_{self._memory_table}_timestamp", self._memory_table, "timestamp DESC"), ] @@ -748,7 +768,7 @@ def _events_query( return sql, tuple(params) -def _event_insert_params(event_record: EventRecord) -> "tuple[Any, ...]": +def _event_insert_params(event_record: StoredEvent) -> "tuple[Any, ...]": return ( event_record["id"], event_record["app_name"], @@ -760,14 +780,14 @@ def _event_insert_params(event_record: EventRecord) -> "tuple[Any, ...]": ) -def _session_record_from_row(row: Any) -> SessionRecord: - return SessionRecord( +def _session_record_from_row(row: Any) -> StoredSession: + return StoredSession( id=row[0], app_name=row[1], user_id=row[2], state=_json_dict(row[3]), create_time=row[4], update_time=row[5] ) -def _event_record_from_row(row: Any) -> EventRecord: - return EventRecord( +def _event_record_from_row(row: Any) -> StoredEvent: + return StoredEvent( id=row[0], app_name=row[1], user_id=row[2], @@ -778,21 +798,23 @@ def _event_record_from_row(row: Any) -> EventRecord: ) -def _memory_record_from_row(row: Any) -> "MemoryRecord": +def _memory_record_from_row(row: Any) -> "StoredMemory": return cast( - "MemoryRecord", + "StoredMemory", { "id": row[0], "session_id": row[1], "app_name": row[2], "user_id": row[3], - "event_id": row[4], - "author": row[5], - "timestamp": row[6], - "content_json": _json_dict(row[7]), - "content_text": row[8], - "metadata_json": _json_dict(row[9]) if row[9] is not None else None, - "inserted_at": row[10], + "scope": row[4], + "event_id": row[5], + "author": row[6], + "timestamp": row[7], + "content_json": _json_dict(row[8]), + "content_text": row[9], + "metadata_json": _json_dict(row[10]) if row[10] is not None else None, + "inserted_at": row[11], + "embedding": None, }, ) @@ -848,3 +870,13 @@ def _escape_sql_literal(value: str) -> str: def _raise_session_not_found(session_id: str) -> None: msg = f"Session {session_id} not found during append_event_and_update_state." raise ValueError(msg) + + +def _build_mssql_scope_where( + app_name: str, user_id: str, scope_filter: Literal["all", "user", "app"] +) -> tuple[str, tuple[Any, ...]]: + if scope_filter == "all": + return "app_name = %s AND ((scope = 'user' AND user_id = %s) OR scope = 'app')", (app_name, user_id) + if scope_filter == "user": + return "app_name = %s AND scope = 'user' AND user_id = %s", (app_name, user_id) + return "app_name = %s AND scope = 'app'", (app_name,) diff --git a/sqlspec/adapters/pymysql/adk/store.py b/sqlspec/adapters/pymysql/adk/store.py index 70d67b9ce..7b71d55a0 100644 --- a/sqlspec/adapters/pymysql/adk/store.py +++ b/sqlspec/adapters/pymysql/adk/store.py @@ -2,13 +2,13 @@ import re from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast import pymysql from typing_extensions import NotRequired from sqlspec.config import ADKConfig -from sqlspec.extensions.adk import BaseSyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.utils.serializers import from_json, to_json @@ -16,7 +16,7 @@ from datetime import datetime, timedelta from sqlspec.adapters.pymysql.config import PyMysqlConfig - from sqlspec.extensions.adk import MemoryRecord + from sqlspec.extensions.adk import StoredMemory __all__ = ("PyMysqlADKConfig", "PyMysqlADKMemoryStore", "PyMysqlADKStore") @@ -70,13 +70,13 @@ def create_tables(self) -> None: def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new session.""" return _create_session(self, session_id, app_name, user_id, state, owner_id) def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get session by scoped identifiers.""" return _select_session(self, app_name, user_id, session_id, renew_for=renew_for) @@ -84,7 +84,7 @@ def update_session_state(self, app_name: str, user_id: str, session_id: str, sta """Update session state.""" _update_session_state(self, app_name, user_id, session_id, state) - def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": """List sessions for an app.""" return _list_sessions(self, app_name, user_id) @@ -92,13 +92,13 @@ def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: """Delete session and associated events.""" _delete_session(self, app_name, user_id, session_id) - def append_event(self, event_record: EventRecord) -> None: + def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session.""" _append_event(self, event_record) def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -106,7 +106,7 @@ def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically append an event and update session + scoped state.""" return _append_event_and_update_state( self, event_record, app_name, user_id, session_id, state, app_state=app_state, user_state=user_state @@ -119,7 +119,7 @@ def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Get events for a session.""" return _select_events(self, app_name, user_id, session_id, after_timestamp, limit) @@ -220,23 +220,28 @@ def create_tables(self) -> None: self._create_tables() - def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with deduplication.""" return self._insert_memory_entries(entries, owner_id) def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries by text query.""" - return self._search_entries(query, app_name, user_id, limit) + return self._search_entries(query, app_name, user_id, limit, scope_filter) def delete_entries_by_session(self, session_id: str) -> int: """Delete all memory entries for a specific session.""" return self._delete_entries_by_session(session_id) - def delete_entries_older_than(self, days: int) -> int: + def delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: """Delete memory entries older than specified days.""" - return self._delete_entries_older_than(days) + return self._delete_entries_older_than(days, app_name, scope) def _memory_table_ddl(self) -> str: adk_config = _adk_config(self._config) @@ -259,6 +264,7 @@ def _memory_table_ddl(self) -> str: session_id VARCHAR(128) NOT NULL, app_name VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, + scope VARCHAR(16) NOT NULL DEFAULT 'user', event_id VARCHAR(128) NOT NULL UNIQUE, author VARCHAR(256){owner_id_line}, timestamp TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), @@ -266,7 +272,8 @@ def _memory_table_ddl(self) -> str: content_text TEXT NOT NULL, metadata_json JSON, inserted_at TIMESTAMP(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6), - INDEX idx_{self._memory_table}_app_user_time (app_name, user_id, timestamp), + INDEX idx_{self._memory_table}_app_scope_user_time (app_name, scope, user_id, timestamp), + INDEX idx_{self._memory_table}_scope (app_name, scope), INDEX idx_{self._memory_table}_session (session_id){fts_index}{fk_constraint} ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci{table_options} """ @@ -281,7 +288,7 @@ def _create_tables(self) -> None: with self._config.provide_session() as driver: driver.execute_script(self._memory_table_ddl()) - def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def _insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -293,17 +300,17 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec if self._owner_id_column_name: sql = f""" INSERT IGNORE INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) + ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """ else: sql = f""" INSERT IGNORE INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author, + id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at - ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) + ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """ with self._config.provide_connection() as conn: @@ -317,6 +324,7 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], owner_id, @@ -332,6 +340,7 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec entry["session_id"], entry["app_name"], entry["user_id"], + entry.get("scope", "user"), entry["event_id"], entry["author"], entry["timestamp"], @@ -348,8 +357,13 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec return inserted_count def _search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -358,23 +372,24 @@ def _search_entries( return [] limit_value = limit or self._max_results + where_scope, scope_params = _build_mysql_scope_where(app_name, user_id, scope_filter) if self._use_fts: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = %s AND user_id = %s + WHERE {where_scope} AND MATCH(content_text) AGAINST (%s IN NATURAL LANGUAGE MODE) ORDER BY timestamp DESC LIMIT %s """ - params = (app_name, user_id, query, limit_value) + params = (*scope_params, query, limit_value) else: sql = f""" SELECT * FROM {self._memory_table} - WHERE app_name = %s AND user_id = %s AND content_text LIKE %s + WHERE {where_scope} AND content_text LIKE %s ORDER BY timestamp DESC LIMIT %s """ - params = (app_name, user_id, f"%{query}%", limit_value) + params = (*scope_params, f"%{query}%", limit_value) with self._config.provide_connection() as conn: cursor = conn.cursor() @@ -385,7 +400,12 @@ def _search_entries( finally: cursor.close() - return [cast("MemoryRecord", dict(zip(columns, row, strict=False))) for row in rows] + records: list[StoredMemory] = [] + for row in rows: + record = cast("StoredMemory", dict(zip(columns, row, strict=False))) + record["embedding"] = None + records.append(record) + return records def _delete_entries_by_session(self, session_id: str) -> int: if not self._enabled: @@ -402,19 +422,25 @@ def _delete_entries_by_session(self, session_id: str) -> int: finally: cursor.close() - def _delete_entries_older_than(self, days: int) -> int: + def _delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) - sql = f""" - DELETE FROM {self._memory_table} - WHERE inserted_at < (UTC_TIMESTAMP(6) - INTERVAL %s DAY) - """ + clauses = ["inserted_at < (UTC_TIMESTAMP(6) - INTERVAL %s DAY)"] + params: list[Any] = [days] + if app_name is not None: + clauses.append("app_name = %s") + params.append(app_name) + if scope is not None: + clauses.append("scope = %s") + params.append(scope) + where_sql = " AND ".join(clauses) + sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" with self._config.provide_connection() as conn: cursor = conn.cursor() try: - cursor.execute(sql, (days,)) + cursor.execute(sql, tuple(params)) conn.commit() return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0 finally: @@ -437,7 +463,7 @@ def _create_session( user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None, -) -> SessionRecord: +) -> StoredSession: params: tuple[Any, ...] if store._owner_id_column_name: sql = f""" @@ -469,7 +495,7 @@ def _create_session( def _select_session( store: PyMysqlADKStore, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None -) -> "SessionRecord | None": +) -> "StoredSession | None": try: with store._config.provide_connection() as conn: cursor = conn.cursor() @@ -520,7 +546,7 @@ def _update_session_state( conn.commit() -def _list_sessions(store: PyMysqlADKStore, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": +def _list_sessions(store: PyMysqlADKStore, app_name: str, user_id: str | None = None) -> "list[StoredSession]": if user_id is None: sql = f""" SELECT id, app_name, user_id, state, create_time, update_time @@ -564,7 +590,7 @@ def _delete_session(store: PyMysqlADKStore, app_name: str, user_id: str, session conn.commit() -def _append_event(store: PyMysqlADKStore, event_record: EventRecord) -> None: +def _append_event(store: PyMysqlADKStore, event_record: StoredEvent) -> None: sql = f""" INSERT INTO {store._events_table} ( id, app_name, user_id, session_id, invocation_id, timestamp, event_data @@ -581,7 +607,7 @@ def _append_event(store: PyMysqlADKStore, event_record: EventRecord) -> None: def _append_event_and_update_state( store: PyMysqlADKStore, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -589,7 +615,7 @@ def _append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, -) -> SessionRecord: +) -> StoredSession: insert_sql = f""" INSERT INTO {store._events_table} ( id, app_name, user_id, session_id, invocation_id, timestamp, event_data @@ -649,7 +675,7 @@ def _select_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, -) -> "list[EventRecord]": +) -> "list[StoredEvent]": if limit == 0: return [] @@ -830,14 +856,14 @@ def _json_dict(value: Any) -> "dict[str, Any]": return cast("dict[str, Any]", value) -def _session_record_from_row(row: Any) -> SessionRecord: - return SessionRecord( +def _session_record_from_row(row: Any) -> StoredSession: + return StoredSession( id=row[0], app_name=row[1], user_id=row[2], state=_json_dict(row[3]), create_time=row[4], update_time=row[5] ) -def _event_record_from_row(row: Any) -> EventRecord: - return EventRecord( +def _event_record_from_row(row: Any) -> StoredEvent: + return StoredEvent( id=row[0], app_name=row[1], user_id=row[2], @@ -848,7 +874,7 @@ def _event_record_from_row(row: Any) -> EventRecord: ) -def _event_insert_params(event_record: EventRecord) -> "tuple[Any, ...]": +def _event_insert_params(event_record: StoredEvent) -> "tuple[Any, ...]": return ( event_record["id"], event_record["app_name"], @@ -972,3 +998,13 @@ def _mysql_upsert_metadata_sql(metadata_table: str) -> str: def _raise_session_not_found(session_id: str) -> None: msg = f"Session {session_id} not found during append_event_and_update_state." raise ValueError(msg) + + +def _build_mysql_scope_where( + app_name: str, user_id: str, scope_filter: Literal["all", "user", "app"] +) -> tuple[str, tuple[Any, ...]]: + if scope_filter == "all": + return "app_name = %s AND ((scope = 'user' AND user_id = %s) OR scope = 'app')", (app_name, user_id) + if scope_filter == "user": + return "app_name = %s AND scope = 'user' AND user_id = %s", (app_name, user_id) + return "app_name = %s AND scope = 'app'", (app_name,) diff --git a/sqlspec/adapters/spanner/adk/store.py b/sqlspec/adapters/spanner/adk/store.py index 730892c7c..3ccf45899 100644 --- a/sqlspec/adapters/spanner/adk/store.py +++ b/sqlspec/adapters/spanner/adk/store.py @@ -2,7 +2,7 @@ from collections.abc import Iterable, Mapping from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, ClassVar, Final, Protocol, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Protocol, cast from google.api_core.exceptions import NotFound from google.cloud.spanner_v1 import param_types @@ -11,7 +11,7 @@ from sqlspec.adapters.spanner.config import SpannerSyncConfig from sqlspec.config import ADKConfig from sqlspec.exceptions import OperationalError -from sqlspec.extensions.adk import BaseSyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.protocols import SpannerParamTypesProtocol from sqlspec.utils.serializers import from_json, to_json @@ -20,7 +20,7 @@ from google.cloud.spanner_v1.database import Database from google.cloud.spanner_v1.transaction import Transaction - from sqlspec.extensions.adk import MemoryRecord + from sqlspec.extensions.adk import StoredMemory __all__ = ("SpannerADKConfig", "SpannerADKRetentionConfig", "SpannerSyncADKMemoryStore", "SpannerSyncADKStore") @@ -91,13 +91,13 @@ def create_tables(self) -> None: def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new session.""" return self._create_session(session_id, app_name, user_id, state, owner_id) def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get session by ID.""" try: return self._get_session(app_name, user_id, session_id, renew_for=renew_for) @@ -110,7 +110,7 @@ def update_session_state(self, app_name: str, user_id: str, session_id: str, sta """Update session state.""" self._update_session_state(app_name, user_id, session_id, state) - def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": """List sessions for an app.""" try: return self._list_sessions(app_name, user_id) @@ -123,13 +123,13 @@ def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: """Delete session and associated events.""" self._delete_session(app_name, user_id, session_id) - def append_event(self, event_record: EventRecord) -> None: + def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session.""" self._append_event(event_record) def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -137,7 +137,7 @@ def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically append an event and update the session's durable state.""" return self._append_event_and_update_state( event_record, app_name, user_id, session_id, state, app_state=app_state, user_state=user_state @@ -150,7 +150,7 @@ def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Get events for a session.""" try: return self._get_events(app_name, user_id, session_id, after_timestamp, limit) @@ -260,7 +260,7 @@ def _decode_json(self, raw: Any) -> Any: def _create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: state_json = to_json(state) params: dict[str, Any] = {"id": session_id, "app_name": app_name, "user_id": user_id, "state": state_json} columns = "id, app_name, user_id, state, create_time, update_time" @@ -289,7 +289,7 @@ def _create_session( def _get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": if renew_for is not None and self._calculate_expires_at(renew_for) is not None: update_sql = f""" UPDATE {self._session_table} @@ -333,7 +333,7 @@ def _get_session( row = rows[0] state_value = self._decode_state(row[3]) - record: SessionRecord = { + record: StoredSession = { "id": row[0], "app_name": row[1], "user_id": row[2], @@ -366,7 +366,7 @@ def _update_session_state(self, app_name: str, user_id: str, session_id: str, st ) ]) - def _list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[SessionRecord]": + def _list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[StoredSession]": sql = f""" SELECT id, app_name, user_id, state, create_time, update_time{", " + self._owner_id_column_name if self._owner_id_column_name else ""} FROM {self._session_table} @@ -383,10 +383,10 @@ def _list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[S sql = f"{sql} ORDER BY update_time DESC" rows = self._run_read(sql, params, types) - records: list[SessionRecord] = [] + records: list[StoredSession] = [] for row in rows: state_value = self._decode_state(row[3]) - record: SessionRecord = { + record: StoredSession = { "id": row[0], "app_name": row[1], "user_id": row[2], @@ -413,7 +413,7 @@ def _delete_session(self, app_name: str, user_id: str, session_id: str) -> None: def _append_event_and_update_state( self, - event_record: "EventRecord", + event_record: "StoredEvent", app_name: str, user_id: str, session_id: str, @@ -421,12 +421,12 @@ def _append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically insert an event and update session state in one transaction. Both the event INSERT and the session state UPDATE execute within a single Spanner transaction so they succeed or fail together. A follow-up - single-use read returns the SessionRecord; we can't capture update_time + single-use read returns the StoredSession; we can't capture update_time inside the write txn because PENDING_COMMIT_TIMESTAMP() only materialises on commit. @@ -506,7 +506,7 @@ def _append_event_and_update_state( raise ValueError(msg) return record - def _insert_event(self, event_record: "EventRecord") -> None: + def _insert_event(self, event_record: "StoredEvent") -> None: event_params: dict[str, Any] = { "id": event_record["id"], "session_id": event_record["session_id"], @@ -527,7 +527,7 @@ def _get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": if limit == 0: return [] @@ -568,7 +568,7 @@ def _get_events( for row in rows ] - def _append_event(self, event_record: EventRecord) -> None: + def _append_event(self, event_record: StoredEvent) -> None: """Synchronous implementation of append_event.""" self._insert_event(event_record) @@ -781,23 +781,28 @@ def create_tables(self) -> None: self._create_tables() - def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with deduplication.""" return self._insert_memory_entries(entries, owner_id) def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries by text query.""" - return self._search_entries(query, app_name, user_id, limit) + return self._search_entries(query, app_name, user_id, limit, scope_filter) def delete_entries_by_session(self, session_id: str) -> int: """Delete all memory entries for a specific session.""" return self._delete_entries_by_session(session_id) - def delete_entries_older_than(self, days: int) -> int: + def delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: """Delete memory entries older than specified days.""" - return self._delete_entries_older_than(days) + return self._delete_entries_older_than(days, app_name, scope) def _database(self) -> "Database": return self._config.get_database() @@ -829,6 +834,7 @@ def _memory_param_types(self, include_owner: bool) -> "dict[str, Any]": "session_id": SPANNER_PARAM_TYPES.STRING, "app_name": SPANNER_PARAM_TYPES.STRING, "user_id": SPANNER_PARAM_TYPES.STRING, + "scope": SPANNER_PARAM_TYPES.STRING, "event_id": SPANNER_PARAM_TYPES.STRING, "author": SPANNER_PARAM_TYPES.STRING, "timestamp": SPANNER_PARAM_TYPES.TIMESTAMP, @@ -888,6 +894,7 @@ def _memory_table_ddl(self) -> "list[str]": session_id STRING(128) NOT NULL, app_name STRING(128) NOT NULL, user_id STRING(128) NOT NULL, + scope STRING(16) NOT NULL, event_id STRING(128) NOT NULL, author STRING(256){owner_line}, timestamp TIMESTAMP NOT NULL OPTIONS (allow_commit_timestamp=true), @@ -898,10 +905,11 @@ def _memory_table_ddl(self) -> "list[str]": ) {pk}{options}{self._memory_row_deletion_policy} """ - app_user_idx = f"CREATE INDEX idx_{self._memory_table}_app_user_time ON {self._memory_table}(app_name, user_id, timestamp DESC)" + app_scope_user_idx = f"CREATE INDEX idx_{self._memory_table}_app_scope_user_time ON {self._memory_table}(app_name, scope, user_id, timestamp DESC)" + scope_idx = f"CREATE INDEX idx_{self._memory_table}_scope ON {self._memory_table}(app_name, scope)" session_idx = f"CREATE INDEX idx_{self._memory_table}_session ON {self._memory_table}(session_id)" - statements = [table_sql, app_user_idx, session_idx] + statements = [table_sql, app_scope_user_idx, scope_idx, session_idx] if fts_index: statements.append(fts_index) return statements @@ -917,12 +925,13 @@ def _drop_memory_table_sql(self) -> "list[str]": statements.append(f"DROP SEARCH INDEX idx_{self._memory_table}_fts") statements.extend([ f"DROP INDEX idx_{self._memory_table}_session", - f"DROP INDEX idx_{self._memory_table}_app_user_time", + f"DROP INDEX idx_{self._memory_table}_app_scope_user_time", + f"DROP INDEX idx_{self._memory_table}_scope", f"DROP TABLE {self._memory_table}", ]) return statements - def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def _insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -938,10 +947,10 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec insert_sql = f""" INSERT INTO {self._memory_table} ( - id, session_id, app_name, user_id, event_id, author{owner_column}, + id, session_id, app_name, user_id, scope, event_id, author{owner_column}, timestamp, content_json, content_text, metadata_json, inserted_at ) VALUES ( - @id, @session_id, @app_name, @user_id, @event_id, @author{owner_param}, + @id, @session_id, @app_name, @user_id, @scope, @event_id, @author{owner_param}, @timestamp, @content_json, @content_text, @metadata_json, @inserted_at ) """ @@ -954,6 +963,7 @@ def _insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "objec "session_id": entry["session_id"], "app_name": entry["app_name"], "user_id": entry["user_id"], + "scope": entry.get("scope", "user"), "event_id": entry["event_id"], "author": entry["author"], "timestamp": entry["timestamp"], @@ -977,8 +987,13 @@ def _event_exists(self, event_id: str) -> bool: return bool(rows) def _search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -986,49 +1001,43 @@ def _search_entries( effective_limit = limit if limit is not None else self._max_results if self._use_fts: - return self._search_entries_fts(query, app_name, user_id, effective_limit) - return self._search_entries_simple(query, app_name, user_id, effective_limit) + return self._search_entries_fts(query, app_name, user_id, effective_limit, scope_filter) + return self._search_entries_simple(query, app_name, user_id, effective_limit, scope_filter) - def _search_entries_fts(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": + def _search_entries_fts( + self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "list[StoredMemory]": + where_scope, scope_params, scope_types = _build_spanner_scope_where(app_name, user_id, scope_filter) sql = f""" - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM {self._memory_table} - WHERE app_name = @app_name - AND user_id = @user_id + WHERE {where_scope} AND SEARCH(content_tokens, @query) ORDER BY timestamp DESC LIMIT @limit """ - params = {"app_name": app_name, "user_id": user_id, "query": query, "limit": limit} - types = { - "app_name": SPANNER_PARAM_TYPES.STRING, - "user_id": SPANNER_PARAM_TYPES.STRING, - "query": SPANNER_PARAM_TYPES.STRING, - "limit": SPANNER_PARAM_TYPES.INT64, - } + params = {**scope_params, "query": query, "limit": limit} + types = {**scope_types, "query": SPANNER_PARAM_TYPES.STRING, "limit": SPANNER_PARAM_TYPES.INT64} rows = self._run_read(sql, params, types) return self._rows_to_records(rows) - def _search_entries_simple(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": + def _search_entries_simple( + self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "list[StoredMemory]": + where_scope, scope_params, scope_types = _build_spanner_scope_where(app_name, user_id, scope_filter) sql = f""" - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM {self._memory_table} - WHERE app_name = @app_name - AND user_id = @user_id + WHERE {where_scope} AND LOWER(content_text) LIKE @pattern ORDER BY timestamp DESC LIMIT @limit """ pattern = f"%{query.lower()}%" - params = {"app_name": app_name, "user_id": user_id, "pattern": pattern, "limit": limit} - types = { - "app_name": SPANNER_PARAM_TYPES.STRING, - "user_id": SPANNER_PARAM_TYPES.STRING, - "pattern": SPANNER_PARAM_TYPES.STRING, - "limit": SPANNER_PARAM_TYPES.INT64, - } + params = {**scope_params, "pattern": pattern, "limit": limit} + types = {**scope_types, "pattern": SPANNER_PARAM_TYPES.STRING, "limit": SPANNER_PARAM_TYPES.INT64} rows = self._run_read(sql, params, types) return self._rows_to_records(rows) @@ -1038,27 +1047,39 @@ def _delete_entries_by_session(self, session_id: str) -> int: types = {"session_id": SPANNER_PARAM_TYPES.STRING} return self._execute_update(sql, params, types) - def _delete_entries_older_than(self, days: int) -> int: + def _delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: cutoff = datetime.now(timezone.utc) - timedelta(days=days) - sql = f"DELETE FROM {self._memory_table} WHERE inserted_at < @cutoff" - params = {"cutoff": cutoff} - types = {"cutoff": SPANNER_PARAM_TYPES.TIMESTAMP} + clauses = ["inserted_at < @cutoff"] + params: dict[str, Any] = {"cutoff": cutoff} + types: dict[str, Any] = {"cutoff": SPANNER_PARAM_TYPES.TIMESTAMP} + if app_name is not None: + clauses.append("app_name = @app_name") + params["app_name"] = app_name + types["app_name"] = SPANNER_PARAM_TYPES.STRING + if scope is not None: + clauses.append("scope = @scope") + params["scope"] = scope + types["scope"] = SPANNER_PARAM_TYPES.STRING + where_sql = " AND ".join(clauses) + sql = f"DELETE FROM {self._memory_table} WHERE {where_sql}" return self._execute_update(sql, params, types) - def _rows_to_records(self, rows: "list[Any]") -> "list[MemoryRecord]": + def _rows_to_records(self, rows: "list[Any]") -> "list[StoredMemory]": return [ { "id": row[0], "session_id": row[1], "app_name": row[2], "user_id": row[3], - "event_id": row[4], - "author": row[5], - "timestamp": row[6], - "content_json": self._decode_json(row[7]), - "content_text": row[8], - "metadata_json": self._decode_json(row[9]), - "inserted_at": row[10], + "scope": row[4], + "event_id": row[5], + "author": row[6], + "timestamp": row[7], + "content_json": self._decode_json(row[8]), + "content_text": row[9], + "metadata_json": self._decode_json(row[10]), + "inserted_at": row[11], + "embedding": None, } for row in rows ] @@ -1194,3 +1215,22 @@ class _SpannerReadProtocol(Protocol): def execute_sql( self, sql: str, params: "dict[str, Any] | None" = None, param_types: "dict[str, Any] | None" = None ) -> Iterable[Any]: ... + + +def _build_spanner_scope_where( + app_name: str, user_id: str, scope_filter: Literal["all", "user", "app"] +) -> tuple[str, dict[str, Any], dict[str, Any]]: + if scope_filter == "all": + where = "app_name = @app_name AND ((scope = 'user' AND user_id = @user_id) OR scope = 'app')" + params = {"app_name": app_name, "user_id": user_id} + types = {"app_name": SPANNER_PARAM_TYPES.STRING, "user_id": SPANNER_PARAM_TYPES.STRING} + return where, params, types + if scope_filter == "user": + where = "app_name = @app_name AND scope = 'user' AND user_id = @user_id" + params = {"app_name": app_name, "user_id": user_id} + types = {"app_name": SPANNER_PARAM_TYPES.STRING, "user_id": SPANNER_PARAM_TYPES.STRING} + return where, params, types + where = "app_name = @app_name AND scope = 'app'" + params = {"app_name": app_name} + types = {"app_name": SPANNER_PARAM_TYPES.STRING} + return where, params, types diff --git a/sqlspec/adapters/sqlite/adk/store.py b/sqlspec/adapters/sqlite/adk/store.py index ce379a921..1d1fe660f 100644 --- a/sqlspec/adapters/sqlite/adk/store.py +++ b/sqlspec/adapters/sqlite/adk/store.py @@ -11,7 +11,7 @@ from sqlspec.adapters.sqlite.config import _render_pragmas from sqlspec.config import ADKConfig from sqlspec.exceptions import ImproperConfigurationError -from sqlspec.extensions.adk import BaseSyncADKStore, EventRecord, SessionRecord +from sqlspec.extensions.adk import BaseSyncADKStore, StoredEvent, StoredSession from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.utils.logging import get_logger from sqlspec.utils.serializers import from_json, to_json @@ -20,7 +20,7 @@ import logging from sqlspec.adapters.sqlite.config import SqliteConfig - from sqlspec.extensions.adk import MemoryRecord + from sqlspec.extensions.adk import StoredMemory __all__ = ("SqliteADKConfig", "SqliteADKMemoryStore", "SqliteADKStore") @@ -98,7 +98,7 @@ def create_tables(self) -> None: def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> SessionRecord: + ) -> StoredSession: """Create a new session. Args: @@ -136,13 +136,13 @@ def create_session( conn.execute(sql, params) conn.commit() - return SessionRecord( + return StoredSession( id=session_id, app_name=app_name, user_id=user_id, state=state, create_time=now, update_time=now ) def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get session by ID. Args: @@ -187,7 +187,7 @@ def get_session( if row is None: return None - return SessionRecord( + return StoredSession( id=row[0], app_name=row[1], user_id=row[2], @@ -224,7 +224,7 @@ def update_session_state(self, app_name: str, user_id: str, session_id: str, sta conn.execute(sql, (state_json, now_julian, app_name, user_id, session_id)) conn.commit() - def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[StoredSession]": """List sessions for an app, optionally filtered by user. Args: @@ -259,7 +259,7 @@ def list_sessions(self, app_name: str, user_id: str | None = None) -> "list[Sess rows = cursor.fetchall() return [ - SessionRecord( + StoredSession( id=row[0], app_name=row[1], user_id=row[2], @@ -290,7 +290,7 @@ def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: conn.execute(sql, (app_name, user_id, session_id)) conn.commit() - def append_event(self, event_record: EventRecord) -> None: + def append_event(self, event_record: StoredEvent) -> None: """Append an event to a session. Args: @@ -324,7 +324,7 @@ def append_event(self, event_record: EventRecord) -> None: def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -332,11 +332,11 @@ def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> SessionRecord: + ) -> StoredSession: """Atomically append an event and update the session's durable state. Inserts the event and updates the session state + update_time in a - single transaction, returning the updated SessionRecord via RETURNING. + single transaction, returning the updated StoredSession via RETURNING. Args: event_record: Event record to store. @@ -418,7 +418,7 @@ def append_event_and_update_state( msg = f"Session {session_id} not found during append_event_and_update_state." raise ValueError(msg) - return SessionRecord( + return StoredSession( id=row[0], app_name=row[1], user_id=row[2], @@ -434,7 +434,7 @@ def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Get events for a session. Args: @@ -475,7 +475,7 @@ def get_events( rows = cursor.fetchall() return [ - EventRecord( + StoredEvent( id=row[0], app_name=row[1], user_id=row[2], @@ -780,7 +780,7 @@ def create_tables(self) -> None: self._enable_foreign_keys(driver.connection) driver.execute_script(self._memory_table_ddl()) - def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with deduplication.""" """Bulk insert memory entries with deduplication. @@ -813,20 +813,22 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object inserted_at_julian = _datetime_to_julian(entry["inserted_at"]) content_json_str = to_json(entry["content_json"]) metadata_json_str = to_json(entry["metadata_json"]) if entry["metadata_json"] else None + scope = entry.get("scope", "user") if self._owner_id_column_name: sql = f""" INSERT OR IGNORE INTO {self._memory_table} - (id, session_id, app_name, user_id, event_id, author, + (id, session_id, app_name, user_id, scope, event_id, author, {self._owner_id_column_name}, timestamp, content_json, content_text, metadata_json, inserted_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """ params: tuple[Any, ...] = ( entry["id"], entry["session_id"], entry["app_name"], entry["user_id"], + scope, entry["event_id"], entry["author"], owner_id, @@ -839,15 +841,16 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object else: sql = f""" INSERT OR IGNORE INTO {self._memory_table} - (id, session_id, app_name, user_id, event_id, author, + (id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """ params = ( entry["id"], entry["session_id"], entry["app_name"], entry["user_id"], + scope, entry["event_id"], entry["author"], timestamp_julian, @@ -866,23 +869,14 @@ def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object return inserted_count def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries by text query.""" - """Search memory entries by text query. - - Args: - query: Text query to search for. - app_name: Application name to filter by. - user_id: User ID to filter by. - limit: Maximum number of results (defaults to max_results config). - - Returns: - List of matching memory records ordered by relevance/timestamp. - - Raises: - RuntimeError: If memory store is disabled. - """ if not self._enabled: msg = "Memory store is disabled" raise RuntimeError(msg) @@ -891,21 +885,13 @@ def search_entries( if self._use_fts: try: - return self._search_entries_fts(query, app_name, user_id, effective_limit) + return self._search_entries_fts(query, app_name, user_id, effective_limit, scope_filter) except Exception as exc: # pragma: no cover logger.warning("FTS search failed; falling back to simple search: %s", exc) - return self._search_entries_simple(query, app_name, user_id, effective_limit) + return self._search_entries_simple(query, app_name, user_id, effective_limit, scope_filter) def delete_entries_by_session(self, session_id: str) -> int: """Delete all memory entries for a specific session.""" - """Delete all memory entries for a specific session. - - Args: - session_id: Session ID to delete entries for. - - Returns: - Number of entries deleted. - """ sql = f"DELETE FROM {self._memory_table} WHERE session_id = ?" with self._config.provide_connection() as conn: @@ -916,25 +902,22 @@ def delete_entries_by_session(self, session_id: str) -> int: return deleted_count - def delete_entries_older_than(self, days: int) -> int: + def delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: """Delete memory entries older than specified days.""" - """Delete memory entries older than specified days. - - Used for TTL cleanup operations. - - Args: - days: Number of days to retain entries. - - Returns: - Number of entries deleted. - """ cutoff_julian = _datetime_to_julian(datetime.now(timezone.utc)) - days sql = f"DELETE FROM {self._memory_table} WHERE inserted_at < ?" + params: list[Any] = [cutoff_julian] + if app_name is not None: + sql += " AND app_name = ?" + params.append(app_name) + if scope is not None: + sql += " AND scope = ?" + params.append(scope) with self._config.provide_connection() as conn: self._enable_foreign_keys(conn) - cursor = conn.execute(sql, (cutoff_julian,)) + cursor = conn.execute(sql, tuple(params)) deleted_count = cursor.rowcount conn.commit() @@ -982,6 +965,7 @@ def _memory_table_ddl(self) -> str: session_id TEXT NOT NULL, app_name TEXT NOT NULL, user_id TEXT NOT NULL, + scope TEXT NOT NULL DEFAULT 'user', event_id TEXT NOT NULL UNIQUE, author TEXT{owner_id_line}, timestamp REAL NOT NULL, @@ -991,8 +975,11 @@ def _memory_table_ddl(self) -> str: inserted_at REAL NOT NULL ); - CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_user_time - ON {self._memory_table}(app_name, user_id, timestamp DESC); + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_app_scope_user_time + ON {self._memory_table}(app_name, scope, user_id, timestamp DESC); + + CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_scope + ON {self._memory_table}(app_name, scope); CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_session ON {self._memory_table}(session_id); @@ -1018,37 +1005,41 @@ def _enable_foreign_keys(self, connection: Any) -> None: """ connection.execute("PRAGMA foreign_keys = ON") - def _search_entries_fts(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": + def _search_entries_fts( + self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "list[StoredMemory]": + where_scope, scope_params = _build_sqlite_scope_clause("m.", app_name, user_id, scope_filter) sql = f""" - SELECT m.id, m.session_id, m.app_name, m.user_id, m.event_id, m.author, + SELECT m.id, m.session_id, m.app_name, m.user_id, m.scope, m.event_id, m.author, m.timestamp, m.content_json, m.content_text, m.metadata_json, m.inserted_at FROM {self._memory_table} m JOIN {self._memory_table}_fts fts ON m.rowid = fts.rowid - WHERE m.app_name = ? - AND m.user_id = ? + WHERE {where_scope} AND fts.content_text MATCH ? ORDER BY m.timestamp DESC LIMIT ? """ - params: tuple[Any, ...] = (app_name, user_id, query, limit) + params = (*scope_params, query, limit) return self._fetch_records(sql, params) - def _search_entries_simple(self, query: str, app_name: str, user_id: str, limit: int) -> "list[MemoryRecord]": + def _search_entries_simple( + self, query: str, app_name: str, user_id: str, limit: int, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "list[StoredMemory]": + where_scope, scope_params = _build_sqlite_scope_clause("", app_name, user_id, scope_filter) sql = f""" - SELECT id, session_id, app_name, user_id, event_id, author, + SELECT id, session_id, app_name, user_id, scope, event_id, author, timestamp, content_json, content_text, metadata_json, inserted_at FROM {self._memory_table} - WHERE app_name = ? - AND user_id = ? + WHERE {where_scope} AND content_text LIKE ? ORDER BY timestamp DESC LIMIT ? """ pattern = f"%{query}%" - params = (app_name, user_id, pattern, limit) + params = (*scope_params, pattern, limit) return self._fetch_records(sql, params) - def _fetch_records(self, sql: str, params: "tuple[Any, ...]") -> "list[MemoryRecord]": + def _fetch_records(self, sql: str, params: "tuple[Any, ...]") -> "list[StoredMemory]": with self._config.provide_connection() as conn: self._enable_foreign_keys(conn) cursor = conn.execute(sql, params) @@ -1059,13 +1050,15 @@ def _fetch_records(self, sql: str, params: "tuple[Any, ...]") -> "list[MemoryRec "session_id": row[1], "app_name": row[2], "user_id": row[3], - "event_id": row[4], - "author": row[5], - "timestamp": _julian_to_datetime(row[6]), - "content_json": from_json(row[7]) if row[7] else {}, - "content_text": row[8], - "metadata_json": from_json(row[9]) if row[9] else None, - "inserted_at": _julian_to_datetime(row[10]), + "scope": row[4], + "event_id": row[5], + "author": row[6], + "timestamp": _julian_to_datetime(row[7]), + "content_json": from_json(row[8]) if row[8] else {}, + "content_text": row[9], + "metadata_json": from_json(row[10]) if row[10] else None, + "inserted_at": _julian_to_datetime(row[11]), + "embedding": None, } for row in rows ] @@ -1153,3 +1146,16 @@ def _julian_to_datetime(julian: float) -> datetime: days_since_epoch = julian - JULIAN_EPOCH timestamp = days_since_epoch * SECONDS_PER_DAY return datetime.fromtimestamp(timestamp, tz=timezone.utc) + + +def _build_sqlite_scope_clause( + prefix: str, app_name: str, user_id: str, scope_filter: "Literal['all', 'user', 'app']" +) -> "tuple[str, tuple[Any, ...]]": + if scope_filter == "all": + return ( + f"{prefix}app_name = ? AND (({prefix}scope = 'user' AND {prefix}user_id = ?) OR {prefix}scope = 'app')", + (app_name, user_id), + ) + if scope_filter == "user": + return f"{prefix}app_name = ? AND {prefix}scope = 'user' AND {prefix}user_id = ?", (app_name, user_id) + return f"{prefix}app_name = ? AND {prefix}scope = 'app'", (app_name,) diff --git a/sqlspec/extensions/adk/__init__.py b/sqlspec/extensions/adk/__init__.py index 78b4d2a4c..d0df11ded 100644 --- a/sqlspec/extensions/adk/__init__.py +++ b/sqlspec/extensions/adk/__init__.py @@ -15,70 +15,44 @@ - BaseSyncADKMemoryStore: Base class for sync memory store implementations - BaseAsyncADKArtifactStore: Base class for async artifact metadata stores - BaseSyncADKArtifactStore: Base class for sync artifact metadata stores - - SessionRecord: TypedDict for session database records - - EventRecord: TypedDict for event database records - - MemoryRecord: TypedDict for memory database records - - ArtifactRecord: TypedDict for artifact metadata database records - -Example (with extension_config): - from sqlspec.adapters.asyncpg import AsyncpgConfig - from sqlspec.adapters.asyncpg.adk import AsyncpgADKStore - from sqlspec.extensions.adk import SQLSpecSessionService - - config = AsyncpgConfig( - connection_config={"dsn": "postgresql://..."}, - extension_config={ - "adk": { - "session_table": "my_sessions", - "events_table": "my_events", - "owner_id_column": "tenant_id INTEGER REFERENCES tenants(id)" - } - } - ) - - store = AsyncpgADKStore(config) - await store.ensure_tables() - - service = SQLSpecSessionService(store) - session = await service.create_session( - app_name="my_app", - user_id="user123", - state={"key": "value"} - ) + - StoredSession: TypedDict for session database records + - StoredEvent: TypedDict for event database records + - StoredMemory: TypedDict for memory database records + - StoredArtifact: TypedDict for artifact metadata database records """ from sqlspec.config import ADKConfig -from sqlspec.extensions.adk._types import EventRecord, SessionRecord +from sqlspec.extensions.adk._types import StoredEvent, StoredSession from sqlspec.extensions.adk.artifact import ( - ArtifactRecord, BaseAsyncADKArtifactStore, BaseSyncADKArtifactStore, SQLSpecArtifactService, + StoredArtifact, ) from sqlspec.extensions.adk.memory import ( BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore, - MemoryRecord, SQLSpecMemoryService, SQLSpecSyncMemoryService, + StoredMemory, ) from sqlspec.extensions.adk.service import SQLSpecSessionService from sqlspec.extensions.adk.store import BaseAsyncADKStore, BaseSyncADKStore __all__ = ( "ADKConfig", - "ArtifactRecord", "BaseAsyncADKArtifactStore", "BaseAsyncADKMemoryStore", "BaseAsyncADKStore", "BaseSyncADKArtifactStore", "BaseSyncADKMemoryStore", "BaseSyncADKStore", - "EventRecord", - "MemoryRecord", "SQLSpecArtifactService", "SQLSpecMemoryService", "SQLSpecSessionService", "SQLSpecSyncMemoryService", - "SessionRecord", + "StoredArtifact", + "StoredEvent", + "StoredMemory", + "StoredSession", ) diff --git a/sqlspec/extensions/adk/_types.py b/sqlspec/extensions/adk/_types.py index 476995708..62d7e64d1 100644 --- a/sqlspec/extensions/adk/_types.py +++ b/sqlspec/extensions/adk/_types.py @@ -7,10 +7,10 @@ from datetime import datetime from typing import Any, TypedDict -__all__ = ("EventRecord", "SessionRecord") +__all__ = ("StoredEvent", "StoredSession") -class SessionRecord(TypedDict): +class StoredSession(TypedDict): """Database record for a session. Represents the schema for sessions stored in the database. @@ -24,7 +24,7 @@ class SessionRecord(TypedDict): update_time: datetime -class EventRecord(TypedDict): +class StoredEvent(TypedDict): """Database record for an event. Stores the full ADK Event as a single JSON blob (``event_data``) alongside diff --git a/sqlspec/extensions/adk/artifact/__init__.py b/sqlspec/extensions/adk/artifact/__init__.py index 417f9d274..dea479286 100644 --- a/sqlspec/extensions/adk/artifact/__init__.py +++ b/sqlspec/extensions/adk/artifact/__init__.py @@ -8,11 +8,11 @@ - SQLSpecArtifactService: Main service implementing BaseArtifactService - BaseAsyncADKArtifactStore: Base class for async artifact metadata stores - BaseSyncADKArtifactStore: Base class for sync artifact metadata stores - - ArtifactRecord: TypedDict for artifact metadata database records + - StoredArtifact: TypedDict for artifact metadata database records """ -from sqlspec.extensions.adk.artifact._types import ArtifactRecord +from sqlspec.extensions.adk.artifact._types import StoredArtifact from sqlspec.extensions.adk.artifact.service import SQLSpecArtifactService from sqlspec.extensions.adk.artifact.store import BaseAsyncADKArtifactStore, BaseSyncADKArtifactStore -__all__ = ("ArtifactRecord", "BaseAsyncADKArtifactStore", "BaseSyncADKArtifactStore", "SQLSpecArtifactService") +__all__ = ("BaseAsyncADKArtifactStore", "BaseSyncADKArtifactStore", "SQLSpecArtifactService", "StoredArtifact") diff --git a/sqlspec/extensions/adk/artifact/_types.py b/sqlspec/extensions/adk/artifact/_types.py index dcedffcf6..810123975 100644 --- a/sqlspec/extensions/adk/artifact/_types.py +++ b/sqlspec/extensions/adk/artifact/_types.py @@ -7,10 +7,10 @@ from datetime import datetime from typing import Any, TypedDict -__all__ = ("ArtifactRecord",) +__all__ = ("StoredArtifact",) -class ArtifactRecord(TypedDict): +class StoredArtifact(TypedDict): """Database record for an artifact version. Represents the schema for artifact metadata stored in the database. diff --git a/sqlspec/extensions/adk/artifact/service.py b/sqlspec/extensions/adk/artifact/service.py index 439485706..f9c88d26d 100644 --- a/sqlspec/extensions/adk/artifact/service.py +++ b/sqlspec/extensions/adk/artifact/service.py @@ -17,7 +17,7 @@ from google.adk.artifacts.base_artifact_service import BaseArtifactService -from sqlspec.extensions.adk.artifact._types import ArtifactRecord +from sqlspec.extensions.adk.artifact._types import StoredArtifact from sqlspec.storage.registry import StorageRegistry, storage_registry from sqlspec.utils.logging import get_logger, log_with_context @@ -154,7 +154,7 @@ def _deserialize_artifact(data: bytes) -> "types.Part": return types.Part.model_validate(parsed) -def _record_to_artifact_version(record: "ArtifactRecord") -> "ArtifactVersion": +def _record_to_artifact_version(record: "StoredArtifact") -> "ArtifactVersion": """Convert a database artifact record to an ADK ArtifactVersion. Args: @@ -266,7 +266,7 @@ async def save_artifact( # Insert metadata row from datetime import datetime, timezone - record = ArtifactRecord( + record = StoredArtifact( app_name=app_name, user_id=user_id, session_id=session_id, diff --git a/sqlspec/extensions/adk/artifact/store.py b/sqlspec/extensions/adk/artifact/store.py index fcac2a28c..056821bcf 100644 --- a/sqlspec/extensions/adk/artifact/store.py +++ b/sqlspec/extensions/adk/artifact/store.py @@ -19,7 +19,7 @@ if TYPE_CHECKING: from sqlspec.config import DatabaseConfigProtocol - from sqlspec.extensions.adk.artifact._types import ArtifactRecord + from sqlspec.extensions.adk.artifact._types import StoredArtifact __all__ = ("BaseAsyncADKArtifactStore", "BaseSyncADKArtifactStore") @@ -71,7 +71,7 @@ class BaseAsyncADKArtifactStore(_ADKArtifactStoreCommon[ConfigT], ABC): __slots__ = () @abstractmethod - async def insert_artifact(self, record: "ArtifactRecord") -> None: + async def insert_artifact(self, record: "StoredArtifact") -> None: """Insert an artifact version metadata row. Args: @@ -81,7 +81,7 @@ async def insert_artifact(self, record: "ArtifactRecord") -> None: @abstractmethod async def get_artifact( self, app_name: str, user_id: str, filename: str, session_id: "str | None" = None, version: "int | None" = None - ) -> "ArtifactRecord | None": + ) -> "StoredArtifact | None": """Get a specific artifact version's metadata. When ``version`` is None, returns the latest version. @@ -117,7 +117,7 @@ async def list_artifact_keys(self, app_name: str, user_id: str, session_id: "str @abstractmethod async def list_artifact_versions( self, app_name: str, user_id: str, filename: str, session_id: "str | None" = None - ) -> "list[ArtifactRecord]": + ) -> "list[StoredArtifact]": """List all version records for an artifact, ordered by version ascending. Args: @@ -133,7 +133,7 @@ async def list_artifact_versions( @abstractmethod async def delete_artifact( self, app_name: str, user_id: str, filename: str, session_id: "str | None" = None - ) -> "list[ArtifactRecord]": + ) -> "list[StoredArtifact]": """Delete all version records for an artifact and return them. The caller uses the returned records to clean up content from @@ -197,7 +197,7 @@ class BaseSyncADKArtifactStore(_ADKArtifactStoreCommon[ConfigT], ABC): __slots__ = () @abstractmethod - def insert_artifact(self, record: "ArtifactRecord") -> None: + def insert_artifact(self, record: "StoredArtifact") -> None: """Insert an artifact version metadata row. Args: @@ -207,7 +207,7 @@ def insert_artifact(self, record: "ArtifactRecord") -> None: @abstractmethod def get_artifact( self, app_name: str, user_id: str, filename: str, session_id: "str | None" = None, version: "int | None" = None - ) -> "ArtifactRecord | None": + ) -> "StoredArtifact | None": """Get a specific artifact version's metadata. When ``version`` is None, returns the latest version. @@ -239,7 +239,7 @@ def list_artifact_keys(self, app_name: str, user_id: str, session_id: "str | Non @abstractmethod def list_artifact_versions( self, app_name: str, user_id: str, filename: str, session_id: "str | None" = None - ) -> "list[ArtifactRecord]": + ) -> "list[StoredArtifact]": """List all version records for an artifact, ordered by version ascending. Args: @@ -255,7 +255,7 @@ def list_artifact_versions( @abstractmethod def delete_artifact( self, app_name: str, user_id: str, filename: str, session_id: "str | None" = None - ) -> "list[ArtifactRecord]": + ) -> "list[StoredArtifact]": """Delete all version records for an artifact and return them. Args: diff --git a/sqlspec/extensions/adk/converters.py b/sqlspec/extensions/adk/converters.py index f7b34303e..daa87cfbd 100644 --- a/sqlspec/extensions/adk/converters.py +++ b/sqlspec/extensions/adk/converters.py @@ -16,7 +16,7 @@ from google.adk.events.event import Event from google.adk.sessions import Session -from sqlspec.extensions.adk._types import EventRecord, SessionRecord +from sqlspec.extensions.adk._types import StoredEvent, StoredSession __all__ = ( "compute_update_marker", @@ -35,16 +35,16 @@ # --------------------------------------------------------------------------- -def session_to_record(session: "Session") -> SessionRecord: +def session_to_record(session: "Session") -> StoredSession: """Convert ADK Session to database record. Args: session: ADK Session object. Returns: - SessionRecord for database storage. + StoredSession for database storage. """ - return SessionRecord( + return StoredSession( id=session.id, app_name=session.app_name, user_id=session.user_id, @@ -74,7 +74,7 @@ def compute_update_marker(update_time: "datetime") -> str: return update_time.isoformat(timespec="microseconds") -def record_to_session(record: SessionRecord, events: "list[EventRecord]") -> "Session": +def record_to_session(record: StoredSession, events: "list[StoredEvent]") -> "Session": """Convert database record to ADK Session. Sets ``_storage_update_marker`` so the service layer can detect @@ -106,7 +106,7 @@ def record_to_session(record: SessionRecord, events: "list[EventRecord]") -> "Se # --------------------------------------------------------------------------- -def event_to_record(event: "Event", app_name: str, user_id: str, session_id: str) -> EventRecord: +def event_to_record(event: "Event", app_name: str, user_id: str, session_id: str) -> StoredEvent: """Convert ADK Event to database record using full-event JSON storage. The entire Event is serialized into ``event_data`` via Pydantic's @@ -120,10 +120,10 @@ def event_to_record(event: "Event", app_name: str, user_id: str, session_id: str session_id: ID of the parent session. Returns: - EventRecord for database storage. + StoredEvent for database storage. """ event_data = _normalize_event_data(event.model_dump(exclude_none=True, mode="json")) - return EventRecord( + return StoredEvent( id=event.id, app_name=app_name, user_id=user_id, @@ -134,7 +134,7 @@ def event_to_record(event: "Event", app_name: str, user_id: str, session_id: str ) -def record_to_event(record: "EventRecord") -> "Event": +def record_to_event(record: "StoredEvent") -> "Event": """Convert database record to ADK Event. Reconstruction is lossless for valid ADK payloads: the full Event is diff --git a/sqlspec/extensions/adk/memory/__init__.py b/sqlspec/extensions/adk/memory/__init__.py index 28ff3880d..c99655f93 100644 --- a/sqlspec/extensions/adk/memory/__init__.py +++ b/sqlspec/extensions/adk/memory/__init__.py @@ -9,46 +9,13 @@ - SQLSpecSyncMemoryService: Sync service for sync adapters - BaseAsyncADKMemoryStore: Base class for async database store implementations - BaseSyncADKMemoryStore: Base class for sync database store implementations - - MemoryRecord: TypedDict for memory database records + - StoredMemory: TypedDict for memory database records - extract_content_text: Helper to extract searchable text from Content - session_to_memory_records: Convert Session to memory records - record_to_memory_entry: Convert database record to MemoryEntry - -Example (async): - from sqlspec.adapters.asyncpg import AsyncpgConfig - from sqlspec.adapters.asyncpg.adk import AsyncpgADKMemoryStore - from sqlspec.extensions.adk.memory import SQLSpecMemoryService - - config = AsyncpgConfig( - connection_config={"dsn": "postgresql://..."}, - extension_config={ - "adk": { - "memory_table": "adk_memory_entries", - "memory_use_fts": True, - "memory_max_results": 50, - } - } - ) - - store = AsyncpgADKMemoryStore(config) - await store.ensure_tables() - - service = SQLSpecMemoryService(store) - - # Store completed session as memories - await service.add_session_to_memory(completed_session) - - # Search memories - response = await service.search_memory( - app_name="my_app", - user_id="user123", - query="previous discussion about Python" - ) - for entry in response.memories: - print(entry.content) """ -from sqlspec.extensions.adk.memory._types import MemoryRecord +from sqlspec.extensions.adk.memory._types import StoredMemory from sqlspec.extensions.adk.memory.converters import ( extract_content_text, record_to_memory_entry, @@ -60,9 +27,9 @@ __all__ = ( "BaseAsyncADKMemoryStore", "BaseSyncADKMemoryStore", - "MemoryRecord", "SQLSpecMemoryService", "SQLSpecSyncMemoryService", + "StoredMemory", "extract_content_text", "record_to_memory_entry", "session_to_memory_records", diff --git a/sqlspec/extensions/adk/memory/_types.py b/sqlspec/extensions/adk/memory/_types.py index 7f55ccb2f..b01dc4ca9 100644 --- a/sqlspec/extensions/adk/memory/_types.py +++ b/sqlspec/extensions/adk/memory/_types.py @@ -5,12 +5,15 @@ """ from datetime import datetime -from typing import Any, TypedDict +from typing import TYPE_CHECKING, Any, TypedDict -__all__ = ("MemoryRecord",) +if TYPE_CHECKING: + from collections.abc import Sequence +__all__ = ("StoredMemory",) -class MemoryRecord(TypedDict): + +class StoredMemory(TypedDict): """Database record for a memory entry. Represents the schema for memory entries stored in the database. @@ -21,6 +24,7 @@ class MemoryRecord(TypedDict): session_id: str app_name: str user_id: str + scope: str event_id: str author: "str | None" timestamp: datetime @@ -28,3 +32,4 @@ class MemoryRecord(TypedDict): content_text: str metadata_json: "dict[str, Any] | None" inserted_at: datetime + embedding: "Sequence[float] | None" diff --git a/sqlspec/extensions/adk/memory/converters.py b/sqlspec/extensions/adk/memory/converters.py index c488e1257..541894e5b 100644 --- a/sqlspec/extensions/adk/memory/converters.py +++ b/sqlspec/extensions/adk/memory/converters.py @@ -7,7 +7,7 @@ from datetime import datetime, timezone from typing import TYPE_CHECKING, Any -from sqlspec.extensions.adk.memory._types import MemoryRecord +from sqlspec.extensions.adk.memory._types import StoredMemory from sqlspec.utils.logging import get_logger from sqlspec.utils.uuids import uuid4 @@ -58,17 +58,20 @@ def extract_content_text(content: "types.Content") -> str: return " ".join(parts_text) -def event_to_memory_record(event: "Event", session_id: str, app_name: str, user_id: str) -> "MemoryRecord | None": - """Convert an ADK Event to a memory record. +def event_to_memory_record( + event: "Event", session_id: str, app_name: str, user_id: str, scope: str = "user" +) -> "StoredMemory | None": + """Convert an ADK Event to a stored memory record. Args: event: ADK Event object. session_id: ID of the parent session. app_name: Name of the application. user_id: ID of the user. + scope: Visibility scope ('user' or 'app'). Returns: - MemoryRecord for database storage, or None if event has no content. + StoredMemory for database storage, or None if event has no content. """ if event.content is None: return None @@ -79,15 +82,16 @@ def event_to_memory_record(event: "Event", session_id: str, app_name: str, user_ content_dict = event.content.model_dump(exclude_none=True, mode="json") - custom_metadata = event.custom_metadata or None + custom_metadata = dict(event.custom_metadata) if event.custom_metadata else None now = datetime.now(timezone.utc) - return MemoryRecord( + return StoredMemory( id=str(uuid4()), session_id=session_id, app_name=app_name, user_id=user_id, + scope=scope, event_id=event.id, author=event.author, timestamp=datetime.fromtimestamp(event.timestamp, tz=timezone.utc), @@ -95,12 +99,17 @@ def event_to_memory_record(event: "Event", session_id: str, app_name: str, user_ content_text=content_text, metadata_json=custom_metadata, inserted_at=now, + embedding=None, ) def memory_entry_to_record( - entry: "MemoryEntry", app_name: str, user_id: str, extra_metadata: "dict[str, Any] | None" = None -) -> "MemoryRecord | None": + entry: "MemoryEntry", + app_name: str, + user_id: str, + extra_metadata: "dict[str, Any] | None" = None, + scope: str = "user", +) -> "StoredMemory | None": """Convert an ADK MemoryEntry to a database record. Serializes the entry's ``content`` to ``content_json``, extracts text @@ -113,9 +122,10 @@ def memory_entry_to_record( user_id: ID of the user. extra_metadata: Optional call-level metadata to merge with the entry's own ``custom_metadata``. + scope: Visibility scope ('user' or 'app'). Returns: - MemoryRecord for database storage, or None if entry has no + StoredMemory for database storage, or None if entry has no indexable content. """ content_text = extract_content_text(entry.content) @@ -143,11 +153,12 @@ def memory_entry_to_record( except (ValueError, TypeError): timestamp = now - return MemoryRecord( + return StoredMemory( id=entry.id or str(uuid4()), session_id="", app_name=app_name, user_id=user_id, + scope=scope, event_id="", author=entry.author or "", timestamp=timestamp, @@ -155,29 +166,31 @@ def memory_entry_to_record( content_text=content_text, metadata_json=merged_metadata, inserted_at=now, + embedding=None, ) -def session_to_memory_records(session: "Session") -> list["MemoryRecord"]: - """Convert a completed ADK Session to a list of memory records. +def session_to_memory_records(session: "Session", scope: str = "user") -> list["StoredMemory"]: + """Convert a completed ADK Session to a list of stored memory records. Extracts all events with content from the session and converts them to memory records for storage. Args: session: ADK Session object with events. + scope: Visibility scope ('user' or 'app'). Returns: - List of MemoryRecord objects for database storage. + List of StoredMemory objects for database storage. """ - records: list[MemoryRecord] = [] + records: list[StoredMemory] = [] if not session.events: return records for event in session.events: record = event_to_memory_record( - event=event, session_id=session.id, app_name=session.app_name, user_id=session.user_id + event=event, session_id=session.id, app_name=session.app_name, user_id=session.user_id, scope=scope ) if record is not None: records.append(record) @@ -185,11 +198,11 @@ def session_to_memory_records(session: "Session") -> list["MemoryRecord"]: return records -def record_to_memory_entry(record: "MemoryRecord") -> "MemoryEntry": +def record_to_memory_entry(record: "StoredMemory") -> "MemoryEntry": """Convert a database record to an ADK MemoryEntry. - Preserves ``id`` and ``custom_metadata`` fields that were previously - dropped on readback. + Preserves ``id``, ``custom_metadata``, and propagates ``scope`` + into the entry's custom metadata. Args: record: Memory database record. @@ -204,16 +217,20 @@ def record_to_memory_entry(record: "MemoryRecord") -> "MemoryEntry": timestamp_str = record["timestamp"].isoformat() if record["timestamp"] else None + custom_metadata = dict(record["metadata_json"] or {}) + if "scope" not in custom_metadata and record.get("scope"): + custom_metadata["scope"] = record["scope"] + return MemoryEntry( id=record["id"], content=content, author=record["author"], timestamp=timestamp_str, - custom_metadata=record["metadata_json"] or {}, + custom_metadata=custom_metadata, ) -def records_to_memory_entries(records: list["MemoryRecord"]) -> list["Any"]: +def records_to_memory_entries(records: list["StoredMemory"]) -> list["Any"]: """Convert a list of database records to ADK MemoryEntry objects. Args: diff --git a/sqlspec/extensions/adk/memory/service.py b/sqlspec/extensions/adk/memory/service.py index cf0972d3e..c07f7d8be 100644 --- a/sqlspec/extensions/adk/memory/service.py +++ b/sqlspec/extensions/adk/memory/service.py @@ -1,7 +1,7 @@ """SQLSpec-backed memory service for Google ADK.""" import inspect -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Literal, cast from google.adk.memory.base_memory_service import BaseMemoryService, SearchMemoryResponse @@ -54,7 +54,7 @@ def store(self) -> "BaseAsyncADKMemoryStore | BaseSyncADKMemoryStore": """Return the database store.""" return self._store - async def add_session_to_memory(self, session: "Session") -> None: + 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 @@ -66,8 +66,9 @@ async def add_session_to_memory(self, session: "Session") -> None: Args: session: Completed ADK Session with events. + scope: Visibility scope ('user' or 'app'). """ - records = session_to_memory_records(session) + records = session_to_memory_records(session, scope=scope) if not records: logger.debug( @@ -88,10 +89,11 @@ async def add_events_to_memory( events: "Sequence[Event]", session_id: "str | None" = None, custom_metadata: "Mapping[str, object] | None" = None, + scope: str = "user", ) -> None: """Add an explicit list of events to the memory service. - Same Event-to-MemoryRecord extraction logic as + Same Event-to-StoredMemory extraction logic as ``add_session_to_memory``, but operates on a sequence of Events directly (no Session wrapper needed). @@ -102,7 +104,8 @@ async def add_events_to_memory( session_id: Optional session ID for memory scope/partitioning. If None, memory entries are user-scoped only. custom_metadata: Optional portable metadata stored in - ``MemoryRecord.metadata_json``. + ``StoredMemory.metadata_json``. + scope: Visibility scope ('user' or 'app'). """ from sqlspec.extensions.adk.memory.converters import event_to_memory_record @@ -110,7 +113,7 @@ async def add_events_to_memory( records = [] for event in events: record = event_to_memory_record( - event=event, session_id=session_id or "", app_name=app_name, user_id=user_id + event=event, session_id=session_id or "", app_name=app_name, user_id=user_id, scope=scope ) if record is not None: if metadata_dict: @@ -135,6 +138,7 @@ async def add_memory( user_id: str, memories: "Sequence[MemoryEntry]", custom_metadata: "Mapping[str, object] | None" = None, + scope: str = "user", ) -> None: """Add explicit memory items directly to the memory service. @@ -149,12 +153,13 @@ async def add_memory( memories: Explicit memory items to add. custom_metadata: Optional portable metadata for memory writes. Merged with each entry's ``custom_metadata``. + scope: Visibility scope ('user' or 'app'). """ call_metadata = dict(custom_metadata) if custom_metadata else {} records = [] for entry in memories: record = memory_entry_to_record( - entry=entry, app_name=app_name, user_id=user_id, extra_metadata=call_metadata + entry=entry, app_name=app_name, user_id=user_id, extra_metadata=call_metadata, scope=scope ) if record is not None: records.append(record) @@ -172,7 +177,9 @@ async def add_memory( user_id, ) - async def search_memory(self, *, app_name: str, user_id: str, query: str) -> "SearchMemoryResponse": + async def search_memory( + self, *, app_name: str, user_id: str, query: str, scope_filter: Literal["all", "user", "app"] = "all" + ) -> "SearchMemoryResponse": """Search memory entries by text query. Uses the store's configured search strategy (simple ILIKE or FTS). @@ -181,11 +188,14 @@ async def search_memory(self, *, app_name: str, user_id: str, query: str) -> "Se app_name: Name of the application. user_id: ID of the user. query: Text query to search for. + scope_filter: Scope filter ('all', 'user', 'app'). Defaults to 'all'. Returns: SearchMemoryResponse with memories: List[MemoryEntry]. """ - records = await self._call_store("search_entries", query=query, app_name=app_name, user_id=user_id) + records = await self._call_store( + "search_entries", query=query, app_name=app_name, user_id=user_id, scope_filter=scope_filter + ) memories = records_to_memory_entries(records) @@ -230,7 +240,7 @@ def store(self) -> "BaseSyncADKMemoryStore": """Return the database store.""" return self._store - def add_session_to_memory(self, session: "Session") -> None: + 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 @@ -238,8 +248,9 @@ def add_session_to_memory(self, session: "Session") -> None: Args: session: Completed ADK Session with events. + scope: Visibility scope ('user' or 'app'). """ - records = session_to_memory_records(session) + records = session_to_memory_records(session, scope=scope) if not records: logger.debug( @@ -252,18 +263,21 @@ def add_session_to_memory(self, session: "Session") -> None: "Stored %d memory entries for session %s (total events: %d)", inserted_count, session.id, len(records) ) - def search_memory(self, *, app_name: str, user_id: str, query: str) -> list["MemoryEntry"]: + def search_memory( + self, *, app_name: str, user_id: str, query: str, scope_filter: Literal["all", "user", "app"] = "all" + ) -> list["MemoryEntry"]: """Search memory entries by text query. Args: app_name: Name of the application. user_id: ID of the user. query: Text query to search for. + scope_filter: Scope filter ('all', 'user', 'app'). Defaults to 'all'. Returns: List of MemoryEntry objects. """ - records = self._store.search_entries(query=query, app_name=app_name, user_id=user_id) + records = self._store.search_entries(query=query, app_name=app_name, user_id=user_id, scope_filter=scope_filter) memories = records_to_memory_entries(records) diff --git a/sqlspec/extensions/adk/memory/store.py b/sqlspec/extensions/adk/memory/store.py index 044289226..e92849d1e 100644 --- a/sqlspec/extensions/adk/memory/store.py +++ b/sqlspec/extensions/adk/memory/store.py @@ -2,19 +2,21 @@ import logging from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Final, Generic, TypeVar, cast +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 ensure_table_name, owner_id_column_name, reset_drop_sql -from sqlspec.migrations.schema import SchemaTarget, ensure_schema_async, ensure_schema_sync +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 from sqlspec.config import DatabaseConfigProtocol - from sqlspec.extensions.adk.memory._types import MemoryRecord + from sqlspec.extensions.adk.memory._types import StoredMemory __all__ = ("BaseAsyncADKMemoryStore", "BaseSyncADKMemoryStore") @@ -58,7 +60,9 @@ def __init__(self, config: ConfigT) -> None: self._owner_id_column_name: str | None = ( owner_id_column_name(self._owner_id_column_ddl) if self._owner_id_column_ddl else None ) - ensure_table_name(self._memory_table) + + def _store_config_from_extension(self) -> _ADKMemoryStoreConfig: + return _adk_memory_store_config(self._config) @property def config(self) -> ConfigT: @@ -66,14 +70,14 @@ def config(self) -> ConfigT: return self._config @property - def memory_table(self) -> str: - """Return the memory table name.""" - return self._memory_table + def is_enabled(self) -> bool: + """Return whether memory storage is enabled.""" + return self._enabled @property - def enabled(self) -> bool: - """Return whether memory store is enabled.""" - return self._enabled + def memory_table(self) -> str: + """Return the configured memory table name.""" + return self._memory_table @property def use_fts(self) -> bool: @@ -82,53 +86,33 @@ def use_fts(self) -> bool: @property def max_results(self) -> int: - """Return the max search results limit.""" + """Return the default maximum results for search.""" return self._max_results @property - def owner_id_column_ddl(self) -> "str | None": - """Return the full owner ID column DDL (or None if not configured).""" + 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 owner ID column name only (or None if not configured).""" + 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 _store_config_from_extension(self) -> "_ADKMemoryStoreConfig": - """Extract ADK memory configuration from config.extension_config. - - Returns: - Dict with memory_table, use_fts, max_results, and optionally owner_id_column. - """ - return _adk_memory_store_config(self._config) - - def _schema_management_flags(self) -> "tuple[bool, bool]": - """Return automatic-management and missing-table creation flags.""" - extension_config = cast("dict[str, Any]", self._config.extension_config) - settings = cast("dict[str, Any]", extension_config.get("adk", {})) - return bool(settings.get("manage_schema", True)), bool(settings.get("create_schema", True)) - - def _schema_target(self, ddl: "str | list[str]") -> SchemaTarget: - """Build a target from the canonical memory table DDL.""" - statement_config = getattr(self._config, "statement_config", None) - dialect = getattr(statement_config, "dialect", None) - script = ddl if isinstance(ddl, str) else ";\n".join(ddl) - return SchemaTarget.from_ddl(self._memory_table, script, dialect=dialect) - - def _reset_drop_memory_table_sql(self) -> "list[str]": - """Return memory drops needed before recreating the clean-break schema.""" - return reset_drop_sql( - list(self._drop_memory_table_sql()), ADK_RESET_MEMORY_TABLES, self._drop_memory_sql_for_table - ) - - def _drop_memory_sql_for_table(self, table_name: str) -> "list[str]": + def _drop_sql_for_table(self, table_name: str) -> list[str]: current_table = self._memory_table self._memory_table = table_name try: @@ -136,23 +120,30 @@ def _drop_memory_sql_for_table(self, table_name: str) -> "list[str]": finally: self._memory_table = current_table - def _log_memory_table_created(self) -> None: - log_with_context( - logger, - logging.DEBUG, - "adk.memory.table.ready", - db_system=resolve_db_system(type(self).__name__), - memory_table=self._memory_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 _log_memory_table_skipped(self) -> None: + 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, - "adk.memory.table.skipped", + event, + table_name=self._memory_table, db_system=resolve_db_system(type(self).__name__), - memory_table=self._memory_table, - reason="disabled", + **kwargs, ) @@ -161,19 +152,6 @@ class BaseAsyncADKMemoryStore(_ADKMemoryStoreCommon[ConfigT], ABC): Implements storage operations for Google ADK memory entries using SQLSpec database adapters with async/await. - - This abstract base class provides common functionality for all database-specific - memory store implementations including: - - Connection management via SQLSpec configs - - Table name validation - - Memory entry CRUD operations - - Text search with optional full-text search support - - Subclasses must implement dialect-specific SQL queries and will be created - in each adapter directory. - - Args: - config: SQLSpec database configuration with extension_config["adk"] settings. """ __slots__ = () @@ -186,73 +164,97 @@ async def create_tables(self) -> None: """ raise NotImplementedError - async def prepare_schema_async(self, driver: Any) -> None: - """Prepare adapter-specific schema decisions with an asynchronous driver.""" + async def drop_tables(self) -> None: + """Drop the memory table and indexes if they exist. - async def ensure_tables(self) -> None: - """Create tables when enabled and emit a standardized log entry.""" + Should drop all dialect-specific objects (tables, indexes, FTS virtual tables, triggers). + """ + statements = self._drop_memory_table_sql() + session_context = self._config.provide_session() + async with cast("Any", session_context) as driver: + for statement in statements: + await driver.execute(statement) + async def ensure_tables(self) -> None: + """Create tables and emit a standardized log entry.""" if not self._enabled: - self._log_memory_table_skipped() + log_with_context( + logger, + logging.DEBUG, + "adk.memory.table.skipped", + memory_table=self._memory_table, + reason="disabled", + db_system=resolve_db_system(type(self).__name__), + ) return + + manage_schema, _create_schema = self._schema_management_flags() if self.create_schema_enabled: await self.create_tables() - await self.reconcile_schema(assume_existing=self.create_schema_enabled) - self._log_memory_table_created() + if manage_schema: + await self.reconcile_schema(assume_existing=self.create_schema_enabled) + log_with_context( + logger, + logging.DEBUG, + "adk.memory.table.ready", + memory_table=self._memory_table, + db_system=resolve_db_system(type(self).__name__), + ) async def reconcile_schema(self, *, assume_existing: bool = False) -> None: - """Apply additive memory-table changes from canonical adapter DDL. - - Args: - assume_existing: Skip table discovery after adapter-level creation. - """ + """Apply additive ADK memory table changes from canonical adapter DDL.""" manage_schema, create_schema = self._schema_management_flags() - if not manage_schema: + if not manage_schema or not self._enabled: return - target = self._schema_target(await self._memory_table_ddl()) + statement_config = getattr(self._config, "statement_config", None) + dialect = getattr(statement_config, "dialect", None) + ddl = await self._memory_table_ddl() + ddl_str = ddl if isinstance(ddl, str) else ";\n".join(ddl) + target = SchemaTarget.from_ddl(self._memory_table, ddl_str, dialect=dialect) session_context = self._config.provide_session() - async with cast("Any", session_context) as driver: - await ensure_schema_async( - driver, [target], manage_schema=True, create_schema=create_schema, assume_existing=assume_existing - ) + if hasattr(session_context, "__aenter__"): + async with cast("Any", session_context) as driver: + await ensure_schema_async( + driver, [target], manage_schema=True, create_schema=create_schema, assume_existing=assume_existing + ) + return + await async_(_reconcile_adk_schema_sync)( + self._config, [target], create_schema=create_schema, assume_existing=assume_existing + ) @abstractmethod - async def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: + async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: """Bulk insert memory entries with deduplication. - Uses UPSERT pattern to skip duplicates based on event_id. - Args: - entries: List of memory records to insert. + entries: List of stored memory records to insert. owner_id: Optional owner ID value for owner_id_column (if configured). Returns: Number of entries actually inserted (excludes duplicates). - - Raises: - RuntimeError: If memory store is disabled. """ raise NotImplementedError @abstractmethod async def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": """Search memory entries by text query. - Uses the configured search strategy (simple ILIKE or FTS). - Args: query: Text query to search for. app_name: Application name to filter by. user_id: User ID to filter by. limit: Maximum number of results (defaults to max_results config). + scope_filter: Scope filter ('all', 'user', 'app'). Returns: List of matching memory records ordered by relevance/timestamp. - - Raises: - RuntimeError: If memory store is disabled. """ raise NotImplementedError @@ -269,13 +271,15 @@ async def delete_entries_by_session(self, session_id: str) -> int: raise NotImplementedError @abstractmethod - async def delete_entries_older_than(self, days: int) -> int: + async def delete_entries_older_than( + self, days: int, app_name: "str | None" = None, scope: "str | None" = None + ) -> int: """Delete memory entries older than specified days. - Used for TTL cleanup operations. - Args: days: Number of days to retain entries. + app_name: Optional application name to scope deletion. + scope: Optional scope ('user' or 'app') to scope deletion. Returns: Number of entries deleted. @@ -284,162 +288,103 @@ async def delete_entries_older_than(self, days: int) -> int: @abstractmethod async def _memory_table_ddl(self) -> "str | list[str]": - """Get the CREATE TABLE SQL for the memory table. - - Returns: - SQL statement(s) to create the memory table with indexes. - """ + """Get the CREATE TABLE SQL for the memory table.""" raise NotImplementedError @abstractmethod def _drop_memory_table_sql(self) -> "list[str]": - """Get the DROP TABLE SQL statements for this database dialect. - - Returns: - List of SQL statements to drop the memory table and indexes. - """ + """Get the DROP TABLE SQL statements for this database dialect.""" raise NotImplementedError class BaseSyncADKMemoryStore(_ADKMemoryStoreCommon[ConfigT], ABC): - """Base class for sync SQLSpec-backed ADK memory stores. - - Implements storage operations for Google ADK memory entries using - SQLSpec database adapters with synchronous execution. - - This abstract base class provides common functionality for sync database-specific - memory store implementations including: - - Connection management via SQLSpec configs - - Table name validation - - Memory entry CRUD operations - - Text search with optional full-text search support - - Subclasses must implement dialect-specific SQL queries and will be created - in each adapter directory. - - Args: - config: SQLSpec database configuration with extension_config["adk"] settings. - """ + """Base class for sync SQLSpec-backed ADK memory stores.""" __slots__ = () @abstractmethod 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. - """ + """Create the memory table and indexes if they don't exist.""" raise NotImplementedError - def prepare_schema_sync(self, driver: Any) -> None: - """Prepare adapter-specific schema decisions with a synchronous driver.""" + def drop_tables(self) -> None: + """Drop the memory table and indexes if they exist.""" + statements = self._drop_memory_table_sql() + with cast("Any", self._config.provide_session()) as driver: + for statement in statements: + driver.execute(statement) def ensure_tables(self) -> None: - """Create tables when enabled and emit a standardized log entry.""" - + """Create tables and emit a standardized log entry.""" if not self._enabled: - self._log_memory_table_skipped() + log_with_context( + logger, + logging.DEBUG, + "adk.memory.table.skipped", + memory_table=self._memory_table, + reason="disabled", + db_system=resolve_db_system(type(self).__name__), + ) return + + manage_schema, _create_schema = self._schema_management_flags() if self.create_schema_enabled: self.create_tables() - self.reconcile_schema(assume_existing=self.create_schema_enabled) - self._log_memory_table_created() + if manage_schema: + self.reconcile_schema(assume_existing=self.create_schema_enabled) + log_with_context( + logger, + logging.DEBUG, + "adk.memory.table.ready", + memory_table=self._memory_table, + db_system=resolve_db_system(type(self).__name__), + ) def reconcile_schema(self, *, assume_existing: bool = False) -> None: - """Apply additive memory-table changes from canonical adapter DDL. - - Args: - assume_existing: Skip table discovery after adapter-level creation. - """ + """Apply additive ADK memory table changes from canonical adapter DDL.""" manage_schema, create_schema = self._schema_management_flags() - if not manage_schema: + if not manage_schema or not self._enabled: return - target = self._schema_target(self._memory_table_ddl()) - with cast("Any", self._config.provide_session()) as driver: - ensure_schema_sync( - driver, [target], manage_schema=True, create_schema=create_schema, assume_existing=assume_existing - ) + statement_config = getattr(self._config, "statement_config", None) + dialect = getattr(statement_config, "dialect", None) + ddl = self._memory_table_ddl() + ddl_str = ddl if isinstance(ddl, str) else ";\n".join(ddl) + target = SchemaTarget.from_ddl(self._memory_table, ddl_str, dialect=dialect) + _reconcile_adk_schema_sync(self._config, [target], create_schema=create_schema, assume_existing=assume_existing) @abstractmethod - def insert_memory_entries(self, entries: "list[MemoryRecord]", owner_id: "object | None" = None) -> int: - """Bulk insert memory entries with deduplication. - - Uses UPSERT pattern to skip duplicates based on event_id. - - Args: - entries: List of memory records to insert. - owner_id: Optional owner ID value for owner_id_column (if configured). - - Returns: - Number of entries actually inserted (excludes duplicates). - - Raises: - RuntimeError: If memory store is disabled. - """ + def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int: + """Bulk insert memory entries with deduplication.""" raise NotImplementedError @abstractmethod def search_entries( - self, query: str, app_name: str, user_id: str, limit: "int | None" = None - ) -> "list[MemoryRecord]": - """Search memory entries by text query. - - Uses the configured search strategy (simple ILIKE or FTS). - - Args: - query: Text query to search for. - app_name: Application name to filter by. - user_id: User ID to filter by. - limit: Maximum number of results (defaults to max_results config). - - Returns: - List of matching memory records ordered by relevance/timestamp. - - Raises: - RuntimeError: If memory store is disabled. - """ + self, + query: str, + app_name: str, + user_id: str, + limit: "int | None" = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> "list[StoredMemory]": + """Search memory entries by text query.""" raise NotImplementedError @abstractmethod def delete_entries_by_session(self, session_id: str) -> int: - """Delete all memory entries for a specific session. - - Args: - session_id: Session ID to delete entries for. - - Returns: - Number of entries deleted. - """ + """Delete all memory entries for a specific session.""" raise NotImplementedError @abstractmethod - def delete_entries_older_than(self, days: int) -> int: - """Delete memory entries older than specified days. - - Used for TTL cleanup operations. - - Args: - days: Number of days to retain entries. - - Returns: - Number of entries deleted. - """ + def delete_entries_older_than(self, days: int, app_name: "str | None" = None, scope: "str | None" = None) -> int: + """Delete memory entries older than specified days.""" raise NotImplementedError @abstractmethod def _memory_table_ddl(self) -> "str | list[str]": - """Get the CREATE TABLE SQL for the memory table. - - Returns: - SQL statement(s) to create the memory table with indexes. - """ + """Get the CREATE TABLE SQL for the memory table.""" raise NotImplementedError @abstractmethod def _drop_memory_table_sql(self) -> "list[str]": - """Get the DROP TABLE SQL statements for this database dialect. - - Returns: - List of SQL statements to drop the memory table and indexes. - """ + """Get the DROP TABLE SQL statements for this database dialect.""" raise NotImplementedError diff --git a/sqlspec/extensions/adk/migrations/0001_create_adk_tables.py b/sqlspec/extensions/adk/migrations/0001_create_adk_tables.py index 8381a35c5..b279eb998 100644 --- a/sqlspec/extensions/adk/migrations/0001_create_adk_tables.py +++ b/sqlspec/extensions/adk/migrations/0001_create_adk_tables.py @@ -1,22 +1,135 @@ -"""No-op migration: superseded by 0002_reset_adk_tables. +"""Create ADK tables migration. -This file used to create the legacy ADK ``sessions`` / ``events`` tables. The -ADK 2.0 clean break replaces that schema in 0002. 0001 is retained as a no-op -so installs that already applied it keep their tracking-table row; fresh -installs run it as a no-op and proceed to 0002. +Creates the canonical ADK tables: +- adk_session +- adk_event +- adk_app_state +- adk_user_state +- adk_internal_metadata +- adk_memory (when memory extension enabled) """ -from typing import TYPE_CHECKING +import inspect +import logging +from typing import TYPE_CHECKING, NoReturn, cast + +from sqlspec.exceptions import SQLSpecError +from sqlspec.extensions.adk._config_utils import ( + _adk_adapter_store_class, + _adk_memory_migration_enabled, + _adk_memory_migration_store_class, +) +from sqlspec.extensions.adk.store import BaseAsyncADKStore, BaseSyncADKStore +from sqlspec.utils.logging import get_logger, log_with_context if TYPE_CHECKING: + from collections.abc import Awaitable + + from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore from sqlspec.migrations.context import MigrationContext __all__ = ("down", "up") +logger = get_logger("sqlspec.migrations.adk.create") + async def up(context: "MigrationContext | None" = None) -> "list[str]": - return [] + if context is None or context.config is None: + _raise_missing_config() + + store_class = _get_store_class(context) + store_instance = store_class(config=context.config) + await _prepare_schema(store_instance, context) + + statements: list[str] = [ + await _resolve_sql(store_instance._sessions_table_ddl()), # pyright: ignore[reportPrivateUsage] + await _resolve_sql(store_instance._events_table_ddl()), # pyright: ignore[reportPrivateUsage] + await _resolve_sql(store_instance._app_states_table_ddl()), # pyright: ignore[reportPrivateUsage] + await _resolve_sql(store_instance._user_states_table_ddl()), # pyright: ignore[reportPrivateUsage] + await _resolve_sql(store_instance._metadata_table_ddl()), # pyright: ignore[reportPrivateUsage] + ] + + memory_store_class = _get_memory_store_class(context) + if _is_memory_enabled(context) and memory_store_class is not None: + memory_store = memory_store_class(config=context.config) + await _prepare_schema(memory_store, context) + memory_sql = memory_store._memory_table_ddl() # pyright: ignore[reportPrivateUsage] + if inspect.isawaitable(memory_sql): + memory_sql = await memory_sql + if isinstance(memory_sql, list): + statements.extend(memory_sql) + else: + statements.append(memory_sql) + log_with_context( + logger, logging.DEBUG, "adk.migration.create.memory.create", table_name=memory_store.memory_table + ) + + return statements async def down(context: "MigrationContext | None" = None) -> "list[str]": - return [] + if context is None or context.config is None: + _raise_missing_config() + + statements: list[str] = [] + store_class = _get_store_class(context) + store_instance = store_class(config=context.config) + + if _is_memory_enabled(context): + memory_store_class = _get_memory_store_class(context) + if memory_store_class is not None: + memory_store = memory_store_class(config=context.config) + statements.extend(memory_store._reset_drop_memory_table_sql()) # pyright: ignore[reportPrivateUsage] + + statements.extend(store_instance._reset_drop_tables_sql()) # pyright: ignore[reportPrivateUsage] + + return statements + + +def _raise_missing_config() -> NoReturn: + msg = "Migration context must have a config to determine store class" + raise SQLSpecError(msg) + + +def _get_store_class(context: "MigrationContext | None") -> "type[BaseAsyncADKStore | BaseSyncADKStore]": + if not context or not context.config: + _raise_missing_config() + return cast("type[BaseAsyncADKStore | BaseSyncADKStore]", _adk_adapter_store_class(context.config, "ADKStore")) + + +async def _resolve_sql(value: "str | Awaitable[str]") -> str: + if inspect.isawaitable(value): + return await value + return value + + +async def _prepare_schema( + store: "BaseAsyncADKStore | BaseSyncADKStore | BaseAsyncADKMemoryStore | BaseSyncADKMemoryStore", + context: "MigrationContext", +) -> None: + driver = getattr(context, "driver", None) + if driver is None: + return + if isinstance(store, BaseAsyncADKStore): + await store.prepare_schema_async(driver) + return + if isinstance(store, BaseSyncADKStore): + store.prepare_schema_sync(driver) + + +def _get_memory_store_class( + context: "MigrationContext | None", +) -> "type[BaseAsyncADKMemoryStore | BaseSyncADKMemoryStore] | None": + if not context or not context.config: + return None + store_class = _adk_memory_migration_store_class(context.config) + if store_class is None: + log_with_context(logger, logging.DEBUG, "adk.migration.create.memory_store.missing") + return None + return cast("type[BaseAsyncADKMemoryStore | BaseSyncADKMemoryStore]", store_class) + + +def _is_memory_enabled(context: "MigrationContext | None") -> bool: + if not context or not context.config: + return False + return _adk_memory_migration_enabled(context.config) diff --git a/sqlspec/extensions/adk/migrations/0002_reset_adk_tables.py b/sqlspec/extensions/adk/migrations/0002_reset_adk_tables.py index f9906dd2e..ca467d2df 100644 --- a/sqlspec/extensions/adk/migrations/0002_reset_adk_tables.py +++ b/sqlspec/extensions/adk/migrations/0002_reset_adk_tables.py @@ -17,13 +17,13 @@ _adk_memory_migration_enabled, _adk_memory_migration_store_class, ) +from sqlspec.extensions.adk.store import BaseAsyncADKStore, BaseSyncADKStore from sqlspec.utils.logging import get_logger, log_with_context if TYPE_CHECKING: from collections.abc import Awaitable from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore, BaseSyncADKMemoryStore - from sqlspec.extensions.adk.store import BaseAsyncADKStore, BaseSyncADKStore from sqlspec.migrations.context import MigrationContext __all__ = ("down", "up") @@ -117,10 +117,11 @@ async def _prepare_schema( driver = getattr(context, "driver", None) if driver is None: return - if getattr(context, "is_async_driver", False): - await cast("BaseAsyncADKStore | BaseAsyncADKMemoryStore", store).prepare_schema_async(driver) + if isinstance(store, BaseAsyncADKStore): + await store.prepare_schema_async(driver) return - cast("BaseSyncADKStore | BaseSyncADKMemoryStore", store).prepare_schema_sync(driver) + if isinstance(store, BaseSyncADKStore): + store.prepare_schema_sync(driver) def _get_memory_store_class( diff --git a/sqlspec/extensions/adk/store.py b/sqlspec/extensions/adk/store.py index eef896b44..79d3c768e 100644 --- a/sqlspec/extensions/adk/store.py +++ b/sqlspec/extensions/adk/store.py @@ -17,7 +17,7 @@ from collections.abc import Callable from sqlspec.config import DatabaseConfigProtocol - from sqlspec.extensions.adk._types import EventRecord, SessionRecord + from sqlspec.extensions.adk._types import StoredEvent, StoredSession __all__ = ("BaseAsyncADKStore", "BaseSyncADKStore") @@ -259,7 +259,7 @@ async def prepare_schema_async(self, driver: Any) -> None: async def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> "SessionRecord": + ) -> "StoredSession": """Create a new session. Args: @@ -277,7 +277,7 @@ async def create_session( @abstractmethod async def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get a session. Args: @@ -304,7 +304,7 @@ async def update_session_state(self, app_name: str, user_id: str, session_id: st raise NotImplementedError @abstractmethod - async def list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[SessionRecord]": + async def list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[StoredSession]": """List all sessions for an app, optionally filtered by user. Args: @@ -328,18 +328,18 @@ async def delete_session(self, app_name: str, user_id: str, session_id: str) -> raise NotImplementedError @abstractmethod - async def append_event(self, event_record: "EventRecord") -> None: + async def append_event(self, event_record: "StoredEvent") -> None: """Append an event to a session. Args: - event_record: Event record to store. + event_record: Event record to insert. """ raise NotImplementedError @abstractmethod async def append_event_and_update_state( self, - event_record: "EventRecord", + event_record: "StoredEvent", app_name: str, user_id: str, session_id: str, @@ -347,16 +347,10 @@ async def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> "SessionRecord": + ) -> "StoredSession": """Atomically append an event and update the session's durable state. This is the authoritative durable write boundary for post-creation - session mutations. The event insert, session state update, and the - optional scoped-state upserts must succeed together or fail together, - and the updated session record is returned in the same round-trip so - callers don't need a follow-up read. - - When ``app_state`` is provided (non-None), it is a full merged app-scoped snapshot to replace/upsert for ``app_name``. When ``user_state`` is provided, it is a full merged user-scoped snapshot to replace/upsert for ``(app_name, user_id)``. ``None`` means that scope @@ -375,7 +369,7 @@ async def append_event_and_update_state( upsert atomically, or ``None`` when untouched. Returns: - The updated SessionRecord reflecting the new state and update_time. + The updated StoredSession reflecting the new state and update_time. Raises: ValueError: If the session row no longer exists at update time @@ -391,7 +385,7 @@ async def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Get events for a session. Args: @@ -676,14 +670,14 @@ def prepare_schema_sync(self, driver: Any) -> None: @abstractmethod def create_session( self, session_id: str, app_name: str, user_id: str, state: "dict[str, Any]", owner_id: "Any | None" = None - ) -> "SessionRecord": + ) -> "StoredSession": """Create a new session.""" raise NotImplementedError @abstractmethod def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: "int | timedelta | None" = None - ) -> "SessionRecord | None": + ) -> "StoredSession | None": """Get a session.""" raise NotImplementedError @@ -693,7 +687,7 @@ def update_session_state(self, app_name: str, user_id: str, session_id: str, sta raise NotImplementedError @abstractmethod - def list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[SessionRecord]": + def list_sessions(self, app_name: str, user_id: "str | None" = None) -> "list[StoredSession]": """List all sessions for an app, optionally filtered by user.""" raise NotImplementedError @@ -703,14 +697,14 @@ def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: raise NotImplementedError @abstractmethod - def append_event(self, event_record: "EventRecord") -> None: + def append_event(self, event_record: "StoredEvent") -> None: """Append an event to a session.""" raise NotImplementedError @abstractmethod def append_event_and_update_state( self, - event_record: "EventRecord", + event_record: "StoredEvent", app_name: str, user_id: str, session_id: str, @@ -718,7 +712,7 @@ def append_event_and_update_state( *, app_state: "dict[str, Any] | None" = None, user_state: "dict[str, Any] | None" = None, - ) -> "SessionRecord": + ) -> "StoredSession": """Atomically append an event and update the session's durable state.""" raise NotImplementedError @@ -730,7 +724,7 @@ def get_events( session_id: str, after_timestamp: "datetime | None" = None, limit: "int | None" = None, - ) -> "list[EventRecord]": + ) -> "list[StoredEvent]": """Get events for a session.""" raise NotImplementedError diff --git a/tests/integration/adapters/_shared/adk_behaviors.py b/tests/integration/adapters/_shared/adk_behaviors.py index 50da8f8f1..d8e3ddb6c 100644 --- a/tests/integration/adapters/_shared/adk_behaviors.py +++ b/tests/integration/adapters/_shared/adk_behaviors.py @@ -5,7 +5,7 @@ from inspect import isawaitable from typing import Any, TypeVar -from sqlspec.extensions.adk import EventRecord +from sqlspec.extensions.adk import StoredEvent T = TypeVar("T") @@ -20,7 +20,7 @@ async def _aclose(config: Any) -> None: await _resolve(config.close_pool()) -def _event(app_name: str, user_id: str, session_id: str, index: int, when: datetime) -> EventRecord: +def _event(app_name: str, user_id: str, session_id: str, index: int, when: datetime) -> StoredEvent: return { "id": f"event-{session_id}-{index}", "app_name": app_name, diff --git a/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_memory_store.py b/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_memory_store.py index cd97a092c..a21c093ef 100644 --- a/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_memory_store.py +++ b/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_memory_store.py @@ -9,18 +9,21 @@ from sqlspec.adapters.duckdb.adk import DuckdbADKMemoryStore from sqlspec.adapters.duckdb.config import DuckDBConfig -from sqlspec.extensions.adk import MemoryRecord +from sqlspec.extensions.adk import StoredMemory pytestmark = [pytest.mark.duckdb, pytest.mark.integration] -def _build_record(*, session_id: str, event_id: str, content_text: str, inserted_at: datetime) -> MemoryRecord: +def _build_record( + *, session_id: str, event_id: str, content_text: str, inserted_at: datetime, scope: str = "user" +) -> StoredMemory: now = datetime.now(timezone.utc) - return MemoryRecord( + return StoredMemory( id=str(uuid4()), session_id=session_id, app_name="app", user_id="user", + scope=scope, event_id=event_id, author="user", timestamp=now, @@ -28,6 +31,7 @@ def _build_record(*, session_id: str, event_id: str, content_text: str, inserted content_text=content_text, metadata_json=None, inserted_at=inserted_at, + embedding=None, ) @@ -120,3 +124,75 @@ def test_duckdb_memory_store_fts_search_uses_bm25_path(tmp_path: Path) -> None: assert len(results) == 1 assert results[0]["event_id"] == "evt-fts-1" assert results[0]["content_json"] == {"text": "espresso roast"} + + +def test_duckdb_memory_store_scoped_search_combined_default(tmp_path: Path) -> None: + """Default search recall returns both user-scoped and app-scoped memories.""" + store = _build_store(tmp_path) + + now = datetime.now(timezone.utc) + record_user = _build_record( + session_id="s1", event_id="evt-u1", content_text="project architecture guideline", inserted_at=now, scope="user" + ) + record_app = _build_record( + session_id="s2", event_id="evt-a1", content_text="company architecture standard", inserted_at=now, scope="app" + ) + record_other_user = _build_record( + session_id="s3", + event_id="evt-other", + content_text="other user architecture note", + inserted_at=now, + scope="user", + ) + record_other_user["user_id"] = "other_user" + + store.insert_memory_entries([record_user, record_app, record_other_user]) + + results = store.search_entries(query="architecture", app_name="app", user_id="user") + event_ids = {r["event_id"] for r in results} + assert event_ids == {"evt-u1", "evt-a1"} + assert "evt-other" not in event_ids + + +def test_duckdb_memory_store_explicit_scope_filters(tmp_path: Path) -> None: + """Explicit scope filters restrict results to only user or only app memories.""" + store = _build_store(tmp_path) + + now = datetime.now(timezone.utc) + record_user = _build_record( + session_id="s1", event_id="evt-u1", content_text="scoped query plan", inserted_at=now, scope="user" + ) + record_app = _build_record( + session_id="s2", event_id="evt-a1", content_text="scoped release plan", inserted_at=now, scope="app" + ) + store.insert_memory_entries([record_user, record_app]) + + user_only = store.search_entries(query="scoped", app_name="app", user_id="user", scope_filter="user") + assert len(user_only) == 1 + assert user_only[0]["event_id"] == "evt-u1" + + app_only = store.search_entries(query="scoped", app_name="app", user_id="user", scope_filter="app") + assert len(app_only) == 1 + assert app_only[0]["event_id"] == "evt-a1" + + +def test_duckdb_memory_store_scoped_retention(tmp_path: Path) -> None: + """Scoped retention deletes entries matching app_name and scope filters.""" + store = _build_store(tmp_path) + + now = datetime.now(timezone.utc) + old = now - timedelta(days=40) + record_user_old = _build_record( + session_id="s1", event_id="evt-uo", content_text="old user memo", inserted_at=old, scope="user" + ) + record_app_old = _build_record( + session_id="s2", event_id="evt-ao", content_text="old app guideline", inserted_at=old, scope="app" + ) + store.insert_memory_entries([record_user_old, record_app_old]) + + deleted = store.delete_entries_older_than(30, app_name="app", scope="user") + assert deleted == 1 + + remaining = store.search_entries(query="old", app_name="app", user_id="user") + assert len(remaining) == 1 + assert remaining[0]["event_id"] == "evt-ao" diff --git a/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_store.py b/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_store.py index d2d515fb7..bb652dbea 100644 --- a/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_store.py +++ b/tests/integration/adapters/duckdb/duckdb/extensions/adk/test_store.py @@ -15,7 +15,7 @@ from sqlspec.adapters.duckdb.adk import DuckdbADKStore from sqlspec.adapters.duckdb.config import DuckDBConfig -from sqlspec.extensions.adk import EventRecord +from sqlspec.extensions.adk import StoredEvent pytestmark = [pytest.mark.duckdb, pytest.mark.integration] @@ -52,7 +52,7 @@ def test_event_with_optional_fields(duckdb_adk_store: DuckdbADKStore) -> None: session_id = "session-008" duckdb_adk_store.create_session(session_id, "test-app", "user-008", {}) - event_record: EventRecord = { + event_record: StoredEvent = { "id": "event-full", "app_name": "test-app", "user_id": "user-008", @@ -96,7 +96,7 @@ def test_event_ordering_by_timestamp(duckdb_adk_store: DuckdbADKStore) -> None: t2 = datetime.now(timezone.utc) t3 = datetime.now(timezone.utc) - ev_middle: EventRecord = { + ev_middle: StoredEvent = { "id": "event-middle", "app_name": "test-app", "user_id": "user-009", @@ -105,7 +105,7 @@ def test_event_ordering_by_timestamp(duckdb_adk_store: DuckdbADKStore) -> None: "timestamp": t2, "event_data": {"id": "event-middle", "app_name": "test-app", "user_id": "user-009"}, } - ev_last: EventRecord = { + ev_last: StoredEvent = { "id": "event-last", "app_name": "test-app", "user_id": "user-009", @@ -114,7 +114,7 @@ def test_event_ordering_by_timestamp(duckdb_adk_store: DuckdbADKStore) -> None: "timestamp": t3, "event_data": {"id": "event-last", "app_name": "test-app", "user_id": "user-009"}, } - ev_first: EventRecord = { + ev_first: StoredEvent = { "id": "event-first", "app_name": "test-app", "user_id": "user-009", @@ -163,7 +163,7 @@ def test_event_data_round_trip(duckdb_adk_store: DuckdbADKStore) -> None: session_id = "session-json-rt" duckdb_adk_store.create_session(session_id, "test-app", "user-012", {}) - event_record: EventRecord = { + event_record: StoredEvent = { "id": "event-json", "app_name": "test-app", "user_id": "user-012", diff --git a/tests/integration/adapters/mysql/aiomysql/extensions/adk/test_store.py b/tests/integration/adapters/mysql/aiomysql/extensions/adk/test_store.py index 92b7b0095..2a5db435f 100644 --- a/tests/integration/adapters/mysql/aiomysql/extensions/adk/test_store.py +++ b/tests/integration/adapters/mysql/aiomysql/extensions/adk/test_store.py @@ -13,7 +13,7 @@ from sqlspec.adapters.aiomysql._typing import AiomysqlCursor from sqlspec.adapters.aiomysql.adk import AiomysqlADKStore -from sqlspec.extensions.adk import EventRecord +from sqlspec.extensions.adk import StoredEvent pytestmark = [pytest.mark.xdist_group("mysql"), pytest.mark.aiomysql, pytest.mark.integration] @@ -76,7 +76,7 @@ async def test_timestamp_precision(aiomysql_adk_store: AiomysqlADKStore) -> None assert hasattr(created["create_time"], "microsecond") event_time = datetime.now(timezone.utc) - event: EventRecord = { + event: StoredEvent = { "id": "event-micro", "app_name": app_name, "user_id": user_id, diff --git a/tests/integration/adapters/mysql/asyncmy/extensions/adk/test_store.py b/tests/integration/adapters/mysql/asyncmy/extensions/adk/test_store.py index 5acaf1931..4ad39f60c 100644 --- a/tests/integration/adapters/mysql/asyncmy/extensions/adk/test_store.py +++ b/tests/integration/adapters/mysql/asyncmy/extensions/adk/test_store.py @@ -12,7 +12,7 @@ import pytest from sqlspec.adapters.asyncmy.adk import AsyncmyADKStore -from sqlspec.extensions.adk import EventRecord +from sqlspec.extensions.adk import StoredEvent pytestmark = [pytest.mark.xdist_group("mysql"), pytest.mark.asyncmy, pytest.mark.integration] @@ -75,7 +75,7 @@ async def test_timestamp_precision(asyncmy_adk_store: AsyncmyADKStore) -> None: assert hasattr(created["create_time"], "microsecond") event_time = datetime.now(timezone.utc) - event: EventRecord = { + event: StoredEvent = { "id": "event-micro", "app_name": app_name, "user_id": user_id, diff --git a/tests/integration/adapters/mysql/mysqlconnector/extensions/adk/test_store.py b/tests/integration/adapters/mysql/mysqlconnector/extensions/adk/test_store.py index 4fc5dd4c6..e1f88c49b 100644 --- a/tests/integration/adapters/mysql/mysqlconnector/extensions/adk/test_store.py +++ b/tests/integration/adapters/mysql/mysqlconnector/extensions/adk/test_store.py @@ -13,7 +13,7 @@ import pytest from sqlspec.adapters.mysqlconnector.adk import MysqlConnectorAsyncADKStore -from sqlspec.extensions.adk import EventRecord +from sqlspec.extensions.adk import StoredEvent pytestmark = [pytest.mark.xdist_group("mysql"), pytest.mark.mysql_connector, pytest.mark.integration] @@ -75,7 +75,7 @@ async def test_timestamp_precision(mysqlconnector_adk_store: MysqlConnectorAsync assert hasattr(created["create_time"], "microsecond") event_time = datetime.now(timezone.utc) - event: EventRecord = { + event: StoredEvent = { "id": "event-micro", "app_name": app_name, "user_id": user_id, diff --git a/tests/integration/adapters/oracle/oracledb/extensions/adk/test_oracle_specific.py b/tests/integration/adapters/oracle/oracledb/extensions/adk/test_oracle_specific.py index 154cd2924..4c2e1a1ad 100644 --- a/tests/integration/adapters/oracle/oracledb/extensions/adk/test_oracle_specific.py +++ b/tests/integration/adapters/oracle/oracledb/extensions/adk/test_oracle_specific.py @@ -11,7 +11,7 @@ from sqlspec.adapters.oracledb import OracleAsyncConfig, OracleSyncConfig from sqlspec.adapters.oracledb.adk import OracleAsyncADKStore, OracleSyncADKStore -from sqlspec.extensions.adk import EventRecord +from sqlspec.extensions.adk import StoredEvent pytestmark = [pytest.mark.xdist_group("oracle"), pytest.mark.oracledb, pytest.mark.integration] @@ -23,9 +23,9 @@ def _unique_session_id(prefix: str) -> str: def _event_record( *, event_id: str, app_name: str, user_id: str, session_id: str, invocation_id: str, event_data: dict[str, Any] -) -> EventRecord: - """Return a clean-break EventRecord for Oracle store tests.""" - return EventRecord( +) -> StoredEvent: + """Return a clean-break StoredEvent for Oracle store tests.""" + return StoredEvent( id=event_id, app_name=app_name, user_id=user_id, @@ -321,7 +321,7 @@ async def test_state_lob_deserialization_sync(oracle_sync_store: "OracleSyncADKS async def test_event_record_clean_break_contract(oracle_async_store: "OracleAsyncADKStore") -> None: - """Test the clean-break EventRecord contract with append_event.""" + """Test the clean-break StoredEvent contract with append_event.""" session_id = _unique_session_id("event-contract") app_name = "test-app" user_id = "user-123" diff --git a/tests/integration/adapters/spanner/spanner/extensions/adk/test_adk_store.py b/tests/integration/adapters/spanner/spanner/extensions/adk/test_adk_store.py index 3098c077b..b46d9b66f 100644 --- a/tests/integration/adapters/spanner/spanner/extensions/adk/test_adk_store.py +++ b/tests/integration/adapters/spanner/spanner/extensions/adk/test_adk_store.py @@ -11,7 +11,7 @@ import pytest -from sqlspec.extensions.adk import EventRecord +from sqlspec.extensions.adk import StoredEvent pytestmark = [pytest.mark.spanner, pytest.mark.integration] @@ -66,7 +66,7 @@ def test_create_and_list_events(spanner_adk_store: Any) -> None: spanner_adk_store.delete_session("app", "user", session_id) spanner_adk_store.create_session(session_id, "app", "user", {"x": 1}) - event_one: EventRecord = { + event_one: StoredEvent = { "id": "event-1", "app_name": "app", "user_id": "user", @@ -75,7 +75,7 @@ def test_create_and_list_events(spanner_adk_store: Any) -> None: "timestamp": datetime.now(timezone.utc), "event_data": {"id": "event-1", "author": "user", "content": {"msg": "hi"}}, } - event_two: EventRecord = { + event_two: StoredEvent = { "id": "event-2", "app_name": "app", "user_id": "user", diff --git a/tests/integration/adapters/sqlite/adbc/extensions/adk/test_memory_store.py b/tests/integration/adapters/sqlite/adbc/extensions/adk/test_memory_store.py index 9d389bb25..22cc30ebc 100644 --- a/tests/integration/adapters/sqlite/adbc/extensions/adk/test_memory_store.py +++ b/tests/integration/adapters/sqlite/adbc/extensions/adk/test_memory_store.py @@ -8,18 +8,21 @@ from sqlspec.adapters.adbc import AdbcConfig from sqlspec.adapters.adbc.adk import AdbcADKMemoryStore -from sqlspec.extensions.adk import MemoryRecord +from sqlspec.extensions.adk import StoredMemory pytestmark = [pytest.mark.xdist_group("sqlite"), pytest.mark.adbc, pytest.mark.integration] -def _build_record(*, session_id: str, event_id: str, content_text: str, inserted_at: datetime) -> MemoryRecord: +def _build_record( + *, session_id: str, event_id: str, content_text: str, inserted_at: datetime, scope: str = "user" +) -> StoredMemory: now = datetime.now(timezone.utc) - return MemoryRecord( + return StoredMemory( id=str(uuid4()), session_id=session_id, app_name="app", user_id="user", + scope=scope, event_id=event_id, author="user", timestamp=now, @@ -27,6 +30,7 @@ def _build_record(*, session_id: str, event_id: str, content_text: str, inserted content_text=content_text, metadata_json=None, inserted_at=inserted_at, + embedding=None, ) diff --git a/tests/integration/adapters/sqlite/aiosqlite/extensions/adk/test_memory_store.py b/tests/integration/adapters/sqlite/aiosqlite/extensions/adk/test_memory_store.py index 3ce5f04d9..fd6c0d099 100644 --- a/tests/integration/adapters/sqlite/aiosqlite/extensions/adk/test_memory_store.py +++ b/tests/integration/adapters/sqlite/aiosqlite/extensions/adk/test_memory_store.py @@ -8,18 +8,21 @@ from sqlspec.adapters.aiosqlite import AiosqliteConfig from sqlspec.adapters.aiosqlite.adk import AiosqliteADKMemoryStore -from sqlspec.extensions.adk import MemoryRecord +from sqlspec.extensions.adk import StoredMemory pytestmark = pytest.mark.xdist_group("sqlite") -def _build_record(*, session_id: str, event_id: str, content_text: str, inserted_at: datetime) -> MemoryRecord: +def _build_record( + *, session_id: str, event_id: str, content_text: str, inserted_at: datetime, scope: str = "user" +) -> StoredMemory: now = datetime.now(timezone.utc) - return MemoryRecord( + return StoredMemory( id=str(uuid4()), session_id=session_id, app_name="app", user_id="user", + scope=scope, event_id=event_id, author="user", timestamp=now, @@ -27,6 +30,7 @@ def _build_record(*, session_id: str, event_id: str, content_text: str, inserted content_text=content_text, metadata_json=None, inserted_at=inserted_at, + embedding=None, ) @@ -145,4 +149,93 @@ async def test_aiosqlite_memory_store_delete_older_than() -> None: assert len(remaining) == 1 assert remaining[0]["event_id"] == "evt-2" + +async def test_aiosqlite_memory_store_scoped_search_combined_default() -> None: + """Default search recall returns both user-scoped and app-scoped memories.""" + with tempfile.NamedTemporaryFile(suffix=".db") as tmp: + config = AiosqliteConfig(connection_config={"database": tmp.name}) + store = AiosqliteADKMemoryStore(config) + await store.create_tables() + + now = datetime.now(timezone.utc) + record_user = _build_record( + session_id="s1", + event_id="evt-u1", + content_text="project architecture guideline", + inserted_at=now, + scope="user", + ) + record_app = _build_record( + session_id="s2", + event_id="evt-a1", + content_text="company architecture standard", + inserted_at=now, + scope="app", + ) + record_other_user = _build_record( + session_id="s3", + event_id="evt-other", + content_text="other user architecture note", + inserted_at=now, + scope="user", + ) + record_other_user["user_id"] = "other_user" + + await store.insert_memory_entries([record_user, record_app, record_other_user]) + + results = await store.search_entries(query="architecture", app_name="app", user_id="user") + event_ids = {r["event_id"] for r in results} + assert event_ids == {"evt-u1", "evt-a1"} + assert "evt-other" not in event_ids + + +async def test_aiosqlite_memory_store_explicit_scope_filters() -> None: + """Explicit scope filters restrict results to only user or only app memories.""" + with tempfile.NamedTemporaryFile(suffix=".db") as tmp: + config = AiosqliteConfig(connection_config={"database": tmp.name}) + store = AiosqliteADKMemoryStore(config) + await store.create_tables() + + now = datetime.now(timezone.utc) + record_user = _build_record( + session_id="s1", event_id="evt-u1", content_text="scoped query plan", inserted_at=now, scope="user" + ) + record_app = _build_record( + session_id="s2", event_id="evt-a1", content_text="scoped release plan", inserted_at=now, scope="app" + ) + await store.insert_memory_entries([record_user, record_app]) + + user_only = await store.search_entries(query="scoped", app_name="app", user_id="user", scope_filter="user") + assert len(user_only) == 1 + assert user_only[0]["event_id"] == "evt-u1" + + app_only = await store.search_entries(query="scoped", app_name="app", user_id="user", scope_filter="app") + assert len(app_only) == 1 + assert app_only[0]["event_id"] == "evt-a1" + + +async def test_aiosqlite_memory_store_scoped_retention() -> None: + """Scoped retention deletes entries matching app_name and scope filters.""" + with tempfile.NamedTemporaryFile(suffix=".db") as tmp: + config = AiosqliteConfig(connection_config={"database": tmp.name}) + store = AiosqliteADKMemoryStore(config) + await store.create_tables() + + now = datetime.now(timezone.utc) + old = now - timedelta(days=40) + record_user_old = _build_record( + session_id="s1", event_id="evt-uo", content_text="old user memo", inserted_at=old, scope="user" + ) + record_app_old = _build_record( + session_id="s2", event_id="evt-ao", content_text="old app guideline", inserted_at=old, scope="app" + ) + await store.insert_memory_entries([record_user_old, record_app_old]) + + deleted = await store.delete_entries_older_than(30, app_name="app", scope="user") + assert deleted == 1 + + remaining = await store.search_entries(query="old", app_name="app", user_id="user") + assert len(remaining) == 1 + assert remaining[0]["event_id"] == "evt-ao" + await config.close_pool() diff --git a/tests/integration/adapters/sqlite/sqlite/extensions/adk/test_memory_store.py b/tests/integration/adapters/sqlite/sqlite/extensions/adk/test_memory_store.py index cf271c192..fef8dba86 100644 --- a/tests/integration/adapters/sqlite/sqlite/extensions/adk/test_memory_store.py +++ b/tests/integration/adapters/sqlite/sqlite/extensions/adk/test_memory_store.py @@ -8,18 +8,21 @@ from sqlspec.adapters.sqlite import SqliteConfig from sqlspec.adapters.sqlite.adk import SqliteADKMemoryStore -from sqlspec.extensions.adk import MemoryRecord +from sqlspec.extensions.adk import StoredMemory pytestmark = pytest.mark.xdist_group("sqlite") -def _build_record(*, session_id: str, event_id: str, content_text: str, inserted_at: datetime) -> MemoryRecord: +def _build_record( + *, session_id: str, event_id: str, content_text: str, inserted_at: datetime, scope: str = "user" +) -> StoredMemory: now = datetime.now(timezone.utc) - return MemoryRecord( + return StoredMemory( id=str(uuid4()), session_id=session_id, app_name="app", user_id="user", + scope=scope, event_id=event_id, author="user", timestamp=now, @@ -27,6 +30,7 @@ def _build_record(*, session_id: str, event_id: str, content_text: str, inserted content_text=content_text, metadata_json=None, inserted_at=inserted_at, + embedding=None, ) @@ -136,3 +140,92 @@ def test_sqlite_memory_store_delete_older_than() -> None: remaining = store.search_entries(query="new", app_name="app", user_id="user") assert len(remaining) == 1 assert remaining[0]["event_id"] == "evt-2" + + +def test_sqlite_memory_store_scoped_search_combined_default() -> None: + """Default search recall returns both user-scoped and app-scoped memories.""" + with tempfile.NamedTemporaryFile(suffix=".db") as tmp: + config = SqliteConfig(connection_config={"database": tmp.name}) + store = SqliteADKMemoryStore(config) + store.create_tables() + + now = datetime.now(timezone.utc) + record_user = _build_record( + session_id="s1", + event_id="evt-u1", + content_text="project architecture guideline", + inserted_at=now, + scope="user", + ) + record_app = _build_record( + session_id="s2", + event_id="evt-a1", + content_text="company architecture standard", + inserted_at=now, + scope="app", + ) + record_other_user = _build_record( + session_id="s3", + event_id="evt-other", + content_text="other user architecture note", + inserted_at=now, + scope="user", + ) + record_other_user["user_id"] = "other_user" + + store.insert_memory_entries([record_user, record_app, record_other_user]) + + results = store.search_entries(query="architecture", app_name="app", user_id="user") + event_ids = {r["event_id"] for r in results} + assert event_ids == {"evt-u1", "evt-a1"} + assert "evt-other" not in event_ids + + +def test_sqlite_memory_store_explicit_scope_filters() -> None: + """Explicit scope filters restrict results to only user or only app memories.""" + with tempfile.NamedTemporaryFile(suffix=".db") as tmp: + config = SqliteConfig(connection_config={"database": tmp.name}) + store = SqliteADKMemoryStore(config) + store.create_tables() + + now = datetime.now(timezone.utc) + record_user = _build_record( + session_id="s1", event_id="evt-u1", content_text="scoped query plan", inserted_at=now, scope="user" + ) + record_app = _build_record( + session_id="s2", event_id="evt-a1", content_text="scoped release plan", inserted_at=now, scope="app" + ) + store.insert_memory_entries([record_user, record_app]) + + user_only = store.search_entries(query="scoped", app_name="app", user_id="user", scope_filter="user") + assert len(user_only) == 1 + assert user_only[0]["event_id"] == "evt-u1" + + app_only = store.search_entries(query="scoped", app_name="app", user_id="user", scope_filter="app") + assert len(app_only) == 1 + assert app_only[0]["event_id"] == "evt-a1" + + +def test_sqlite_memory_store_scoped_retention() -> None: + """Scoped retention deletes entries matching app_name and scope filters.""" + with tempfile.NamedTemporaryFile(suffix=".db") as tmp: + config = SqliteConfig(connection_config={"database": tmp.name}) + store = SqliteADKMemoryStore(config) + store.create_tables() + + now = datetime.now(timezone.utc) + old = now - timedelta(days=40) + record_user_old = _build_record( + session_id="s1", event_id="evt-uo", content_text="old user memo", inserted_at=old, scope="user" + ) + record_app_old = _build_record( + session_id="s2", event_id="evt-ao", content_text="old app guideline", inserted_at=old, scope="app" + ) + store.insert_memory_entries([record_user_old, record_app_old]) + + deleted = store.delete_entries_older_than(30, app_name="app", scope="user") + assert deleted == 1 + + remaining = store.search_entries(query="old", app_name="app", user_id="user") + assert len(remaining) == 1 + assert remaining[0]["event_id"] == "evt-ao" diff --git a/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py b/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py index 0184e4332..bfec4c7e7 100644 --- a/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py +++ b/tests/unit/adapters/test_oracledb/test_oracle_adk_store.py @@ -247,6 +247,7 @@ async def test_oracle_async_adk_memory_rows_to_records_deserializes_json_fields( "session-1", "app", "user", + "user", "event-1", "assistant", timestamp, @@ -264,6 +265,7 @@ async def test_oracle_async_adk_memory_rows_to_records_deserializes_json_fields( "session_id": "session-1", "app_name": "app", "user_id": "user", + "scope": "user", "event_id": "event-1", "author": "assistant", "timestamp": timestamp, @@ -271,6 +273,7 @@ async def test_oracle_async_adk_memory_rows_to_records_deserializes_json_fields( "content_text": "hello", "metadata_json": {"source": "unit"}, "inserted_at": timestamp, + "embedding": None, } ] @@ -284,6 +287,7 @@ def test_oracle_sync_adk_memory_rows_to_records_deserializes_json_fields() -> No "session-2", "app", "user", + "user", "event-2", "user", timestamp, @@ -298,6 +302,8 @@ def test_oracle_sync_adk_memory_rows_to_records_deserializes_json_fields() -> No assert records[0]["content_json"] == {"text": "sync"} assert records[0]["metadata_json"] == {"source": "unit"} assert records[0]["content_text"] == "sync" + assert records[0]["scope"] == "user" + assert records[0]["embedding"] is None def _sync_store_with_driver() -> "tuple[Any, MagicMock, MagicMock]": diff --git a/tests/unit/adapters/test_spanner/test_adk_store.py b/tests/unit/adapters/test_spanner/test_adk_store.py index bc94caf69..a7ea74454 100644 --- a/tests/unit/adapters/test_spanner/test_adk_store.py +++ b/tests/unit/adapters/test_spanner/test_adk_store.py @@ -16,7 +16,7 @@ SpannerSyncADKStore, ) from sqlspec.config import ADKConfig -from sqlspec.extensions.adk import EventRecord, MemoryRecord +from sqlspec.extensions.adk import StoredEvent, StoredMemory def _mock_config(adk_config: dict[str, object] | None = None) -> MagicMock: @@ -57,7 +57,7 @@ def test_insert_event_preserves_event_record_timestamp() -> None: """Spanner stores the ADK event timestamp, not the commit timestamp.""" store = SpannerSyncADKStore(_mock_config()) timestamp = datetime(2026, 5, 10, 12, 0, tzinfo=timezone.utc) - event: EventRecord = { + event: StoredEvent = { "id": "event-1", "app_name": "app", "user_id": "u1", @@ -82,7 +82,7 @@ def test_append_event_and_update_state_preserves_event_record_timestamp() -> Non """Atomic append uses the ADK event timestamp while session update uses commit time.""" store = SpannerSyncADKStore(_mock_config()) timestamp = datetime(2026, 5, 10, 12, 0, tzinfo=timezone.utc) - event: EventRecord = { + event: StoredEvent = { "id": "event-1", "app_name": "app", "user_id": "u1", @@ -172,11 +172,12 @@ def test_spanner_session_store_drops_expiration_indexes_before_tables() -> None: def test_spanner_memory_insert_entries_writes_clean_break_record() -> None: store = SpannerSyncADKMemoryStore(_mock_config()) timestamp = datetime(2026, 5, 10, 12, 0, tzinfo=timezone.utc) - entry: MemoryRecord = { + entry: StoredMemory = { "id": "memory-1", "session_id": "session-1", "app_name": "app", "user_id": "user", + "scope": "user", "event_id": "event-1", "author": "assistant", "timestamp": timestamp, @@ -184,6 +185,7 @@ def test_spanner_memory_insert_entries_writes_clean_break_record() -> None: "content_text": "hello", "metadata_json": {"source": "unit"}, "inserted_at": timestamp, + "embedding": None, } with patch.object(store, "_event_exists", return_value=False), patch.object(store, "_run_write") as run_write: @@ -208,6 +210,7 @@ def test_spanner_memory_rows_to_records_decodes_json_fields() -> None: "session-1", "app", "user", + "user", "event-1", "assistant", timestamp, @@ -221,6 +224,8 @@ def test_spanner_memory_rows_to_records_decodes_json_fields() -> None: assert records[0]["content_json"] == {"text": "hello"} assert records[0]["metadata_json"] == {"source": "unit"} assert records[0]["content_text"] == "hello" + assert records[0]["scope"] == "user" + assert records[0]["embedding"] is None def test_spanner_reset_drop_tables_filters_absent_tables() -> None: @@ -242,7 +247,8 @@ def test_spanner_memory_reset_drop_tables_filters_absent_tables_and_indexes() -> assert statements == [ "DROP INDEX idx_adk_memory_entries_session", - "DROP INDEX idx_adk_memory_entries_app_user_time", + "DROP INDEX idx_adk_memory_entries_app_scope_user_time", + "DROP INDEX idx_adk_memory_entries_scope", "DROP TABLE adk_memory_entries", ] diff --git a/tests/unit/extensions/test_adk/test_converters.py b/tests/unit/extensions/test_adk/test_converters.py index b8c8d160b..7aed86ec0 100644 --- a/tests/unit/extensions/test_adk/test_converters.py +++ b/tests/unit/extensions/test_adk/test_converters.py @@ -1,7 +1,7 @@ """Unit tests for ADK session/event converters and scoped state helpers. Tests the NEW contract specified in Chapter 1 of the ADK Clean-Break Overhaul: -- EventRecord has exactly 7 keys (id, app_name, user_id, session_id, invocation_id, timestamp, event_data) +- StoredEvent has exactly 7 keys (id, app_name, user_id, session_id, invocation_id, timestamp, event_data) - event_to_record takes (event, app_name, user_id, session_id) - record_to_event uses Event.model_validate for full round-trip fidelity - filter_temp_state, split_scoped_state, merge_scoped_state for scoped state handling @@ -247,7 +247,7 @@ def test_compute_update_marker_normalizes_aware_datetime_to_utc() -> None: def test_event_to_record_clean_break_keys() -> None: - """EventRecord has exactly the clean-break indexed fields plus event_data.""" + """StoredEvent has exactly the clean-break indexed fields plus event_data.""" event = _make_event() record = event_to_record(event, "app", "u1", "session-1") assert set(record.keys()) == {"id", "app_name", "user_id", "session_id", "invocation_id", "timestamp", "event_data"} @@ -496,9 +496,9 @@ def test_session_to_record_includes_required_fields() -> None: def test_record_to_session_with_events_round_trip() -> None: """Sessions with events reconstruct correctly using record_to_session.""" - from sqlspec.extensions.adk._types import SessionRecord + from sqlspec.extensions.adk._types import StoredSession - session_record = SessionRecord( + session_record = StoredSession( id="s1", app_name="app", user_id="u1", @@ -521,9 +521,9 @@ def test_record_to_session_with_events_round_trip() -> None: def test_record_to_session_empty_events() -> None: """Sessions without events reconstruct with empty events list.""" - from sqlspec.extensions.adk._types import SessionRecord + from sqlspec.extensions.adk._types import StoredSession - session_record = SessionRecord( + session_record = StoredSession( id="s2", app_name="app", user_id="u2", diff --git a/tests/unit/extensions/test_adk/test_service.py b/tests/unit/extensions/test_adk/test_service.py index ad1e4c040..eed2e948b 100644 --- a/tests/unit/extensions/test_adk/test_service.py +++ b/tests/unit/extensions/test_adk/test_service.py @@ -23,7 +23,7 @@ from google.adk.sessions.base_session_service import GetSessionConfig from google.adk.sessions.session import Session -from sqlspec.extensions.adk._types import EventRecord, SessionRecord +from sqlspec.extensions.adk._types import StoredEvent, StoredSession from sqlspec.extensions.adk.service import SQLSpecSessionService # --------------------------------------------------------------------------- @@ -86,7 +86,7 @@ async def append_event_and_update_state( self.app_state = dict(app_state) if user_state is not None: self.user_state = dict(user_state) - # Return the updated SessionRecord — caller no longer needs a follow-up get_session(). + # Return the updated StoredSession — caller no longer needs a follow-up get_session(). updated = dict(self._session_record) updated["state"] = state updated["update_time"] = datetime.now(timezone.utc) @@ -164,7 +164,7 @@ def __init__(self) -> None: self.app_state: dict[str, Any] = {} self.user_state: dict[str, Any] = {} self.create_session_calls: list[dict[str, Any]] = [] - self._session_record = SessionRecord( + self._session_record = StoredSession( id="s1", app_name="app", user_id="u1", @@ -175,7 +175,7 @@ def __init__(self) -> None: def create_session( self, session_id: str, app_name: str, user_id: str, state: dict[str, Any], owner_id: Any | None = None - ) -> SessionRecord: + ) -> StoredSession: self.create_session_calls.append({ "session_id": session_id, "app_name": app_name, @@ -183,7 +183,7 @@ def create_session( "state": state, "owner_id": owner_id, }) - self._session_record = SessionRecord( + self._session_record = StoredSession( id=session_id, app_name=app_name, user_id=user_id, @@ -195,7 +195,7 @@ def create_session( def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: Any | None = None - ) -> SessionRecord | None: + ) -> StoredSession | None: if self._session_record["app_name"] != app_name or self._session_record["user_id"] != user_id: return None if self._session_record["id"] != session_id: @@ -204,7 +204,7 @@ def get_session( def get_events( self, app_name: str, user_id: str, session_id: str, after_timestamp: Any = None, limit: Any = None - ) -> list[EventRecord]: + ) -> list[StoredEvent]: return [] def get_app_state(self, app_name: str) -> dict[str, Any]: @@ -221,19 +221,19 @@ def upsert_user_state(self, app_name: str, user_id: str, state: dict[str, Any]) def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, state: dict[str, Any], app_state: dict[str, Any] | None = None, user_state: dict[str, Any] | None = None, - ) -> SessionRecord: + ) -> StoredSession: if app_state is not None: self.app_state = dict(app_state) if user_state is not None: self.user_state = dict(user_state) - self._session_record = SessionRecord( + self._session_record = StoredSession( id=self._session_record["id"], app_name=self._session_record["app_name"], user_id=self._session_record["user_id"], @@ -243,7 +243,7 @@ def append_event_and_update_state( ) return self._session_record - def list_sessions(self, app_name: str, user_id: str | None = None) -> list[SessionRecord]: + def list_sessions(self, app_name: str, user_id: str | None = None) -> list[StoredSession]: return [self._session_record] def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: @@ -330,7 +330,7 @@ async def test_append_event_uses_returned_record_no_extra_get_session() -> None: """The post-append get_session round-trip must be eliminated. With the atomic-return contract, append_event_and_update_state returns the - updated SessionRecord; the service must use it directly instead of calling + updated StoredSession; the service must use it directly instead of calling get_session again. We expect exactly 1 get_session call (the stale-check BEFORE the append) per append_event. """ diff --git a/tests/unit/extensions/test_adk/test_store_common_mixins.py b/tests/unit/extensions/test_adk/test_store_common_mixins.py index 16d9d4c7f..954852dea 100644 --- a/tests/unit/extensions/test_adk/test_store_common_mixins.py +++ b/tests/unit/extensions/test_adk/test_store_common_mixins.py @@ -4,7 +4,7 @@ from sqlspec.adapters.sqlite import SqliteConfig from sqlspec.adapters.sqlite.adk import SqliteADKMemoryStore, SqliteADKStore -from sqlspec.extensions.adk.artifact._types import ArtifactRecord +from sqlspec.extensions.adk.artifact._types import StoredArtifact from sqlspec.extensions.adk.artifact.store import ( BaseAsyncADKArtifactStore, BaseSyncADKArtifactStore, @@ -17,12 +17,12 @@ class _ConcreteArtifactStore(BaseSyncADKArtifactStore[SqliteConfig]): __slots__ = () - def insert_artifact(self, record: ArtifactRecord) -> None: + def insert_artifact(self, record: StoredArtifact) -> None: return None def get_artifact( self, app_name: str, user_id: str, filename: str, session_id: str | None = None, version: int | None = None - ) -> ArtifactRecord | None: + ) -> StoredArtifact | None: return None def list_artifact_keys(self, app_name: str, user_id: str, session_id: str | None = None) -> list[str]: @@ -30,12 +30,12 @@ def list_artifact_keys(self, app_name: str, user_id: str, session_id: str | None def list_artifact_versions( self, app_name: str, user_id: str, filename: str, session_id: str | None = None - ) -> list[ArtifactRecord]: + ) -> list[StoredArtifact]: return [] def delete_artifact( self, app_name: str, user_id: str, filename: str, session_id: str | None = None - ) -> list[ArtifactRecord]: + ) -> list[StoredArtifact]: return [] def get_next_version(self, app_name: str, user_id: str, filename: str, session_id: str | None = None) -> int: diff --git a/tests/unit/extensions/test_adk/test_store_config.py b/tests/unit/extensions/test_adk/test_store_config.py index f721663d2..932b59988 100644 --- a/tests/unit/extensions/test_adk/test_store_config.py +++ b/tests/unit/extensions/test_adk/test_store_config.py @@ -4,14 +4,14 @@ import importlib import logging from datetime import datetime -from typing import Any +from typing import Any, Literal import pytest -from sqlspec.extensions.adk import EventRecord, SessionRecord -from sqlspec.extensions.adk.artifact._types import ArtifactRecord +from sqlspec.extensions.adk import StoredEvent, StoredSession +from sqlspec.extensions.adk.artifact._types import StoredArtifact from sqlspec.extensions.adk.artifact.store import BaseSyncADKArtifactStore -from sqlspec.extensions.adk.memory import MemoryRecord +from sqlspec.extensions.adk.memory import StoredMemory from sqlspec.extensions.adk.memory import store as memory_store_module from sqlspec.extensions.adk.memory.store import BaseSyncADKMemoryStore from sqlspec.extensions.adk.store import BaseAsyncADKStore, BaseSyncADKStore @@ -33,8 +33,8 @@ def provide_session(self) -> str: class _AsyncSessionStore(BaseAsyncADKStore[Any]): async def create_session( self, session_id: str, app_name: str, user_id: str, state: dict[str, Any], owner_id: Any | None = None - ) -> SessionRecord: - return SessionRecord( + ) -> StoredSession: + return StoredSession( id=session_id, app_name=app_name, user_id=user_id, @@ -45,24 +45,24 @@ async def create_session( async def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: Any | None = None - ) -> SessionRecord | None: + ) -> StoredSession | None: return None async def update_session_state(self, app_name: str, user_id: str, session_id: str, state: dict[str, Any]) -> None: return None - async def list_sessions(self, app_name: str, user_id: str | None = None) -> list[SessionRecord]: + async def list_sessions(self, app_name: str, user_id: str | None = None) -> list[StoredSession]: return [] async def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: return None - async def append_event(self, event_record: EventRecord) -> None: + async def append_event(self, event_record: StoredEvent) -> None: return None async def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -70,7 +70,7 @@ async def append_event_and_update_state( *, app_state: dict[str, Any] | None = None, user_state: dict[str, Any] | None = None, - ) -> SessionRecord: + ) -> StoredSession: return await self.create_session(session_id, app_name, user_id, state) async def get_events( @@ -80,7 +80,7 @@ async def get_events( session_id: str, after_timestamp: datetime | None = None, limit: int | None = None, - ) -> list[EventRecord]: + ) -> list[StoredEvent]: return [] async def delete_expired_events(self, before: datetime) -> int: @@ -151,8 +151,8 @@ def __init__(self, config: _Config) -> None: def create_session( self, session_id: str, app_name: str, user_id: str, state: dict[str, Any], owner_id: Any | None = None - ) -> SessionRecord: - return SessionRecord( + ) -> StoredSession: + return StoredSession( id=session_id, app_name=app_name, user_id=user_id, @@ -163,24 +163,24 @@ def create_session( def get_session( self, app_name: str, user_id: str, session_id: str, *, renew_for: Any | None = None - ) -> SessionRecord | None: + ) -> StoredSession | None: return None def update_session_state(self, app_name: str, user_id: str, session_id: str, state: dict[str, Any]) -> None: return None - def list_sessions(self, app_name: str, user_id: str | None = None) -> list[SessionRecord]: + def list_sessions(self, app_name: str, user_id: str | None = None) -> list[StoredSession]: return [] def delete_session(self, app_name: str, user_id: str, session_id: str) -> None: return None - def append_event(self, event_record: EventRecord) -> None: + def append_event(self, event_record: StoredEvent) -> None: return None def append_event_and_update_state( self, - event_record: EventRecord, + event_record: StoredEvent, app_name: str, user_id: str, session_id: str, @@ -188,7 +188,7 @@ def append_event_and_update_state( *, app_state: dict[str, Any] | None = None, user_state: dict[str, Any] | None = None, - ) -> SessionRecord: + ) -> StoredSession: return self.create_session(session_id, app_name, user_id, state) def get_events( @@ -198,7 +198,7 @@ def get_events( session_id: str, after_timestamp: datetime | None = None, limit: int | None = None, - ) -> list[EventRecord]: + ) -> list[StoredEvent]: return [] def delete_expired_events(self, before: datetime) -> int: @@ -270,16 +270,23 @@ def __init__(self, config: _Config) -> None: def create_tables(self) -> None: self.create_tables_called = True - def insert_memory_entries(self, entries: list[MemoryRecord], owner_id: object | None = None) -> int: + def insert_memory_entries(self, entries: list[StoredMemory], owner_id: object | None = None) -> int: return len(entries) - def search_entries(self, query: str, app_name: str, user_id: str, limit: int | None = None) -> list[MemoryRecord]: + def search_entries( + self, + query: str, + app_name: str, + user_id: str, + limit: int | None = None, + scope_filter: Literal["all", "user", "app"] = "all", + ) -> list[StoredMemory]: return [] def delete_entries_by_session(self, session_id: str) -> int: return 0 - def delete_entries_older_than(self, days: int) -> int: + def delete_entries_older_than(self, days: int, app_name: str | None = None, scope: str | None = None) -> int: return 0 def _memory_table_ddl(self) -> str | list[str]: @@ -290,12 +297,12 @@ def _drop_memory_table_sql(self) -> list[str]: class _SyncArtifactStore(BaseSyncADKArtifactStore[Any]): - def insert_artifact(self, record: ArtifactRecord) -> None: + def insert_artifact(self, record: StoredArtifact) -> None: return None def get_artifact( self, app_name: str, user_id: str, filename: str, session_id: str | None = None, version: int | None = None - ) -> ArtifactRecord | None: + ) -> StoredArtifact | None: return None def list_artifact_keys(self, app_name: str, user_id: str, session_id: str | None = None) -> list[str]: @@ -303,12 +310,12 @@ def list_artifact_keys(self, app_name: str, user_id: str, session_id: str | None def list_artifact_versions( self, app_name: str, user_id: str, filename: str, session_id: str | None = None - ) -> list[ArtifactRecord]: + ) -> list[StoredArtifact]: return [] def delete_artifact( self, app_name: str, user_id: str, filename: str, session_id: str | None = None - ) -> list[ArtifactRecord]: + ) -> list[StoredArtifact]: return [] def get_next_version(self, app_name: str, user_id: str, filename: str, session_id: str | None = None) -> int: diff --git a/tests/unit/extensions/test_adk/test_store_instantiation.py b/tests/unit/extensions/test_adk/test_store_instantiation.py index cf6ca6ece..f93944edf 100644 --- a/tests/unit/extensions/test_adk/test_store_instantiation.py +++ b/tests/unit/extensions/test_adk/test_store_instantiation.py @@ -118,14 +118,14 @@ def test_store_method_signatures_match_base_contract(class_path: str) -> None: """Every shipped concrete store keeps the base store method signatures.""" cls = _load_class(class_path) - if issubclass(cls, BaseAsyncADKStore): - base: type = BaseAsyncADKStore - elif issubclass(cls, BaseSyncADKStore): - base = BaseSyncADKStore - elif issubclass(cls, BaseAsyncADKMemoryStore): - base = BaseAsyncADKMemoryStore - else: + if issubclass(cls, BaseAsyncADKMemoryStore): + base: type = BaseAsyncADKMemoryStore + elif issubclass(cls, BaseSyncADKMemoryStore): base = BaseSyncADKMemoryStore + elif issubclass(cls, BaseAsyncADKStore): + base = BaseAsyncADKStore + else: + base = BaseSyncADKStore for method_name in base.__abstractmethods__: base_signature = inspect.signature(getattr(base, method_name)) diff --git a/tests/unit/utils/test_mypyc_smoke.py b/tests/unit/utils/test_mypyc_smoke.py index c9de78bb9..ecbe50f51 100644 --- a/tests/unit/utils/test_mypyc_smoke.py +++ b/tests/unit/utils/test_mypyc_smoke.py @@ -169,7 +169,7 @@ def test_smoke_runner_skips_optional_adk_dependency(monkeypatch: MonkeyPatch) -> monkeypatch.setattr( module, "SMOKE_IMPORTS", - (module.SmokeImport("adk_record_types", "sqlspec.extensions.adk._types", "SessionRecord", True, "google.adk"),), + (module.SmokeImport("adk_record_types", "sqlspec.extensions.adk._types", "StoredSession", True, "google.adk"),), ) def import_missing_optional_dependency(name: str) -> ModuleType: @@ -183,7 +183,7 @@ def import_missing_optional_dependency(name: str) -> ModuleType: { "name": "adk_record_types", "module": "sqlspec.extensions.adk._types", - "attribute": "SessionRecord", + "attribute": "StoredSession", "imported": False, "compiled": False, "compiled_required": True, @@ -200,7 +200,7 @@ def test_smoke_runner_skips_missing_optional_parent_package(monkeypatch: MonkeyP monkeypatch.setattr( module, "SMOKE_IMPORTS", - (module.SmokeImport("adk_record_types", "sqlspec.extensions.adk._types", "SessionRecord", True, "google.adk"),), + (module.SmokeImport("adk_record_types", "sqlspec.extensions.adk._types", "StoredSession", True, "google.adk"),), ) def import_missing_optional_parent(name: str) -> ModuleType: diff --git a/tools/scripts/mypyc_smoke.py b/tools/scripts/mypyc_smoke.py index b2d80846e..05205d3fd 100644 --- a/tools/scripts/mypyc_smoke.py +++ b/tools/scripts/mypyc_smoke.py @@ -52,7 +52,7 @@ class SmokeImport(NamedTuple): SmokeImport("event_payload", "sqlspec.extensions.events._payload", "encode_notify_payload", True), SmokeImport("event_channel", "sqlspec.extensions.events._channel", "SyncEventChannel", True), SmokeImport("event_queue", "sqlspec.extensions.events._queue", "SyncTableEventQueue", True), - SmokeImport("adk_record_types", "sqlspec.extensions.adk._types", "SessionRecord", True, "google.adk"), + SmokeImport("adk_record_types", "sqlspec.extensions.adk._types", "StoredSession", True, "google.adk"), SmokeImport("migration_runner", "sqlspec.migrations.runner", "SyncMigrationRunner", True), SmokeImport("sqlite_type_converter", "sqlspec.adapters.sqlite.type_converter", "register_type_handlers", True), )