Agent Openai Memory

Stateful Memory with OpenAI Agents SDK Sessions

This template uses OpenAI Agents SDK Sessions with AsyncDatabricksSession to persist conversation history to a Databricks Lakebase instance.

How Sessions Work

Sessions automatically manage conversation history for multi-turn interactions:

1. Before each run: The session retrieves prior conversation history and prepends it to input
2. During the run: New items (user messages, responses, tool calls) are generated
3. After each run: All new items are automatically stored in the session

This eliminates the need to manually manage conversation state between runs.

Key Concepts

How This Template Uses Sessions

Session Creation (agent_server/agent.py)

from databricks_openai.agents import AsyncDatabricksSession

session = AsyncDatabricksSession(
    session_id=get_session_id(request),
    instance_name=LAKEBASE_INSTANCE_NAME,
)

result = await Runner.run(agent, messages, session=session)

Session ID Extraction (agent_server/agent.py)

The session_id is extracted from custom_inputs or auto-generated:

def get_session_id(request: ResponsesAgentRequest) -> str:
    if hasattr(request, "custom_inputs") and request.custom_inputs:
        if "session_id" in request.custom_inputs:
            return request.custom_inputs["session_id"]
    return str(uuid7())

Lakebase Instance Resolution (agent_server/utils.py)

The LAKEBASE_INSTANCE_NAME env var can be either an instance name or a hostname. The resolve_lakebase_instance_name() function handles both cases:

_LAKEBASE_INSTANCE_NAME_RAW = os.environ.get("LAKEBASE_INSTANCE_NAME")
LAKEBASE_INSTANCE_NAME = resolve_lakebase_instance_name(_LAKEBASE_INSTANCE_NAME_RAW)

Prerequisites

1. Dependency: databricks-openai[memory] must be in pyproject.toml (already included)

2. Lakebase instance: You need a Databricks Lakebase instance. See the lakebase-setup skill for creating and configuring one.

3. Environment variable: Set LAKEBASE_INSTANCE_NAME in your .env file:

   LAKEBASE_INSTANCE_NAME=<your-lakebase-instance-name>
   

Configuration Files

databricks.yml (Lakebase Resource)

Add the Lakebase database resource to your app:

resources:
  apps:
    agent_openai_agents_sdk_short_term_memory:
      name: "your-app-name"
      source_code_path: ./

      resources:
        # ... other resources (experiment, etc.) ...

        # Lakebase instance for session storage
        - name: 'database'
          database:
            instance_name: '<your-lakebase-instance-name>'
            database_name: 'postgres'
            permission: 'CAN_CONNECT_AND_CREATE'

app.yaml (Environment Variables)

The LAKEBASE_INSTANCE_NAME env var is resolved from the database resource at deploy time:

env:
  - name: LAKEBASE_INSTANCE_NAME
    valueFrom: "database"

.env (Local Development)

LAKEBASE_INSTANCE_NAME=<your-lakebase-instance-name>

Testing Sessions

Test Multi-Turn Conversation Locally

# Start the server
uv run start-app

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

# Note the session_id from custom_outputs in the response

# Second message - continues the same session
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": "<session_id from previous response>"}
  }'

Test Streaming

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

Troubleshooting

Next Steps

• Configure Lakebase: see lakebase-setup skill
• Test locally: see run-locally skill
• Deploy: see deploy skill