Skip to content

Commit 383df5e

Browse files
committed
feat(protocols): add new protocols for enhanced functionality
- Introduced several new protocols including HasAddListenerProtocol, HasArrowStoreProtocol, HasErrorsProtocol, HasNameProtocol, HasNotifiesProtocol, HasRowcountProtocol, HasSqlStateProtocol, HasSqliteErrorProtocol, HasValueProtocol, and HasTypeCodeProtocol. - Updated existing protocols to improve type safety and functionality. - Enhanced ReadableProtocol to accept an optional size parameter for the read method. refactor(storage): simplify row count handling in SyncStoragePipeline - Removed unnecessary checks for has_arrow_table_stats when determining row counts in SyncStoragePipeline methods. - Streamlined the logic for processing rows in both synchronous and asynchronous storage pipelines. fix(module_loader): improve import_string error handling - Refined the import_string function to handle attribute access more gracefully, reducing the likelihood of unhandled exceptions. refactor(schema): optimize type checking and field access - Improved type guards for dataclass and attrs instances, enhancing performance and readability. - Simplified field access logic in extract_dataclass_fields and related functions. test(adk): add integration tests for AioSQLite and SQLite memory stores - Implemented comprehensive tests for memory store operations including insert, search, delete by session, and delete older than. - Ensured that memory records are correctly handled and deduplicated. test(unit): add unit tests for ADK memory converters - Created tests for event and session conversion functions to validate correct behavior and output. chore(tests): remove redundant has_attr tests - Eliminated outdated tests for has_attr function, focusing on more relevant type guard tests.
1 parent 245d0ba commit 383df5e

File tree

97 files changed

+3989
-635
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

97 files changed

+3989
-635
lines changed

docs/changelog.rst

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,21 @@ SQLSpec Changelog
1010
Recent Updates
1111
==============
1212

13+
ADK Memory Store
14+
----------------
15+
16+
- Added ``SQLSpecMemoryService`` and ``SQLSpecSyncMemoryService`` for SQLSpec-backed ADK memory storage.
17+
- Implemented adapter-specific memory stores with optional full-text search (`memory_use_fts`) and simple fallback search.
18+
- Extended ADK migrations to include memory tables with configurable ``include_memory_migration`` toggles.
19+
- Added CLI commands for memory cleanup and verification (`sqlspec adk memory cleanup/verify`).
20+
21+
Driver Layer Compilation
22+
------------------------
23+
24+
- Compiled driver base classes and mixins with mypyc to reduce dispatch overhead in the execution pipeline.
25+
- Replaced dynamic ``getattr`` patterns with protocol-driven access for mypyc compatibility.
26+
- Added driver protocols and updated mypyc build configuration to include driver modules.
27+
1328
Database Event Channels
1429
-----------------------
1530

docs/examples/extensions/adk/litestar_aiosqlite.py

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
"""Expose SQLSpec-backed ADK sessions through a Litestar endpoint."""
1+
"""Expose SQLSpec-backed ADK sessions and memory through Litestar endpoints."""
22

33
import asyncio
44
from typing import Any
@@ -7,18 +7,25 @@
77

88
from sqlspec.adapters.aiosqlite import AiosqliteConfig
99
from sqlspec.adapters.aiosqlite.adk import AiosqliteADKStore
10+
from sqlspec.adapters.aiosqlite.adk.memory_store import AiosqliteADKMemoryStore
1011
from sqlspec.extensions.adk import SQLSpecSessionService
12+
from sqlspec.extensions.adk.memory import SQLSpecMemoryService
1113

1214
config = AiosqliteConfig(connection_config={"database": ":memory:"})
1315
service: "SQLSpecSessionService | None" = None
16+
memory_service: "SQLSpecMemoryService | None" = None
1417

1518

1619
async def startup() -> None:
1720
"""Initialize the ADK store when the app boots."""
1821
global service
22+
global memory_service
1923
store = AiosqliteADKStore(config)
24+
memory_store = AiosqliteADKMemoryStore(config)
2025
await store.create_tables()
26+
await memory_store.create_tables()
2127
service = SQLSpecSessionService(store)
28+
memory_service = SQLSpecMemoryService(memory_store)
2229

2330

2431
@get("/sessions")
@@ -29,7 +36,15 @@ async def list_sessions() -> "dict[str, Any]":
2936
return {"count": len(sessions.sessions)}
3037

3138

32-
app = Litestar(route_handlers=[list_sessions], on_startup=[startup])
39+
@get("/memories")
40+
async def list_memories(query: str = "demo") -> "dict[str, Any]":
41+
"""Return memory count for a query string."""
42+
assert memory_service is not None
43+
response = await memory_service.search_memory(app_name="docs", user_id="demo", query=query)
44+
return {"count": len(response.memories)}
45+
46+
47+
app = Litestar(route_handlers=[list_sessions, list_memories], on_startup=[startup])
3348

3449

3550
def main() -> None:

docs/examples/extensions/adk/litestar_aiosqlite.rst

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
ADK + Litestar Endpoint
22
=======================
33

4-
Initialize ``SQLSpecSessionService`` inside Litestar and expose a ``/sessions`` endpoint backed by
5-
AioSQLite.
4+
Initialize ``SQLSpecSessionService`` and ``SQLSpecMemoryService`` inside Litestar and expose
5+
``/sessions`` plus ``/memories`` endpoints backed by AioSQLite.
66

77
.. code-block:: console
88
Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
"""Run an ADK agent with SQLSpec-backed session and memory services (AioSQLite)."""
2+
3+
import asyncio
4+
5+
from google.adk.agents.llm_agent import LlmAgent
6+
from google.adk.apps.app import App
7+
from google.adk.runners import Runner
8+
from google.genai import types
9+
10+
from sqlspec.adapters.aiosqlite import AiosqliteConfig
11+
from sqlspec.adapters.aiosqlite.adk import AiosqliteADKStore
12+
from sqlspec.adapters.aiosqlite.adk.memory_store import AiosqliteADKMemoryStore
13+
from sqlspec.extensions.adk import SQLSpecSessionService
14+
from sqlspec.extensions.adk.memory import SQLSpecMemoryService
15+
16+
__all__ = ("main",)
17+
18+
19+
async def main() -> None:
20+
"""Run a single ADK turn, then persist memory and search it."""
21+
config = AiosqliteConfig(
22+
connection_config={"database": ":memory:"}, extension_config={"adk": {"memory_use_fts": False}}
23+
)
24+
session_store = AiosqliteADKStore(config)
25+
memory_store = AiosqliteADKMemoryStore(config)
26+
await session_store.create_tables()
27+
await memory_store.create_tables()
28+
29+
session_service = SQLSpecSessionService(session_store)
30+
memory_service = SQLSpecMemoryService(memory_store)
31+
32+
agent = LlmAgent(name="sqlspec_agent", model="gemini-2.5-flash", instruction="Answer briefly.")
33+
app = App(name="sqlspec_demo", root_agent=agent)
34+
runner = Runner(app=app, session_service=session_service, memory_service=memory_service)
35+
36+
session_id = "session-1"
37+
user_id = "demo-user"
38+
await session_service.create_session(app_name=app.name, user_id=user_id, session_id=session_id)
39+
40+
new_message = types.UserContent(parts=[types.Part(text="Remember I like espresso.")])
41+
async for _event in runner.run_async(user_id=user_id, session_id=session_id, new_message=new_message):
42+
pass
43+
44+
session = await session_service.get_session(app_name=app.name, user_id=user_id, session_id=session_id)
45+
if session:
46+
await memory_service.add_session_to_memory(session)
47+
48+
response = await memory_service.search_memory(app_name=app.name, user_id=user_id, query="espresso")
49+
print({"memories": len(response.memories)})
50+
51+
52+
if __name__ == "__main__":
53+
asyncio.run(main())
Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
ADK Runner with Memory (AioSQLite)
2+
==================================
3+
4+
Run an ADK ``Runner`` with SQLSpec-backed session and memory services, then
5+
persist memories from the completed session.
6+
7+
This example requires Google ADK credentials (for example, a configured API key)
8+
and network access to the model provider.
9+
10+
.. code-block:: console
11+
12+
uv run python docs/examples/extensions/adk/runner_memory_aiosqlite.py
13+
14+
Source
15+
------
16+
17+
.. literalinclude:: runner_memory_aiosqlite.py
18+
:language: python
19+
:linenos:

