diff --git a/litellm/llms/bedrock_mantle/passthrough/transformation.py b/litellm/llms/bedrock_mantle/passthrough/transformation.py index 1393ac7c6e7..e6b831efa57 100644 --- a/litellm/llms/bedrock_mantle/passthrough/transformation.py +++ b/litellm/llms/bedrock_mantle/passthrough/transformation.py @@ -1,12 +1,19 @@ from collections.abc import Mapping -from typing import Final, Literal +from typing import TYPE_CHECKING, Final, Literal, Optional +from httpx import Response + +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig from litellm.llms.bedrock_mantle.common_utils import ( MANTLE_HOST_RE, resolve_mantle_bearer_token, resolve_mantle_region, ) +from litellm.types.utils import LlmProviders + +if TYPE_CHECKING: + from litellm.types.utils import CostResponseTypes class BedrockMantlePassthroughConfig(BedrockPassthroughConfig): @@ -42,3 +49,23 @@ class BedrockMantlePassthroughConfig(BedrockPassthroughConfig): def get_bedrock_bearer_token(self, litellm_params: Mapping[str, object]) -> str | None: api_key: Final = litellm_params.get("api_key") return resolve_mantle_bearer_token(api_key if isinstance(api_key, str) else None) + + def logging_non_streaming_response( + self, + model: str, + custom_llm_provider: str, + httpx_response: Response, + request_data: dict, # mutable-ok: mirrors the inherited BedrockPassthroughConfig signature + logging_obj: Logging, + endpoint: str, + ) -> Optional["CostResponseTypes"]: + is_converse: Final = "invoke" not in endpoint and "converse" in endpoint + shape_provider: Final = LlmProviders.BEDROCK.value if is_converse else custom_llm_provider + return super().logging_non_streaming_response( + model=model, + custom_llm_provider=shape_provider, + httpx_response=httpx_response, + request_data=request_data, + logging_obj=logging_obj, + endpoint=endpoint, + ) diff --git a/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py b/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py index b7f9e492e14..8c6eda605ca 100644 --- a/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py @@ -14,6 +14,7 @@ from litellm.utils import ProviderConfigManager MANTLE_API_BASE = "https://bedrock-mantle.us-east-2.api.aws" INVOKE_ENDPOINT = "model/us.openai.gpt-5.6-sol/invoke" +CONVERSE_ENDPOINT = "model/us.openai.gpt-5.6-sol/converse" REQUEST_BODY = {"messages": [{"role": "user", "content": "say pong"}], "max_completion_tokens": 64} @@ -150,3 +151,47 @@ def test_invoke_passthrough_route_reaches_bedrock_runtime_for_a_mantle_deploymen assert str(sent["url"]) == f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}" assert sent["headers"]["Authorization"] == f"Bearer {expected_bearer}" assert json.loads(sent["content"]) == REQUEST_BODY + + +def _logged_model_response(endpoint, body): + request = httpx.Request("POST", f"https://bedrock-runtime.us-east-1.amazonaws.com/{endpoint}") + return BedrockMantlePassthroughConfig().logging_non_streaming_response( + model="us.openai.gpt-5.6-sol", + custom_llm_provider="bedrock_mantle", + httpx_response=httpx.Response(200, json=body, request=request), + request_data={"messages": [{"role": "user", "content": [{"text": "say pong"}]}]}, + logging_obj=MagicMock(), + endpoint=endpoint, + ) + + +def test_converse_logging_parses_the_converse_response_shape(): + result = _logged_model_response( + CONVERSE_ENDPOINT, + { + "metrics": {"latencyMs": 800.0}, + "output": {"message": {"content": [{"text": "pong"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 8, "outputTokens": 5, "totalTokens": 13}, + }, + ) + assert result.choices[0].message.content == "pong" + assert result.usage.prompt_tokens == 8 + assert result.usage.completion_tokens == 5 + + +def test_invoke_logging_parses_the_openai_chat_response_shape(): + result = _logged_model_response( + INVOKE_ENDPOINT, + { + "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "pong", "role": "assistant"}}], + "created": 1787677792, + "id": "chatcmpl-regression", + "model": "us.openai.gpt-5.6-sol", + "object": "chat.completion", + "usage": {"completion_tokens": 5, "prompt_tokens": 8, "total_tokens": 13}, + }, + ) + assert result.choices[0].message.content == "pong" + assert result.usage.prompt_tokens == 8 + assert result.usage.completion_tokens == 5