scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

init commit (#16200)

Browse source
This commit is contained in:
Anthony Ivan 2025-11-04 10:58:03 +08:00 • committed by GitHub
parent 232d1558dd
commit e5e9523958
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 300 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/litellm_core_utils/prompt_templates/common_utils.py
View file

@ -94,6 +94,15 @@ def handle_messages_with_content_list_to_str_conversion(
return messages
def strip_name_from_message(message: AllMessageValues, allowed_name_roles: List[str] = ["user"]) -> AllMessageValues:
"""
Removes 'name' from message
"""
msg_copy = message.copy()
if msg_copy.get("role") not in allowed_name_roles:
msg_copy.pop("name", None) # type: ignore
return msg_copy
def strip_name_from_messages(
messages: List[AllMessageValues], allowed_name_roles: List[str] = ["user"]
) -> List[AllMessageValues]:

33
litellm/llms/databricks/chat/transformation.py
View file

@ -26,7 +26,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
_should_convert_tool_call_to_json_mode,
)
from litellm.litellm_core_utils.prompt_templates.common_utils import (
strip_name_from_messages,
strip_name_from_message
)
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.types.llms.anthropic import AllAnthropicToolsValues
@ -332,8 +332,11 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
_message = message.model_dump(exclude_none=True)
else:
_message = message
_message = strip_name_from_message(_message, allowed_name_roles=["user"])
# Move message-level cache_control into a content block when content is a string.
if "cache_control" in _message and isinstance(_message.get("content"), str):
_message = self._move_cache_control_into_string_content_block(_message)
new_messages.append(_message)
new_messages = strip_name_from_messages(new_messages)
if is_async:
return super()._transform_messages(
@ -344,6 +347,32 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
messages=new_messages, model=model, is_async=cast(Literal[False], False)
)
def _move_cache_control_into_string_content_block(self, message: AllMessageValues) -> AllMessageValues:
"""
Moves message-level cache_control into a content block when content is a string.
Transforms:
{"role": "user", "content": "text", "cache_control": {...}}
Into:
{"role": "user", "content": [{"type": "text", "text": "text", "cache_control": {...}}]}
This is required for Anthropic's prompt caching API when cache_control is specified
at the message level but content is a simple string (not already an array of content blocks).
"""
content = message.get("content")
# Create new message with cache_control moved into content block
transformed_message = cast(dict[str, Any], message.copy())
cache_control = transformed_message.pop("cache_control")
transformed_message["content"] = [
{
"type": "text",
"text": content,
"cache_control": cache_control,
}
]
return cast(AllMessageValues, transformed_message)
@staticmethod
def extract_content_str(
content: Optional[AllDatabricksContentValues],

293
tests/llm_translation/test_databricks.py
View file

@ -54,7 +54,7 @@ def mock_chat_response() -> Dict[str, Any]:
}
def mock_chat_response_claude_prompt_caching() -> Dict[str, Any]:
def mock_chat_response_anthropic_prompt_caching() -> Dict[str, Any]:
return {
"id": "msg_01234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ",
"object": "chat.completion",
@ -65,7 +65,7 @@ def mock_chat_response_claude_prompt_caching() -> Dict[str, Any]:
"index": 0,
"message": {
"role": "assistant",
"content": "The text you've provided consists entirely of the phrase \"example text\" repeated many times without any variation or additional content. There is no specific information, narrative, argument, or structured content to explain. This appears to be placeholder or filler text that would typically be replaced with actual content in a final document.",
"content": "I notice that you've provided a repetitive text that simply repeats \"example text\" many times rather than actual content to summarize. \n\nTo provide you with a meaningful summary, I would need:\n- Actual substantive text with real information, arguments, or narrative\n- Content that has key points, themes, or conclusions to extract\n- Material with varying ideas or concepts to synthesize\n\nCould you please share the actual text you'd like me to summarize? I'm ready to help once you provide content with real information to work with.",
"refusal": None,
"function_call": None,
"tool_calls": None,
@ -76,20 +76,26 @@ def mock_chat_response_claude_prompt_caching() -> Dict[str, Any]:
"logprobs": None,
}
],
"usage": {
"prompt_tokens": 1556,
"completion_tokens": 65,
"total_tokens": 1621,
"usage": {
"completion_tokens": 117,
"prompt_tokens": 1549,
"total_tokens": 1666,
"completion_tokens_details": None,
"prompt_tokens_details": None,
"prompt_tokens_details": {
"audio_tokens": None,
"cached_tokens": 0,
"text_tokens": None,
"image_tokens": None,
"cache_creation_tokens": 1545
},
"cache_read_input_tokens": 0,
"cache_creation_input_tokens": 1552,
"cache_creation_input_tokens": 1545
},
"service_tier": None,
"system_fingerprint": None,
}
def mock_chat_response_claude_prompt_caching_repeat() -> Dict[str, Any]:
def mock_chat_response_anthropic_prompt_caching_not_enough_tokens() -> Dict[str, Any]:
return {
"id": "msg_01234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ",
"object": "chat.completion",
@ -100,7 +106,7 @@ def mock_chat_response_claude_prompt_caching_repeat() -> Dict[str, Any]:
"index": 0,
"message": {
"role": "assistant",
"content": "The text you've provided consists entirely of the phrase \"example text\" repeated many times without any variation or additional content. There is no specific information, narrative, argument, or structured content to explain. This appears to be placeholder or filler text that would typically be replaced with actual content in a final document.",
"content": "I notice that you've provided a repetitive text that simply repeats \"example text\" many times rather than actual content to summarize. \n\nTo provide you with a meaningful summary, I would need:\n- Actual substantive text with real information, arguments, or narrative\n- Content that has key points, themes, or conclusions to extract\n- Material with varying ideas or concepts to synthesize\n\nCould you please share the actual text you'd like me to summarize? I'm ready to help once you provide content with real information to work with.",
"refusal": None,
"function_call": None,
"tool_calls": None,
@ -111,26 +117,73 @@ def mock_chat_response_claude_prompt_caching_repeat() -> Dict[str, Any]:
"logprobs": None,
}
],
"usage": {
"prompt_tokens": 1556,
"completion_tokens": 65,
"total_tokens": 1621,
"usage": {
"completion_tokens": 117,
"prompt_tokens": 1549,
"total_tokens": 1666,
"completion_tokens_details": None,
"prompt_tokens_details": None,
"cache_read_input_tokens": 1552,
"cache_creation_input_tokens": 0,
"prompt_tokens_details": {
"audio_tokens": None,
"cached_tokens": 0,
"text_tokens": None,
"image_tokens": None,
"cache_creation_tokens": 0
},
"cache_read_input_tokens": 0,
"cache_creation_input_tokens": 0
},
"service_tier": None,
"system_fingerprint": None,
}
def mock_chat_response_anthropic_prompt_caching_repeat() -> Dict[str, Any]:
return {
"id": "msg_01234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ",
"object": "chat.completion",
"created": 1761118943,
"model": "claude-3-7-sonnet", # Mock model name for testing
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": "I notice that you've provided a repetitive text that simply repeats \"example text\" many times rather than actual content to summarize. \n\nTo provide you with a meaningful summary, I would need:\n- Actual substantive text with real information, arguments, or narrative\n- Content that has key points, themes, or conclusions to extract\n- Material with varying ideas or concepts to synthesize\n\nCould you please share the actual text you'd like me to summarize? I'm ready to help once you provide content with real information to work with.",
"refusal": None,
"function_call": None,
"tool_calls": None,
"annotations": None,
"audio": None,
},
"finish_reason": "stop",
"logprobs": None,
}
],
"usage": {
"completion_tokens": 117,
"prompt_tokens": 1549,
"total_tokens": 1666,
"completion_tokens_details": None,
"prompt_tokens_details": {
"audio_tokens": None,
"cached_tokens": 0,
"text_tokens": None,
"image_tokens": None,
"cache_creation_tokens": 1545
},
"cache_read_input_tokens": 1545,
"cache_creation_input_tokens": 0
},
"service_tier": None,
"system_fingerprint": None,
}
def mock_chat_response_nonclaude_prompt_caching() -> Dict[str, Any]:
def mock_chat_response_nonanthropic_prompt_caching() -> Dict[str, Any]:
return {
"id": "msg_01234567890ABCDEFGHIJKLMNOPQRSTUVWXYZ",
"object": "chat.completion",
"created": 1761119150,
"model": "gpt-oss-20b",
"model": "gpt-oss-20b", # Mock model nama for testing
"choices": [
{
"index": 0,
@ -881,7 +934,7 @@ async def test_databricks_embeddings(sync_mode, monkeypatch):
)
def test_completion_with_prompt_caching_claude_model(monkeypatch):
def test_completion_with_prompt_caching_anthropic_model(monkeypatch):
base_url = "https://my.workspace.cloud.databricks.com/serving-endpoints"
api_key = "dapimykey"
monkeypatch.setenv("DATABRICKS_API_BASE", base_url)
@ -890,7 +943,7 @@ def test_completion_with_prompt_caching_claude_model(monkeypatch):
sync_handler = HTTPHandler()
mock_response = Mock(spec=httpx.Response)
mock_response.status_code = 200
mock_response.json.return_value = mock_chat_response_claude_prompt_caching()
mock_response.json.return_value = mock_chat_response_anthropic_prompt_caching()
mock_text = 'example text' * 512
messages = [
@ -936,13 +989,13 @@ def test_completion_with_prompt_caching_claude_model(monkeypatch):
# Check the response object returned from litellm.completion()
assert 'claude-3-7-sonnet' in response['model']
assert response['usage']['cache_read_input_tokens'] == 0
assert response['usage']['cache_creation_input_tokens'] == 1552
assert response['usage']['prompt_tokens'] == 1556
assert response['usage']['completion_tokens'] == 65
assert response['usage']['total_tokens'] == 1621
assert response['usage']['cache_creation_input_tokens'] == 1545
assert response['usage']['prompt_tokens'] == 1549
assert response['usage']['completion_tokens'] == 117
assert response['usage']['total_tokens'] == 1666
def test_completion_with_prompt_caching_claude_model_repeat(monkeypatch):
def test_completion_with_prompt_caching_anthropic_model_repeat(monkeypatch):
base_url = "https://my.workspace.cloud.databricks.com/serving-endpoints"
api_key = "dapimykey"
monkeypatch.setenv("DATABRICKS_API_BASE", base_url)
@ -951,7 +1004,7 @@ def test_completion_with_prompt_caching_claude_model_repeat(monkeypatch):
sync_handler = HTTPHandler()
mock_response = Mock(spec=httpx.Response)
mock_response.status_code = 200
mock_response.json.return_value = mock_chat_response_claude_prompt_caching_repeat()
mock_response.json.return_value = mock_chat_response_anthropic_prompt_caching_repeat()
mock_text = 'example text' * 512
messages = [
@ -998,14 +1051,14 @@ def test_completion_with_prompt_caching_claude_model_repeat(monkeypatch):
# TODO: add test for entire expected output schema in the future
# Check the response object returned from litellm.completion()
assert 'claude-3-7-sonnet' in response['model']
assert response['usage']['cache_read_input_tokens'] == 1552
assert response['usage']['cache_read_input_tokens'] == 1545
assert response['usage']['cache_creation_input_tokens'] == 0
assert response['usage']['prompt_tokens'] == 1556
assert response['usage']['completion_tokens'] == 65
assert response['usage']['total_tokens'] == 1621
assert response['usage']['prompt_tokens'] == 1549
assert response['usage']['completion_tokens'] == 117
assert response['usage']['total_tokens'] == 1666
def test_completion_with_prompt_caching_nonclaude_model(monkeypatch):
def test_completion_with_prompt_caching_nonanthropic_model(monkeypatch):
base_url = "https://my.workspace.cloud.databricks.com/serving-endpoints"
api_key = "dapimykey"
monkeypatch.setenv("DATABRICKS_API_BASE", base_url)
@ -1014,7 +1067,7 @@ def test_completion_with_prompt_caching_nonclaude_model(monkeypatch):
sync_handler = HTTPHandler()
mock_response = Mock(spec=httpx.Response)
mock_response.status_code = 200
mock_response.json.return_value = mock_chat_response_nonclaude_prompt_caching()
mock_response.json.return_value = mock_chat_response_nonanthropic_prompt_caching()
mock_text = 'example text' * 512
messages = [
@ -1148,3 +1201,177 @@ def test_databricks_anthropic_function_call_with_no_schema(model, monkeypatch):
assert len(response.choices[0].message.tool_calls) == 1
assert response.choices[0].message.tool_calls[0].function.name == "get_current_weather"
def test_databricks_anthropic_user_string_content_cache_injection(monkeypatch):
base_url = "https://my.workspace.cloud.databricks.com/serving-endpoints"
api_key = "dapimykey"
monkeypatch.setenv("DATABRICKS_API_BASE", base_url)
monkeypatch.setenv("DATABRICKS_API_KEY", api_key)
sync_handler = HTTPHandler()
mock_response = Mock(spec=httpx.Response)
mock_response.status_code = 200
mock_response.json.return_value = mock_chat_response_anthropic_prompt_caching()
mock_text = 'example text' * 512
messages = [
{
"role": "system",
"content": "You are an expert summarizer."
},
{
"role": "user",
"content": mock_text
}
]
cache_control_injection_points = [
{
"location": "message",
"role": "user"
}
]
with patch.object(HTTPHandler, "post", return_value=mock_response) as mock_post:
response = litellm.completion(
model="databricks/databricks-claude-3-7-sonnet",
messages=messages,
client=sync_handler,
temperature=0.5,
cache_control_injection_points=cache_control_injection_points,
extraparam="testpassingextraparam",
)
assert (
mock_post.call_args.kwargs["headers"]["Content-Type"] == "application/json"
)
assert (
mock_post.call_args.kwargs["headers"]["Authorization"]
== f"Bearer {api_key}"
)
assert mock_post.call_args.kwargs["url"] == f"{base_url}/chat/completions"
assert mock_post.call_args.kwargs["stream"] == False
# TODO: add test for entire expected output schema in the future
# Check the response object returned from litellm.completion()
assert 'claude-3-7-sonnet' in response['model']
assert response['usage']['cache_read_input_tokens'] == 0
assert response['usage']['cache_creation_input_tokens'] == 1545
assert response['usage']['prompt_tokens'] == 1549
assert response['usage']['completion_tokens'] == 117
assert response['usage']['total_tokens'] == 1666
def test_databricks_anthropic_system_string_content_cache_injection(monkeypatch):
base_url = "https://my.workspace.cloud.databricks.com/serving-endpoints"
api_key = "dapimykey"
monkeypatch.setenv("DATABRICKS_API_BASE", base_url)
monkeypatch.setenv("DATABRICKS_API_KEY", api_key)
sync_handler = HTTPHandler()
mock_response = Mock(spec=httpx.Response)
mock_response.status_code = 200
mock_response.json.return_value = mock_chat_response_anthropic_prompt_caching()
mock_text = 'example text' * 512
messages = [
{
"role": "system",
"content": mock_text
},
{
"role": "user",
"content": "You are an expert summarizer."
}
]
cache_control_injection_points = [
{
"location": "message",
"role": "system"
}
]
with patch.object(HTTPHandler, "post", return_value=mock_response) as mock_post:
response = litellm.completion(
model="databricks/databricks-claude-3-7-sonnet",
messages=messages,
client=sync_handler,
temperature=0.5,
cache_control_injection_points=cache_control_injection_points,
extraparam="testpassingextraparam",
)
assert (
mock_post.call_args.kwargs["headers"]["Content-Type"] == "application/json"
)
assert (
mock_post.call_args.kwargs["headers"]["Authorization"]
== f"Bearer {api_key}"
)
assert mock_post.call_args.kwargs["url"] == f"{base_url}/chat/completions"
assert mock_post.call_args.kwargs["stream"] == False
# TODO: add test for entire expected output schema in the future
# Check the response object returned from litellm.completion()
assert 'claude-3-7-sonnet' in response['model']
assert response['usage']['cache_read_input_tokens'] == 0
assert response['usage']['cache_creation_input_tokens'] == 1545
assert response['usage']['prompt_tokens'] == 1549
assert response['usage']['completion_tokens'] == 117
assert response['usage']['total_tokens'] == 1666
def test_databricks_anthropic_system_string_content_cache_injection_not_enough_tokens(monkeypatch):
base_url = "https://my.workspace.cloud.databricks.com/serving-endpoints"
api_key = "dapimykey"
monkeypatch.setenv("DATABRICKS_API_BASE", base_url)
monkeypatch.setenv("DATABRICKS_API_KEY", api_key)
sync_handler = HTTPHandler()
mock_response = Mock(spec=httpx.Response)
mock_response.status_code = 200
mock_response.json.return_value = mock_chat_response_anthropic_prompt_caching_not_enough_tokens()
mock_text = 'example text' * 512
messages = [
{
"role": "system",
"content": "You are a helpful assistant that explains the content of the given text."
},
{
"role": "user",
"content": mock_text
}
]
cache_control_injection_points = [
{
"location": "message",
"role": "system"
}
]
with patch.object(HTTPHandler, "post", return_value=mock_response) as mock_post:
response = litellm.completion(
model="databricks/databricks-claude-3-7-sonnet",
messages=messages,
client=sync_handler,
temperature=0.5,
cache_control_injection_points=cache_control_injection_points,
extraparam="testpassingextraparam",
)
assert (
mock_post.call_args.kwargs["headers"]["Content-Type"] == "application/json"
)
assert (
mock_post.call_args.kwargs["headers"]["Authorization"]
== f"Bearer {api_key}"
)
assert mock_post.call_args.kwargs["url"] == f"{base_url}/chat/completions"
assert mock_post.call_args.kwargs["stream"] == False
# TODO: add test for entire expected output schema in the future
# Check the response object returned from litellm.completion()
assert 'claude-3-7-sonnet' in response['model']
assert response['usage']['cache_read_input_tokens'] == 0
assert response['usage']['cache_creation_input_tokens'] == 0
assert response['usage']['prompt_tokens'] == 1549
assert response['usage']['completion_tokens'] == 117
assert response['usage']['total_tokens'] == 1666