fix(adk): repair PostgreSQL vector memory search (#726)

.pre-commit-config.yaml

@@ -25,7 +25,7 @@ repos:
       - id: mixed-line-ending
       - id: trailing-whitespace
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.16.2"
+    rev: "v0.16.3"
     hooks:
       - id: ruff-check
         args: ["--fix"]

Makefile

@@ -113,7 +113,7 @@ release:                                           ## Bump version and create re
 	@make docs
 	@make clean
 	@make build
-	@uv run bump-my-version bump $(bump)
+	@uv run bump-my-version bump --current-version "$$(uv version --short)" $(bump)
 	@uv lock --upgrade-package sqlspec >/dev/null 2>&1
 	@echo "${OK} Release complete 🎉"
 
@@ -133,7 +133,7 @@ pre-release:                                       ## Start a pre-release: make
 	@echo "${INFO} Preparing pre-release $(version)... 🧪"
 	@make clean
 	@make build
-	@uv run bump-my-version bump --new-version $(version) pre
+	@uv run bump-my-version bump --current-version "$$(uv version --short)" --new-version $(version) pre
 	@uv lock --upgrade-package sqlspec >/dev/null 2>&1
 	@echo "${OK} Pre-release $(version) complete 🧪"
 	@echo ""

docs/changelog.rst

@@ -9,6 +9,28 @@ important operational fixes.
 Recent Updates
 ==============
 
+v0.62.1 - PostgreSQL ADK memory and migration fixes
+---------------------------------------------------
+
+**Fixed:**
+
+* ADK now saves all memory vectors on PostgreSQL, including null values. Vector
+  and hybrid searches pass query values with a portable ``float8[]`` cast.
+  Asyncpg and psycopg no longer need optional pgvector codecs for these tasks.
+* BM25 search now checks for ``pg_textsearch``. It no longer treats ParadeDB's
+  ``pg_search`` as the same feature. The ADK migration enables
+  ``pg_textsearch`` before it creates a BM25 index. Table checks and searches
+  never try to install the extension.
+* Migration commands now set the SQL dialect before they build the context.
+  Calls made from Python no longer capture an empty statement config.
+
+**Upgrade notes:**
+
+* BM25 needs a server that ships ``pg_textsearch``. The role that runs the
+  migration must be able to run ``CREATE EXTENSION``. AlloyDB offers BM25 on
+  PostgreSQL 17 and 18. Use an ``alloydbsuperuser`` role or ask an administrator
+  to install the extension first.
+
 v0.62.0 - ADK session paging and retention, migrations, and event payloads
 ------------------------------------------------------------------------------
 

docs/extensions/adk/migrations.rst

@@ -92,8 +92,8 @@ per-feature switches:
    )
 
 ``enable_sessions=False`` suppresses the session, event, state, and metadata
-DDL. ``enable_memory=False`` suppresses the memory table DDL and the PostgreSQL
-vector extension statement described below. Both default to ``True``.
+DDL. ``enable_memory=False`` suppresses the memory table DDL and its PostgreSQL
+extension statements described below. Both default to ``True``.
 
 PostgreSQL pgvector Requirement
 ===============================
@@ -126,6 +126,21 @@ attempts extension installation, so it runs no repeated startup privilege
 check. Deployments that rely on that path instead of versioned migrations must
 pre-provision pgvector themselves.
 
+PostgreSQL BM25 Requirement
+===========================
+
+When ``enable_bm25=True``, the same migration runs
+``CREATE EXTENSION IF NOT EXISTS pg_textsearch`` before creating the BM25
+index. This is the extension used by AlloyDB for PostgreSQL 17 and 18 as well
+as PostgreSQL installations that package ``pg_textsearch``. ParadeDB's
+``pg_search`` extension does not satisfy this requirement.
+
+As with pgvector, the server must provide the extension and the migration role
+must be allowed to enable it. On managed AlloyDB, run the migration as a role
+with the documented ``alloydbsuperuser`` privileges or have an administrator
+pre-provision ``pg_textsearch``. The idempotent migration statement is then a
+no-op. Automatic table reconciliation never attempts to enable it.
+
 Clean-Break Migration Notes
 ============================
 