docs/examples/index.rst

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,8 @@ Extensions
128128
- Create an ADK session, append events, and fetch the transcript using SQLSpec’s AioSQLite store.
129129
* - ``extensions/adk/litestar_aiosqlite.py``
130130
- Wire ``SQLSpecSessionService`` into Litestar and expose a simple ``/sessions`` endpoint.
131+
* - ``extensions/adk/runner_memory_aiosqlite.py``
132+
- Run an ADK ``Runner`` with SQLSpec-backed session + memory services, then query stored memories.
131133

132134
Shared Utilities
133135
----------------
@@ -150,6 +152,7 @@ Shared Utilities
150152
loaders/sql_files
151153
extensions/adk/basic_aiosqlite
152154
extensions/adk/litestar_aiosqlite
155+
extensions/adk/runner_memory_aiosqlite
153156
frameworks/fastapi/aiosqlite_app
154157
frameworks/fastapi/sqlite_app
155158
frameworks/starlette/aiosqlite_app

docs/extensions/adk/api.rst

Lines changed: 84 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,8 +56,54 @@ SQLSpecSessionService
5656
:doc:`/examples/extensions/adk/litestar_aiosqlite`
5757
Web framework integration using Litestar
5858

59-
Base Store Classes
60-
==================
59+
Memory Service
60+
==============
61+
62+
SQLSpecMemoryService
63+
--------------------
64+
65+
.. autoclass:: sqlspec.extensions.adk.memory.SQLSpecMemoryService
66+
:show-inheritance:
67+
68+
SQLSpec-backed implementation of Google ADK's ``BaseMemoryService``.
69+
70+
This service persists memories extracted from completed sessions and exposes
71+
search capabilities via adapter-specific stores.
72+
73+
**Attributes:**
74+
75+
.. attribute:: store
76+
:no-index:
77+
78+
The database store implementation (e.g., ``AsyncpgADKMemoryStore``).
79+
80+
**Example:**
81+
82+
.. code-block:: python
83+
84+
from sqlspec.adapters.asyncpg.adk.memory_store import AsyncpgADKMemoryStore
85+
from sqlspec.extensions.adk.memory import SQLSpecMemoryService
86+
87+
store = AsyncpgADKMemoryStore(config)
88+
await store.create_tables()
89+
memory_service = SQLSpecMemoryService(store)
90+
91+
.. seealso::
92+
93+
:doc:`/examples/extensions/adk/runner_memory_aiosqlite`
94+
ADK Runner example with SQLSpec-backed memory service
95+
96+
SQLSpecSyncMemoryService
97+
------------------------
98+
99+
.. autoclass:: sqlspec.extensions.adk.memory.SQLSpecSyncMemoryService
100+
:show-inheritance:
101+
102+
Sync memory service for sync adapters (SQLite/DuckDB). This class does not
103+
inherit from ADK's async ``BaseMemoryService`` but mirrors the async API.
104+
105+
Session Store Base Classes
106+
==========================
61107

