diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 7290bd7d528..b64bd095293 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -627,17 +627,22 @@ class CheckBatchCost: return None @classmethod - def _get_input_file_id(cls, job: "ManagedObjectRow") -> Optional[str]: + def _stored_batch(cls, job: "ManagedObjectRow") -> Optional["LiteLLMBatch"]: from litellm.types.utils import LiteLLMBatch parsed_file_object: Final = cls._parse_file_object(job) if parsed_file_object is None: return None try: - return LiteLLMBatch.model_validate(parsed_file_object).input_file_id + return LiteLLMBatch.model_validate(parsed_file_object) except Exception: return None + @classmethod + def _get_input_file_id(cls, job: "ManagedObjectRow") -> Optional[str]: + stored_batch: Final = cls._stored_batch(job) + return stored_batch.input_file_id if stored_batch is not None else None + @classmethod def _get_stored_model_id(cls, job: "ManagedObjectRow") -> Optional[str]: """The deployment that created the batch, written into file_object by the managed @@ -1105,6 +1110,8 @@ class CheckBatchCost: - if not, return False - if so, return True """ + from litellm.proxy.openai_files_endpoints.common_utils import created_batch_endpoint + try: from litellm.integrations.prometheus import PrometheusLogger @@ -1188,6 +1195,7 @@ class CheckBatchCost: if self._is_batch_gone_at_provider(e, batch_id) and self._batch_deployment_exists(model_id): await self._retire_job(job, f"batch {batch_id} no longer exists at the provider") continue + response.endpoint = created_batch_endpoint(response, self._stored_batch(job)) ## RETRIEVE THE BATCH JOB OUTPUT FILE if response.status in PROVIDER_TERMINAL_BATCH_STATUSES and response.output_file_id is not None: diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 81057a1ee23..ba9b9f30a2c 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -315,7 +315,10 @@ class BedrockBatchesHandler: # provider uses, so model_list / env / role-assumption configs are # honored. We instantiate BedrockBatchesConfig (which extends # BaseAWSLLM) lazily to avoid a circular import at module load. - from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig + from litellm.llms.bedrock.batches.transformation import ( + BedrockBatchesConfig, + openai_endpoint_for_bedrock_job_model, + ) creds: Final = BedrockBatchesConfig().resolve_credentials(AwsAuthParams.model_validate(kwargs), region) @@ -404,7 +407,7 @@ class BedrockBatchesHandler: request_counts=_record_counts_from_response(response), metadata=openai_batch_metadata, completion_window="24h", - endpoint="/v1/chat/completions", + endpoint=openai_endpoint_for_bedrock_job_model(response.get("modelId")), input_file_id=input_uri, output_file_id=output_file_uri if openai_status == "completed" else None, ) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 3ad8049ec26..52847c71efc 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final, Literal, cast from httpx import Headers, Response from pydantic import TypeAdapter, ValidationError +from litellm.exceptions import ModelNotMappedError from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix, is_bedrock_arn from litellm.litellm_core_utils.cloud_storage_security import ( BEDROCK_MANAGED_S3_BATCH_PREFIX, @@ -32,6 +33,7 @@ from litellm.types.utils import LiteLLMBatch, LlmProviders, Usage from ..base_aws_llm import BaseAWSLLM from ..common_utils import ( CommonBatchFilesUtils, + get_cached_model_info, merge_bedrock_aws_request_params, resolve_s3_bucket_owner, resolve_s3_encryption_key_id, @@ -82,6 +84,21 @@ def _validate_bedrock_tags(raw_tags: object) -> list[BedrockTag]: ) from e +def openai_endpoint_for_bedrock_job_model( + model_id: str | None, +) -> Literal["/v1/chat/completions", "/v1/embeddings"]: + if not model_id: + return "/v1/chat/completions" + try: + model_info: Final = get_cached_model_info()(model=model_id, custom_llm_provider="bedrock") + except ModelNotMappedError: + return "/v1/chat/completions" + is_embedding_model: Final = ( + model_info.get("mode") == "embedding" or model_info.get("output_vector_size") is not None + ) + return "/v1/embeddings" if is_embedding_model else "/v1/chat/completions" + + def titan_embedding_usage_from_batch_output(model_output: Mapping[str, object]) -> Usage | None: """Titan embedding batch lines report usage as a top-level inputTextTokenCount, not a usage block.""" if "embedding" not in model_output and "embeddingsByType" not in model_output: @@ -639,7 +656,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): return LiteLLMBatch( id=job_arn, object="batch", - endpoint="/v1/chat/completions", + endpoint=openai_endpoint_for_bedrock_job_model(response_data.get("modelId")), errors=errors, input_file_id=input_file_id, completion_window="24h", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 891e247747e..b60f6cfb8f0 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -16994,6 +16994,7 @@ "max_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 }, @@ -17004,6 +17005,7 @@ "max_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/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index ead1cd8e2ca..b749dcd5780 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -52,6 +52,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( # noqa: F401 # apply_team_provider_credentials, authorize_model_for_key, batch_cost_poller_is_active, + created_batch_endpoint, decode_model_from_file_id, encode_batch_response_ids, encode_file_id_with_model, @@ -635,6 +636,7 @@ async def retrieve_batch( proxy_logging_obj=proxy_logging_obj, user_api_key_dict=user_api_key_dict, ) + stored_batch: Final = response if executed_batch and response is None: raise batch_error(404, f"No batch found with id '{batch_id}'.") @@ -779,6 +781,8 @@ async def retrieve_batch( **data, ) + response.endpoint = created_batch_endpoint(response, stored_batch) + # FIX: Update the database with the latest state from provider await update_batch_in_database( batch_id=batch_id, @@ -1221,6 +1225,14 @@ async def cancel_batch( managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") from litellm.proxy.proxy_server import prisma_client + db_batch_object, stored_batch = await get_batch_from_database( + batch_id=batch_id, + unified_batch_id=unified_batch_id, + managed_files_obj=managed_files_obj, + prisma_client=prisma_client, + verbose_proxy_logger=verbose_proxy_logger, + ) + response.endpoint = created_batch_endpoint(response, stored_batch) await update_batch_in_database( batch_id=batch_id, unified_batch_id=unified_batch_id, @@ -1228,6 +1240,7 @@ async def cancel_batch( managed_files_obj=managed_files_obj, prisma_client=prisma_client, verbose_proxy_logger=verbose_proxy_logger, + db_batch_object=db_batch_object, operation="cancel", user_api_key_dict=user_api_key_dict, ) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 48306cc8d7a..4f2a7277e4c 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1480,6 +1480,13 @@ def batch_cost_poller_is_active() -> bool: return False +def created_batch_endpoint(provider_batch: "LiteLLMBatch", stored_batch: "LiteLLMBatch | None") -> str: + 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: """Whether a "completed" batch may be retired from cost recovery. diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 891e247747e..b60f6cfb8f0 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -16994,6 +16994,7 @@ "max_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 }, @@ -17004,6 +17005,7 @@ "max_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/integration/spend/test_batch_observability.py b/tests/integration/spend/test_batch_observability.py index ca9c26ff513..bd3462febd0 100644 --- a/tests/integration/spend/test_batch_observability.py +++ b/tests/integration/spend/test_batch_observability.py @@ -2,12 +2,16 @@ from __future__ import annotations import json import uuid +from dataclasses import dataclass from hashlib import sha256 +from pathlib import Path from typing import Final +import httpx import pytest -from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows +from integration._support.process import owned_proxy from integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse from pydantic import JsonValue @@ -59,11 +63,11 @@ def _failed_line(index: int) -> str: ) -def _batch(status: str, *, files_ready: bool) -> dict[str, JsonValue]: +def _batch(status: str, *, files_ready: bool, endpoint: str = "/v1/chat/completions") -> dict[str, JsonValue]: return { "id": "batch-$REQUEST_ID", "object": "batch", - "endpoint": "/v1/chat/completions", + "endpoint": endpoint, "errors": None, "input_file_id": "file-in-$REQUEST_ID", "completion_window": "24h", @@ -79,7 +83,7 @@ def _batch(status: str, *, files_ready: bool) -> dict[str, JsonValue]: } -def _provider_routes() -> RoutedResponse: +def _provider_routes(created_endpoint: str = "/v1/chat/completions") -> RoutedResponse: output_lines: Final = ( _successful_line(1, REASONING_TOKENS[0]), _failed_line(2), @@ -102,7 +106,8 @@ def _provider_routes() -> RoutedResponse: }, ), "POST /batches": JsonResponse( - content_type="application/json", body=_batch("validating", files_ready=False) + content_type="application/json", + body=_batch("validating", files_ready=False, endpoint=created_endpoint), ), "GET /batches/batch-$REQUEST_ID": JsonResponse( content_type="application/json", body=_batch("completed", files_ready=True) @@ -150,44 +155,53 @@ def _metadata(row: dict[str, JsonValue]) -> dict[str, JsonValue]: return object_value(JSON_OBJECT.validate_json(value) if isinstance(value, str) else value) +@dataclass(frozen=True, slots=True) +class CompletedBatch: + key: str + batch_id: str + retrieval: httpx.Response + + +def _complete_a_batch(gateway: Gateway, scenario: Scenario, created_endpoint: str, *, managed: bool) -> CompletedBatch: + scenario_id: Final = f"batch-observability-{uuid.uuid4().hex[:12]}" + handle: Final = register_scenario(scenario_id, _provider_routes(created_endpoint)) + scenario.cleanups.callback(delete_scenario, handle) + model_name: Final = scenario.model(api_base=handle.api_base()) + key: Final = scenario.key(models=[model_name]) + file_response: Final = gateway.request_multipart( + "/v1/files", + {"purpose": "batch", "target_model_names" if managed else "model": model_name}, + {"file": ("in.jsonl", _input_file(model_name), "application/jsonl")}, + key=key, + ) + assert file_response.status_code == 200, file_response.text + input_file_id: Final = string_value(JSON_OBJECT.validate_json(file_response.content)["id"]) + batch_response: Final = gateway.request( + "POST", + "/v1/batches", + {"input_file_id": input_file_id, "endpoint": created_endpoint, "completion_window": "24h"} + | ({} if managed else {"model": model_name}), + key=key, + ) + assert batch_response.status_code == 200, batch_response.text + batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"]) + retrieval: Final = eventually( + lambda: gateway.request("GET", f"/v1/batches/{batch_id}", key=key), + lambda response: response.status_code == 200 and response.json()["status"] == "completed", + seconds=30, + ) + assert retrieval.status_code == 200, retrieval.text + return CompletedBatch(key=key, batch_id=batch_id, retrieval=retrieval) + + @pytest.mark.covers("spend.batches.retrieval_row_aggregates_reasoning_tokens_and_per_request_counts") def test_batch_retrieval_row_sums_reasoning_tokens_and_counts_output_and_error_file_failures( gateway: Gateway, ) -> None: with gateway.scenario() as scenario: - scenario_id: Final = f"batch-observability-{uuid.uuid4().hex[:12]}" - handle: Final = register_scenario(scenario_id, _provider_routes()) - scenario.cleanups.callback(delete_scenario, handle) - model_name: Final = scenario.model(api_base=handle.api_base()) - key: Final = scenario.key(models=[model_name]) - file_response: Final = gateway.request_multipart( - "/v1/files", - {"purpose": "batch", "model": model_name}, - {"file": ("in.jsonl", _input_file(model_name), "application/jsonl")}, - key=key, - ) - assert file_response.status_code == 200, file_response.text - input_file_id: Final = string_value(JSON_OBJECT.validate_json(file_response.content)["id"]) - batch_response: Final = gateway.request( - "POST", - "/v1/batches", - { - "input_file_id": input_file_id, - "endpoint": "/v1/chat/completions", - "completion_window": "24h", - "model": model_name, - }, - key=key, - ) - assert batch_response.status_code == 200, batch_response.text - batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"]) - retrieval: Final = eventually( - lambda: gateway.request("GET", f"/v1/batches/{batch_id}", key=key), - lambda response: response.status_code == 200 and response.json()["status"] == "completed", - seconds=30, - ) - assert retrieval.status_code == 200, retrieval.text - rows: Final = eventually(lambda: _retrieval_rows(key), lambda values: len(values) == 1, seconds=70) + batch: Final = _complete_a_batch(gateway, scenario, "/v1/chat/completions", managed=False) + retrieval: Final = batch.retrieval + rows: Final = eventually(lambda: _retrieval_rows(batch.key), lambda values: len(values) == 1, seconds=70) row: Final = rows[0] metadata: Final = _metadata(row) usage: Final = object_value(metadata["usage_object"]) @@ -199,3 +213,19 @@ def test_batch_retrieval_row_sums_reasoning_tokens_and_counts_output_and_error_f ) assert metadata["batch_successful_requests"] == 2, (retrieval.text, metadata) assert metadata["batch_failed_requests"] == 1 + ERROR_FILE_FAILURES, (retrieval.text, metadata) + + +@pytest.mark.timeout(240) +def test_managed_batch_keeps_its_created_endpoint_when_the_provider_reports_another( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + owned_proxy(gateway, tmp_path, {"PROXY_BATCH_POLLING_ENABLED": "false"}) as candidate, + candidate.scenario() as scenario, + ): + batch: Final = _complete_a_batch(candidate, scenario, "/v1/embeddings", managed=True) + stored: Final = candidate.request("GET", f"/v1/batches/{batch.batch_id}", key=batch.key) + assert (batch.retrieval.json()["endpoint"], stored.json()["endpoint"]) == ( + "/v1/embeddings", + "/v1/embeddings", + ), stored.text diff --git a/tests/unit/llms/bedrock/batches/test_handler.py b/tests/unit/llms/bedrock/batches/test_handler.py index d920eabcdf6..3c105725e10 100644 --- a/tests/unit/llms/bedrock/batches/test_handler.py +++ b/tests/unit/llms/bedrock/batches/test_handler.py @@ -165,6 +165,24 @@ def test_handle_model_invocation_job_status_completed(patched_boto3): assert batch.metadata["output_s3_uri"] == OUTPUT_PREFIX +@pytest.mark.parametrize( + "job_model_id,expected_endpoint", + [ + ("amazon.titan-embed-text-v2:0", "/v1/embeddings"), + ("anthropic.claude-3-5-sonnet-20240620-v1:0", "/v1/chat/completions"), + ], +) +def test_job_status_endpoint_follows_the_job_models_mode( + patched_boto3: tuple[MagicMock, MagicMock], job_model_id: str, expected_endpoint: str +) -> None: + fake_client, _ = patched_boto3 + fake_client.get_model_invocation_job.return_value = {**_fake_boto3_response(), "modelId": job_model_id} + + batch: Final = BedrockBatchesHandler.handle_model_invocation_job_status(batch_id=JOB_ARN) + + assert batch.endpoint == expected_endpoint + + @pytest.mark.parametrize("success_count,error_count", [(100, 0), (86, 14)]) def test_completed_job_maps_provider_record_counts(patched_boto3, success_count, error_count): fake_client, _ = patched_boto3 diff --git a/tests/unit/llms/bedrock/batches/test_transformation.py b/tests/unit/llms/bedrock/batches/test_transformation.py index 5e987239c54..b81100d6b95 100644 --- a/tests/unit/llms/bedrock/batches/test_transformation.py +++ b/tests/unit/llms/bedrock/batches/test_transformation.py @@ -14,11 +14,14 @@ 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 typing import Final from unittest.mock import MagicMock, patch import httpx import pytest +import litellm from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig from litellm.types.utils import LlmProviders @@ -805,6 +808,49 @@ def test_retrieve_response_no_errors_when_no_message(config): assert out.errors is None +@pytest.mark.parametrize( + "job_model_id,expected_endpoint", + [ + ("amazon.titan-embed-text-v2:0", "/v1/embeddings"), + ("arn:aws:bedrock:us-west-2::foundation-model/amazon.titan-embed-text-v2:0", "/v1/embeddings"), + ("anthropic.claude-3-5-sonnet-20240620-v1:0", "/v1/chat/completions"), + ("a-model-the-cost-map-does-not-list", "/v1/chat/completions"), + (None, "/v1/chat/completions"), + ], +) +def test_retrieve_response_endpoint_follows_the_job_models_mode( + config: BedrockBatchesConfig, job_model_id: str | None, expected_endpoint: str +) -> None: + out: Final = config.transform_retrieve_batch_response( + model=None, + raw_response=_raw({"jobArn": ARN, "status": "InProgress", "modelId": job_model_id}), + logging_obj=MagicMock(), + litellm_params={}, + ) + 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, job_model_id: str +) -> None: + monkeypatch.setitem(litellm.model_cost, job_model_id, {**litellm.model_cost[job_model_id], "mode": "batch"}) + out: Final = config.transform_retrieve_batch_response( + model=None, + raw_response=_raw({"jobArn": ARN, "status": "InProgress", "modelId": job_model_id}), + logging_obj=MagicMock(), + litellm_params={}, + ) + assert out.endpoint == "/v1/embeddings" + + def test_retrieve_response_enriches_metadata(config): body = { "jobArn": ARN, diff --git a/tests/unit/proxy/batches_endpoints/test_endpoints.py b/tests/unit/proxy/batches_endpoints/test_endpoints.py index 914a28dac71..e7057019863 100644 --- a/tests/unit/proxy/batches_endpoints/test_endpoints.py +++ b/tests/unit/proxy/batches_endpoints/test_endpoints.py @@ -36,6 +36,7 @@ from contextlib import ExitStack from dataclasses import dataclass from datetime import datetime, timedelta, timezone from pathlib import Path +from collections.abc import Callable from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -1841,6 +1842,66 @@ async def test_retrieve__db_non_terminal_state_syncs_with_provider(retrieve_harn assert retrieve_harness.update_batch_in_db.call_count == 1 +def _hook_rewriting_the_batch_id_to(batch_id: str) -> Callable[..., LiteLLMBatch]: + def hook(*, response: LiteLLMBatch, **_: object) -> LiteLLMBatch: + response.id = batch_id + return response + + return hook + + +@pytest.mark.asyncio +async def test_retrieve__provider_sync_keeps_the_endpoint_the_batch_was_created_for(retrieve_harness): + stored: Final = make_batch(id="batch-unified-blob", status="validating").model_copy(update={"endpoint": "/v1/embeddings"}) + provider_batch: Final = make_batch(id="batch-raw-xyz", status="in_progress") + retrieve_harness.get_batch_from_db.return_value = (MagicMock(status="validating"), stored) + retrieve_harness.router.aretrieve_batch.return_value = provider_batch + retrieve_harness.logging.post_call_success_hook.side_effect = _hook_rewriting_the_batch_id_to("batch-unified-blob") + + with patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID): + resp: Final = await call_retrieve(retrieve_harness, "batch-unified-blob") + + assert (resp.endpoint, resp.status) == ("/v1/embeddings", "in_progress") + assert retrieve_harness.update_batch_in_db.call_args.kwargs["response"].endpoint == "/v1/embeddings" + assert retrieve_harness.logging.post_call_success_hook.call_args.kwargs["response"].endpoint == "/v1/embeddings" + assert (provider_batch.id, provider_batch.endpoint) == ("batch-unified-blob", "/v1/embeddings") + assert resp._hidden_params["unified_batch_id"] == UNIFIED_BATCH_ID + + +@pytest.mark.asyncio +async def test_retrieve__a_saved_batch_without_an_endpoint_keeps_the_providers(retrieve_harness): + stored: Final = make_batch(id="batch-unified-blob", status="validating").model_copy(update={"endpoint": ""}) + provider_batch: Final = 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: Final = 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 +async def test_retrieve__an_embeddings_endpoint_from_the_provider_corrects_a_saved_chat_endpoint(retrieve_harness): + stored: Final = make_batch(id="batch-unified-blob", status="validating").model_copy( + update={"endpoint": "/v1/chat/completions"} + ) + provider_batch: Final = 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: Final = 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): @@ -2595,6 +2656,27 @@ async def test_cancel__db_write_receives_caller_auth(cancel_harness): assert cancel_harness.update_batch_in_db.call_args.kwargs["user_api_key_dict"] is caller +@pytest.mark.asyncio +async def test_cancel__keeps_the_endpoint_the_batch_was_created_for(cancel_harness): + stored_row: Final = MagicMock(status="in_progress") + stored: Final = make_batch(id="batch-unified-blob", status="in_progress").model_copy(update={"endpoint": "/v1/embeddings"}) + provider_batch: Final = make_batch(id="batch-raw-xyz", status="cancelling") + cancel_harness.router.acancel_batch.return_value = provider_batch + cancel_harness.logging.post_call_success_hook.side_effect = _hook_rewriting_the_batch_id_to("batch-unified-blob") + + with ( + patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID), + patch.object(endpoints, "get_batch_from_database", AsyncMock(return_value=(stored_row, stored))), + ): + resp: Final = await call_cancel(cancel_harness, "batch-unified-blob") + + assert (resp.endpoint, resp.status) == ("/v1/embeddings", "cancelling") + written: Final = cancel_harness.update_batch_in_db.call_args.kwargs + assert (written["response"].endpoint, written["db_batch_object"]) == ("/v1/embeddings", stored_row) + assert cancel_harness.logging.post_call_success_hook.call_args.kwargs["response"].endpoint == "/v1/embeddings" + assert (provider_batch.id, provider_batch.endpoint) == ("batch-unified-blob", "/v1/embeddings") + + @pytest.mark.asyncio async def test_cancel__unified_missing_model_id_400(cancel_harness): # unified id with no model_id segment -> get_model_id returns None -> 400. diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index aff145f1e27..ef1471906dd 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -3662,6 +3662,30 @@ class TestMultiPodBatchCostClaim: assert row.batch_processed is True await self._run_deletion_guard(prisma, _CLAIM_OUTPUT_FILE_ID) + @pytest.mark.asyncio + @pytest.mark.parametrize("has_batch_processed_column", (True, False)) + async def test_billing_a_batch_keeps_the_endpoint_it_was_created_for(self, has_batch_processed_column: bool) -> None: + from litellm.types.utils import LiteLLMBatch + + row: Final = _FakeManagedObjectRow() + row.file_object = LiteLLMBatch( + id="batch-456", + status="in_progress", + input_file_id="file-input-1", + endpoint="/v1/embeddings", + completion_window="24h", + created_at=1, + object="batch", + ).model_dump_json() + instance: Final = self._instance(self._prisma(row, []), self._router()) + instance._has_batch_processed_column = has_batch_processed_column + + with self._billing_patches([]): + await instance.check_batch_cost() + + assert row.status == "complete" + assert json.loads(row.file_object)["endpoint"] == "/v1/embeddings" + @pytest.mark.asyncio async def test_schema_without_batch_processed_still_bills(self): """Older schemas have no column to claim, so they keep the pre-fix behavior instead