db-openai: Add AsyncDatabricksSession class to support Session Protocol for Stateful Conversation Management via SQLAlchemy engine by jennsun · Pull Request #316 · databricks/databricks-ai-bridge

jennsun · 2026-02-03T21:32:50Z

Adds AsyncDatabricksSession, a session storage implementation for the OpenAI Agents SDK that persists conversation history to Databricks Lakebase.

This class subclasses OpenAI's SQLAlchemySession original code to inherit all SQL logic while adding Lakebase-specific features:

Automatic OAuth token rotation via SQLAlchemy's do_connect event
Instance name resolution and username inference from _LakebasePoolBase

More on Session protocol:
https://openai.github.io/openai-agents-python/ref/memory/session/#agents.memory.session.Session

Using SQLAlchemy's default QueuePool (pool_size=5, max_overflow=10) for connection pooling

Usage:

from databricks_openai.agents import AsyncDatabricksSession

  session = AsyncDatabricksSession(
  session_id=session_id,
        instance_name=LAKEBASE_INSTANCE_NAME,
    )
   result = Runner.run_streamed(agent, input=messages, session=session)

Example queries:

curl -X POST http://localhost:8000/invocations \
    -H "Content-Type: application/json" \
    -d '{"input": [{"role": "user", "content": "Hello I live in SF!"}]}'

returns responses with session id:

{"object":"response","output":[{"type":"message","id":"__fake_id__","content":[{"annotations":[],"text":"Hi! What part of San Francisco are you in, and what are you looking for—recommendations (food/coffee, parks, things to do), help planning a day, or something else?","type":"output_text","logprobs":[]}],"role":"assistant","status":"completed","provider_data":{"model":"databricks-gpt-5-2","response_id":"chatcmpl-D5Ki6f7TKBNrVVfuxDL0Lu16YCQjz"}}],"custom_outputs":{"session_id":"fd57ff2c-1d66-4da3-ba28-3216d4e6d86e"}}

follow-up stateful question:

curl -X POST http://localhost:8000/invocations \
    -H "Content-Type: application/json" \
    -d '{
        "input": [{"role": "user", "content": "What city did I say I live in?"}],
        "custom_inputs": {"session_id": "fd57ff2c-1d66-4da3-ba28-3216d4e6d86e"}
    }'

gives us:

{"object":"response","output":[{"type":"message","id":"__fake_id__","content":[{"annotations":[],"text":"You said you live in SF (San Francisco).","type":"output_text","logprobs":[]}],"role":"assistant","status":"completed","provider_data":{"model":"databricks-gpt-5-2","response_id":"chatcmpl-D5KiakiUvMW2weIeLXUuPesxyNpwv"}}],"custom_outputs":{"session_id":"fd57ff2c-1d66-4da3-ba28-3216d4e6d86e"}}

testing:
unit + integration tests

sample agent: OpenAI MemorySession Stateful Agent Example
sample app: https://eng-ml-agent-platform.staging.cloud.databricks.com/apps/j-openai-stateful?o=2850744067564480

integrations/openai/src/databricks_openai/agents/session.py

jennsun · 2026-02-04T16:59:17Z

integrations/openai/src/databricks_openai/agents/session.py

