OpenAI SDK Stateful Agent Example by jennsun · Pull Request #109 · databricks/app-templates

jennsun · 2026-02-04T00:09:10Z

OpenAI AsyncDatabricksSession Stateful Agent Example
using session protocol class implemented in databricks/databricks-ai-bridge#316

agent-openai-agents-sdk-stateful-memory/pyproject.toml

agent-openai-agents-sdk-short-term-memory/agent_server/agent.py

…ol for Stateful Conversation Management via SQLAlchemy engine (#316) Adds AsyncDatabricksSession, a session storage implementation for the OpenAI Agents SDK that persists conversation history to Databricks Lakebase. This class subclasses [OpenAI's SQLAlchemySession](https://openai.github.io/openai-agents-python/sessions/sqlalchemy_session/) [original code](https://github.com/openai/openai-agents-python/blob/main/src/agents/extensions/memory/sqlalchemy_session.py) 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](databricks/app-templates#109) sample app: https://eng-ml-agent-platform.staging.cloud.databricks.com/apps/j-openai-stateful?o=2850744067564480 <img width="1014" height="886" alt="image" src="https://github.com/user-attachments/assets/b53a3567-7cc2-4967-b8b0-ace481940f2a" />

jennsun · 2026-02-11T02:12:12Z

agent-openai-agents-sdk-short-term-memory/pyproject.toml

+    "openai-agents>=0.4.1",
+    "python-dotenv",
+    "uuid-utils>=0.10.0",
+    "databricks-openai[memory]>=0.11.1",


new package

dhruv0811 · 2026-02-11T19:26:06Z

agent-openai-agents-sdk-stateful-memory/agent_server/agent.py

+)
+
+# Lakebase instance name for persistent session storage
+LAKEBASE_INSTANCE_NAME = os.environ.get("LAKEBASE_INSTANCE_NAME", "lakebase")


Do we want to add any validation here like we did in the langgraph agent.py for lakebase instance name?

app-templates/agent-langgraph-long-term-memory/agent_server/agent.py

Line 48 in bc598b8

if not _LAKEBASE_INSTANCE_NAME_RAW:

dhruv0811

Looks good overall!

Do we want to add this template to the sync script as well? I can see that the skills are synced, maybe the changes to the sync script didn't get added to this PR?

agent-openai-agents-sdk-stateful-memory/databricks.yml

agent-openai-agents-sdk-stateful-memory/app.yaml

agent-openai-agents-sdk-stateful-memory/README.md

bbqiu · 2026-02-12T18:28:52Z

agent-openai-agents-sdk-stateful-memory/databricks.yml

+      name: /Users/${workspace.current_user.userName}/${bundle.name}-${bundle.target}
+
+  apps:
+    agent_openai_agents_sdk_stateful_memory:


should we also have a lakebase instance in here?

additionally, can we also add a databricks.yml for the long and short term langgraph examples?

bbqiu · 2026-02-12T18:32:51Z

agent-openai-agents-sdk-short-term-memory/agent_server/agent.py