62108
BaseAsyncADKStore
63109
------------
@@ -170,6 +216,35 @@ BaseSyncADKStore
170216
store = SqliteADKStore(config)
171217
store.create_tables()
172218
219+
Memory Store Base Classes
220+
=========================
221+
222+
BaseAsyncADKMemoryStore
223+
-----------------------
224+
225+
.. autoclass:: sqlspec.extensions.adk.memory.BaseAsyncADKMemoryStore
226+
:show-inheritance:
227+
228+
Abstract base class for async SQLSpec-backed ADK memory stores.
229+
230+
**Abstract Methods:**
231+
232+
- :meth:`create_tables`
233+
- :meth:`insert_memory_entries`
234+
- :meth:`search_entries`
235+
- :meth:`delete_entries_by_session`
236+
- :meth:`delete_entries_older_than`
237+
- :meth:`_get_create_memory_table_sql`
238+
- :meth:`_get_drop_memory_table_sql`
239+
240+
BaseSyncADKMemoryStore
241+
----------------------
242+
243+
.. autoclass:: sqlspec.extensions.adk.memory.BaseSyncADKMemoryStore
244+
:show-inheritance:
245+
246+
Abstract base class for sync SQLSpec-backed ADK memory stores.
247+
173248
Type Definitions
174249
================
175250

@@ -218,6 +293,13 @@ SessionRecord
218293
219294
from datetime import datetime, timezone
220295
296+
MemoryRecord
297+
------------
298+
299+
.. autoclass:: sqlspec.extensions.adk.memory._types.MemoryRecord
300+
301+
TypedDict representing a memory database record.
302+
221303
record: SessionRecord = {
222304
"id": "550e8400-e29b-41d4-a716-446655440000",
223305
"app_name": "weather_agent",

docs/extensions/adk/index.rst

Lines changed: 23 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -25,17 +25,18 @@ Google ADK Extension
2525
migrations
2626
schema
2727

28-
Session and event storage for the Google Agent Development Kit (ADK) using SQLSpec database adapters.
28+
Session, event, and memory storage for the Google Agent Development Kit (ADK) using SQLSpec database adapters.
2929

3030
Overview
3131
========
3232

33-
The SQLSpec ADK extension provides persistent storage for `Google Agent Development Kit <https://github.com/google/genai>`_ sessions and events, enabling stateful AI agent applications with database-backed conversation history.
33+
The SQLSpec ADK extension provides persistent storage for `Google Agent Development Kit <https://github.com/google/genai>`_ sessions, events, and long-term memory entries, enabling stateful AI agent applications with database-backed conversation history and recall.
3434

3535
This extension implements ADK's ``BaseSessionService`` protocol, allowing AI agents to store and retrieve:
3636

3737
- **Session State**: Persistent conversation context and application state
3838
- **Event History**: Complete record of user/assistant interactions
39+
- **Long-term Memory**: Searchable memory entries extracted from completed sessions
3940
- **Multi-User Support**: Isolated sessions per application and user
4041
- **Type-Safe Storage**: Full type safety with TypedDicts and validated records
4142

@@ -149,25 +150,30 @@ The extension follows a layered architecture:
149150
└──────────┬──────────┘
150151
151152
┌──────────▼──────────┐
152-
│ SQLSpecSessionService│ ← Implements BaseSessionService
153-
└──────────┬──────────┘
154-
155-
┌──────────▼──────────┐
156-
│ Store Implementation│ ← AsyncpgADKStore, SqliteADKStore, etc.
157-
└──────────┬──────────┘
158-
159-
┌──────────▼──────────┐
160-
│ SQLSpec Config │ ← AsyncpgConfig, SqliteConfig, etc.
153+
┌─────────────────────┐
154+
│ ADK Runner │
161155
└──────────┬──────────┘
162156
163-
┌──────────▼──────────┐
164-
│ Database │
165-
└─────────────────────┘
157+
┌──────────▼──────────┐ ┌────────────────────┐
158+
│ SQLSpecSessionService│ │ SQLSpecMemoryService│
159+
└──────────┬──────────┘ └──────────┬─────────┘
160+
│ │
161+
┌──────────▼──────────┐ ┌─────────▼─────────┐
162+
│ Session Store │ │ Memory Store │
163+
└──────────┬──────────┘ └─────────┬─────────┘
164+
│ │
165+
┌──────────▼──────────┐ ┌─────────▼─────────┐
166+
│ SQLSpec Config │ │ SQLSpec Config │
167+
└──────────┬──────────┘ └─────────┬─────────┘
168+
│ │
169+
┌──────────▼──────────┐ ┌─────────▼─────────┐
170+
│ Database │ │ Database │
171+
└─────────────────────┘ └───────────────────┘
166172
167173
**Layers:**
168174

169-
1. **Service Layer** (``SQLSpecSessionService``): Implements ADK's ``BaseSessionService`` protocol
170-
2. **Store Layer** (``BaseAsyncADKStore``): Abstract database operations for each adapter
175+
1. **Service Layer** (``SQLSpecSessionService`` / ``SQLSpecMemoryService``): Implements ADK service protocols
176+
2. **Store Layer** (``BaseAsyncADKStore`` / ``BaseAsyncADKMemoryStore``): Abstract database operations per adapter
171177
3. **Config Layer** (SQLSpec): Connection pooling and resource management
172178
4. **Database Layer**: Physical storage with database-specific optimizations
173179

@@ -178,6 +184,7 @@ New curated examples live in the :doc:`examples catalog </examples/index>`:
178184

179185
* :doc:`/examples/extensions/adk/basic_aiosqlite` – create a session, append two events, and read the transcript using AioSQLite storage.
180186
* :doc:`/examples/extensions/adk/litestar_aiosqlite` – initialize ``SQLSpecSessionService`` inside a Litestar app and expose a ``/sessions`` route.
187+
* :doc:`/examples/extensions/adk/runner_memory_aiosqlite` – run an ADK ``Runner`` with SQLSpec-backed memory and search stored memories.
181188

182189
Use Cases
183190
=========

docs/extensions/adk/migrations.rst

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,8 @@ Setting Up Migrations
5555
"adk": {
5656
"session_table": "adk_sessions",
5757
"events_table": "adk_events",
58+
"memory_table": "adk_memory_entries",
59+
"memory_use_fts": True,
5860
"owner_id_column": "account_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE"
5961
}
6062
},
@@ -76,6 +78,13 @@ Setting Up Migrations
7678
``owner_id_column`` configuration when creating tables. The column is added to
7779
the sessions table DDL if specified in ``extension_config["adk"]["owner_id_column"]``.
7880