+        # ensuring fresh tokens are injected via do_connect event.
+        engine = create_async_engine(
+            url,
+            pool_recycle=DEFAULT_POOL_RECYCLE_SECONDS,


Connection pooling happens here via SQLAlchemy's default QueuePool (SQLAlchemy defaults to QueuePool) and pool_recycle=2700 setting ensures connections are recycled every 45 minutes (before the 50-min token cache expires), at which point the do_connect event injects a fresh token.

bbqiu · 2026-02-04T19:39:33Z

integrations/openai/src/databricks_openai/agents/session.py

@@ -0,0 +1,234 @@
+"""


do we want this to be importable from databricks_openai.agents?

sure, I'll make this importable so instead of:

from databricks_openai.agents.session import MemorySession

import path will look like:

from databricks_openai.agents import MemorySession

bbqiu · 2026-02-04T19:43:25Z

integrations/openai/src/databricks_openai/agents/session.py

+DEFAULT_DATABASE = "databricks_postgres"
+
+
+class _LakebaseCredentials(_LakebasePoolBase):


does it make sense to rename _LakebasePoolBase?

also could you remind me why we didn't have the cache lock as an attribute of the _LakebasePoolBase instance?

we had separate locks (sync vs async) for the sync vs async LakebasePools we implemented, so each of the subclasses adds their own cache lock

I'll rename _LakebasePoolBase to be _LakebaseBase as it's a more generic class for resolving lakebase host/username/token caching - the actual pooling logic is actually implemented in the subclasses LakebasePool/AsyncLakebasePool

bbqiu · 2026-02-04T19:46:38Z

integrations/openai/src/databricks_openai/agents/session.py

+            return token
+
+
+class MemorySession(SQLAlchemySession):


pls correct me if i'm wrong, but i think this is async only?

can we leave some clarification about this in the docstrings / throw a helpful error if someone tries to run this synchronously

yes it's async only - taking a look at the source code all of these classes implement the async interface (since they follow the session protocol):
https://github.com/openai/openai-agents-python/blob/main/src/agents/memory/session.py

i'll cover this in unit tests/rename to AsyncDatabricksSession to make it clearer

bbqiu · 2026-02-04T19:50:31Z

integrations/openai/src/databricks_openai/agents/session.py

+        )
+
+        # Attach event to inject Lakebase token before each connection
+        # Note: do_connect fires on sync_engine even for async operations


we are creating an async engine right? is this an old comment

yes, but the async engine is a wrapper around a sync engine:

there is not yet an “async” version of a SQLAlchemy event handler

"Events can be registered at the instance level (e.g. a specific AsyncEngine instance) by associating the event with the sync attribute that refers to the proxied object. For example to register the PoolEvents.connect() event against an AsyncEngine instance, use its AsyncEngine.sync_engine attribute as target."
link: https://docs.sqlalchemy.org/en/20/orm/extensions/asyncio.html#using-events-with-the-asyncio-extension

bbqiu · 2026-02-04T19:51:06Z

integrations/openai/src/databricks_openai/agents/session.py

+        return self._credentials.username
+
+    @property
+    def connection_url(self) -> str:


nit: do we wanna reuse this in _create_engine func above?

bbqiu · 2026-02-04T19:58:30Z

integrations/openai/src/databricks_openai/agents/session.py

+            token_cache_duration_seconds=token_cache_duration_seconds,
+        )
+
+        engine = self._create_engine(**engine_kwargs)


iiuc, this can potentially be run on every new conversation right?

is there any way to reuse an engine across sessions? if not, should we try to make these operations async via event loops

…e LakebasePoolBase to LakebaseBase

bbqiu · 2026-02-05T23:49:28Z

integrations/openai/src/databricks_openai/agents/session.py

+            return token
+
+
+class AsyncDatabricksSession(SQLAlchemySession):


from talking to the research team working on DBRA, they actually have a very similar snippet as us to manage a SQLAlchemy connection to lakebase: https://sourcegraph.prod.databricks-corp.com/databricks-eng/universe/-/blob/research/aroll/app/aroll_app/db/connection.py?L162-182

would it make sense for us to further abstract this by providing a similar AsyncLakebaseSQLAlchemy / LakebaseSQLAlchemy class?

discussed offline but I'll refactor such that:

in db-ai-bridge we add AsyncLakebaseSqlAlchemy support for creating SQLAlchemy engines (similar to https://sourcegraph.prod.databricks-corp.com/databricks-eng/universe/-/blob/research/aroll/app/aroll_app/db/connection.py?L185)

in db-openai we subclass SQLAlchemySession in AsyncDatabricksSession (session protocol is async + specific to openai agents sdk) and pass in engine from AsyncLakebaseSqlAlchemy

this will create much cleaner separation of concerns for future frameworks to reuse any sqlalchemy engines etc!

…n that returns engines to manage connections to db

jennsun · 2026-02-06T18:53:58Z

integrations/openai/src/databricks_openai/agents/session.py

+    # Class-level cache for AsyncLakebaseSQLAlchemy instances, keyed by instance_name.
+    # This allows multiple AsyncDatabricksSession instances to share a single engine/pool.
+    _lakebase_sql_alchemy_cache: dict[str, AsyncLakebaseSQLAlchemy] = {}
+    _lakebase_sql_alchemy_cache_lock = Lock()


thoughts on the class-level cache for AsyncLakebaseSQLAlchemy engines keyed by instance_name?

this is so we reuse a single SQLAlchemy engine / pool per Lakebase instance, avoiding repeated pool creation, TCP handshakes, and auth setup.

sessions are still created per Runner.run(), but engines are shared

this approach looks good to me to minimize IO. two comments:

we may want to include a param for a func for customers to customize the cache key. currently, diff engine kwargs for the same instance name will be ignored

let's also call this out in the docstring and add a param to optionally disable this engine caching

the best case would be include engine kwargs + instance name in the cache key

sounds good - going to create cache key that takes into consideration both instance name + engine kwards, as well as ability to not cache the engines (but defaults to caching)

tests/databricks_ai_bridge/test_lakebase.py

bbqiu · 2026-02-06T23:24:07Z

src/databricks_ai_bridge/lakebase.py

@@ -6,11 +6,15 @@
 import uuid


(ok for followup PR) we should probably think about separating this file into a few separate ones since it's getting quite long

bbqiu · 2026-02-06T23:24:59Z

integrations/openai/tests/unit_tests/test_clients.py

                    model="databricks-claude-3-7-sonnet",
                    messages=[{"role": "user", "content": "hi"}],
-                    tools=tools,
+                    tools=cast(Any, tools),


did we delete this change from the diff? i think we still need it cc @fanzeyi who ran into a bug that that was fixing earlier

context: #274 (comment)

added this back here and included unit tests to make sure the non-list inputs are handled gracefully!

bbqiu · 2026-02-06T23:25:32Z

integrations/openai/pyproject.toml

 ]

+[project.optional-dependencies]
+memory = [


can we update the CI job for this memory extra too

good catch!

integrations/openai/src/databricks_openai/agents/__init__.py

bbqiu

overall LGTM, please address all comments and this looks ready to merge!

…QLAlchemy Tests

… includes it

integrations/openai/src/databricks_openai/agents/session.py

bbqiu

lgtm, feel free to merge after addressing comments!

…ched_engine

OpenAI AsyncDatabricksSession Stateful Agent Example using session protocol class implemented in databricks/databricks-ai-bridge#316 * openai agents stateful example * add session id to outputs * update example w/ asyncdatabrickssession * package release agent updates * use uuid7 for example * pr review updates * add openai agent memory skill * add to openai templates sync script * run python sync skills * databricks yml and use chatcontext convo id * sanitize mcp tool output items https://github.com/databricks/app-templates/pull/119/changes * deduplicate input logic * update sanitize mcp handler to be more defensive * rename from agent-openai-agents-sdk-stateful-memory to agent-openai-agents-sdk-short-term-memory

jennsun added 3 commits February 3, 2026 13:32

memorysession subclassing SQLAlchemySession

3351a09

fix type check issues

1e7b165

update tests

9ac71d9

jennsun changed the title ~~memorysession subclassing SQLAlchemySession~~ OpenAI: Add MemorySession class to support Session Protocol for Stateful Conversation Management Feb 4, 2026

jennsun marked this pull request as ready for review February 4, 2026 00:13

jennsun mentioned this pull request Feb 4, 2026

OpenAI SDK Add MemorySession that follows the Session protocol #274

Closed

jennsun commented Feb 4, 2026

View reviewed changes

integrations/openai/src/databricks_openai/agents/session.py Show resolved Hide resolved

jennsun requested a review from bbqiu February 4, 2026 01:49

enable connection pooling using SQLAlchemy's default QueuePool

1433da0

jennsun commented Feb 4, 2026

View reviewed changes

bbqiu reviewed Feb 4, 2026

View reviewed changes

jennsun added 4 commits February 4, 2026 16:33

update to asyncdatabrickssession, class-level cache for engine, renam…

d1afb16

…e LakebasePoolBase to LakebaseBase

update dependnecies needed in openai package

d7c4739

format

5aa6d1a

integrations/openai typecheck

bbb0b0e

jennsun mentioned this pull request Feb 5, 2026

OpenAI SDK Stateful Agent Example databricks/app-templates#109

Merged

jennsun added 2 commits February 4, 2026 17:20

Linting / typechecking for integrations/openai

e1ff83c

add dependency to test dev

b9b2cf8

jennsun requested a review from bbqiu February 5, 2026 19:49

bbqiu reviewed Feb 5, 2026

View reviewed changes

jennsun added 2 commits February 5, 2026 17:51

refactor AsyncDatabricksSession to use AsyncLakebaseSqlAlchemy sessio…

04b7e46

…n that returns engines to manage connections to db

ruff lint

00280db

jennsun requested a review from bbqiu February 6, 2026 02:04

jennsun commented Feb 6, 2026

View reviewed changes

bbqiu reviewed Feb 6, 2026

View reviewed changes

tests/databricks_ai_bridge/test_lakebase.py Show resolved Hide resolved

bbqiu reviewed Feb 6, 2026

View reviewed changes

integrations/openai/src/databricks_openai/agents/__init__.py Outdated Show resolved Hide resolved

bbqiu reviewed Feb 6, 2026

View reviewed changes

jennsun added 2 commits February 9, 2026 11:56

openai integration cleanup - ensure imports available, AsyncLakebaseS…

ec0c1c9

…QLAlchemy Tests

lint

be3a9cb

jennsun requested review from bbqiu and fanzeyi February 9, 2026 21:11

remove redundant sqlalchemy import since ai-bridge[memory] dependency…

6135b57

… includes it

bbqiu reviewed Feb 9, 2026

View reviewed changes

integrations/openai/src/databricks_openai/agents/session.py Outdated Show resolved Hide resolved

bbqiu approved these changes Feb 9, 2026

View reviewed changes

jennsun added 2 commits February 9, 2026 16:27

cache using both instances name + engine kwargs, can tune with use_ca…

0ce9d42

…ched_engine

lint fix

62bd09f

jennsun changed the title ~~OpenAI: Add MemorySession class to support Session Protocol for Stateful Conversation Management~~ db-openai: Add AsyncDatabricksSession class to support Session Protocol for Stateful Conversation Management via SQLAlchemy engine Feb 10, 2026

jennsun merged commit 4fbc9ea into main Feb 10, 2026
38 checks passed

jennsun deleted the openai-lakebasesqlalchemysession branch February 10, 2026 00:48

smurching mentioned this pull request Feb 10, 2026

Release v0.14.0: Version bump for databricks-ai-bridge, databricks-langchain, databricks-openai, databricks-mcp #319

Merged

		DEFAULT_DATABASE = "databricks_postgres"


		class _LakebaseCredentials(_LakebasePoolBase):

		return token


		class AsyncDatabricksSession(SQLAlchemySession):

Conversation

jennsun commented Feb 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jennsun Feb 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jennsun Feb 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bbqiu Feb 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

bbqiu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

bbqiu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jennsun commented Feb 3, 2026 •

edited

Loading

jennsun Feb 6, 2026 •

edited

Loading

jennsun Feb 6, 2026 •

edited

Loading

bbqiu Feb 6, 2026 •

edited

Loading