Add Async Stateful Classes for Langgraph memory: AsyncLakebasePool, AsyncCheckpointSaver, AsyncDatabricksStore (#255)

This PR adds async versions of the Lakebase connection pool (AsyncLakebasePool), LangGraph checkpoint saver (AsyncCheckpointSaver) and LangGraph PostgresStore(AsyncDatabricksStore) so they can be used in async workflows such as async agent servers and streaming applications

Changes
**Core Library (databricks-ai-bridge)**
src/databricks_ai_bridge/lakebase.py
Added AsyncLakebasePool class that wraps [psycopg_pool.AsyncConnectionPool](https://www.psycopg.org/psycopg3/docs/api/pool.html#psycopg_pool.AsyncConnectionPool)

**LangChain (databricks-langchain)**
**integrations/langchain/src/databricks_langchain/checkpoint.py**
Added AsyncCheckpointSaver class extending [langgraph.checkpoint.postgres.aio.AsyncPostgresSaver](https://pypi.org/project/langgraph-checkpoint-postgres/)

**integrations/langchain/src/databricks_langchain/store.py**
Added AsyncDatabricksStore class extending [AsyncBatchedBaseStore](https://reference.langchain.com/python/langgraph/store/#langgraph.store.postgres.AsyncPostgresStore)
Uses AsyncLakebasePool for connection pooling
Creates short-lived AsyncPostgresStore instances for each operation

Added unit tests for new classes 

Used in agent on apps langgraph examples:
short-term: bbqiu/agent-on-app-proto#39
long-term: bbqiu/agent-on-app-proto#42

Author: jennsun

integrations/langchain/src/databricks_langchain/__init__.py

@@ -18,20 +18,22 @@
 from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit, UnityCatalogTool
 
 from databricks_langchain.chat_models import ChatDatabricks
-from databricks_langchain.checkpoint import CheckpointSaver
+from databricks_langchain.checkpoint import AsyncCheckpointSaver, CheckpointSaver
 from databricks_langchain.embeddings import DatabricksEmbeddings
 from databricks_langchain.genie import GenieAgent
 from databricks_langchain.multi_server_mcp_client import (
     DatabricksMCPServer,
     DatabricksMultiServerMCPClient,
     MCPServer,
 )
-from databricks_langchain.store import DatabricksStore
+from databricks_langchain.store import AsyncDatabricksStore, DatabricksStore
 from databricks_langchain.vector_search_retriever_tool import VectorSearchRetrieverTool
 from databricks_langchain.vectorstores import DatabricksVectorSearch
 
 # Expose all integrations to users under databricks-langchain
 __all__ = [
+    "AsyncCheckpointSaver",
+    "AsyncDatabricksStore",
     "ChatDatabricks",
     "CheckpointSaver",
     "DatabricksEmbeddings",

integrations/langchain/src/databricks_langchain/checkpoint.py

@@ -3,12 +3,14 @@
 from databricks.sdk import WorkspaceClient
 
 try:
-    from databricks_ai_bridge.lakebase import LakebasePool
+    from databricks_ai_bridge.lakebase import AsyncLakebasePool, LakebasePool
     from langgraph.checkpoint.postgres import PostgresSaver
+    from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
 
     _checkpoint_imports_available = True
 except ImportError:
     PostgresSaver = object
+    AsyncPostgresSaver = object
     _checkpoint_imports_available = False
 
 
@@ -48,3 +50,42 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         """Exit context manager and close the connection pool."""
         self._lakebase.close()
         return False
+
+
+class AsyncCheckpointSaver(AsyncPostgresSaver):
+    """
+    Async LangGraph PostgresSaver using a Lakebase connection pool.
+
+    instance_name: Name of Lakebase Instance
+    """
+
+    def __init__(
+        self,
+        *,
+        instance_name: str,
+        workspace_client: WorkspaceClient | None = None,
+        **pool_kwargs: object,
+    ) -> None:
+        # Lazy imports
+        if not _checkpoint_imports_available:
+            raise ImportError(
+                "AsyncCheckpointSaver requires databricks-langchain[memory]. "
+                "Please install with: pip install databricks-langchain[memory]"
+            )
+
+        self._lakebase: AsyncLakebasePool = AsyncLakebasePool(
+            instance_name=instance_name,
+            workspace_client=workspace_client,
+            **dict(pool_kwargs),
+        )
+        super().__init__(self._lakebase.pool)
+
+    async def __aenter__(self):
+        """Enter async context manager and open the connection pool."""
+        await self._lakebase.open()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Exit async context manager and close the connection pool."""
+        await self._lakebase.close()
+        return False

integrations/langchain/src/databricks_langchain/store.py

@@ -6,22 +6,27 @@
 from databricks.sdk import WorkspaceClient
 
 try:
-    from databricks_ai_bridge.lakebase import LakebasePool
+    from databricks_ai_bridge.lakebase import AsyncLakebasePool, LakebasePool
     from langgraph.store.base import BaseStore, Op, Result
-    from langgraph.store.postgres import PostgresStore
+    from langgraph.store.base.batch import AsyncBatchedBaseStore
+    from langgraph.store.postgres import AsyncPostgresStore, PostgresStore
+
+    from databricks_langchain.embeddings import DatabricksEmbeddings
 
     _store_imports_available = True
 except ImportError:
     LakebasePool = object
+    AsyncLakebasePool = object
     PostgresStore = object
+    AsyncPostgresStore = object
     BaseStore = object
+    AsyncBatchedBaseStore = object
     Item = object
     Op = object
     Result = object
+    DatabricksEmbeddings = object
     _store_imports_available = False
 
-from databricks_langchain.embeddings import DatabricksEmbeddings
-
 
 class DatabricksStore(BaseStore):
     """Provides APIs for working with long-term memory on Databricks using Lakebase.
@@ -135,3 +140,121 @@ async def abatch(self, ops: Iterable[Op]) -> list[Result]:
         would need async-compatible connection pooling.
         """
         return self.batch(ops)
+
+
+class AsyncDatabricksStore(AsyncBatchedBaseStore):
+    """Async version of DatabricksStore for working with long-term memory on Databricks.
+
+    Extends LangGraph AsyncBatchedBaseStore interface using Databricks Lakebase
+    for async connection pooling, with semantic search capabilities via DatabricksEmbeddings.
+
+    Operations borrow a connection from the async pool, create a short-lived AsyncPostgresStore,
+    execute the operation, and return the connection to the pool.
+    """
+
+    def __init__(
+        self,
+        *,
+        instance_name: str,
+        workspace_client: Optional[WorkspaceClient] = None,
+        embedding_endpoint: Optional[str] = None,
+        embedding_dims: Optional[int] = None,
+        embedding_fields: Optional[List[str]] = None,
+        embeddings: Optional[DatabricksEmbeddings] = None,
+        **pool_kwargs: Any,
+    ) -> None:
+        """Initialize AsyncDatabricksStore with embedding support.
+
+        Args:
+            instance_name: The name of the Lakebase instance to connect to.
+            workspace_client: Optional Databricks WorkspaceClient for authentication.
+            embedding_endpoint: Name of the Databricks Model Serving endpoint for embeddings
+                (e.g., "databricks-gte-large-en"). If provided, enables semantic search.
+            embedding_dims: Dimension of the embedding vectors (e.g., 1024 for gte-large-en,
+                1536 for OpenAI-compatible models). Required if embedding_endpoint is set.
+            embedding_fields: List of JSON paths to vectorize. Defaults to ["$"] which
+                vectorizes the entire JSON value.
+            embeddings: Optional pre-configured DatabricksEmbeddings instance. If provided,
+                takes precedence over embedding_endpoint.
+            **pool_kwargs: Additional keyword arguments passed to AsyncLakebasePool.
+        """
+        if not _store_imports_available:
+            raise ImportError(
+                "AsyncDatabricksStore requires databricks-langchain[memory]. "
+                "Install with: pip install 'databricks-langchain[memory]'"
+            )
+
+        super().__init__()
+
+        self._lakebase: AsyncLakebasePool = AsyncLakebasePool(
+            instance_name=instance_name,
+            workspace_client=workspace_client,
+            **pool_kwargs,
+        )
+
+        # Initialize embeddings and index configuration for semantic search
+        self.embeddings: Optional[DatabricksEmbeddings] = None
+        self.index_config: Optional[dict] = None
+
+        if embeddings is not None:
+            # Use pre-configured embeddings instance
+            if embedding_endpoint is not None:
+                warnings.warn(
+                    "Both 'embeddings' and 'embedding_endpoint' were specified. "
+                    "Using the provided 'embeddings' instance.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+            self.embeddings = embeddings
+            if embedding_dims is None:
+                raise ValueError("embedding_dims is required when providing an embeddings instance")
+            self.index_config = {
+                "dims": embedding_dims,
+                "embed": self.embeddings,
+                "fields": embedding_fields or ["$"],
+            }
+        elif embedding_endpoint is not None:
+            # Create embeddings from endpoint configuration
+            if embedding_dims is None:
+                raise ValueError("embedding_dims is required when embedding_endpoint is specified")
+            self.embeddings = DatabricksEmbeddings(endpoint=embedding_endpoint)
+            self.index_config = {
+                "dims": embedding_dims,
+                "embed": self.embeddings,
+                "fields": embedding_fields or ["$"],
+            }
+
+    async def _with_store(self, fn, *args, **kwargs):
+        """
+        Borrow an async connection, create a short-lived AsyncPostgresStore with index config,
+        call fn(store), then return the connection to the pool.
+        """
+        async with self._lakebase.connection() as conn:
+            if self.index_config is not None:
+                store = AsyncPostgresStore(conn=conn, index=self.index_config)
+            else:
+                store = AsyncPostgresStore(conn=conn)
+            return await fn(store, *args, **kwargs)
+
+    async def setup(self) -> None:
+        """Instantiate the store, setting up necessary persistent storage."""
+        return await self._with_store(lambda s: s.setup())
+
+    async def abatch(self, ops: Iterable[Op]) -> list[Result]:
+        """Execute a batch of operations asynchronously.
+
+        This is the core method required by AsyncBatchedBaseStore. All other async operations
+        (aget, aput, asearch, adelete, alist_namespaces) are inherited from AsyncBatchedBaseStore
+        and internally call this abatch() method.
+        """
+        return await self._with_store(lambda s: s.abatch(ops))
+
+    async def __aenter__(self):
+        """Enter async context manager and open the connection pool."""
+        await self._lakebase.open()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Exit async context manager and close the connection pool."""
+        await self._lakebase.close()
+        return False

integrations/langchain/tests/unit_tests/test_checkpoint.py

@@ -10,7 +10,7 @@
 
 from databricks_ai_bridge import lakebase
 
-from databricks_langchain import CheckpointSaver
+from databricks_langchain import AsyncCheckpointSaver, CheckpointSaver
 
 
 class TestConnectionPool:
@@ -66,3 +66,106 @@ def test_checkpoint_saver_configures_lakebase(monkeypatch):
 
     with saver._lakebase.connection() as conn:
         assert conn == "lake-conn"
+
+
+class TestAsyncConnectionPool:
+    """Mock async connection pool for testing."""
+
+    def __init__(self, connection_value="async-conn"):
+        self.connection_value = connection_value
+        self.conninfo = ""
+        self._opened = False
+        self._closed = False
+
+    def __call__(
+        self,
+        *,
+        conninfo,
+        connection_class=None,
+        **kwargs,
+    ):
+        self.conninfo = conninfo
+        return self
+
+    def connection(self):
+        class _AsyncCtx:
+            def __init__(self, outer):
+                self.outer = outer
+
+            async def __aenter__(self):
+                return self.outer.connection_value
+
+            async def __aexit__(self, exc_type, exc, tb):
+                pass
+
+        return _AsyncCtx(self)
+
+    async def open(self):
+        self._opened = True
+
+    async def close(self):
+        self._closed = True
+
+
+@pytest.mark.asyncio
+async def test_async_checkpoint_saver_configures_lakebase(monkeypatch):
+    test_pool = TestAsyncConnectionPool(connection_value="async-lake-conn")
+    monkeypatch.setattr(lakebase, "AsyncConnectionPool", test_pool)
+
+    workspace = MagicMock()
+    workspace.database.generate_database_credential.return_value = MagicMock(token="stub-token")
+    workspace.database.get_database_instance.return_value.read_write_dns = "db-host"
+    workspace.current_service_principal.me.side_effect = RuntimeError("no sp")
+    workspace.current_user.me.return_value = MagicMock(user_name="test@databricks.com")
+
+    saver = AsyncCheckpointSaver(
+        instance_name="lakebase-instance",
+        workspace_client=workspace,
+    )
+
+    assert (
+        test_pool.conninfo
+        == "dbname=databricks_postgres user=test@databricks.com host=db-host port=5432 sslmode=require"
+    )
+
+    assert saver._lakebase.pool == test_pool
+
+
+@pytest.mark.asyncio
+async def test_async_checkpoint_saver_context_manager(monkeypatch):
+    test_pool = TestAsyncConnectionPool(connection_value="async-lake-conn")
+    monkeypatch.setattr(lakebase, "AsyncConnectionPool", test_pool)
+
+    workspace = MagicMock()
+    workspace.database.generate_database_credential.return_value = MagicMock(token="stub-token")
+    workspace.database.get_database_instance.return_value.read_write_dns = "db-host"
+    workspace.current_service_principal.me.side_effect = RuntimeError("no sp")
+    workspace.current_user.me.return_value = MagicMock(user_name="test@databricks.com")
+
+    async with AsyncCheckpointSaver(
+        instance_name="lakebase-instance",
+        workspace_client=workspace,
+    ) as saver:
+        assert test_pool._opened
+        assert saver._lakebase.pool == test_pool
+
+    assert test_pool._closed
+
+
+@pytest.mark.asyncio
+async def test_async_checkpoint_saver_connection(monkeypatch):
+    test_pool = TestAsyncConnectionPool(connection_value="async-lake-conn")
+    monkeypatch.setattr(lakebase, "AsyncConnectionPool", test_pool)
+
+    workspace = MagicMock()
+    workspace.database.generate_database_credential.return_value = MagicMock(token="stub-token")
+    workspace.database.get_database_instance.return_value.read_write_dns = "db-host"
+    workspace.current_service_principal.me.side_effect = RuntimeError("no sp")
+    workspace.current_user.me.return_value = MagicMock(user_name="test@databricks.com")
+
+    async with AsyncCheckpointSaver(
+        instance_name="lakebase-instance",
+        workspace_client=workspace,
+    ) as saver:
+        async with saver._lakebase.connection() as conn:
+            assert conn == "async-lake-conn"