Add DBRX Anthropic w/ thinking + response_format support (#9744)

* feat(databricks/chat/): add anthropic w/ reasoning content support via databricks

Allows user to call claude-3-7-sonnet with thinking via databricks

* refactor: refactor choices transformation + add unit testing

* fix(databricks/chat/transformation.py): support thinking blocks on databricks response streaming

* feat(databricks/chat/transformation.py): support response_format for claude models

* fix(databricks/chat/transformation.py): correctly handle response_format={"type": "text"}

* feat(databricks/chat/transformation.py): support 'reasoning_effort' param mapping for anthropic

* fix: fix ruff errors

* fix: fix linting error

* test: update test

* fix(databricks/chat/transformation.py): handle json mode output parsing

* fix(databricks/chat/transformation.py): handle json mode on streaming

* test: update test

* test: update dbrx testing

* test: update testing

* fix(base_model_iterator.py): handle non-json chunk

* test: update tests

* fix: fix ruff check

* fix: fix databricks config import

* fix: handle _tool = none

* test: skip invalid test
This commit is contained in:
Krish Dholakia 2025-04-04 22:13:32 -07:00 committed by GitHub
parent e3b231bc11
commit 5099aac1a5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 872 additions and 340 deletions

View file

@ -3,6 +3,7 @@ Utility functions for base LLM classes.
"""
import copy
import json
from abc import ABC, abstractmethod
from typing import List, Optional, Type, Union
@ -10,8 +11,8 @@ from openai.lib import _parsing, _pydantic
from pydantic import BaseModel
from litellm._logging import verbose_logger
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ProviderSpecificModelInfo
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
from litellm.types.utils import Message, ProviderSpecificModelInfo
class BaseLLMModelInfo(ABC):
@ -55,6 +56,32 @@ class BaseLLMModelInfo(ABC):
pass
def _convert_tool_response_to_message(
tool_calls: List[ChatCompletionToolCallChunk],
) -> Optional[Message]:
"""
In JSON mode, Anthropic API returns JSON schema as a tool call, we need to convert it to a message to follow the OpenAI format
"""
## HANDLE JSON MODE - anthropic returns single function call
json_mode_content_str: Optional[str] = tool_calls[0]["function"].get("arguments")
try:
if json_mode_content_str is not None:
args = json.loads(json_mode_content_str)
if isinstance(args, dict) and (values := args.get("values")) is not None:
_message = Message(content=json.dumps(values))
return _message
else:
# a lot of the times the `values` key is not present in the tool response
# relevant issue: https://github.com/BerriAI/litellm/issues/6741
_message = Message(content=json.dumps(args))
return _message
except json.JSONDecodeError:
# json decode error does occur, return the original tool response str
return Message(content=json_mode_content_str)
return None
def _dict_to_response_format_helper(
response_format: dict, ref_template: Optional[str] = None
) -> dict: