diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 995c45b925..57a9857dd6 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -1045,28 +1045,44 @@ def _handle_async_invoke_status( # Transform response to a LiteLLMBatch object from litellm.types.utils import LiteLLMBatch + # Normalize status to lowercase (AWS returns 'Completed', 'Failed', etc.) + aws_status_raw = status_response.get("status", "") + aws_status_lower = aws_status_raw.lower() + # Map AWS status values to LiteLLM expected values + status_mapping = { + "completed": "completed", + "failed": "failed", + "inprogress": "in_progress", + "in_progress": "in_progress", + } + normalized_status = status_mapping.get(aws_status_lower, aws_status_lower) + + # Get output S3 URI safely + output_s3_uri = "" + try: + output_s3_uri = status_response["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"] + except (KeyError, TypeError): + pass + + # Use BedrockBatchesConfig's timestamp parsing method (expects raw AWS status string) + from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig + created_at, in_progress_at, completed_at, failed_at, _, _ = BedrockBatchesConfig()._parse_timestamps_and_status(status_response, aws_status_raw) result = LiteLLMBatch( id=status_response["invocationArn"], object="batch", - status=status_response["status"], - created_at=status_response["submitTime"], - in_progress_at=status_response["lastModifiedTime"], - completed_at=status_response.get("endTime"), - failed_at=( - status_response.get("endTime") - if status_response["status"] == "failed" - else None - ), + status=normalized_status, + created_at=created_at, + in_progress_at=in_progress_at, + completed_at=completed_at, + failed_at=failed_at, request_counts=BatchRequestCounts( total=1, - completed=1 if status_response["status"] == "completed" else 0, - failed=1 if status_response["status"] == "failed" else 0, + completed=1 if normalized_status == "completed" else 0, + failed=1 if normalized_status == "failed" else 0, ), metadata=dict( **{ - "output_file_id": status_response["outputDataConfig"][ - "s3OutputDataConfig" - ]["s3Uri"], + "output_file_id": output_s3_uri, "failure_message": status_response.get("failureMessage") or "", "model_arn": status_response["modelArn"], } diff --git a/litellm/llms/bedrock/embed/amazon_nova_transformation.py b/litellm/llms/bedrock/embed/amazon_nova_transformation.py index 97652175a9..ada49d0ff2 100644 --- a/litellm/llms/bedrock/embed/amazon_nova_transformation.py +++ b/litellm/llms/bedrock/embed/amazon_nova_transformation.py @@ -83,6 +83,8 @@ class AmazonNovaEmbeddingConfig: # Start with inference_params (user-provided params) embedding_params = inference_params.copy() + embedding_params.pop("output_s3_uri", None) + # Map OpenAI dimensions to embeddingDimension if provided if "dimensions" in embedding_params: embedding_params["embeddingDimension"] = embedding_params.pop("dimensions") diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index be2bfcd70e..c9eea516eb 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -603,22 +603,39 @@ class BedrockEmbedding(BaseAWSLLM): aws_region_name=aws_region_name, ) - # Construct the status check URL - status_url = f"{endpoint_url}/async-invoke/{invocation_arn}" - # Prepare headers + from urllib.parse import quote + + # Encode the ARN for use in URL path + encoded_arn = quote(invocation_arn, safe="") + status_url = f"{endpoint_url.rstrip('/')}/async-invoke/{encoded_arn}" + + # Prepare headers for GET request headers = {"Content-Type": "application/json"} - # Get AWS signed headers - prepped = self.get_request_headers( # type: ignore - credentials=credentials, - aws_region_name=aws_region_name, - extra_headers=None, - endpoint_url=status_url, - data="", # GET request, no body + # Use AWSRequest directly for GET requests (get_request_headers hardcodes POST) + try: + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError( + "Missing boto3 to call bedrock. Run 'pip install boto3'." + ) + + # Create AWSRequest with GET method and encoded URL + request = AWSRequest( + method="GET", + url=status_url, + data=None, # GET request, no body headers=headers, - api_key=None, ) + + # Sign the request - SigV4Auth will create canonical string from request URL + sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name) + sigv4.add_auth(request) + + # Prepare the request + prepped = request.prepare() # LOGGING if logging_obj is not None: