tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-07 19:03:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #23164 from BerriAI/litellm_oss_staging_03_09_2026

Browse Source 
oss staging 03/09/2026
This commit is contained in:
Sameer Kankute
2026-03-10 17:45:17 +05:30
committed by GitHub
parent d0d1291c15 72fc2f77cf
commit b211b67cea
8 changed files with 490 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/litellm_core_utils/duration_parser.py
+4 -6
View File
@@ -64,12 +64,10 @@ def duration_in_seconds(duration: str) -> int:
now = time.time()
current_time = datetime.fromtimestamp(now)
if current_time.month == 12:
target_year = current_time.year + 1
target_month = 1
else:
target_year = current_time.year
target_month = current_time.month + value
# Calculate target month and year, handling overflow past December
total_months = current_time.month - 1 + value # 0-indexed months
target_year = current_time.year + total_months // 12
target_month = total_months % 12 + 1 # back to 1-indexed
# Determine the day to set for next month
target_day = current_time.day
litellm/llms/fireworks_ai/chat/transformation.py
+4 -1
View File
@@ -426,8 +426,11 @@ class FireworksAIConfig(OpenAIGPTConfig):
"FIREWORKS_ACCOUNT_ID is not set. Please set the environment variable, to query Fireworks AI's `/models` endpoint."
)
base = api_base.rstrip("/")
if base.endswith("/v1"):
base = base[: -len("/v1")]
response = litellm.module_level_client.get(
url=f"{api_base}/v1/accounts/{account_id}/models",
url=f"{base}/v1/accounts/{account_id}/models",
headers={"Authorization": f"Bearer {api_key}"},
)
litellm/llms/sagemaker/completion/handler.py
+20 -36
View File
@@ -583,35 +583,17 @@ class SagemakerLLM(BaseAWSLLM):
### BOTO3 INIT
import boto3
# pop aws_secret_access_key, aws_access_key_id, aws_region_name from kwargs, since completion calls fail with them
aws_secret_access_key = optional_params.pop("aws_secret_access_key", None)
aws_access_key_id = optional_params.pop("aws_access_key_id", None)
aws_region_name = optional_params.pop("aws_region_name", None)
# Use _load_credentials to support role assumption (aws_role_name, aws_session_name)
credentials, aws_region_name = self._load_credentials(optional_params)
if aws_access_key_id is not None:
# uses auth params passed to completion
# aws_access_key_id is not None, assume user is trying to auth using litellm.completion
client = boto3.client(
service_name="sagemaker-runtime",
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
region_name=aws_region_name,
)
else:
# aws_access_key_id is None, assume user is trying to auth using env variables
# boto3 automaticaly reads env variables
# we need to read region name from env
# I assume majority of users use .env for auth
region_name = (
get_secret("AWS_REGION_NAME")
or aws_region_name # get region from config file if specified
or "us-west-2" # default to us-west-2 if region not specified
)
client = boto3.client(
service_name="sagemaker-runtime",
region_name=region_name,
)
# Create boto3 session with the loaded credentials
session = boto3.Session(
aws_access_key_id=credentials.access_key,
aws_secret_access_key=credentials.secret_key,
aws_session_token=credentials.token,
region_name=aws_region_name,
)
client = session.client(service_name="sagemaker-runtime")
# pop streaming if it's in the optional params as 'stream' raises an error with sagemaker
inference_params = deepcopy(optional_params)
@@ -628,7 +610,9 @@ class SagemakerLLM(BaseAWSLLM):
#### EMBEDDING LOGIC
# Transform request based on model type
provider_config = SagemakerEmbeddingConfig.get_model_config(model)
request_data = provider_config.transform_embedding_request(model, input, optional_params, {})
request_data = provider_config.transform_embedding_request(
model, input, optional_params, {}
)
data = json.dumps(request_data).encode("utf-8")
## LOGGING
@@ -673,19 +657,19 @@ class SagemakerLLM(BaseAWSLLM):
)
print_verbose(f"raw model_response: {response}")
# Transform response based on model type
from httpx import Response as HttpxResponse
# Create a mock httpx Response object for the transformation
mock_response = HttpxResponse(
status_code=200,
content=json.dumps(response).encode('utf-8'),
headers={"content-type": "application/json"}
content=json.dumps(response).encode("utf-8"),
headers={"content-type": "application/json"},
)
model_response = EmbeddingResponse()
# Use the request_data that was already transformed above
return provider_config.transform_embedding_response(
model=model,
@@ -695,5 +679,5 @@ class SagemakerLLM(BaseAWSLLM):
api_key=None,
request_data=request_data,
optional_params=optional_params,
litellm_params=litellm_params or {}
litellm_params=litellm_params or {},
)
litellm/responses/litellm_completion_transformation/transformation.py
+35
View File
@@ -384,6 +384,41 @@ class LiteLLMCompletionResponsesConfig:
if call_id_raw:
existing_tool_call_ids.add(str(call_id_raw))
#########################################################
# Merge consecutive function_call items into a single assistant
# message. Anthropic requires that all tool_use blocks appear in
# ONE assistant message immediately followed by the tool_result
# blocks. Without this merging, each function_call creates its own
# assistant message, producing back-to-back assistant messages that
# Anthropic rejects with "tool_use ids were found without
# tool_result blocks immediately after".
#########################################################
if messages:
last_msg = messages[-1]
last_role = (
last_msg.get("role")
if isinstance(last_msg, dict)
else getattr(last_msg, "role", None)
)
if last_role == "assistant":
for new_msg in chat_completion_messages:
new_role = (
new_msg.get("role")
if isinstance(new_msg, dict)
else getattr(new_msg, "role", None)
)
if new_role == "assistant":
new_tcs = (
new_msg.get("tool_calls")
if isinstance(new_msg, dict)
else getattr(new_msg, "tool_calls", None)
) or []
for tc in new_tcs:
LiteLLMCompletionResponsesConfig._add_tool_call_to_assistant(
last_msg, tc
)
continue
#########################################################
# If Input Item is a Tool Call Output, add it to the tool_call_output_messages list
# preserving the ordering of tool call outputs. Some models require the tool
litellm/router_strategy/lowest_latency.py
+5 -3
View File
@@ -490,20 +490,22 @@ class LowestLatencyLoggingHandler(CustomLogger):
# get average latency or average ttft (depending on streaming/non-streaming)
total: float = 0.0
if (
use_ttft = (
request_kwargs is not None
and request_kwargs.get("stream", None) is not None
and request_kwargs["stream"] is True
and len(item_ttft_latency) > 0
):
)
if use_ttft:
for _call_latency in item_ttft_latency:
if isinstance(_call_latency, float):
total += _call_latency
item_latency = total / len(item_ttft_latency)
else:
for _call_latency in item_latency:
if isinstance(_call_latency, float):
total += _call_latency
item_latency = total / len(item_latency)
item_latency = total / len(item_latency)
# -------------- #
# Debugging Logic
tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py
+54
View File
@@ -110,6 +110,60 @@ def test_get_supported_openai_params_reasoning_effort():
assert "reasoning_effort" not in unsupported_params
@pytest.mark.parametrize(
"api_base, expected_url_prefix",
[
(
"https://api.fireworks.ai/inference/v1",
"https://api.fireworks.ai/inference/v1/accounts/",
),
(
"https://api.fireworks.ai/inference/v1/",
"https://api.fireworks.ai/inference/v1/accounts/",
),
(
"https://custom-host.example.com/v1",
"https://custom-host.example.com/v1/accounts/",
),
(
"https://custom-host.example.com/api",
"https://custom-host.example.com/api/v1/accounts/",
),
],
ids=["default", "trailing-slash", "custom-with-v1", "custom-without-v1"],
)
def test_get_models_url_no_double_v1(api_base, expected_url_prefix):
"""Ensure get_models never produces a /v1/v1/ URL segment (fixes #23106)."""
config = FireworksAIConfig()
account_id = "fireworks"
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"models": [{"name": "accounts/fireworks/models/llama-v3-70b"}]
}
with (
patch("litellm.module_level_client.get", return_value=mock_response) as mock_get,
patch(
"litellm.llms.fireworks_ai.chat.transformation.get_secret_str",
side_effect=lambda key: {
"FIREWORKS_API_KEY": "test-key",
"FIREWORKS_API_BASE": api_base,
"FIREWORKS_ACCOUNT_ID": account_id,
}.get(key),
),
):
result = config.get_models(api_key="test-key", api_base=api_base)
called_url = mock_get.call_args.kwargs.get("url") or mock_get.call_args[1].get("url", "")
assert "/v1/v1/" not in called_url, f"Double /v1/ detected in URL: {called_url}"
assert called_url.startswith(expected_url_prefix), (
f"URL {called_url} does not start with {expected_url_prefix}"
)
assert result == ["fireworks_ai/accounts/fireworks/models/llama-v3-70b"]
def test_transform_messages_helper_removes_provider_specific_fields():
"""
Test that _transform_messages_helper removes provider_specific_fields from messages.
tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_role_assumption.py
+243
View File
@@ -0,0 +1,243 @@
"""
Test cases for SageMaker embedding role assumption support
This module tests that the SageMaker embedding handler properly supports
AWS IAM role assumption via aws_role_name and aws_session_name parameters,
matching the behavior of the completion handler.
"""
import json
import os
import sys
from datetime import timezone
from unittest.mock import MagicMock, call, patch
sys.path.insert(0, os.path.abspath("../../../../.."))
from botocore.credentials import Credentials
from litellm.llms.sagemaker.completion.handler import SagemakerLLM
from litellm.types.utils import EmbeddingResponse
class TestSagemakerEmbeddingRoleAssumption:
"""Test that SageMaker embedding supports role assumption like completion does"""
def setup_method(self):
self.sagemaker_llm = SagemakerLLM()
def test_embedding_uses_load_credentials(self):
"""
Test that embedding() calls _load_credentials() to support role assumption.
This ensures aws_role_name and aws_session_name parameters are properly handled.
"""
# Mock credentials that would be returned after role assumption
mock_credentials = Credentials(
access_key="assumed-access-key",
secret_key="assumed-secret-key",
token="assumed-session-token",
)
# Mock the SageMaker client response
mock_sagemaker_client = MagicMock()
mock_sagemaker_client.invoke_endpoint.return_value = {
"Body": MagicMock(
read=MagicMock(return_value=json.dumps({"embedding": [[0.1, 0.2, 0.3]]}).encode())
)
}
# Mock boto3.Session to return our mock client
mock_session = MagicMock()
mock_session.client.return_value = mock_sagemaker_client
with patch.object(
self.sagemaker_llm, "_load_credentials", return_value=(mock_credentials, "us-east-1")
) as mock_load_creds, patch("boto3.Session", return_value=mock_session):
# Create mock logging object
mock_logging = MagicMock()
optional_params = {
"aws_role_name": "arn:aws:iam::123456789012:role/TestRole",
"aws_session_name": "test-session",
}
self.sagemaker_llm.embedding(
model="test-endpoint",
input=["hello world"],
model_response=EmbeddingResponse(),
print_verbose=print,
encoding=None,
logging_obj=mock_logging,
optional_params=optional_params,
)
# Verify _load_credentials was called with the optional_params
mock_load_creds.assert_called_once()
# Verify boto3.Session was created with the assumed credentials
mock_session_calls = mock_session.client.call_args_list
assert len(mock_session_calls) == 1
assert mock_session_calls[0] == call(service_name="sagemaker-runtime")
def test_embedding_role_assumption_with_sts(self):
"""
Test the full role assumption flow for embeddings, similar to completion.
Verifies that STS assume_role is called when aws_role_name is provided.
"""
# Mock the STS client for role assumption
mock_sts_client = MagicMock()
# Mock the STS response with proper expiration handling
mock_expiry = MagicMock()
mock_expiry.tzinfo = timezone.utc
time_diff = MagicMock()
time_diff.total_seconds.return_value = 3600
mock_expiry.__sub__ = MagicMock(return_value=time_diff)
mock_sts_response = {
"Credentials": {
"AccessKeyId": "assumed-access-key",
"SecretAccessKey": "assumed-secret-key",
"SessionToken": "assumed-session-token",
"Expiration": mock_expiry,
}
}
mock_sts_client.assume_role.return_value = mock_sts_response
# Mock the SageMaker client response
mock_sagemaker_client = MagicMock()
mock_sagemaker_client.invoke_endpoint.return_value = {
"Body": MagicMock(
read=MagicMock(return_value=json.dumps({"embedding": [[0.1, 0.2, 0.3]]}).encode())
)
}
# Mock boto3.Session for SageMaker client creation
mock_session = MagicMock()
mock_session.client.return_value = mock_sagemaker_client
def mock_boto3_client(service_name, **kwargs):
if service_name == "sts":
return mock_sts_client
return mock_sagemaker_client
with patch("boto3.client", side_effect=mock_boto3_client), \
patch("boto3.Session", return_value=mock_session):
mock_logging = MagicMock()
optional_params = {
"aws_role_name": "arn:aws:iam::123456789012:role/CrossAccountRole",
"aws_session_name": "litellm-embedding-session",
"aws_region_name": "us-east-1",
}
self.sagemaker_llm.embedding(
model="test-endpoint",
input=["hello world"],
model_response=EmbeddingResponse(),
print_verbose=print,
encoding=None,
logging_obj=mock_logging,
optional_params=optional_params,
)
# Verify STS assume_role was called with correct parameters
mock_sts_client.assume_role.assert_called_once()
call_args = mock_sts_client.assume_role.call_args
assert call_args[1]["RoleArn"] == "arn:aws:iam::123456789012:role/CrossAccountRole"
assert call_args[1]["RoleSessionName"] == "litellm-embedding-session"
def test_embedding_without_role_assumption(self):
"""
Test that embedding works without role assumption when aws_role_name is not provided.
Should use default credentials from environment/instance profile.
"""
# Mock the SageMaker client response
mock_sagemaker_client = MagicMock()
mock_sagemaker_client.invoke_endpoint.return_value = {
"Body": MagicMock(
read=MagicMock(return_value=json.dumps({"embedding": [[0.1, 0.2, 0.3]]}).encode())
)
}
mock_session = MagicMock()
mock_session.client.return_value = mock_sagemaker_client
# Mock credentials returned from environment
mock_credentials = Credentials(
access_key="env-access-key",
secret_key="env-secret-key",
token=None,
)
with patch.object(
self.sagemaker_llm, "_load_credentials", return_value=(mock_credentials, "us-west-2")
), patch("boto3.Session", return_value=mock_session):
mock_logging = MagicMock()
# No aws_role_name provided
optional_params = {
"aws_region_name": "us-west-2",
}
result = self.sagemaker_llm.embedding(
model="test-endpoint",
input=["hello world"],
model_response=EmbeddingResponse(),
print_verbose=print,
encoding=None,
logging_obj=mock_logging,
optional_params=optional_params,
)
# Should still work and return embeddings
assert result is not None
def test_embedding_session_created_with_assumed_credentials(self):
"""
Test that boto3.Session is created with the credentials from role assumption.
This verifies the credentials flow from _load_credentials to the SageMaker client.
"""
mock_credentials = Credentials(
access_key="assumed-key",
secret_key="assumed-secret",
token="assumed-token",
)
mock_sagemaker_client = MagicMock()
mock_sagemaker_client.invoke_endpoint.return_value = {
"Body": MagicMock(
read=MagicMock(return_value=json.dumps({"embedding": [[0.1, 0.2, 0.3]]}).encode())
)
}
with patch.object(
self.sagemaker_llm, "_load_credentials", return_value=(mock_credentials, "us-east-1")
), patch("boto3.Session") as mock_session_class:
mock_session = MagicMock()
mock_session.client.return_value = mock_sagemaker_client
mock_session_class.return_value = mock_session
mock_logging = MagicMock()
self.sagemaker_llm.embedding(
model="test-endpoint",
input=["hello world"],
model_response=EmbeddingResponse(),
print_verbose=print,
encoding=None,
logging_obj=mock_logging,
optional_params={},
)
# Verify Session was created with the assumed credentials
mock_session_class.assert_called_once_with(
aws_access_key_id="assumed-key",
aws_secret_access_key="assumed-secret",
aws_session_token="assumed-token",
region_name="us-east-1",
)
tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py
+125
View File
@@ -1774,3 +1774,128 @@ class TestStreamingIDConsistency:
# Verify it matches the cached ID
assert iterator._cached_item_id is not None
assert iterator._cached_item_id == text_done_id
def test_parallel_tool_calls_merged_into_single_assistant_message(self):
"""
Regression test: multi-turn parallel tool calls via the Responses API must
produce a single assistant message with all tool_calls, not one assistant
message per function_call item.
When the model responds with two parallel tool calls (e.g. get_weather for
SF and NYC), the next Responses API request includes two consecutive
function_call items followed by two function_call_output items.
Without the fix each function_call becomes its own assistant message,
producing back-to-back assistant messages that Anthropic/Vertex AI rejects:
"tool_use ids were found without tool_result blocks immediately after".
"""
input_items = [
{"type": "message", "role": "user", "content": "Weather in SF and NYC?"},
# Two parallel tool calls from the previous assistant response
{
"type": "function_call",
"call_id": "toolu_01",
"name": "get_weather",
"arguments": '{"city": "SF"}',
},
{
"type": "function_call",
"call_id": "toolu_02",
"name": "get_weather",
"arguments": '{"city": "NYC"}',
},
# Tool results
{"type": "function_call_output", "call_id": "toolu_01", "output": "72°F"},
{"type": "function_call_output", "call_id": "toolu_02", "output": "55°F"},
]
messages = LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message(
input=input_items
)
roles = [
m.get("role") if isinstance(m, dict) else getattr(m, "role", None)
for m in messages
]
# Must not have two consecutive assistant messages
for i in range(len(roles) - 1):
assert not (
roles[i] == "assistant" and roles[i + 1] == "assistant"
), f"Consecutive assistant messages at indices {i} and {i+1}: {roles}"
# The single assistant message must contain BOTH tool_calls
assistant_messages = [
m for m in messages
if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None))
== "assistant"
]
assert len(assistant_messages) == 1, (
f"Expected 1 assistant message, got {len(assistant_messages)}"
)
assistant_msg = assistant_messages[0]
tool_calls = (
assistant_msg.get("tool_calls")
if isinstance(assistant_msg, dict)
else getattr(assistant_msg, "tool_calls", None)
)
assert tool_calls is not None and len(tool_calls) == 2, (
f"Expected 2 tool_calls in the merged assistant message, got: {tool_calls}"
)
call_ids = [
(tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None))
for tc in tool_calls
]
assert "toolu_01" in call_ids, f"toolu_01 missing from tool_calls: {call_ids}"
assert "toolu_02" in call_ids, f"toolu_02 missing from tool_calls: {call_ids}"
# Both tool messages must be present
tool_messages = [
m for m in messages
if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None))
== "tool"
]
assert len(tool_messages) == 2, (
f"Expected 2 tool messages, got {len(tool_messages)}"
)
def test_single_tool_call_still_works_after_merge_fix(self):
"""
Ensure the parallel-tool-call merging fix does not break the existing
single-tool-call path.
"""
input_items = [
{"type": "message", "role": "user", "content": "Weather in SF?"},
{
"type": "function_call",
"call_id": "toolu_01",
"name": "get_weather",
"arguments": '{"city": "SF"}',
},
{"type": "function_call_output", "call_id": "toolu_01", "output": "72°F"},
]
messages = LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message(
input=input_items
)
roles = [
m.get("role") if isinstance(m, dict) else getattr(m, "role", None)
for m in messages
]
assert "user" in roles
assert "assistant" in roles
assert "tool" in roles
assistant_messages = [m for m in messages if (m.get("role") if isinstance(m, dict) else getattr(m, "role", None)) == "assistant"]
assert len(assistant_messages) == 1
tool_calls = (
assistant_messages[0].get("tool_calls")
if isinstance(assistant_messages[0], dict)
else getattr(assistant_messages[0], "tool_calls", None)
)
assert tool_calls is not None and len(tool_calls) == 1