+    # user_workspace_client = get_user_workspace_client()
+
+    # Create session for persistent conversation history with your Databricks Lakebase instance
+    session = AsyncDatabricksSession(


can we clarify that this example is short term memory only?

we can extend this in the future to include context compression like https://openai.github.io/openai-agents-python/sessions/#openai-responses-compaction-sessions

bbqiu · 2026-02-12T18:34:04Z

agent-openai-agents-sdk-short-term-memory/README.md

+This template uses OpenAI Agents SDK [Sessions](https://openai.github.io/openai-agents-python/sessions/) to automatically manage conversation history. Sessions store conversation history for a specific session, allowing agents to maintain context without requiring explicit manual memory management. This is particularly useful for building chat applications or multi-turn conversations where you want the agent to remember previous interactions.
+
+How it works:
+- **Before each run**: The session retrieves prior conversation history and prepends it to the input


is the convo history that is retrieved deduped with the items that are coming in? bc clients (ours included) will send back in the whole convo

I see, looks like deduplication only happens for tool calls but not duplicate conversation messages, this is called during Runner.run called here in runner: https://github.com/openai/openai-agents-python/blob/main/src/agents/run.py#L484

which uses the deduplicating logic here: https://github.com/openai/openai-agents-python/blob/main/src/agents/run_internal/items.py#L167

should we be updating our message to only pass the latest message to Runner.run then? like

latest_message = request.input[-1].model_dump() result = await Runner.run(agent, [latest_message], session=session)

instead of what we currently do:

messages = [i.model_dump() for i in request.input] result = await Runner.run(agent, messages, session=session)

this would also match up with the openai example where it passes a single message instead of the history when you pass in session since the runner will get the history from the session:
https://openai.github.io/openai-agents-python/sessions/#sqlalchemy-sessions

experiment trace from a session: do you know if mlflow tracing does any deduping work or is this trace represented of the request payload?

added deduplication logic in utils - session's messages is always 1 less than the total client's messages since client message includes all previous turns + new message

agent-openai-agents-sdk-short-term-memory/agent_server/agent.py

…lates/pull/119/changes

bbqiu

lgtm with one nit: let's rename this to ...sdk-short-term-memory to match the langgraph examples

…gents-sdk-short-term-memory

jennsun commented Feb 4, 2026

View reviewed changes

agent-openai-agents-sdk-stateful-memory/pyproject.toml Outdated Show resolved Hide resolved

jennsun mentioned this pull request Feb 4, 2026

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

Merged

bbqiu reviewed Feb 6, 2026

View reviewed changes

agent-openai-agents-sdk-short-term-memory/agent_server/agent.py Show resolved Hide resolved

jennsun requested a review from bbqiu February 11, 2026 02:08

jennsun changed the title ~~OpenAI MemorySession Stateful Agent Example~~ OpenAI Stateful Agent Example Feb 11, 2026

jennsun commented Feb 11, 2026

View reviewed changes

jennsun changed the title ~~OpenAI Stateful Agent Example~~ OpenAI SDK Stateful Agent Example Feb 11, 2026

jennsun requested a review from dhruv0811 February 11, 2026 18:34

dhruv0811 reviewed Feb 11, 2026

View reviewed changes

agent-openai-agents-sdk-stateful-memory/databricks.yml Outdated Show resolved Hide resolved

agent-openai-agents-sdk-stateful-memory/app.yaml Show resolved Hide resolved

agent-openai-agents-sdk-stateful-memory/README.md Outdated Show resolved Hide resolved

jennsun requested a review from dhruv0811 February 11, 2026 21:36

jennsun force-pushed the openai-memorysession-example branch from c77d036 to 80a7ff6 Compare February 11, 2026 23:29

jennsun added 9 commits February 11, 2026 15:30

openai agents stateful example

917f98b

add session id to outputs

48745b9

update example w/ asyncdatabrickssession

62b2010

package release agent updates

d1619bf

use uuid7 for example

d8da483

pr review updates

013443c

add openai agent memory skill

ebbeb19

add to openai templates sync script

283e1a0

run python sync skills

ac6c759

jennsun force-pushed the openai-memorysession-example branch from 80a7ff6 to ac6c759 Compare February 11, 2026 23:31

Merge branch 'main' into openai-memorysession-example

648e15d

bbqiu reviewed Feb 12, 2026

View reviewed changes

agent-openai-agents-sdk-short-term-memory/agent_server/agent.py Show resolved Hide resolved

jennsun added 2 commits February 12, 2026 15:39

databricks yml and use chatcontext convo id

dfcb42f

sanitize mcp tool output items https://github.com/databricks/app-temp…

b33b628

…lates/pull/119/changes

jennsun added 2 commits February 12, 2026 17:28

deduplicate input logic

4e3f787

update sanitize mcp handler to be more defensive

23be8cb

bbqiu approved these changes Feb 13, 2026

View reviewed changes

rename from agent-openai-agents-sdk-stateful-memory to agent-openai-a…

438a38e

…gents-sdk-short-term-memory

jennsun merged commit efc28d1 into databricks:main Feb 13, 2026

Conversation

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

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dhruv0811 left a comment

Choose a reason for hiding this comment

Uh oh!

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!

Uh oh!

bbqiu left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jennsun commented Feb 4, 2026 •

edited

Loading