diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 80351664dfe..4fd35c9d691 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -207,7 +207,10 @@ async def _get_batch_output_file_content_as_dictionary( ) if custom_llm_provider == "vertex_ai": - raise ValueError("Vertex AI does not support file content retrieval") + return await _get_vertex_ai_batch_output_with_custom_id( + batch=batch, + litellm_params=litellm_params, + ) if batch.output_file_id is None: raise ValueError("Output file id is None cannot retrieve file content") @@ -235,6 +238,60 @@ async def _get_batch_output_file_content_as_dictionary( return _get_file_content_as_dictionary(_file_content.content) +async def _get_vertex_ai_batch_output_with_custom_id( + batch: Batch, + litellm_params: Optional[dict] = None, +) -> List[dict]: + """ + Vertex AI batch outputs do not reliably include custom_id. + Map output lines back to custom_id using keyField (preferred). + """ + from litellm.files.main import afile_content + + if batch.output_file_id is None: + raise ValueError("Output file id is None cannot retrieve file content") + + if batch.input_file_id is None: + raise ValueError("Input file id is None cannot map vertex batch outputs") + + credentials = _extract_file_access_credentials(litellm_params) + + output_content = await afile_content( + file_id=batch.output_file_id, + custom_llm_provider="vertex_ai", + **credentials, + ) + output_lines = _get_file_content_as_dictionary(output_content.content) + + input_content = await afile_content( + file_id=batch.input_file_id, + custom_llm_provider="vertex_ai", + **credentials, + ) + input_lines = _get_file_content_as_dictionary(input_content.content) + + if len(output_lines) != len(input_lines): + raise ValueError( + "Vertex AI batch output line count does not match input line count" + ) + + output_has_keys = all( + (line.get("custom_id") or line.get("key")) for line in output_lines + ) + if not output_has_keys: + raise ValueError( + "Vertex AI batch output is missing custom_id/key; cannot safely map results" + ) + + mapped_lines = [] + for line in output_lines: + key = line.get("custom_id") or line.get("key") + if key is not None: + line["custom_id"] = key + mapped_lines.append(line) + return mapped_lines + + def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict: """ Extract credentials from litellm_params for file access operations. @@ -405,4 +462,4 @@ def _batch_response_was_successful(batch_job_output_file: dict) -> bool: Check if the batch job response status == 200 """ _response: dict = batch_job_output_file.get("response", None) or {} - return _response.get("status_code", None) == 200 \ No newline at end of file + return _response.get("status_code", None) == 200 diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index 7cb06fea9e2..5b08b693721 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -41,6 +41,7 @@ class VertexAIBatchTransformation: return VertexAIBatchPredictionJob( inputConfig=input_config, outputConfig=output_config, + instanceConfig={"keyField": "custom_id"}, model=model, displayName=request_display_name, ) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index f0493cd6be9..40e5702b061 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -243,7 +243,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): litellm_params={}, cached_content=None, ) - vertex_jsonl_content.append({"request": vertex_request_body}) + vertex_line: Dict[str, Any] = {"request": vertex_request_body} + custom_id = _openai_jsonl_content.get("custom_id") + if custom_id is not None: + vertex_line["custom_id"] = custom_id + vertex_jsonl_content.append(vertex_line) return vertex_jsonl_content def transform_create_file_request( diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 190e680b7b9..3a1a90c0b0f 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -2,6 +2,7 @@ from enum import Enum from typing import Any, Dict, List, Literal, Optional, Union from typing_extensions import ( + NotRequired, Required, TypedDict, ) @@ -568,6 +569,13 @@ class InputConfig(TypedDict): gcsSource: GcsSource +class InstanceConfig(TypedDict, total=False): + instanceType: str + keyField: str + includedFields: List[str] + excludedFields: List[str] + + class GcsDestination(TypedDict): outputUriPrefix: str @@ -630,6 +638,7 @@ class VertexAIBatchPredictionJob(TypedDict): displayName: str model: str inputConfig: InputConfig + instanceConfig: NotRequired[InstanceConfig] outputConfig: OutputConfig diff --git a/tests/batches_tests/test_batch_custom_id_mapping_live.py b/tests/batches_tests/test_batch_custom_id_mapping_live.py new file mode 100644 index 00000000000..c028a8aaa9d --- /dev/null +++ b/tests/batches_tests/test_batch_custom_id_mapping_live.py @@ -0,0 +1,235 @@ +import asyncio +import json +import os +import tempfile +from typing import List, Optional + +import pytest + +import litellm + + +def _write_batch_jsonl(model: str) -> str: + records = [ + { + "custom_id": "request-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": model, + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello world!"}, + ], + "max_tokens": 10, + }, + }, + { + "custom_id": "request-2", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": model, + "messages": [ + {"role": "system", "content": "You are an unhelpful assistant."}, + {"role": "user", "content": "Hello world!"}, + ], + "max_tokens": 10, + }, + }, + { + "custom_id": "request-3", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": model, + "messages": [ + {"role": "system", "content": "Answer with a single word."}, + {"role": "user", "content": "Hi"}, + ], + "max_tokens": 5, + }, + }, + ] + + fd, path = tempfile.mkstemp(suffix=".jsonl") + os.close(fd) + with open(path, "w", encoding="utf-8") as f: + for record in records: + f.write(json.dumps(record) + "\n") + return path + + +def _load_vertex_ai_credentials_from_env() -> None: + os.environ["GCS_FLUSH_INTERVAL"] = "1" + + private_key_id = os.environ.get("GCS_PRIVATE_KEY_ID", "") + private_key = os.environ.get("GCS_PRIVATE_KEY", "").replace("\\n", "\n") + + if not private_key_id or not private_key: + return + + service_account_key_data = { + "private_key_id": private_key_id, + "private_key": private_key, + } + + with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_file: + json.dump(service_account_key_data, temp_file, indent=2) + + abs_path = os.path.abspath(temp_file.name) + os.environ["GCS_PATH_SERVICE_ACCOUNT"] = abs_path + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = abs_path + + +def _normalize_credential_path(env_key: str) -> Optional[str]: + path = os.getenv(env_key) + if not path: + return None + abs_path = os.path.abspath(path) + if not os.path.exists(abs_path): + return None + os.environ[env_key] = abs_path + return abs_path + + +def _extract_custom_ids_from_output(lines: List[dict]) -> List[str]: + custom_ids: List[str] = [] + for line in lines: + if "custom_id" in line: + custom_ids.append(line["custom_id"]) + elif "key" in line: + custom_ids.append(line["key"]) + return custom_ids + + +async def _wait_for_batch_completion( + batch_id: str, + provider: str, + timeout_seconds: Optional[int] = None, + poll_seconds: Optional[int] = None, + require_output_file_id: bool = True, +): + timeout_seconds = timeout_seconds or int( + os.getenv("BATCH_WAIT_TIMEOUT_SECONDS", "1800") + ) + poll_seconds = poll_seconds or int(os.getenv("BATCH_WAIT_POLL_SECONDS", "10")) + + start = asyncio.get_event_loop().time() + while True: + batch = await litellm.aretrieve_batch( + batch_id=batch_id, + custom_llm_provider=provider, + litellm_params={ + "litellm_metadata": {"batch_ignore_default_logging": True}, + }, + ) + if batch.status == "completed": + if not require_output_file_id or getattr(batch, "output_file_id", None): + return batch + if batch.status in ["failed", "cancelled", "expired"]: + raise AssertionError(f"Batch ended with status={batch.status}") + if asyncio.get_event_loop().time() - start > timeout_seconds: + raise AssertionError("Timed out waiting for batch completion") + await asyncio.sleep(poll_seconds) + + +@pytest.mark.asyncio +async def test_openai_batch_custom_id_mapping_live(): + api_key = os.getenv("OPENAI_API_KEY") + if not api_key: + pytest.skip("OPENAI_API_KEY not set") + + file_path = _write_batch_jsonl(model="gpt-4o-mini") + + with open(file_path, "rb") as file_handle: + file_obj = await litellm.acreate_file( + file=file_handle, + purpose="batch", + custom_llm_provider="openai", + ) + + create_batch_response = await litellm.acreate_batch( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id=file_obj.id, + custom_llm_provider="openai", + ) + + completed_batch = await _wait_for_batch_completion( + batch_id=create_batch_response.id, + provider="openai", + ) + + assert completed_batch.output_file_id is not None + + output_content = await litellm.afile_content( + file_id=completed_batch.output_file_id, + custom_llm_provider="openai", + ) + + output_lines = [ + json.loads(line) + for line in output_content.content.decode("utf-8").strip().split("\n") + if line.strip() + ] + + custom_ids = _extract_custom_ids_from_output(output_lines) + assert sorted(custom_ids) == ["request-1", "request-2", "request-3"] + + +@pytest.mark.asyncio +async def test_vertex_batch_custom_id_mapping_live(): + if not os.getenv("VERTEXAI_PROJECT") or not os.getenv("VERTEXAI_LOCATION"): + pytest.skip("VERTEXAI_PROJECT or VERTEXAI_LOCATION not set") + if not os.getenv("GCS_BUCKET_NAME"): + pytest.skip("GCS_BUCKET_NAME not set") + + _load_vertex_ai_credentials_from_env() + + normalized_gcs_path = _normalize_credential_path("GCS_PATH_SERVICE_ACCOUNT") + normalized_google_path = _normalize_credential_path( + "GOOGLE_APPLICATION_CREDENTIALS" + ) + if not normalized_gcs_path and not normalized_google_path: + pytest.skip( + "Vertex credentials not set or file not found " + "(GCS_PATH_SERVICE_ACCOUNT/GOOGLE_APPLICATION_CREDENTIALS)" + ) + + file_path = _write_batch_jsonl(model="gemini-2.5-flash-lite") + + with open(file_path, "rb") as file_handle: + file_obj = await litellm.acreate_file( + file=file_handle, + purpose="batch", + custom_llm_provider="vertex_ai", + ) + + create_batch_response = await litellm.acreate_batch( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id=file_obj.id, + custom_llm_provider="vertex_ai", + ) + + completed_batch = await _wait_for_batch_completion( + batch_id=create_batch_response.id, + provider="vertex_ai", + ) + + assert completed_batch.output_file_id is not None + + output_content = await litellm.afile_content( + file_id=completed_batch.output_file_id, + custom_llm_provider="vertex_ai", + ) + + output_lines = [ + json.loads(line) + for line in output_content.content.decode("utf-8").strip().split("\n") + if line.strip() + ] + + custom_ids = _extract_custom_ids_from_output(output_lines) + assert sorted(custom_ids) == ["request-1", "request-2", "request-3"] diff --git a/tests/test_litellm/batches/test_vertex_batch_mapping.py b/tests/test_litellm/batches/test_vertex_batch_mapping.py new file mode 100644 index 00000000000..7787477000c --- /dev/null +++ b/tests/test_litellm/batches/test_vertex_batch_mapping.py @@ -0,0 +1,127 @@ +import json +from types import SimpleNamespace + +import pytest + +from litellm.batches.batch_utils import _get_batch_output_file_content_as_dictionary + + +class DummyBatch(SimpleNamespace): + pass + + +@pytest.mark.asyncio +async def test_should_map_vertex_output_by_key_field(monkeypatch): + """Ensure Vertex output lines with key/custom_id are mapped directly.""" + input_lines = [ + {"custom_id": "id-1", "body": {"model": "m", "messages": []}}, + {"custom_id": "id-2", "body": {"model": "m", "messages": []}}, + ] + output_lines = [ + {"key": "id-2", "response": {"status_code": 200}}, + {"key": "id-1", "response": {"status_code": 200}}, + ] + + content_map = { + "gs://bucket/input.jsonl": "\n".join( + json.dumps(line) for line in input_lines + ).encode("utf-8"), + "gs://bucket/output.jsonl": "\n".join( + json.dumps(line) for line in output_lines + ).encode("utf-8"), + } + + async def fake_afile_content(*, file_id, **_kwargs): + return SimpleNamespace(content=content_map[file_id]) + + monkeypatch.setattr("litellm.files.main.afile_content", fake_afile_content) + + batch = DummyBatch( + input_file_id="gs://bucket/input.jsonl", + output_file_id="gs://bucket/output.jsonl", + ) + + result = await _get_batch_output_file_content_as_dictionary( + batch=batch, + custom_llm_provider="vertex_ai", + ) + + assert [line.get("custom_id") for line in result] == ["id-2", "id-1"] + + +@pytest.mark.asyncio +async def test_should_map_vertex_output_by_input_order(monkeypatch): + """Should error when output lacks key/custom_id fields.""" + input_lines = [ + {"custom_id": "id-1", "body": {"model": "m", "messages": []}}, + {"custom_id": "id-2", "body": {"model": "m", "messages": []}}, + ] + output_lines = [ + {"response": {"status_code": 200}}, + {"response": {"status_code": 200}}, + ] + + content_map = { + "gs://bucket/input.jsonl": "\n".join( + json.dumps(line) for line in input_lines + ).encode("utf-8"), + "gs://bucket/output.jsonl": "\n".join( + json.dumps(line) for line in output_lines + ).encode("utf-8"), + } + + async def fake_afile_content(*, file_id, **_kwargs): + return SimpleNamespace(content=content_map[file_id]) + + monkeypatch.setattr("litellm.files.main.afile_content", fake_afile_content) + + batch = DummyBatch( + input_file_id="gs://bucket/input.jsonl", + output_file_id="gs://bucket/output.jsonl", + ) + + with pytest.raises( + ValueError, + match="Vertex AI batch output is missing custom_id/key", + ): + await _get_batch_output_file_content_as_dictionary( + batch=batch, + custom_llm_provider="vertex_ai", + ) + + +@pytest.mark.asyncio +async def test_should_raise_on_vertex_output_count_mismatch(monkeypatch): + """Guard against mismatched input/output line counts for Vertex.""" + input_lines = [ + {"custom_id": "id-1", "body": {"model": "m", "messages": []}}, + {"custom_id": "id-2", "body": {"model": "m", "messages": []}}, + ] + output_lines = [ + {"response": {"status_code": 200}}, + ] + + content_map = { + "gs://bucket/input.jsonl": "\n".join( + json.dumps(line) for line in input_lines + ).encode("utf-8"), + "gs://bucket/output.jsonl": "\n".join( + json.dumps(line) for line in output_lines + ).encode("utf-8"), + } + + async def fake_afile_content(*, file_id, **_kwargs): + return SimpleNamespace(content=content_map[file_id]) + + monkeypatch.setattr("litellm.files.main.afile_content", fake_afile_content) + + batch = DummyBatch( + input_file_id="gs://bucket/input.jsonl", + output_file_id="gs://bucket/output.jsonl", + ) + + with pytest.raises(ValueError, match="output line count does not match input"): + await _get_batch_output_file_content_as_dictionary( + batch=batch, + custom_llm_provider="vertex_ai", + )