From 09b6ce2bb8f73f623404211955f79a00449c83e1 Mon Sep 17 00:00:00 2001 From: Jenny Date: Wed, 1 Jul 2026 14:56:24 -0700 Subject: [PATCH 1/3] Add memory store kwargs to DatabricksOpenAI conversations.create 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 --- .../src/databricks_openai/utils/clients.py | 114 ++++++++++++++++++ .../openai/tests/unit_tests/test_clients.py | 110 +++++++++++++++++ 2 files changed, 224 insertions(+) diff --git a/integrations/openai/src/databricks_openai/utils/clients.py b/integrations/openai/src/databricks_openai/utils/clients.py index 840b6c97..acc720f2 100644 --- a/integrations/openai/src/databricks_openai/utils/clients.py +++ b/integrations/openai/src/databricks_openai/utils/clients.py @@ -6,6 +6,7 @@ from openai import APIConnectionError, APIStatusError, AsyncOpenAI, OpenAI from openai.resources.chat import AsyncChat, Chat from openai.resources.chat.completions import AsyncCompletions, Completions +from openai.resources.conversations import AsyncConversations, Conversations from openai.resources.responses import AsyncResponses, Responses from typing_extensions import override @@ -322,6 +323,61 @@ def create(self, **kwargs): return response +def _merge_memory_params_into_extra_body( + kwargs: dict, + memory_store_name: str | None, + memory_scope: str | None, + memory_scope_kind: str, +) -> dict: + """Translate memory kwargs into the Databricks-specific conversation body fields. + + The Databricks conversations endpoint accepts ``memory_store`` and ``scope`` fields + that aren't part of the OpenAI API, so they're sent via ``extra_body``. Explicit + kwargs take precedence over the same keys in a caller-provided ``extra_body``. + """ + if memory_store_name is None and memory_scope is None: + return kwargs + extra_body = dict(kwargs.pop("extra_body", None) or {}) + if memory_store_name is not None: + extra_body["memory_store"] = {"name": memory_store_name} + if memory_scope is not None: + extra_body["scope"] = {"kind": memory_scope_kind, "value": memory_scope} + kwargs["extra_body"] = extra_body + return kwargs + + +class DatabricksConversations(Conversations): + """Conversations resource with Databricks agent memory support. + + Adds ``memory_store_name`` and ``memory_scope`` convenience kwargs to ``create``, + which attach the conversation to a Unity Catalog memory store. See + https://docs.databricks.com/aws/en/generative-ai/agent-memory/ for details. + """ + + def create( + self, + *, + memory_store_name: str | None = None, + memory_scope: str | None = None, + memory_scope_kind: str = "user", + **kwargs, + ): + """Create a conversation, optionally backed by a Databricks memory store. + + Args: + memory_store_name: Full three-part Unity Catalog name of the memory store + (e.g. ``"main.default.support_agent_memory"``). + memory_scope: Scope value that partitions memory entries, typically a user id. + memory_scope_kind: Kind of the scope. Defaults to ``"user"``. + **kwargs: Standard OpenAI ``conversations.create`` arguments + (``items``, ``metadata``, ``extra_body``, etc.). + """ + kwargs = _merge_memory_params_into_extra_body( + kwargs, memory_store_name, memory_scope, memory_scope_kind + ) + return super().create(**kwargs) + + class DatabricksOpenAI(OpenAI): """OpenAI client authenticated with Databricks to query LLMs and agents hosted on Databricks. @@ -385,6 +441,13 @@ class DatabricksOpenAI(OpenAI): ... model="apps/my-agent", # Looks up app URL automatically ... input=[{"role": "user", "content": "Hello"}], ... ) + + Example - Create a conversation backed by a Databricks memory store: + >>> client = DatabricksOpenAI(use_ai_gateway=True) + >>> conversation = client.conversations.create( + ... memory_store_name="main.default.support_agent_memory", + ... memory_scope="user-123", + ... ) """ def __init__( @@ -432,6 +495,12 @@ def responses(self) -> Responses: self._databricks_responses = DatabricksResponses(self, self._workspace_client) return self._databricks_responses + @property + def conversations(self) -> Conversations: + if not hasattr(self, "_databricks_conversations"): + self._databricks_conversations = DatabricksConversations(self) + return self._databricks_conversations + class AsyncDatabricksCompletions(AsyncCompletions): """Async completions that conditionally strips 'strict' from tools for non-GPT models.""" @@ -490,6 +559,38 @@ async def create(self, **kwargs): return response +class AsyncDatabricksConversations(AsyncConversations): + """Async conversations resource with Databricks agent memory support. + + Adds ``memory_store_name`` and ``memory_scope`` convenience kwargs to ``create``, + which attach the conversation to a Unity Catalog memory store. See + https://docs.databricks.com/aws/en/generative-ai/agent-memory/ for details. + """ + + async def create( + self, + *, + memory_store_name: str | None = None, + memory_scope: str | None = None, + memory_scope_kind: str = "user", + **kwargs, + ): + """Create a conversation, optionally backed by a Databricks memory store. + + Args: + memory_store_name: Full three-part Unity Catalog name of the memory store + (e.g. ``"main.default.support_agent_memory"``). + memory_scope: Scope value that partitions memory entries, typically a user id. + memory_scope_kind: Kind of the scope. Defaults to ``"user"``. + **kwargs: Standard OpenAI ``conversations.create`` arguments + (``items``, ``metadata``, ``extra_body``, etc.). + """ + kwargs = _merge_memory_params_into_extra_body( + kwargs, memory_store_name, memory_scope, memory_scope_kind + ) + return await super().create(**kwargs) + + class AsyncDatabricksOpenAI(AsyncOpenAI): """Async OpenAI client authenticated with Databricks to query LLMs and agents hosted on Databricks. @@ -553,6 +654,13 @@ class AsyncDatabricksOpenAI(AsyncOpenAI): ... model="apps/my-agent", # Looks up app URL automatically ... input=[{"role": "user", "content": "Hello"}], ... ) + + Example - Create a conversation backed by a Databricks memory store: + >>> client = AsyncDatabricksOpenAI(use_ai_gateway=True) + >>> conversation = await client.conversations.create( + ... memory_store_name="main.default.support_agent_memory", + ... memory_scope="user-123", + ... ) """ def __init__( @@ -598,3 +706,9 @@ def responses(self) -> AsyncResponses: if not hasattr(self, "_databricks_responses"): self._databricks_responses = AsyncDatabricksResponses(self, self._workspace_client) return self._databricks_responses + + @property + def conversations(self) -> AsyncConversations: + if not hasattr(self, "_databricks_conversations"): + self._databricks_conversations = AsyncDatabricksConversations(self) + return self._databricks_conversations diff --git a/integrations/openai/tests/unit_tests/test_clients.py b/integrations/openai/tests/unit_tests/test_clients.py index aab52864..e084e8ff 100644 --- a/integrations/openai/tests/unit_tests/test_clients.py +++ b/integrations/openai/tests/unit_tests/test_clients.py @@ -9,6 +9,7 @@ from openai import APIConnectionError, APIStatusError, AsyncOpenAI, OpenAI from openai._types import NOT_GIVEN, Omit from openai.resources.chat.completions import AsyncCompletions, Completions +from openai.resources.conversations import AsyncConversations, Conversations from openai.resources.responses import AsyncResponses, Responses from databricks_openai import AsyncDatabricksOpenAI, DatabricksOpenAI @@ -806,6 +807,115 @@ def test_falls_back_to_no_token_when_empty_string(self): assert _get_openai_api_key() == "no-token" +class TestDatabricksConversationsMemory: + """Tests for memory store kwargs on conversations.create.""" + + def test_memory_kwargs_translated_to_extra_body(self, mock_workspace_client): + client = DatabricksOpenAI(workspace_client=mock_workspace_client, use_ai_gateway=True) + + with patch.object(Conversations, "create") as mock_create: + mock_create.return_value = MagicMock() + client.conversations.create( + memory_store_name="main.default.support_agent_memory", + memory_scope="user-123", + ) + + call_kwargs = mock_create.call_args.kwargs + assert call_kwargs["extra_body"] == { + "memory_store": {"name": "main.default.support_agent_memory"}, + "scope": {"kind": "user", "value": "user-123"}, + } + assert "memory_store_name" not in call_kwargs + assert "memory_scope" not in call_kwargs + + def test_memory_scope_kind_override(self, mock_workspace_client): + client = DatabricksOpenAI(workspace_client=mock_workspace_client, use_ai_gateway=True) + + with patch.object(Conversations, "create") as mock_create: + mock_create.return_value = MagicMock() + client.conversations.create( + memory_store_name="main.default.support_agent_memory", + memory_scope="acct-42", + memory_scope_kind="account", + ) + + extra_body = mock_create.call_args.kwargs["extra_body"] + assert extra_body["scope"] == {"kind": "account", "value": "acct-42"} + + def test_memory_kwargs_merge_with_existing_extra_body(self, mock_workspace_client): + client = DatabricksOpenAI(workspace_client=mock_workspace_client, use_ai_gateway=True) + + with patch.object(Conversations, "create") as mock_create: + mock_create.return_value = MagicMock() + client.conversations.create( + memory_store_name="main.default.support_agent_memory", + memory_scope="user-123", + extra_body={"custom_field": "value", "scope": {"kind": "user", "value": "old"}}, + ) + + extra_body = mock_create.call_args.kwargs["extra_body"] + assert extra_body["custom_field"] == "value" + # Explicit kwargs take precedence over extra_body entries + assert extra_body["scope"] == {"kind": "user", "value": "user-123"} + + def test_create_without_memory_kwargs_passes_through(self, mock_workspace_client): + client = DatabricksOpenAI(workspace_client=mock_workspace_client, use_ai_gateway=True) + + with patch.object(Conversations, "create") as mock_create: + mock_create.return_value = MagicMock() + client.conversations.create(metadata={"topic": "support"}) + + call_kwargs = mock_create.call_args.kwargs + assert call_kwargs == {"metadata": {"topic": "support"}} + + def test_raw_extra_body_still_works(self, mock_workspace_client): + """The original extra_body-based usage keeps working unchanged.""" + client = DatabricksOpenAI(workspace_client=mock_workspace_client, use_ai_gateway=True) + + with patch.object(Conversations, "create") as mock_create: + mock_create.return_value = MagicMock() + client.conversations.create( + extra_body={ + "memory_store": {"name": "main.default.support_agent_memory"}, + "scope": {"kind": "user", "value": "user-123"}, + }, + ) + + extra_body = mock_create.call_args.kwargs["extra_body"] + assert extra_body["memory_store"] == {"name": "main.default.support_agent_memory"} + assert extra_body["scope"] == {"kind": "user", "value": "user-123"} + + def test_conversations_items_resource_still_accessible(self, mock_workspace_client): + client = DatabricksOpenAI(workspace_client=mock_workspace_client, use_ai_gateway=True) + assert client.conversations.items is not None + + @pytest.mark.asyncio + async def test_async_memory_kwargs_translated_to_extra_body(self, mock_workspace_client): + client = AsyncDatabricksOpenAI(workspace_client=mock_workspace_client, use_ai_gateway=True) + + with patch.object(AsyncConversations, "create", new_callable=AsyncMock) as mock_create: + await client.conversations.create( + memory_store_name="main.default.support_agent_memory", + memory_scope="user-123", + ) + + call_kwargs = mock_create.call_args.kwargs + assert call_kwargs["extra_body"] == { + "memory_store": {"name": "main.default.support_agent_memory"}, + "scope": {"kind": "user", "value": "user-123"}, + } + + @pytest.mark.asyncio + async def test_async_create_without_memory_kwargs_passes_through(self, mock_workspace_client): + client = AsyncDatabricksOpenAI(workspace_client=mock_workspace_client, use_ai_gateway=True) + + with patch.object(AsyncConversations, "create", new_callable=AsyncMock) as mock_create: + await client.conversations.create(metadata={"topic": "support"}) + + call_kwargs = mock_create.call_args.kwargs + assert call_kwargs == {"metadata": {"topic": "support"}} + + class TestDatabricksOpenAIWithGateway: """Tests for AI Gateway routing in DatabricksOpenAI and AsyncDatabricksOpenAI.""" From ef7df3e935cdf304fddc4441245ac38485e237f7 Mon Sep 17 00:00:00 2001 From: Jenny Date: Wed, 1 Jul 2026 15:04:51 -0700 Subject: [PATCH 2/3] Rename conversation memory kwargs to store/scope; accept WorkspaceClient 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 --- .../src/databricks_openai/utils/clients.py | 75 ++++++++++--------- .../openai/tests/unit_tests/test_clients.py | 39 +++++++--- 2 files changed, 68 insertions(+), 46 deletions(-) diff --git a/integrations/openai/src/databricks_openai/utils/clients.py b/integrations/openai/src/databricks_openai/utils/clients.py index acc720f2..7dd31d03 100644 --- a/integrations/openai/src/databricks_openai/utils/clients.py +++ b/integrations/openai/src/databricks_openai/utils/clients.py @@ -323,11 +323,18 @@ def create(self, **kwargs): return response +def _resolve_scope_value(scope: "str | WorkspaceClient") -> str: + """Resolve a scope value, extracting the authenticated user's id from a WorkspaceClient.""" + if isinstance(scope, WorkspaceClient): + return scope.current_user.me().user_name + return scope + + def _merge_memory_params_into_extra_body( kwargs: dict, - memory_store_name: str | None, - memory_scope: str | None, - memory_scope_kind: str, + store: str | None, + scope: "str | WorkspaceClient | None", + scope_kind: str, ) -> dict: """Translate memory kwargs into the Databricks-specific conversation body fields. @@ -335,13 +342,13 @@ def _merge_memory_params_into_extra_body( that aren't part of the OpenAI API, so they're sent via ``extra_body``. Explicit kwargs take precedence over the same keys in a caller-provided ``extra_body``. """ - if memory_store_name is None and memory_scope is None: + if store is None and scope is None: return kwargs extra_body = dict(kwargs.pop("extra_body", None) or {}) - if memory_store_name is not None: - extra_body["memory_store"] = {"name": memory_store_name} - if memory_scope is not None: - extra_body["scope"] = {"kind": memory_scope_kind, "value": memory_scope} + if store is not None: + extra_body["memory_store"] = {"name": store} + if scope is not None: + extra_body["scope"] = {"kind": scope_kind, "value": _resolve_scope_value(scope)} kwargs["extra_body"] = extra_body return kwargs @@ -349,32 +356,32 @@ def _merge_memory_params_into_extra_body( class DatabricksConversations(Conversations): """Conversations resource with Databricks agent memory support. - Adds ``memory_store_name`` and ``memory_scope`` convenience kwargs to ``create``, - which attach the conversation to a Unity Catalog memory store. See + Adds ``store`` and ``scope`` convenience kwargs to ``create``, which attach the + conversation to a Unity Catalog memory store. See https://docs.databricks.com/aws/en/generative-ai/agent-memory/ for details. """ def create( self, *, - memory_store_name: str | None = None, - memory_scope: str | None = None, - memory_scope_kind: str = "user", + store: str | None = None, + scope: "str | WorkspaceClient | None" = None, + scope_kind: str = "user", **kwargs, ): """Create a conversation, optionally backed by a Databricks memory store. Args: - memory_store_name: Full three-part Unity Catalog name of the memory store + store: Full three-part Unity Catalog name of the memory store (e.g. ``"main.default.support_agent_memory"``). - memory_scope: Scope value that partitions memory entries, typically a user id. - memory_scope_kind: Kind of the scope. Defaults to ``"user"``. + scope: Scope value that partitions memory entries. Pass a ``WorkspaceClient`` + to use its authenticated user's identity, or a string user id for + external (non-Databricks) users. + scope_kind: Kind of the scope. Defaults to ``"user"``. **kwargs: Standard OpenAI ``conversations.create`` arguments (``items``, ``metadata``, ``extra_body``, etc.). """ - kwargs = _merge_memory_params_into_extra_body( - kwargs, memory_store_name, memory_scope, memory_scope_kind - ) + kwargs = _merge_memory_params_into_extra_body(kwargs, store, scope, scope_kind) return super().create(**kwargs) @@ -445,8 +452,8 @@ class DatabricksOpenAI(OpenAI): Example - Create a conversation backed by a Databricks memory store: >>> client = DatabricksOpenAI(use_ai_gateway=True) >>> conversation = client.conversations.create( - ... memory_store_name="main.default.support_agent_memory", - ... memory_scope="user-123", + ... store="main.default.support_agent_memory", + ... scope=user_client, # extracts the user id; pass a string for external users ... ) """ @@ -562,32 +569,32 @@ async def create(self, **kwargs): class AsyncDatabricksConversations(AsyncConversations): """Async conversations resource with Databricks agent memory support. - Adds ``memory_store_name`` and ``memory_scope`` convenience kwargs to ``create``, - which attach the conversation to a Unity Catalog memory store. See + Adds ``store`` and ``scope`` convenience kwargs to ``create``, which attach the + conversation to a Unity Catalog memory store. See https://docs.databricks.com/aws/en/generative-ai/agent-memory/ for details. """ async def create( self, *, - memory_store_name: str | None = None, - memory_scope: str | None = None, - memory_scope_kind: str = "user", + store: str | None = None, + scope: "str | WorkspaceClient | None" = None, + scope_kind: str = "user", **kwargs, ): """Create a conversation, optionally backed by a Databricks memory store. Args: - memory_store_name: Full three-part Unity Catalog name of the memory store + store: Full three-part Unity Catalog name of the memory store (e.g. ``"main.default.support_agent_memory"``). - memory_scope: Scope value that partitions memory entries, typically a user id. - memory_scope_kind: Kind of the scope. Defaults to ``"user"``. + scope: Scope value that partitions memory entries. Pass a ``WorkspaceClient`` + to use its authenticated user's identity, or a string user id for + external (non-Databricks) users. + scope_kind: Kind of the scope. Defaults to ``"user"``. **kwargs: Standard OpenAI ``conversations.create`` arguments (``items``, ``metadata``, ``extra_body``, etc.). """ - kwargs = _merge_memory_params_into_extra_body( - kwargs, memory_store_name, memory_scope, memory_scope_kind - ) + kwargs = _merge_memory_params_into_extra_body(kwargs, store, scope, scope_kind) return await super().create(**kwargs) @@ -658,8 +665,8 @@ class AsyncDatabricksOpenAI(AsyncOpenAI): Example - Create a conversation backed by a Databricks memory store: >>> client = AsyncDatabricksOpenAI(use_ai_gateway=True) >>> conversation = await client.conversations.create( - ... memory_store_name="main.default.support_agent_memory", - ... memory_scope="user-123", + ... store="main.default.support_agent_memory", + ... scope=user_client, # extracts the user id; pass a string for external users ... ) """ diff --git a/integrations/openai/tests/unit_tests/test_clients.py b/integrations/openai/tests/unit_tests/test_clients.py index e084e8ff..5a090635 100644 --- a/integrations/openai/tests/unit_tests/test_clients.py +++ b/integrations/openai/tests/unit_tests/test_clients.py @@ -816,8 +816,8 @@ def test_memory_kwargs_translated_to_extra_body(self, mock_workspace_client): with patch.object(Conversations, "create") as mock_create: mock_create.return_value = MagicMock() client.conversations.create( - memory_store_name="main.default.support_agent_memory", - memory_scope="user-123", + store="main.default.support_agent_memory", + scope="user-123", ) call_kwargs = mock_create.call_args.kwargs @@ -825,18 +825,33 @@ def test_memory_kwargs_translated_to_extra_body(self, mock_workspace_client): "memory_store": {"name": "main.default.support_agent_memory"}, "scope": {"kind": "user", "value": "user-123"}, } - assert "memory_store_name" not in call_kwargs - assert "memory_scope" not in call_kwargs + assert "store" not in call_kwargs + assert "scope" not in call_kwargs - def test_memory_scope_kind_override(self, mock_workspace_client): + def test_workspace_client_scope_extracts_user_id(self, mock_workspace_client): + user_client = MagicMock(spec=WorkspaceClient) + user_client.current_user.me.return_value.user_name = "jenny@databricks.com" client = DatabricksOpenAI(workspace_client=mock_workspace_client, use_ai_gateway=True) with patch.object(Conversations, "create") as mock_create: mock_create.return_value = MagicMock() client.conversations.create( - memory_store_name="main.default.support_agent_memory", - memory_scope="acct-42", - memory_scope_kind="account", + store="main.default.support_agent_memory", + scope=user_client, + ) + + extra_body = mock_create.call_args.kwargs["extra_body"] + assert extra_body["scope"] == {"kind": "user", "value": "jenny@databricks.com"} + + def test_scope_kind_override(self, mock_workspace_client): + client = DatabricksOpenAI(workspace_client=mock_workspace_client, use_ai_gateway=True) + + with patch.object(Conversations, "create") as mock_create: + mock_create.return_value = MagicMock() + client.conversations.create( + store="main.default.support_agent_memory", + scope="acct-42", + scope_kind="account", ) extra_body = mock_create.call_args.kwargs["extra_body"] @@ -848,8 +863,8 @@ def test_memory_kwargs_merge_with_existing_extra_body(self, mock_workspace_clien with patch.object(Conversations, "create") as mock_create: mock_create.return_value = MagicMock() client.conversations.create( - memory_store_name="main.default.support_agent_memory", - memory_scope="user-123", + store="main.default.support_agent_memory", + scope="user-123", extra_body={"custom_field": "value", "scope": {"kind": "user", "value": "old"}}, ) @@ -895,8 +910,8 @@ async def test_async_memory_kwargs_translated_to_extra_body(self, mock_workspace with patch.object(AsyncConversations, "create", new_callable=AsyncMock) as mock_create: await client.conversations.create( - memory_store_name="main.default.support_agent_memory", - memory_scope="user-123", + store="main.default.support_agent_memory", + scope="user-123", ) call_kwargs = mock_create.call_args.kwargs From 9154280018aeca9955e4cdf060ba7fff36b28bb6 Mon Sep 17 00:00:00 2001 From: Jenny Date: Wed, 1 Jul 2026 15:25:04 -0700 Subject: [PATCH 3/3] Fix ty typecheck errors for conversations memory kwargs 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 --- .../openai/src/databricks_openai/utils/clients.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/integrations/openai/src/databricks_openai/utils/clients.py b/integrations/openai/src/databricks_openai/utils/clients.py index 7dd31d03..d29c70d4 100644 --- a/integrations/openai/src/databricks_openai/utils/clients.py +++ b/integrations/openai/src/databricks_openai/utils/clients.py @@ -326,7 +326,13 @@ def create(self, **kwargs): def _resolve_scope_value(scope: "str | WorkspaceClient") -> str: """Resolve a scope value, extracting the authenticated user's id from a WorkspaceClient.""" if isinstance(scope, WorkspaceClient): - return scope.current_user.me().user_name + user_name = scope.current_user.me().user_name + if not user_name: + raise ValueError( + "Could not determine the current user from the provided WorkspaceClient; " + "pass the scope value as a string instead." + ) + return user_name return scope @@ -503,7 +509,7 @@ def responses(self) -> Responses: return self._databricks_responses @property - def conversations(self) -> Conversations: + def conversations(self) -> "DatabricksConversations": if not hasattr(self, "_databricks_conversations"): self._databricks_conversations = DatabricksConversations(self) return self._databricks_conversations @@ -715,7 +721,7 @@ def responses(self) -> AsyncResponses: return self._databricks_responses @property - def conversations(self) -> AsyncConversations: + def conversations(self) -> "AsyncDatabricksConversations": if not hasattr(self, "_databricks_conversations"): self._databricks_conversations = AsyncDatabricksConversations(self) return self._databricks_conversations