Add memory store kwargs to DatabricksOpenAI conversations.create

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,74 @@ 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):
+        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
+
+
+def _merge_memory_params_into_extra_body(
+    kwargs: dict,
+    store: str | None,
+    scope: "str | WorkspaceClient | None",
+    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 store is None and scope is None:
+        return kwargs
+    extra_body = dict(kwargs.pop("extra_body", None) or {})
+    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
+
+
+class DatabricksConversations(Conversations):
+    """Conversations resource with Databricks agent memory support.
+
+    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,
+        *,
+        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:
+            store: Full three-part Unity Catalog name of the memory store
+                (e.g. ``"main.default.support_agent_memory"``).
+            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, store, scope, scope_kind)
+        return super().create(**kwargs)
+
+
 class DatabricksOpenAI(OpenAI):
     """OpenAI client authenticated with Databricks to query LLMs and agents hosted on Databricks.
 
@@ -385,6 +454,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(
+        ...     store="main.default.support_agent_memory",
+        ...     scope=user_client,  # extracts the user id; pass a string for external users
+        ... )
     """
 
     def __init__(
@@ -432,6 +508,12 @@ def responses(self) -> Responses:
             self._databricks_responses = DatabricksResponses(self, self._workspace_client)
         return self._databricks_responses
 
+    @property
+    def conversations(self) -> "DatabricksConversations":
+        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 +572,38 @@ async def create(self, **kwargs):
         return response
 
 
+class AsyncDatabricksConversations(AsyncConversations):
+    """Async conversations resource with Databricks agent memory support.
+
+    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,
+        *,
+        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:
+            store: Full three-part Unity Catalog name of the memory store
+                (e.g. ``"main.default.support_agent_memory"``).
+            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, store, scope, 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 +667,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(
+        ...     store="main.default.support_agent_memory",
+        ...     scope=user_client,  # extracts the user id; pass a string for external users
+        ... )
     """
 
     def __init__(
@@ -598,3 +719,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) -> "AsyncDatabricksConversations":
+        if not hasattr(self, "_databricks_conversations"):
+            self._databricks_conversations = AsyncDatabricksConversations(self)
+        return self._databricks_conversations

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,130 @@ 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(
+                store="main.default.support_agent_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 "store" not in call_kwargs
+            assert "scope" not in call_kwargs
+
+    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(
+                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"]
+            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(
+                store="main.default.support_agent_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(
+                store="main.default.support_agent_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."""