81+
.. note::
82+
83+
**Memory Tables**: ``ext_adk_0001`` also creates the memory table when
84+
``enable_memory`` (default) or ``include_memory_migration`` is set to ``True``.
85+
Set ``include_memory_migration=False`` to skip memory DDL while keeping the
86+
runtime memory service enabled.
87+
7988
**2. Initialize Migration Directory:**
8089

8190
.. code-block:: bash

docs/extensions/adk/quickstart.rst

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,21 @@ Create the session service that implements ADK's ``BaseSessionService`` protocol
100100
service = SQLSpecSessionService(store)
101101
return service
102102
103+
Optional: Initialize Memory Service
104+
===================================
105+
106+
If you want long-term memory search, create a memory store and service alongside the session service:
107+
108+
.. code-block:: python
109+
110+
from sqlspec.adapters.asyncpg.adk.memory_store import AsyncpgADKMemoryStore
111+
from sqlspec.extensions.adk.memory import SQLSpecMemoryService
112+
113+
async def create_memory_service(config):
114+
memory_store = AsyncpgADKMemoryStore(config)
115+
await memory_store.create_tables()
116+
return SQLSpecMemoryService(memory_store)
117+
103118
Step 5: Create a Session
104119
=========================
105120

0 commit comments

Comments
 (0)