From 9b15f639a4a6edc37a6096d258b8e605f94dae4b Mon Sep 17 00:00:00 2001 From: Varad Khonde <72742264+Varad2001@users.noreply.github.com> Date: Mon, 9 Mar 2026 21:31:31 +0530 Subject: [PATCH 1/5] fix(responses): merge parallel function_call items into single assistant message (#23116) --- .../transformation.py | 35 +++++ .../test_litellm_completion_responses.py | 125 ++++++++++++++++++ 2 files changed, 160 insertions(+) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 19845d7c49..d901bbc210 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -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 diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 6d6162437c..a931a9bc93 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -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 From 0d3735f9c0ee8102d43dcba9b39f4028c22d59c2 Mon Sep 17 00:00:00 2001 From: JiangNan <1394485448@qq.com> Date: Tue, 10 Mar 2026 10:46:26 +0800 Subject: [PATCH 2/5] fix: handle month overflow in duration_in_seconds for multi-month durations (#23099) When value > 1 (e.g., "2mo") and current_month + value > 12, target_month exceeds valid range (1-12), causing ValueError in datetime constructor. For example, calling duration_in_seconds("2mo") in November produces target_month=13. Use modular arithmetic to correctly wrap months and increment year. Signed-off-by: JiangNan <1394485448@qq.com> --- litellm/litellm_core_utils/duration_parser.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index 70c28c4e06..6d2b4226ff 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -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 From 3ce37e6c35f41d9e337e9db6c6e58b3616e4699b Mon Sep 17 00:00:00 2001 From: JiangNan <1394485448@qq.com> Date: Tue, 10 Mar 2026 10:49:31 +0800 Subject: [PATCH 3/5] fix: use correct list length when averaging TTFT latency for streaming requests (#23100) In _get_available_deployments, when streaming mode is active the code sums time-to-first-token values from item_ttft_latency but divides by len(item_latency) instead of len(item_ttft_latency). These lists can have different lengths, producing an incorrect average that skews lowest-latency routing decisions for streaming requests. Signed-off-by: JiangNan <1394485448@qq.com> --- litellm/router_strategy/lowest_latency.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 0449a843bd..e09b5c1456 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -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 From 9314963697120c289c489f4decb49bea2e8d55e9 Mon Sep 17 00:00:00 2001 From: zxshen Date: Tue, 10 Mar 2026 10:49:58 +0800 Subject: [PATCH 4/5] fix(fireworks): strip duplicate /v1 from models endpoint URL (#23113) _get_openai_compatible_provider_info already returns an api_base ending in /v1, but get_models prepended another /v1, producing .../inference/v1/v1/accounts/... which 404s. Strip the trailing /v1 from api_base before re-adding it so that both the default and any user-supplied base work correctly. Add parametrized tests covering the default URL, trailing-slash, custom base with /v1, and custom base without /v1. Fixes #23106 --- .../llms/fireworks_ai/chat/transformation.py | 5 +- .../test_fireworks_ai_chat_transformation.py | 54 +++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 7ec32fecc4..c9b6f330ec 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -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}"}, ) diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 8006ffdff1..5d5aaa64c8 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -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. From 33d4b93bee36ed1419f3043ada679e6011fed022 Mon Sep 17 00:00:00 2001 From: jymmi Date: Mon, 9 Mar 2026 20:58:14 -0700 Subject: [PATCH 5/5] fix(sagemaker): Add role assumption support for embedding endpoint (#20435) The SageMaker embedding handler was not using _load_credentials(), which meant aws_role_name and aws_session_name parameters were ignored. This prevented cross-account role assumption for embeddings while it worked for completions. Changes: - Replace direct boto3 client creation with _load_credentials() call - Create boto3.Session with assumed credentials - Add comprehensive unit tests for role assumption This aligns the embedding handler behavior with the completion handler, which already supports role assumption via the BaseAWSLLM.get_credentials() method. Fixes cross-account SageMaker embedding access where users need to assume a role in another account to invoke endpoints. --- litellm/llms/sagemaker/completion/handler.py | 56 ++-- ...est_sagemaker_embedding_role_assumption.py | 243 ++++++++++++++++++ 2 files changed, 263 insertions(+), 36 deletions(-) create mode 100644 tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_role_assumption.py diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index 2a30dc5ef3..efbb218f57 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -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 {}, ) diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_role_assumption.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_role_assumption.py new file mode 100644 index 0000000000..82c84af5e2 --- /dev/null +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_role_assumption.py @@ -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", + )