Skip to content

Add memory store kwargs to DatabricksOpenAI conversations.create#451

Open
jennsun wants to merge 3 commits into
mainfrom
conversations-memory-kwargs
Open

Add memory store kwargs to DatabricksOpenAI conversations.create#451
jennsun wants to merge 3 commits into
mainfrom
conversations-memory-kwargs

Conversation

@jennsun

@jennsun jennsun commented Jul 1, 2026

Copy link
Copy Markdown
Contributor

Summary

Databricks agent memory attaches conversations to a Unity Catalog memory store via memory_store and scope fields on the conversations endpoint. These aren't part of the OpenAI API, so callers previously had to pass them through extra_body:

client = DatabricksOpenAI(use_ai_gateway=True)
conversation = client.conversations.create(
    extra_body={
        "memory_store": {"name": "main.default.support_agent_memory"},
        "scope": {"kind": "user", "value": user_id},
    },
)

This PR adds DatabricksConversations / AsyncDatabricksConversations resources (following the existing DatabricksResponses pattern) that accept first-class kwargs and translate them into extra_body before delegating to the stock OpenAI SDK resource:

conversation = client.conversations.create(
    store="main.default.support_agent_memory",
    scope=user_client,  # SDK extracts the user id; pass a string for external (non-Databricks) users
)
  • scope accepts either a WorkspaceClient (the authenticated user's id is extracted via current_user.me().user_name) or a plain string.
  • scope_kind is exposed for non-user scope kinds (defaults to "user").
  • All standard conversations.create arguments (items, metadata, extra_body, extra_headers, ...) pass through unchanged, so existing extra_body usage keeps working; explicit kwargs win over the same keys in a caller-provided extra_body.
  • client.conversations.items.* and the other conversation methods are inherited untouched from the OpenAI SDK.

No databricks-sdk changes required — the Databricks conversation endpoints are OpenAI-shaped, so this is purely client-side sugar over the gateway's OpenAI-compatible surface.

Test plan

  • New unit tests in tests/unit_tests/test_clients.py covering: kwarg-to-extra_body translation, WorkspaceClient scope extraction, scope kind override, merging with caller-provided extra_body (kwargs take precedence), passthrough when no memory kwargs are given, backward-compatible raw extra_body usage, and the async variants.
  • python -m pytest tests/unit_tests/ — 222 passed.
  • ruff check / ruff format --check clean.

🤖 Generated with Claude Code

jennsun and others added 3 commits July 1, 2026 14:56
Databricks agent memory attaches conversations to a Unity Catalog memory
store via memory_store and scope fields on the conversations endpoint.
These aren't part of the OpenAI API, so callers had to pass them through
extra_body:

    client.conversations.create(
        extra_body={
            "memory_store": {"name": "main.default.support_agent_memory"},
            "scope": {"kind": "user", "value": user_id},
        },
    )

Add DatabricksConversations / AsyncDatabricksConversations resources that
accept first-class kwargs and translate them into extra_body before
delegating to the stock OpenAI SDK resource:

    client.conversations.create(
        memory_store_name="main.default.support_agent_memory",
        memory_scope=user_id,
    )

All standard conversations.create arguments (items, metadata, extra_body)
pass through unchanged, so existing extra_body usage keeps working.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ent scope

conversations.create(store=..., scope=...) instead of memory_store_name /
memory_scope. scope also accepts a WorkspaceClient, from which the
authenticated user's id is extracted; pass a string for external
(non-Databricks) users.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Annotate the conversations properties with the Databricks subclass types so
the store/scope kwargs typecheck (and autocomplete) at call sites, and handle
the str | None type of current_user.me().user_name explicitly.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@jennsun jennsun requested a review from annzhang-db July 2, 2026 21:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant