diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 41c1baffab8..4fab9384032 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -1,6 +1,7 @@ import reprlib from collections.abc import Mapping from dataclasses import dataclass, fields +from types import MappingProxyType from typing import Final, cast from pydantic import TypeAdapter, ValidationError @@ -31,6 +32,23 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset( } ) + +def with_aws_credential_fallbacks( + optional_params: Mapping[str, object], litellm_params: Mapping[str, object] +) -> Mapping[str, object]: + """``optional_params`` with each AWS credential kwarg it leaves unset taken from ``litellm_params``.""" + return MappingProxyType( + { + **optional_params, + **{ + key: litellm_params[key] + for key in AWS_CREDENTIAL_KWARGS_KEYS + if optional_params.get(key) is None and litellm_params.get(key) is not None + }, + } + ) + + PROVIDER_AFFINITY_HEADER_KWARG_KEY: Final = "provider_affinity_header" AGENTIC_LOOP_CONTROL_KWARGS_KEYS: Final = frozenset( {"_agentic_loop_depth", "_agentic_loop_fingerprints", "max_agentic_loops"} diff --git a/litellm/llms/bedrock/chat/__init__.py b/litellm/llms/bedrock/chat/__init__.py index 3c870226640..d30272babc2 100644 --- a/litellm/llms/bedrock/chat/__init__.py +++ b/litellm/llms/bedrock/chat/__init__.py @@ -1,4 +1,3 @@ -from .converse_handler import BedrockConverseLLM from .invoke_handler import ( AmazonAnthropicClaudeStreamDecoder, AmazonDeepSeekR1StreamDecoder, diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py deleted file mode 100644 index a6f9192339f..00000000000 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ /dev/null @@ -1,525 +0,0 @@ -import json -from collections.abc import Iterator, Mapping -from typing import Final - -import httpx - -import litellm -from litellm.anthropic_beta_headers_manager import ( - update_headers_with_filtered_beta, -) -from litellm.litellm_core_utils.get_litellm_params import stored_control_options -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject -from litellm.litellm_core_utils.url_utils import join_api_base -from litellm.llms.base_llm import base_model_iterator -from litellm.llms.custom_httpx.http_handler import ( - AsyncHTTPHandler, - HTTPHandler, - get_async_httpx_client, - get_httpx_client, -) -from litellm.types.utils import ModelResponse -from litellm.utils import CustomStreamWrapper - -from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token, pop_aws_auth_params, run_aws_signing -from ..common_utils import BedrockError, error_response_text, get_all_bedrock_regions -from .invoke_handler import AWSEventStreamDecoder, make_call - - -def make_sync_call( - client: HTTPHandler | None, - api_base: str, - headers: dict, - data: str, - model: str, - messages: list, - logging_obj: LiteLLMLoggingObject, - json_mode: bool | None = False, - fake_stream: bool = False, - stream_chunk_size: int | None = None, - timeout: float | httpx.Timeout | None = None, - optional_params: Mapping[str, object] | None = None, -) -> tuple[base_model_iterator.MockResponseIterator | Iterator[object], httpx.Headers]: - if client is None: - client = get_httpx_client() # Create a new client if none provided - - try: - response: Final = client.post( - api_base, - headers=headers, - data=data, - stream=not fake_stream, - logging_obj=logging_obj, - timeout=timeout, - ) - except httpx.HTTPStatusError as err: - raise BedrockError( - status_code=err.response.status_code, - message=error_response_text(err.response), - headers=err.response.headers, - response=err.response, - ) from err - - if response.status_code != 200: - raise BedrockError( - status_code=response.status_code, - message=str(response.read()), - headers=response.headers, - response=response, - ) - - if fake_stream: - model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response( - model=model, - response=response, - model_response=litellm.ModelResponse(), - stream=True, - logging_obj=logging_obj, - optional_params=dict(optional_params or {}), - api_key="", - data=data, - messages=messages, - encoding=litellm.encoding, - json_mode=json_mode, - ) - completion_stream: base_model_iterator.MockResponseIterator | Iterator[object] = ( - base_model_iterator.MockResponseIterator(model_response=model_response) - ) - else: - decoder: Final = AWSEventStreamDecoder(model=model, json_mode=json_mode) - completion_stream = decoder.iter_bytes( - response.iter_bytes(chunk_size=stream_chunk_size), response_headers=response.headers - ) - - # LOGGING - logging_obj.post_call( - input=messages, - api_key="", - original_response="first stream response received", - additional_args={"complete_input_dict": data}, - ) - - return completion_stream, response.headers - - -class BedrockConverseLLM(BaseAWSLLM): - def __init__(self) -> None: - super().__init__() - - async def async_streaming( - self, - model: str, - messages: list, - api_base: str, - model_response: ModelResponse, - timeout: float | httpx.Timeout | None, - encoding, - logging_obj: LiteLLMLoggingObject, - stream, - optional_params: dict, - litellm_params: dict, - credentials: Credentials | None, - logger_fn=None, - headers={}, - client: AsyncHTTPHandler | None = None, - fake_stream: bool = False, - json_mode: bool | None = False, - api_key: str | None = None, - stream_chunk_size: int | None = None, - ) -> CustomStreamWrapper: - request_data: Final = await litellm.AmazonConverseConfig()._async_transform_request( - model=model, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - headers=headers, - ) - data: Final = json.dumps(request_data) - - prepped: Final = await run_aws_signing( - self.get_request_headers, - credentials=credentials, - aws_region_name=litellm_params.get("aws_region_name") or "us-west-2", - extra_headers=headers, - endpoint_url=api_base, - data=data, - headers=headers, - api_key=api_key, - ) - - ## LOGGING - logging_obj.pre_call( - input=messages, - api_key="", - additional_args={ - "complete_input_dict": data, - "api_base": api_base, - "headers": dict(prepped.headers), - }, - ) - - completion_stream, response_headers = await make_call( - client=client, - api_base=api_base, - headers=dict(prepped.headers), - data=data, - model=model, - messages=messages, - logging_obj=logging_obj, - fake_stream=fake_stream, - json_mode=json_mode, - stream_chunk_size=stream_chunk_size, - timeout=timeout, - optional_params=optional_params, - ) - streaming_response: Final = CustomStreamWrapper( - completion_stream=completion_stream, - model=model, - custom_llm_provider="bedrock", - logging_obj=logging_obj, - _response_headers=response_headers, - ) - return streaming_response - - async def async_completion( - self, - model: str, - messages: list, - api_base: str, - model_response: ModelResponse, - timeout: float | httpx.Timeout | None, - encoding, - logging_obj: LiteLLMLoggingObject, - stream, - optional_params: dict, - litellm_params: dict, - credentials: Credentials | None, - logger_fn=None, - headers: dict = {}, - client: AsyncHTTPHandler | None = None, - api_key: str | None = None, - skip_pre_call_logging: bool = False, - ) -> ModelResponse | CustomStreamWrapper: - request_data: Final = await litellm.AmazonConverseConfig()._async_transform_request( - model=model, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - headers=headers, - ) - data: Final = json.dumps(request_data) - - prepped: Final = await run_aws_signing( - self.get_request_headers, - credentials=credentials, - aws_region_name=litellm_params.get("aws_region_name") or "us-west-2", - extra_headers=headers, - endpoint_url=api_base, - data=data, - headers=headers, - api_key=api_key, - ) - - ## LOGGING - # The Rust path already logged this request's pre_call before handing - # it here, and it only declines before the provider is called, so this - # is the same attempt continuing rather than a second one. - if not skip_pre_call_logging: - logging_obj.pre_call( - input=messages, - api_key="", - additional_args={ - "complete_input_dict": data, - "api_base": api_base, - "headers": prepped.headers, - }, - ) - - headers = dict(prepped.headers) - if client is None or not isinstance(client, AsyncHTTPHandler): - _params: Final = {} - if timeout is not None: - if isinstance(timeout, float) or isinstance(timeout, int): - timeout = httpx.Timeout(timeout) - _params["timeout"] = timeout - client = get_async_httpx_client(params=_params, llm_provider=litellm.LlmProviders.BEDROCK) - else: - client = client - - try: - response: Final = await client.post( - url=api_base, - headers=headers, - data=data, - logging_obj=logging_obj, - ) - response.raise_for_status() - except httpx.HTTPStatusError as err: - error_code: Final = err.response.status_code - raise BedrockError( - status_code=error_code, - message=error_response_text(err.response), - headers=err.response.headers, - response=err.response, - ) - except httpx.TimeoutException: - raise BedrockError(status_code=408, message="Timeout error occurred.") - - transformed_response: Final = litellm.AmazonConverseConfig()._transform_response( - model=model, - response=response, - model_response=model_response, - stream=stream if isinstance(stream, bool) else False, - logging_obj=logging_obj, - api_key="", - data=data, - messages=messages, - optional_params=optional_params, - encoding=encoding, - ) - transformed_response.set_provider_response_headers(response.headers) - return transformed_response - - def completion( - self, - model: str, - messages: list, - api_base: str | None, - custom_prompt_dict: dict, - model_response: ModelResponse, - encoding, - logging_obj: LiteLLMLoggingObject, - optional_params: dict, - acompletion: bool, - timeout: float | httpx.Timeout | None, - litellm_params: dict, - logger_fn=None, - extra_headers: dict | None = None, - client: AsyncHTTPHandler | HTTPHandler | None = None, - api_key: str | None = None, - ): - ## SETUP ## - stream: Final = optional_params.pop("stream", None) - stream_chunk_size: Final = stored_control_options(litellm_params).stream_chunk_size if stream is True else None - unencoded_model_id: Final = optional_params.pop("model_id", None) - fake_stream = optional_params.pop("fake_stream", False) - json_mode: Final = optional_params.get("json_mode") is True - if unencoded_model_id is not None: - modelId = self.encode_model_id(model_id=unencoded_model_id) - else: - # Strip nova spec prefixes before encoding model ID for API URL - _model_for_id = model - _stripped = _model_for_id - for rp in ["bedrock/converse/", "bedrock/", "converse/"]: - if _stripped.startswith(rp): - _stripped = _stripped[len(rp) :] - break - # Strip embedded region prefix (e.g. "bedrock/us-east-1/model" -> "model") - # and capture it so it can be used as aws_region_name below. - _region_from_model: str | None = None - _potential_region: Final = _stripped.split("/", 1)[0] - if _potential_region in get_all_bedrock_regions() and "/" in _stripped: - _region_from_model = _potential_region - _stripped = _stripped.split("/", 1)[1] - _model_for_id = _stripped - for _nova_prefix in ["nova-2/", "nova/"]: - if _stripped.startswith(_nova_prefix): - _model_for_id = _model_for_id.replace(_nova_prefix, "", 1) - break - modelId = self.encode_model_id(model_id=_model_for_id) - # Inject region extracted from model path so _get_aws_region_name picks it up - if _region_from_model is not None and "aws_region_name" not in optional_params: - optional_params["aws_region_name"] = _region_from_model - - fake_stream = litellm.AmazonConverseConfig().should_fake_stream( - fake_stream=fake_stream, - model=model, - stream=stream, - custom_llm_provider="bedrock", - ) - - ### SET REGION NAME ### - aws_region_name: Final = self._get_aws_region_name( - optional_params=optional_params, - model=model, - model_id=unencoded_model_id, - ) - - auth_params: Final = pop_aws_auth_params(optional_params) - aws_bedrock_runtime_endpoint: Final = optional_params.pop("aws_bedrock_runtime_endpoint", None) - optional_params.pop("aws_region_name", None) - - litellm_params["aws_region_name"] = aws_region_name # [DO NOT DELETE] important for async calls - - credentials: Final[Credentials | None] = ( - None - if bedrock_bearer_token(api_key) is not None - else self.resolve_credentials(auth_params, aws_region_name) - ) - - ### SET RUNTIME ENDPOINT ### - endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint( - api_base=api_base, - aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, - aws_region_name=aws_region_name, - ) - if (stream is not None and stream is True) and not fake_stream: - endpoint_url = join_api_base(endpoint_url, f"model/{modelId}/converse-stream") - proxy_endpoint_url = join_api_base(proxy_endpoint_url, f"model/{modelId}/converse-stream") - else: - endpoint_url = join_api_base(endpoint_url, f"model/{modelId}/converse") - proxy_endpoint_url = join_api_base(proxy_endpoint_url, f"model/{modelId}/converse") - - ## COMPLETION CALL - headers = {"Content-Type": "application/json"} - if extra_headers is not None: - headers = {"Content-Type": "application/json", **extra_headers} - - # Filter beta headers in HTTP headers before making the request - headers = update_headers_with_filtered_beta(headers=headers, provider="bedrock_converse") - - ### ROUTING (ASYNC, STREAMING, SYNC) - if acompletion: - if isinstance(client, HTTPHandler): - client = None - if stream is True: - return self.async_streaming( - model=model, - messages=messages, - api_base=proxy_endpoint_url, - model_response=model_response, - encoding=encoding, - logging_obj=logging_obj, - optional_params=optional_params, - stream=True, - litellm_params=litellm_params, - logger_fn=logger_fn, - headers=headers, - timeout=timeout, - client=client, - json_mode=json_mode, - fake_stream=fake_stream, - credentials=credentials, - api_key=api_key, - stream_chunk_size=stream_chunk_size, - ) - ### ASYNC COMPLETION - return self.async_completion( - model=model, - messages=messages, - api_base=proxy_endpoint_url, - model_response=model_response, - encoding=encoding, - logging_obj=logging_obj, - optional_params=optional_params, - stream=stream, - litellm_params=litellm_params, - logger_fn=logger_fn, - headers=headers, - timeout=timeout, - client=client, - credentials=credentials, - api_key=api_key, - ) - - ## TRANSFORMATION ## - - _data: Final = litellm.AmazonConverseConfig()._transform_request( - model=model, - messages=messages, - optional_params=optional_params, - litellm_params=litellm_params, - headers=extra_headers, - ) - data: Final = json.dumps(_data) - - prepped: Final = self.get_request_headers( - credentials=credentials, - aws_region_name=aws_region_name, - extra_headers=extra_headers, - endpoint_url=proxy_endpoint_url, - data=data, - headers=headers, - api_key=api_key, - ) - - ## LOGGING - logging_obj.pre_call( - input=messages, - api_key="", - additional_args={ - "complete_input_dict": data, - "api_base": proxy_endpoint_url, - "headers": prepped.headers, - }, - ) - if client is None or isinstance(client, AsyncHTTPHandler): - _params: Final = {} - if timeout is not None: - if isinstance(timeout, float) or isinstance(timeout, int): - timeout = httpx.Timeout(timeout) - _params["timeout"] = timeout - client = get_httpx_client(_params) - else: - client = client - - if stream is not None and stream is True: - completion_stream, response_headers = make_sync_call( - client=(client if client is not None and isinstance(client, HTTPHandler) else None), - api_base=proxy_endpoint_url, - headers=prepped.headers, - data=data, - model=model, - messages=messages, - logging_obj=logging_obj, - json_mode=json_mode, - fake_stream=fake_stream, - stream_chunk_size=stream_chunk_size, - timeout=timeout, - optional_params=optional_params, - ) - streaming_response: Final = CustomStreamWrapper( - completion_stream=completion_stream, - model=model, - custom_llm_provider="bedrock", - logging_obj=logging_obj, - _response_headers=response_headers, - ) - - return streaming_response - - ### COMPLETION - - try: - response: Final = client.post( - url=proxy_endpoint_url, - headers=prepped.headers, - data=data, - logging_obj=logging_obj, - ) - response.raise_for_status() - except httpx.HTTPStatusError as err: - error_code: Final = err.response.status_code - raise BedrockError( - status_code=error_code, - message=error_response_text(err.response), - headers=err.response.headers, - response=err.response, - ) - except httpx.TimeoutException: - raise BedrockError(status_code=408, message="Timeout error occurred.") - - sync_transformed_response: Final = litellm.AmazonConverseConfig()._transform_response( - model=model, - response=response, - model_response=model_response, - stream=stream if isinstance(stream, bool) else False, - logging_obj=logging_obj, - api_key="", - data=data, - messages=messages, - optional_params=optional_params, - encoding=encoding, - ) - sync_transformed_response.set_provider_response_headers(response.headers) - return sync_transformed_response diff --git a/litellm/llms/bedrock/chat/converse_runtime.py b/litellm/llms/bedrock/chat/converse_runtime.py new file mode 100644 index 00000000000..cae793d57e4 --- /dev/null +++ b/litellm/llms/bedrock/chat/converse_runtime.py @@ -0,0 +1,273 @@ +"""Bedrock Runtime transport for the Converse route, driven by the shared ``BaseLLMHTTPHandler``. + +``AmazonConverseConfig`` stays transport-free because the ``converse_like`` route reuses it against +endpoints that are not AWS and are never SigV4-signed. +""" + +import json +from collections.abc import Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, cast + +import httpx + +from litellm.litellm_core_utils.get_litellm_params import stored_control_options, with_aws_credential_fallbacks +from litellm.litellm_core_utils.logging_utils import track_llm_api_timing +from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS +from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig +from litellm.llms.bedrock.chat.invoke_handler import make_call, make_sync_call +from litellm.llms.bedrock.common_utils import get_all_bedrock_regions +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import ModelResponse +from litellm.utils import CustomStreamWrapper + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer + +_AWS: Final = BaseAWSLLM() +_ROUTE_PREFIXES: Final = ("bedrock/converse/", "bedrock/", "converse/") +_MODEL_SPEC_PREFIXES: Final = ("nova-2/", "nova/") +_PARAMS_SENT_OUTSIDE_THE_BODY: Final = frozenset( + {*AWS_AUTH_PARAM_KEYS, "aws_region_name", "aws_bedrock_runtime_endpoint", "model_id", "stream", "fake_stream"} +) + + +@dataclass(frozen=True, slots=True) +class _ModelPath: + model_id: str + region: str | None + + +def _model_path(model: str) -> _ModelPath: + without_route: Final = next((model[len(prefix) :] for prefix in _ROUTE_PREFIXES if model.startswith(prefix)), model) + head, separator, tail = without_route.partition("/") + region: Final = head if separator and head in get_all_bedrock_regions() else None + after_region: Final = without_route if region is None else tail + model_id: Final = model if region is None else tail + spec: Final = next((prefix for prefix in _MODEL_SPEC_PREFIXES if after_region.startswith(prefix)), None) + return _ModelPath(model_id=model_id if spec is None else model_id.replace(spec, "", 1), region=region) + + +def _explicit_model_id(params: Mapping[str, object]) -> str | None: + model_id: Final = params.get("model_id") + return model_id if isinstance(model_id, str) else None + + +def _region(model: str, params: Mapping[str, object]) -> str: + model_id: Final = _explicit_model_id(params) + path_region: Final = None if model_id is not None else _model_path(model).region + region_params: Final = ( + {**params, "aws_region_name": path_region} + if path_region is not None and "aws_region_name" not in params + else params + ) + return _AWS._get_aws_region_name(optional_params=dict(region_params), model=model, model_id=model_id) + + +def _body_params(optional_params: Mapping[str, object]) -> dict[str, object]: + return {key: value for key, value in optional_params.items() if key not in _PARAMS_SENT_OUTSIDE_THE_BODY} + + +def _signed_body_text(data: Mapping[str, object], signed_json_body: bytes | None) -> str: + return json.dumps(data) if signed_json_body is None else signed_json_body.decode() + + +class AmazonConverseRuntimeConfig(AmazonConverseConfig, SignsRequestsWithAWS): + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict, + litellm_params: dict, + stream: bool | None = None, + ) -> str: + params: Final = with_aws_credential_fallbacks(optional_params, litellm_params) + runtime_endpoint: Final = params.get("aws_bedrock_runtime_endpoint") + _, endpoint_url = _AWS.get_runtime_endpoint( + api_base=api_base, + aws_bedrock_runtime_endpoint=runtime_endpoint if isinstance(runtime_endpoint, str) else None, + aws_region_name=_region(model, params), + ) + model_id: Final = BaseAWSLLM.encode_model_id(_explicit_model_id(params) or _model_path(model).model_id) + action: Final = "converse-stream" if stream is True else "converse" + return join_api_base(endpoint_url, f"model/{model_id}/{action}") + + def sign_request( + self, + headers: dict, + optional_params: dict, + request_data: dict, + api_base: str, + api_key: str | None = None, + model: str | None = None, + stream: bool | None = None, + fake_stream: bool | None = None, + ) -> tuple[dict, bytes | None]: + return _AWS._sign_request( + service_name="bedrock", + headers=headers, + optional_params={**optional_params, "aws_region_name": _region(model or "", optional_params)}, + request_data=request_data, + api_base=api_base, + api_key=api_key, + model=model, + stream=stream, + fake_stream=fake_stream, + ) + + def transform_request( + self, + model: str, + messages: list[AllMessageValues], + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> dict: + return super().transform_request( + model=model, + messages=messages, + optional_params=_body_params(optional_params), + litellm_params=litellm_params, + headers=headers, + ) + + async def async_transform_request( + self, + model: str, + messages: list[AllMessageValues], + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> dict: + return cast( + dict, + await self._async_transform_request( + model=model, + messages=messages, + optional_params=_body_params(optional_params), + litellm_params=litellm_params, + headers=headers, + ), + ) + + @property + def uses_async_transform_request(self) -> bool: + return True + + def transform_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ModelResponse, + logging_obj: "LiteLLMLoggingObj", + request_data: dict, + messages: list[AllMessageValues], + optional_params: dict, + litellm_params: dict, + encoding: "Tokenizer | None", + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ModelResponse: + response: Final = super().transform_response( + model=model, + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data=request_data, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + encoding=encoding, + api_key=api_key, + json_mode=json_mode, + ) + response.set_provider_response_headers(raw_response.headers) + return response + + @property + def has_custom_stream_wrapper(self) -> bool: + return True + + @property + def supports_stream_param_in_request_body(self) -> bool: + return False + + @track_llm_api_timing() + async def get_async_custom_stream_wrapper( + self, + model: str, + custom_llm_provider: str, + logging_obj: "LiteLLMLoggingObj", + api_base: str, + headers: dict, + data: dict, + messages: list, + client: AsyncHTTPHandler | None = None, + json_mode: bool | None = None, + signed_json_body: bytes | None = None, + *, + litellm_params: Mapping[str, object], + timeout: float | httpx.Timeout | None = None, + ) -> CustomStreamWrapper: + completion_stream, response_headers = await make_call( + client=client, + api_base=api_base, + headers=headers, + data=_signed_body_text(data, signed_json_body), + model=model, + messages=messages, + logging_obj=logging_obj, + json_mode=json_mode, + stream_chunk_size=stored_control_options(litellm_params).stream_chunk_size, + timeout=timeout, + ) + return CustomStreamWrapper( + completion_stream=completion_stream, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + _response_headers=response_headers, + ) + + @track_llm_api_timing() + def get_sync_custom_stream_wrapper( + self, + model: str, + custom_llm_provider: str, + logging_obj: "LiteLLMLoggingObj", + api_base: str, + headers: dict, + data: dict, + messages: list, + client: HTTPHandler | AsyncHTTPHandler | None = None, + json_mode: bool | None = None, + signed_json_body: bytes | None = None, + *, + litellm_params: Mapping[str, object], + timeout: float | httpx.Timeout | None = None, + ) -> CustomStreamWrapper: + completion_stream, response_headers = make_sync_call( + client=client if isinstance(client, HTTPHandler) else None, + api_base=api_base, + headers=headers, + data=_signed_body_text(data, signed_json_body), + signed_json_body=signed_json_body, + model=model, + messages=messages, + logging_obj=logging_obj, + json_mode=json_mode, + stream_chunk_size=stored_control_options(litellm_params).stream_chunk_size, + timeout=timeout, + ) + return CustomStreamWrapper( + completion_stream=completion_stream, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + _response_headers=response_headers, + ) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 8c9de272c0d..0412df3deee 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -68,7 +68,7 @@ from litellm.litellm_core_utils.chat_completion_agentic_loop import ( ) from litellm.litellm_core_utils.core_helpers import as_str_mapping, set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value -from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS +from litellm.litellm_core_utils.get_litellm_params import with_aws_credential_fallbacks from litellm.litellm_core_utils.get_llm_provider_logic import get_provider_affinity_header from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header @@ -577,18 +577,6 @@ def _mask_presigned_request_headers(transformed_request: bytes | str | dict) -> } -def _aws_signing_overrides( - optional_params: Mapping[str, object], litellm_params: Mapping[str, object] -) -> Mapping[str, object]: - return MappingProxyType( - { - key: litellm_params[key] - for key in AWS_CREDENTIAL_KWARGS_KEYS - if optional_params.get(key) is None and litellm_params.get(key) is not None - } - ) - - class _PreparedFileContentRequest(NamedTuple): url: str params: dict @@ -957,10 +945,7 @@ class BaseLLMHTTPHandler: "tuple[dict[str, object], bytes | None]", provider_config.sign_request( headers=request_headers, - optional_params={ - **optional_params, - **_aws_signing_overrides(optional_params, litellm_params), - }, + optional_params={**with_aws_credential_fallbacks(optional_params, litellm_params)}, request_data=data, api_base=api_base, api_key=api_key, diff --git a/litellm/main.py b/litellm/main.py index 1d18100d1d5..b7a715fc890 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -64,6 +64,7 @@ if TYPE_CHECKING: from litellm.router import Router from litellm.types.utils import TokenCountResponse +from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.constants import ( AZURE_OPENAI_AUDIO_PROVIDERS, DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, @@ -258,7 +259,7 @@ from .llms.azure.completion.handler import AzureTextCompletion from .llms.azure_ai.anthropic.handler import AzureAnthropicChatCompletion from .llms.azure_ai.anthropic.transformation import azure_anthropic_messages_url from .llms.azure_ai.embed import AzureAIEmbedding -from .llms.bedrock.chat import BedrockConverseLLM +from .llms.bedrock.chat.converse_runtime import AmazonConverseRuntimeConfig from .llms.bedrock.chat.converse_transformation import AmazonConverseConfig from .llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import AmazonAnthropicClaudeConfig from .llms.bedrock.chat.native_output_fallback import NativeOutputAttemptResult, call_with_native_output_fallback @@ -364,7 +365,6 @@ azure_audio_transcriptions: Final = AzureAudioTranscription() huggingface_embed: Final = HuggingFaceEmbedding() predibase_chat_completions: Final = PredibaseChatCompletion() codestral_text_completions: Final = CodestralTextCompletion() -bedrock_converse_chat_completion: Final = BedrockConverseLLM() bedrock_embedding: Final = BedrockEmbedding() bedrock_image_generation = BedrockImageGeneration() bedrock_image_edit = BedrockImageEdit() @@ -4421,10 +4421,8 @@ def _complete_bedrock(ctx: CompletionDispatchContext) -> _CompletionDispatchResu api_base: Final = ctx.api_base api_key: Final = ctx.api_key client: Final = _dispatch_client_http(ctx) - custom_prompt_dict = ctx.custom_prompt_dict headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params - logger_fn: Final = ctx.logger_fn logging: Final = ctx.logging messages: Final = ctx.messages model = ctx.model @@ -4435,8 +4433,6 @@ def _complete_bedrock(ctx: CompletionDispatchContext) -> _CompletionDispatchResu stream: Final = ctx.stream timeout: Final = ctx.timeout - custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict - if "aws_bedrock_client" in optional_params: verbose_logger.warning( "'aws_bedrock_client' is a deprecated param. Please move to another auth method - https://docs.litellm.ai/docs/providers/bedrock#boto3---authentication." @@ -4488,24 +4484,29 @@ def _complete_bedrock(ctx: CompletionDispatchContext) -> _CompletionDispatchResu elif bedrock_route == "converse": model = without_bedrock_route_prefix(model) converse_model: Final = model + converse_headers: Final = update_headers_with_filtered_beta( + headers=dict(headers or {}), provider="bedrock_converse" + ) def converse_attempt(fallback: NativeOutputFallback | None) -> NativeOutputAttemptResult: - return bedrock_converse_chat_completion.completion( + return base_llm_http_handler.completion( model=converse_model, + stream=stream, messages=messages, - custom_prompt_dict=custom_prompt_dict, + acompletion=acompletion, + api_base=api_base, model_response=model_response, optional_params=optional_params if fallback is None else dict(fallback.optional_params), litellm_params=litellm_params if fallback is None else dict(fallback.litellm_params), - logger_fn=logger_fn, - encoding=_get_encoding(), - logging_obj=logging, - extra_headers=headers, # Use merged headers instead of original extra_headers + shared_session=shared_session, + custom_llm_provider="bedrock", timeout=_native_output_attempt_timeout(ctx, fallback), - acompletion=acompletion, - client=client, - api_base=api_base, + headers=dict(converse_headers), + encoding=_get_encoding(), api_key=api_key, + logging_obj=logging, + client=client, + provider_config=AmazonConverseRuntimeConfig(), ) converse_params: Final = _OPTIONAL_PARAMS_ADAPTER.validate_python(optional_params) diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index 640a947f9c5..73acc367a17 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -48,9 +48,9 @@ def s3_object_reader(bucket_name: str) -> SyncBucketObjectReader: import boto3 from botocore.credentials import Credentials - from litellm.main import bedrock_converse_chat_completion + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM - credentials: Final[Credentials] = bedrock_converse_chat_completion.get_credentials() + credentials: Final[Credentials] = BaseAWSLLM().get_credentials() s3_client: Final = boto3.client( "s3", aws_access_key_id=credentials.access_key, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 627164abb93..4dcbbd73036 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1614,8 +1614,7 @@ async def bedrock_proxy_route( detail="bedrock-agent-runtime pass-through is disabled on this proxy.", ) - from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post - from litellm.llms.bedrock.chat import BedrockConverseLLM + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing, sign_aws_json_post aws_region_name: Final = await run_aws_signing(get_secret_str, secret_name="AWS_REGION_NAME") base_target_url: Final = f"https://bedrock-agent-runtime.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" @@ -1632,7 +1631,7 @@ async def bedrock_proxy_route( path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint) ) - bedrock_llm: Final = BedrockConverseLLM() + bedrock_llm: Final = BaseAWSLLM() # Assuming the body contains JSON data, parse it try: data: Final = await _json_request_body(request) diff --git a/tests/unit/litellm_core_utils/test_get_litellm_params.py b/tests/unit/litellm_core_utils/test_get_litellm_params.py index a1184039102..02793e43ca8 100644 --- a/tests/unit/litellm_core_utils/test_get_litellm_params.py +++ b/tests/unit/litellm_core_utils/test_get_litellm_params.py @@ -20,6 +20,7 @@ from litellm.litellm_core_utils.get_litellm_params import ( get_litellm_params, parse_control_options, stored_control_options, + with_aws_credential_fallbacks, ) from litellm.types.litellm_params import ControlOptions from litellm.utils import get_non_default_completion_params @@ -559,3 +560,24 @@ class TestOpenAIWifKeys: params = get_litellm_params() for key in self.THREE_KEYS: assert key not in params + + +def test_aws_credential_fallbacks_only_fill_credentials_optional_params_leave_unset(): + merged = with_aws_credential_fallbacks( + {"temperature": 0.2, "aws_region_name": "us-west-2", "aws_profile_name": None}, + { + "aws_role_name": "arn:aws:iam::000000000000:role/attributed", + "aws_session_name": "user-123", + "aws_region_name": "us-east-1", + "aws_profile_name": "from-deployment", + "api_key": "not-an-aws-param", + }, + ) + + assert dict(merged) == { + "temperature": 0.2, + "aws_region_name": "us-west-2", + "aws_profile_name": "from-deployment", + "aws_role_name": "arn:aws:iam::000000000000:role/attributed", + "aws_session_name": "user-123", + } diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index ba8fb3bf107..a61acb0adc7 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -3541,25 +3541,6 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques assert retry_authorization != first_attempt_headers["Authorization"] -def test_aws_signing_overrides_only_fills_missing_credentials(): - from litellm.llms.custom_httpx.llm_http_handler import _aws_signing_overrides - - overrides = _aws_signing_overrides( - {"temperature": 0.2, "aws_region_name": "us-west-2"}, - { - "aws_role_name": "arn:aws:iam::000000000000:role/attributed", - "aws_session_name": "user-123", - "aws_region_name": "us-east-1", - "api_key": "not-an-aws-param", - }, - ) - - assert dict(overrides) == { - "aws_role_name": "arn:aws:iam::000000000000:role/attributed", - "aws_session_name": "user-123", - } - - class TestServerFulfilledToolsInRequest: """_server_fulfilled_tools_in_request gates the buffered (non-leaking) streaming mode for server-fulfilled tools like headroom_retrieve.""" diff --git a/tests/unit/proxy/common_utils/test_load_config_utils.py b/tests/unit/proxy/common_utils/test_load_config_utils.py index 927e6ee511d..e78c4ba438d 100644 --- a/tests/unit/proxy/common_utils/test_load_config_utils.py +++ b/tests/unit/proxy/common_utils/test_load_config_utils.py @@ -20,7 +20,7 @@ class TestGetFileContentsFromS3: """Test suite for S3 config loading functionality.""" @patch("boto3.client") - @patch("litellm.main.bedrock_converse_chat_completion") + @patch("litellm.llms.bedrock.base_aws_llm.BaseAWSLLM.get_credentials") @patch("yaml.safe_load") def test_get_file_contents_from_s3_no_temp_file_creation( self, mock_yaml_load, mock_bedrock, mock_boto3_client @@ -38,7 +38,7 @@ class TestGetFileContentsFromS3: mock_credentials.access_key = "test_access_key" mock_credentials.secret_key = "test_secret_key" mock_credentials.token = "test_token" - mock_bedrock.get_credentials.return_value = mock_credentials + mock_bedrock.return_value = mock_credentials # Mock S3 client and response mock_s3_client = MagicMock() diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 11b07472eaa..4402ee8bfaa 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -2245,7 +2245,7 @@ class TestBedrockAgentRuntimePassthroughToggle: "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_secret_str", return_value="us-east-1", ), - patch("litellm.llms.bedrock.chat.BedrockConverseLLM", return_value=bedrock_llm), + patch("litellm.llms.bedrock.base_aws_llm.BaseAWSLLM", return_value=bedrock_llm), patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_request_copy", Mock(),