diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7d61e731774..bd38f77fa3b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -17536,6 +17536,7 @@ "max_input_tokens": 512, "mode": "embedding", "output_cost_per_token": 0.0, + "output_vector_size": 1024, "source": "https://aws.amazon.com/bedrock/pricing/", "supports_embedding_image_input": true }, @@ -17545,6 +17546,7 @@ "max_input_tokens": 512, "mode": "embedding", "output_cost_per_token": 0.0, + "output_vector_size": 1024, "source": "https://aws.amazon.com/bedrock/pricing/", "supports_embedding_image_input": true }, diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 1222b7004a5..8223ae34940 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1608,7 +1608,10 @@ def batch_cost_poller_is_active() -> bool: def created_batch_endpoint(provider_batch: "LiteLLMBatch", stored_batch: "LiteLLMBatch | None") -> str: - return stored_batch.endpoint if stored_batch is not None and stored_batch.endpoint else provider_batch.endpoint + if stored_batch is None or not stored_batch.endpoint: + return provider_batch.endpoint + provider_named_the_endpoint: Final = provider_batch.endpoint not in ("", "/v1/chat/completions") + return provider_batch.endpoint if provider_named_the_endpoint else stored_batch.endpoint def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7d61e731774..bd38f77fa3b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -17536,6 +17536,7 @@ "max_input_tokens": 512, "mode": "embedding", "output_cost_per_token": 0.0, + "output_vector_size": 1024, "source": "https://aws.amazon.com/bedrock/pricing/", "supports_embedding_image_input": true }, @@ -17545,6 +17546,7 @@ "max_input_tokens": 512, "mode": "embedding", "output_cost_per_token": 0.0, + "output_vector_size": 1024, "source": "https://aws.amazon.com/bedrock/pricing/", "supports_embedding_image_input": true }, diff --git a/tests/unit/llms/bedrock/batches/test_handler.py b/tests/unit/llms/bedrock/batches/test_handler.py index ce507e99f4b..a04903e85e6 100644 --- a/tests/unit/llms/bedrock/batches/test_handler.py +++ b/tests/unit/llms/bedrock/batches/test_handler.py @@ -178,7 +178,7 @@ def test_job_status_endpoint_follows_the_job_models_mode( fake_client, _ = patched_boto3 fake_client.get_model_invocation_job.return_value = {**_fake_boto3_response(), "modelId": job_model_id} - batch = BedrockBatchesHandler.handle_model_invocation_job_status(batch_id=JOB_ARN) + batch: Final = BedrockBatchesHandler.handle_model_invocation_job_status(batch_id=JOB_ARN) assert batch.endpoint == expected_endpoint diff --git a/tests/unit/llms/bedrock/batches/test_transformation.py b/tests/unit/llms/bedrock/batches/test_transformation.py index e774082668b..7fe4ba1a881 100644 --- a/tests/unit/llms/bedrock/batches/test_transformation.py +++ b/tests/unit/llms/bedrock/batches/test_transformation.py @@ -14,6 +14,7 @@ URL/ARN handling, and the error class. AWS auth/sigv4 is the only external seam we mock; everything else runs for real. """ +from collections.abc import Mapping from datetime import datetime, timezone from typing import Final from unittest.mock import MagicMock, patch @@ -932,7 +933,7 @@ def test_retrieve_response_no_errors_when_no_message(config): def test_retrieve_response_endpoint_follows_the_job_models_mode( config: BedrockBatchesConfig, job_model_id: str | None, expected_endpoint: str ) -> None: - out = config.transform_retrieve_batch_response( + out: Final = config.transform_retrieve_batch_response( model=None, raw_response=_raw({"jobArn": ARN, "status": "InProgress", "modelId": job_model_id}), logging_obj=MagicMock(), @@ -941,12 +942,19 @@ def test_retrieve_response_endpoint_follows_the_job_models_mode( assert out.endpoint == expected_endpoint +def _is_bedrock_embedding_model(cost_map_entry: Mapping[str, object]) -> bool: + return cost_map_entry.get("litellm_provider") == "bedrock" and cost_map_entry.get("mode") == "embedding" + + +@pytest.mark.parametrize( + "job_model_id", + sorted(model_id for model_id, entry in litellm.model_cost.items() if _is_bedrock_embedding_model(entry)), +) def test_retrieve_response_endpoint_survives_a_batch_deployment_overriding_the_models_mode( - config: BedrockBatchesConfig, monkeypatch: pytest.MonkeyPatch + config: BedrockBatchesConfig, monkeypatch: pytest.MonkeyPatch, job_model_id: str ) -> None: - job_model_id = "amazon.titan-embed-text-v2:0" monkeypatch.setitem(litellm.model_cost, job_model_id, {**litellm.model_cost[job_model_id], "mode": "batch"}) - out = config.transform_retrieve_batch_response( + out: Final = config.transform_retrieve_batch_response( model=None, raw_response=_raw({"jobArn": ARN, "status": "InProgress", "modelId": job_model_id}), logging_obj=MagicMock(), diff --git a/tests/unit/proxy/batches_endpoints/test_endpoints.py b/tests/unit/proxy/batches_endpoints/test_endpoints.py index 01c5ef8c3bd..b3e04c4ea94 100644 --- a/tests/unit/proxy/batches_endpoints/test_endpoints.py +++ b/tests/unit/proxy/batches_endpoints/test_endpoints.py @@ -1952,6 +1952,24 @@ async def test_retrieve__a_saved_batch_without_an_endpoint_keeps_the_providers(r assert retrieve_harness.update_batch_in_db.call_args.kwargs["response"].endpoint == "/v1/embeddings" +@pytest.mark.asyncio +async def test_retrieve__an_embeddings_endpoint_from_the_provider_corrects_a_saved_chat_endpoint(retrieve_harness): + stored = make_batch(id="batch-unified-blob", status="validating").model_copy( + update={"endpoint": "/v1/chat/completions"} + ) + provider_batch = make_batch(id="batch-raw-xyz", status="in_progress").model_copy( + update={"endpoint": "/v1/embeddings"} + ) + retrieve_harness.get_batch_from_db.return_value = (MagicMock(status="validating"), stored) + retrieve_harness.router.aretrieve_batch.return_value = provider_batch + + with patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID): + resp = await call_retrieve(retrieve_harness, "batch-unified-blob") + + assert resp.endpoint == "/v1/embeddings" + assert retrieve_harness.update_batch_in_db.call_args.kwargs["response"].endpoint == "/v1/embeddings" + + @pytest.mark.asyncio @pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"]) async def test_retrieve__executed_batch_served_from_db_in_every_status(retrieve_harness, status):