Add memory store kwargs to DatabricksOpenAI conversations.create#451
Open
jennsun wants to merge 3 commits into
Open
Add memory store kwargs to DatabricksOpenAI conversations.create#451jennsun wants to merge 3 commits into
jennsun wants to merge 3 commits into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Databricks agent memory attaches conversations to a Unity Catalog memory store via
memory_storeandscopefields on the conversations endpoint. These aren't part of the OpenAI API, so callers previously had to pass them throughextra_body:This PR adds
DatabricksConversations/AsyncDatabricksConversationsresources (following the existingDatabricksResponsespattern) that accept first-class kwargs and translate them intoextra_bodybefore delegating to the stock OpenAI SDK resource:scopeaccepts either aWorkspaceClient(the authenticated user's id is extracted viacurrent_user.me().user_name) or a plain string.scope_kindis exposed for non-userscope kinds (defaults to"user").conversations.createarguments (items,metadata,extra_body,extra_headers, ...) pass through unchanged, so existingextra_bodyusage keeps working; explicit kwargs win over the same keys in a caller-providedextra_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
tests/unit_tests/test_clients.pycovering: kwarg-to-extra_body translation,WorkspaceClientscope extraction, scope kind override, merging with caller-providedextra_body(kwargs take precedence), passthrough when no memory kwargs are given, backward-compatible rawextra_bodyusage, and the async variants.python -m pytest tests/unit_tests/— 222 passed.ruff check/ruff format --checkclean.🤖 Generated with Claude Code