fix(databricks): keep #/$defs refs in json_schema response_format (#45659)

Backport of #45659 to stable/1.103.x.
Cherry-picked from fa7ad80a62 (main).
Adapted to the line: unit tests ported into tests/test_litellm; the databricks chat wire-test additions are dropped since that file is not on this line.

litellm/llms/databricks/chat/transformation.py

@@ -7,7 +7,7 @@ from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping
 from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload
 
 import httpx
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 
 from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
 from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
@@ -20,6 +20,9 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
     strip_name_from_message,
 )
 from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
+from litellm.llms.base_llm.base_utils import (
+    type_to_response_format_param,  # pyright: ignore[reportUnknownVariableType]  # base_utils helper returns an untyped dict
+)
 from litellm.types.llms.anthropic import AllAnthropicToolsValues
 from litellm.types.llms.databricks import (
     AllDatabricksContentValues,
@@ -58,6 +61,10 @@ from ...anthropic.chat.transformation import (
 from ...openai_like.chat.transformation import OpenAILikeChatConfig
 from ..common_utils import DatabricksBase, DatabricksException
 
+_RESPONSE_FORMAT_ADAPTER: (  # mutable-ok: mirrors the dict return contract of get_json_schema_from_pydantic_object
+    Final[TypeAdapter[dict[str, object] | None]]
+) = TypeAdapter(dict[str, object] | None)
+
 
 def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool:
     """Databricks rejects assistant messages with neither content nor tool calls, e.g. a replayed
@@ -191,6 +198,14 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
     def get_config(cls):
         return super().get_config()
 
+    def get_json_schema_from_pydantic_object(
+        self,
+        response_format: (  # mutable-ok: matches BaseConfig override signature
+            type[BaseModel] | dict[str, object] | None
+        ),
+    ) -> dict[str, object] | None:  # mutable-ok: BaseConfig contract returns a dict
+        return _RESPONSE_FORMAT_ADAPTER.validate_python(type_to_response_format_param(response_format=response_format))
+
     def get_required_params(self) -> list[ProviderField]:
         """For a given provider, return it's required fields with a description"""
         return [

tests/integration/sdk/test_databricks_json_schema_refs_sdk.py

@@ -0,0 +1,186 @@
+import json
+import uuid
+from collections.abc import Iterator, Mapping
+from typing import Final
+
+import pytest
+from integration._support.wire import Reply, Request, wire_server
+from pydantic import BaseModel, ConfigDict, Field, JsonValue
+
+import litellm
+
+_QWEN_BACKEND: Final = "databricks-qwen35-122b-a10b"
+_CLAUDE_BACKEND: Final = "databricks-claude-haiku-4-5"
+_API_KEY: Final = "synthetic-databricks-key"
+_ORDER_JSON: Final = '{"buyer":{"name":"Ada"},"seller":{"name":"Bo"}}'
+_MESSAGES: Final[tuple[Mapping[str, str], ...]] = ({"role": "user", "content": "Return the order as JSON."},)
+
+
+class Person(BaseModel):
+    model_config = ConfigDict(frozen=True)
+    name: str
+
+
+class Order(BaseModel):
+    model_config = ConfigDict(frozen=True)
+    buyer: Person
+    seller: Person
+
+
+class _StrictObject(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+    additional_properties: bool | None = Field(default=None, alias="additionalProperties")
+
+
+class _OrderSchema(_StrictObject):
+    defs: Mapping[str, _StrictObject] = Field(alias="$defs")
+
+
+class _JsonSchema(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+    name: str
+    strict: bool | None = None
+    json_schema: dict[str, JsonValue] = Field(alias="schema")
+
+
+class _ResponseFormat(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+    type: str
+    json_schema: _JsonSchema
+
+
+class _Function(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+    name: str
+    parameters: dict[str, JsonValue]
+
+
+class _Tool(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+    function: _Function
+
+
+class _Outbound(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+    model: str
+    response_format: _ResponseFormat | None = None
+    tools: tuple[_Tool, ...] = ()
+
+
+class _CallerMessage(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+    content: str
+
+
+class _CallerChoice(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+    message: _CallerMessage
+
+
+class _CallerCompletion(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="ignore")
+    choices: tuple[_CallerChoice, ...]
+
+
+def _completion(backend: str) -> Reply:
+    claude: Final = "claude" in backend
+    tool_call: Final[dict[str, JsonValue]] = {
+        "id": "json-tool-call",
+        "type": "function",
+        "function": {"name": "json_tool_call", "arguments": _ORDER_JSON},
+    }
+    message: Final[dict[str, JsonValue]] = (
+        {"role": "assistant", "content": None, "tool_calls": [tool_call]}
+        if claude
+        else {"role": "assistant", "content": _ORDER_JSON}
+    )
+    body: Final[dict[str, JsonValue]] = {
+        "id": f"chatcmpl-{uuid.uuid4().hex}",
+        "object": "chat.completion",
+        "created": 1,
+        "model": backend,
+        "choices": [{"index": 0, "message": message, "finish_reason": "tool_calls" if claude else "stop"}],
+        "usage": {"prompt_tokens": 9, "completion_tokens": 7, "total_tokens": 16},
+    }
+    return Reply(body=json.dumps(body).encode())
+
+
+def _refs(node: JsonValue) -> Iterator[str]:
+    if isinstance(node, list):
+        for item in node:
+            yield from _refs(item)
+        return
+    if not isinstance(node, dict):
+        return
+    for key, value in node.items():
+        if key == "$ref" and isinstance(value, str):
+            yield value
+        else:
+            yield from _refs(value)
+
+
+def _sent_schema(backend: str, received: tuple[Request, ...]) -> dict[str, JsonValue]:
+    assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")]
+    outbound: Final = _Outbound.model_validate_json(received[0].body)
+    assert outbound.model == backend
+    if "claude" in backend:
+        assert outbound.response_format is None
+        assert [tool.function.name for tool in outbound.tools] == ["json_tool_call"]
+        return outbound.tools[0].function.parameters
+    assert outbound.response_format is not None
+    assert (outbound.response_format.type, outbound.response_format.json_schema.name) == ("json_schema", "Order")
+    assert outbound.response_format.json_schema.strict is True
+    return outbound.response_format.json_schema.json_schema
+
+
+def _assert_strict_order_schema(schema: dict[str, JsonValue]) -> None:
+    assert set(_refs(schema)) == {"#/$defs/Person"}, schema
+    parsed: Final = _OrderSchema.model_validate(schema)
+    assert (parsed.additional_properties, parsed.defs["Person"].additional_properties) == (False, False), schema
+
+
+def _caller_order(response: litellm.ModelResponse) -> Order:
+    return Order.model_validate_json(
+        _CallerCompletion.model_validate_json(response.model_dump_json()).choices[0].message.content
+    )
+
+
+_BACKENDS: Final = (
+    pytest.param(_QWEN_BACKEND, id="qwen"),
+    pytest.param(_CLAUDE_BACKEND, id="claude"),
+)
+_EXPECTED_ORDER: Final = Order(buyer=Person(name="Ada"), seller=Person(name="Bo"))
+
+
+@pytest.mark.parametrize("backend", _BACKENDS)
+def test_databricks_pydantic_response_format_sends_strict_defs_refs(backend: str) -> None:
+    with wire_server(lambda _request: _completion(backend)) as wire:
+        response: Final = litellm.completion(
+            model=f"databricks/{backend}",
+            api_base=wire.url,
+            api_key=_API_KEY,
+            messages=[dict(message) for message in _MESSAGES],
+            response_format=Order,
+            num_retries=0,
+        )
+        received: Final = wire.drain()
+    _assert_strict_order_schema(_sent_schema(backend, received))
+    assert isinstance(response, litellm.ModelResponse)
+    assert _caller_order(response) == _EXPECTED_ORDER
+
+
+@pytest.mark.parametrize("backend", _BACKENDS)
+async def test_databricks_pydantic_response_format_sends_strict_defs_refs_async(backend: str) -> None:
+    with wire_server(lambda _request: _completion(backend)) as wire:
+        response: Final = await litellm.acompletion(
+            model=f"databricks/{backend}",
+            api_base=wire.url,
+            api_key=_API_KEY,
+            messages=[dict(message) for message in _MESSAGES],
+            response_format=Order,
+            num_retries=0,
+        )
+        received: Final = wire.drain()
+    _assert_strict_order_schema(_sent_schema(backend, received))
+    assert isinstance(response, litellm.ModelResponse)
+    assert _caller_order(response) == _EXPECTED_ORDER

tests/test_litellm/llms/databricks/chat/test_databricks_chat_transformation.py

@@ -1,8 +1,11 @@
+import copy
 import json
 
 import pytest
 from fastapi.testclient import TestClient
+from pydantic import BaseModel, ConfigDict
 
+from typing import Final
 from unittest.mock import MagicMock, patch
 
 import litellm
@@ -18,6 +21,26 @@ from litellm.llms.databricks.chat.transformation import (
 )
 
 
+JSON_SCHEMA_RESPONSE_FORMAT: Final = {
+    "type": "json_schema",
+    "json_schema": {
+        "name": "P",
+        "strict": True,
+        "schema": {
+            "$defs": {
+                "Person": {
+                    "type": "object",
+                    "properties": {"name": {"type": "string"}},
+                    "required": ["name"],
+                }
+            },
+            "type": "object",
+            "properties": {"p": {"$ref": "#/$defs/Person"}},
+            "required": ["p"],
+        },
+    },
+}
+
 @pytest.fixture()
 def _use_local_model_cost_map(monkeypatch):
     monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
@@ -808,3 +831,68 @@ def test_chunk_parser_surfaces_top_level_reasoning_delta(reasoning_key: str) ->
 
     assert parsed.choices[0].delta.reasoning_content == "We need answer"
     assert parsed.choices[0].delta.content is None
+
+
+@pytest.mark.parametrize(
+    "model",
+    (
+        "databricks-meta-llama-3-3-70b-instruct",
+        "databricks-qwen35-122b-a10b",
+        "databricks-gpt-oss-120b",
+    ),
+)
+def test_databricks_non_claude_json_schema_preserves_refs(model: str) -> None:
+    optional_params: Final = litellm.utils.get_optional_params(
+        model=model,
+        custom_llm_provider="databricks",
+        response_format=copy.deepcopy(JSON_SCHEMA_RESPONSE_FORMAT),
+    )
+
+    assert optional_params["response_format"] == JSON_SCHEMA_RESPONSE_FORMAT
+
+
+def test_databricks_claude_json_schema_preserves_refs_in_tool_parameters() -> None:
+    optional_params: Final = litellm.utils.get_optional_params(
+        model="databricks-claude-haiku-4-5",
+        custom_llm_provider="databricks",
+        response_format=copy.deepcopy(JSON_SCHEMA_RESPONSE_FORMAT),
+    )
+
+    assert "response_format" not in optional_params
+    assert optional_params["tools"] == [
+        {
+            "type": "function",
+            "function": {
+                "name": "json_tool_call",
+                "parameters": JSON_SCHEMA_RESPONSE_FORMAT["json_schema"]["schema"],
+            },
+        }
+    ]
+
+
+def test_databricks_pydantic_json_schema_preserves_nested_refs() -> None:
+    class Person(BaseModel):
+        model_config = ConfigDict(frozen=True)
+
+        name: str
+
+    class Order(BaseModel):
+        model_config = ConfigDict(frozen=True)
+
+        person: Person
+
+    optional_params: Final = litellm.utils.get_optional_params(
+        model="databricks-meta-llama-3-3-70b-instruct",
+        custom_llm_provider="databricks",
+        response_format=Order,
+    )
+    schema: Final = optional_params["response_format"]["json_schema"]["schema"]
+
+    assert schema["properties"]["person"] == {"$ref": "#/$defs/Person"}
+    assert schema["$defs"]["Person"] == {
+        "additionalProperties": False,
+        "properties": {"name": {"title": "Name", "type": "string"}},
+        "required": ["name"],
+        "title": "Person",
+        "type": "object",
+    }