Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
125 changes: 38 additions & 87 deletions .claude/skills/agent-openai-memory/SKILL.md
Original file line number Diff line number Diff line change
@@ -1,32 +1,24 @@
---
name: agent-openai-memory
description: "Add memory capabilities to your agent. Use when: (1) User asks about 'memory', 'state', 'remember', 'conversation history', (2) Want to persist conversations or user preferences, (3) Adding checkpointing or long-term storage."
description: "Add session-based memory to OpenAI Agents SDK agent using AsyncDatabricksSession and Lakebase. Use when: (1) User asks about 'memory', 'state', 'remember', 'conversation history', (2) Want to persist conversations or user preferences, (3) Adding session-based checkpointing."
---

# Stateful Memory with OpenAI Agents SDK Sessions

This template uses OpenAI Agents SDK [Sessions](https://openai.github.io/openai-agents-python/sessions/) with `AsyncDatabricksSession` to persist conversation history to a Databricks Lakebase instance.
Uses `AsyncDatabricksSession` to persist conversation history to Lakebase, enabling multi-turn interactions where the agent remembers prior messages within a session.

## 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.
## Prerequisites

## Key Concepts
1. **Dependency**: `databricks-openai[memory]` in `pyproject.toml` (already included in memory templates)
2. **Lakebase instance**: See **lakebase-setup** skill for creating and configuring one
3. **Environment variable**: Set `LAKEBASE_INSTANCE_NAME` in `.env`:
```bash
LAKEBASE_INSTANCE_NAME=<your-lakebase-instance-name>
```

| Concept | Description |
|---------|-------------|
| **Session** | Stores conversation history for a specific `session_id` |
| **`session_id`** | Unique identifier linking requests to the same conversation |
| **`AsyncDatabricksSession`** | Session implementation backed by Databricks Lakebase |
| **`LAKEBASE_INSTANCE_NAME`** | Environment variable specifying the Lakebase instance |
---

## How This Template Uses Sessions
## Implementation

### Session Creation (`agent_server/agent.py`)

Expand All @@ -43,8 +35,6 @@ 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:

```python
def get_session_id(request: ResponsesAgentRequest) -> str:
if hasattr(request, "custom_inputs") and request.custom_inputs:
Expand All @@ -55,87 +45,57 @@ def get_session_id(request: ResponsesAgentRequest) -> str:

### 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:

```python
_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:
```bash
LAKEBASE_INSTANCE_NAME=<your-lakebase-instance-name>
```

---

## Configuration Files
## Configuration

### databricks.yml (Lakebase Resource)

Add the Lakebase database resource to your app:

```yaml
resources:
apps:
agent_openai_advanced:
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: 'databricks_postgres'
permission: 'CAN_CONNECT_AND_CREATE'
- name: 'database'
database:
instance_name: '<your-lakebase-instance-name>'
database_name: 'databricks_postgres'
permission: 'CAN_CONNECT_AND_CREATE'
```

### databricks.yml config block (Environment Variables)

The `LAKEBASE_INSTANCE_NAME` env var is resolved from the database resource at deploy time. Add to your app's `config.env` in `databricks.yml`:

```yaml
config:
env:
- name: LAKEBASE_INSTANCE_NAME
value_from: "database"
config:
env:
- name: LAKEBASE_INSTANCE_NAME
value_from: "database"
```

### .env (Local Development)
---

```bash
LAKEBASE_INSTANCE_NAME=<your-lakebase-instance-name>
```
## Testing

---
### Verify Lakebase Connectivity

## Testing Sessions
```bash
databricks lakebase instances get <instance-name> --profile <profile>
```

### Test Multi-Turn Conversation Locally
### Test Multi-Turn Conversation

```bash
# Start the server
uv run start-app

# First message - starts a new session
# 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
# Second message -- continues the same session (should remember SF)
curl -X POST http://localhost:8000/invocations \
-H "Content-Type: application/json" \
-d '{
Expand All @@ -144,28 +104,19 @@ curl -X POST http://localhost:8000/invocations \
}'
```

### Test Streaming

```bash
curl -X POST http://localhost:8000/invocations \
-H "Content-Type: application/json" \
-d '{
"input": [{"role": "user", "content": "Hello!"}],
"stream": true
}'
```
If the agent responds with "SF" or "San Francisco", session memory is working.

---

## Troubleshooting

| Issue | Cause | Solution |
|-------|-------|----------|
| **"LAKEBASE_INSTANCE_NAME environment variable is required"** | Missing env var | Set `LAKEBASE_INSTANCE_NAME` in `.env` |
| **SSL connection closed unexpectedly** | Network/instance issue | Verify Lakebase instance is running: `databricks lakebase instances get <name>` |
| **Agent doesn't remember previous messages** | Different session_id | Pass the same `session_id` via `custom_inputs` across requests |
| **"Unable to resolve hostname"** | Hostname doesn't match any instance | Verify the hostname or use the instance name directly |
| **Permission denied** | Missing Lakebase access | Add `database` resource to `databricks.yml` with `CAN_CONNECT_AND_CREATE` |
| Issue | Solution |
|-------|----------|
| "LAKEBASE_INSTANCE_NAME environment variable is required" | Set `LAKEBASE_INSTANCE_NAME` in `.env` |
| SSL connection closed unexpectedly | Verify instance is running: `databricks lakebase instances get <name>` |
| Agent doesn't remember previous messages | Pass same `session_id` via `custom_inputs` across requests |
| "Unable to resolve hostname" | Use instance name directly instead of hostname |
| Permission denied | Add `database` resource to `databricks.yml` with `CAN_CONNECT_AND_CREATE` |

---

Expand Down
Loading