pyproject.toml

@@ -24,7 +24,7 @@ maintainers = [{ name = "Litestar Developers", email = "hello@litestar.dev" }]
 name = "sqlspec"
 readme = "README.md"
 requires-python = ">=3.10, <4.0"
-version = "0.62.0"
+version = "0.62.1"
 
 [project.urls]
 Discord = "https://discord.gg/litestar"
@@ -331,7 +331,7 @@ opt_level = "3"   # Maximum optimization (0-3)
 allow_dirty = true
 commit = false
 commit_args = "--no-verify"
-current_version = "0.61.1"
+current_version = "0.62.1"
 ignore_missing_files = false
 ignore_missing_version = false
 message = "chore(release): bump to v{new_version}"

sqlspec/adapters/asyncmy/adk/store.py

@@ -7,6 +7,7 @@
 import asyncmy
 from typing_extensions import NotRequired
 
+from sqlspec.adapters.asyncmy.core import resolve_rowcount
 from sqlspec.config import ADKConfig
 from sqlspec.extensions.adk import BaseAsyncADKStore, StoredEvent, StoredSession, normalize_session_list_options
 from sqlspec.extensions.adk.memory.store import BaseAsyncADKMemoryStore
@@ -309,7 +310,7 @@ async def delete_expired_events(self, before: "datetime", app_name: "str | None"
             async with self._config.provide_connection() as conn, conn.cursor() as cursor:
                 await cursor.execute(sql, tuple(params))
                 await conn.commit()
-                return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0
+                return resolve_rowcount(cursor)
         except asyncmy.errors.ProgrammingError as exc:  # pyright: ignore[reportAttributeAccessIssue]
             if _is_mysql_table_missing(exc):
                 return 0
@@ -327,7 +328,7 @@ async def delete_idle_sessions(self, updated_before: "datetime", app_name: "str
             async with self._config.provide_connection() as conn, conn.cursor() as cursor:
                 await cursor.execute(sql, tuple(params))
                 await conn.commit()
-                return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0
+                return resolve_rowcount(cursor)
         except asyncmy.errors.ProgrammingError as exc:  # pyright: ignore[reportAttributeAccessIssue]
             if _is_mysql_table_missing(exc):
                 return 0
@@ -345,7 +346,7 @@ async def delete_idle_user_states(self, updated_before: "datetime", app_name: "s
             async with self._config.provide_connection() as conn, conn.cursor() as cursor:
                 await cursor.execute(sql, tuple(params))
                 await conn.commit()
-                return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0
+                return resolve_rowcount(cursor)
         except asyncmy.errors.ProgrammingError as exc:  # pyright: ignore[reportAttributeAccessIssue]
             if _is_mysql_table_missing(exc):
                 return 0
@@ -545,7 +546,7 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "
                             entry["inserted_at"],
                         )
                     await cursor.execute(sql, params)
-                    inserted_count += cursor.rowcount
+                    inserted_count += resolve_rowcount(cursor)
             await conn.commit()
         return inserted_count
 
@@ -608,7 +609,7 @@ async def delete_entries_by_session(self, session_id: str) -> int:
         async with self._config.provide_connection() as conn, conn.cursor() as cursor:
             await cursor.execute(sql, (session_id,))
             await conn.commit()
-            return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0
+            return resolve_rowcount(cursor)
 
     async def delete_entries_older_than(
         self, days: int, app_name: "str | None" = None, scope: "str | None" = None
@@ -632,7 +633,7 @@ async def delete_entries_older_than(
         async with self._config.provide_connection() as conn, conn.cursor() as cursor:
             await cursor.execute(sql, tuple(params))
             await conn.commit()
-            return cursor.rowcount if cursor.rowcount and cursor.rowcount > 0 else 0
+            return resolve_rowcount(cursor)
 
     async def _memory_table_ddl(self) -> str:
         """Get MySQL CREATE TABLE SQL for memory entries."""

sqlspec/adapters/asyncmy/config.py

@@ -325,7 +325,7 @@ async def _create_pool(self) -> "AsyncmyPool":
 
         Future driver_features can be added here if needed.
         """
-        return cast("AsyncmyPool", await asyncmy.create_pool(**_pool_config(self.connection_config)))
+        return await asyncmy.create_pool(**_pool_config(self.connection_config))
 
     async def _ensure_connection(self, connection: "AsyncmyConnection") -> None:
         """Ensure connection callback has been called exactly once for this connection.

sqlspec/adapters/asyncmy/litestar/store.py

@@ -3,6 +3,7 @@
 from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, Any, Final, cast
 
+from sqlspec.adapters.asyncmy.core import resolve_rowcount
 from sqlspec.exceptions import ImproperConfigurationError
 from sqlspec.extensions.litestar.store import BaseSQLSpecStore
 from sqlspec.utils.logging import get_logger
@@ -227,7 +228,7 @@ async def delete_expired(self) -> int:
         async with self._config.provide_connection() as conn, conn.cursor() as cursor:
             await cursor.execute(sql)
             await conn.commit()
-            count: int = cursor.rowcount
+            count = resolve_rowcount(cursor)
             if count > 0:
                 self._log_delete_expired(count)
             return count

sqlspec/adapters/asyncpg/adk/store.py

@@ -599,6 +599,8 @@ async def create_tables(self) -> None:
             return
 
         async with self._config.provide_session() as driver:
+            if self._enable_bm25:
+                self._config._ensure_pg_textsearch_available()
             await driver.execute_script(await self._memory_table_ddl())
 
     async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "object | None" = None) -> int:
@@ -616,9 +618,9 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "
                     sql = f"""
                     INSERT INTO {self._memory_table}
                     (id, session_id, app_name, user_id, scope, event_id, author,
-                     {self._owner_id_column_name}, timestamp, content_json,
+                     {self._owner_id_column_name}, timestamp, embedding, content_json,
                      content_text, metadata_json, inserted_at)
-                    VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)
+                    VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10::float8[]::vector, $11, $12, $13, $14)
                     ON CONFLICT (event_id) DO NOTHING
                     """
                     result = await conn.execute(
@@ -632,6 +634,7 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "
                         entry["author"],
                         owner_id,
                         entry["timestamp"],
+                        entry.get("embedding"),
                         entry["content_json"],
                         entry["content_text"],
                         entry["metadata_json"],
@@ -641,8 +644,8 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "
                     sql = f"""
                     INSERT INTO {self._memory_table}
                     (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, $12)
+                     timestamp, embedding, content_json, content_text, metadata_json, inserted_at)
+                    VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9::float8[]::vector, $10, $11, $12, $13)
                     ON CONFLICT (event_id) DO NOTHING
                     """
                     result = await conn.execute(
@@ -655,13 +658,14 @@ async def insert_memory_entries(self, entries: "list[StoredMemory]", owner_id: "
                         entry["event_id"],
                         entry["author"],
                         entry["timestamp"],
+                        entry.get("embedding"),
                         entry["content_json"],
                         entry["content_text"],
                         entry["metadata_json"],
                         entry["inserted_at"],
                     )
                 try:
-                    inserted_count += int(result.split(" ")[1])
+                    inserted_count += int(result.rsplit(" ", 1)[-1])
                 except (IndexError, ValueError):
                     continue
 
@@ -706,7 +710,7 @@ async def search_entries(
             candidate_limit = max(limit_value * 2, 50)
             sql = f"""
             WITH vector_matches AS (
-                SELECT id, RANK() OVER (ORDER BY embedding <=> {p_vec}) AS rank_vec
+                SELECT id, RANK() OVER (ORDER BY embedding <=> {p_vec}::float8[]::vector) AS rank_vec
                 FROM {self._memory_table}
                 WHERE {where_scope} AND embedding IS NOT NULL
                 LIMIT {p_cand}
@@ -732,7 +736,7 @@ async def search_entries(
             sql = f"""
             SELECT * FROM {self._memory_table}
             WHERE {where_scope} AND embedding IS NOT NULL
-            ORDER BY embedding <=> {p_vec} ASC, timestamp DESC
+            ORDER BY embedding <=> {p_vec}::float8[]::vector ASC, timestamp DESC
             LIMIT {p_lim}
             """
             params = (*scope_params, list(embedding), limit_value)
@@ -759,6 +763,8 @@ async def search_entries(
             params = (*scope_params, f"%{query}%", limit_value)
 
         async with self._config.provide_connection() as conn:
+            if embedding is not None and self._enable_bm25 and query:
+                self._config._ensure_pg_textsearch_available()
             rows = await conn.fetch(sql, *params)
         return [cast("StoredMemory", dict(row)) for row in rows]
 
@@ -821,7 +827,7 @@ async def _memory_table_ddl(self) -> str:
 
         if self._enable_bm25:
             indexes.append(
-                f"CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_bm25 ON {self._memory_table} USING bm25 (content_text);"
+                f"CREATE INDEX IF NOT EXISTS idx_{self._memory_table}_bm25 ON {self._memory_table} USING bm25 (content_text) WITH (text_config='english');"
             )
 
         if self._vector_index_type == "scann":

sqlspec/adapters/asyncpg/config.py

@@ -326,6 +326,8 @@ def __init__(
         self._alloydb_connector: Any | None = None
         self._pgvector_available: bool | None = None
         self._paradedb_available: bool | None = None
+        self._pg_textsearch_available: bool | None = None
+        self._pg_textsearch_probe_error: Exception | None = None
 
         self._validate_connector_config()
 
@@ -452,17 +454,28 @@ async def _init_connection(self, connection: "AsyncpgConnection") -> None:
         if self._pgvector_available is None:
             detected_extensions: set[str] = set()
             extensions = build_postgres_extension_probe_names(self.driver_features)
+            adk_config = self.extension_config.get("adk", {})
+            bm25_enabled = bool(
+                isinstance(adk_config, dict)
+                and adk_config.get("enable_memory", True)
+                and adk_config.get("enable_bm25", False)
+            )
+            if bm25_enabled:
+                extensions.append("pg_textsearch")
             if extensions:
                 try:
                     results = await connection.fetch(
                         "SELECT extname FROM pg_extension WHERE extname = ANY($1::text[])", extensions
                     )
                     detected_extensions = {r["extname"] for r in results}
-                except Exception:
+                except Exception as exc:
                     detected_extensions = set()
+                    if bm25_enabled:
+                        self._pg_textsearch_probe_error = exc
             self.statement_config, self._pgvector_available, self._paradedb_available = (
                 resolve_postgres_extension_state(self.statement_config, self.driver_features, detected_extensions)
             )
+            self._pg_textsearch_available = "pg_textsearch" in detected_extensions if bm25_enabled else False
 
         if self._pgvector_available:
             await register_pgvector_support(connection)
@@ -471,6 +484,12 @@ async def _init_connection(self, connection: "AsyncpgConnection") -> None:
         if self._user_connection_hook is not None:
             await self._user_connection_hook(connection)
 
+    def _ensure_pg_textsearch_available(self) -> None:
+        if self._pg_textsearch_available:
+            return
+        msg = "ADK memory enable_bm25 requires the pg_textsearch PostgreSQL extension"
+        raise ImproperConfigurationError(msg) from self._pg_textsearch_probe_error
+
     async def _close_pool(self) -> None:
         """Close the actual async connection pool and cleanup connectors."""
         if self.connection_instance: