diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index d63482977ed..5ebdadba4b5 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -5,6 +5,7 @@ from typing import Any, List, Literal, Optional, Tuple import httpx import litellm +from litellm.llms.vertex_ai.batches.utils import get_vertex_ai_batch_output_with_custom_id from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.types.llms.openai import Batch @@ -207,7 +208,7 @@ async def _get_batch_output_file_content_as_dictionary( ) if custom_llm_provider == "vertex_ai": - return await _get_vertex_ai_batch_output_with_custom_id( + return await get_vertex_ai_batch_output_with_custom_id( batch=batch, litellm_params=litellm_params, ) @@ -238,45 +239,6 @@ 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") - - 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) - - 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. diff --git a/litellm/llms/vertex_ai/batches/utils.py b/litellm/llms/vertex_ai/batches/utils.py new file mode 100644 index 00000000000..26b1545504f --- /dev/null +++ b/litellm/llms/vertex_ai/batches/utils.py @@ -0,0 +1,47 @@ +from typing import List, Optional + +from litellm.types.llms.openai import Batch + + +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") + + from litellm.batches.batch_utils import ( + _extract_file_access_credentials, + _get_file_content_as_dictionary, + ) + + 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) + + 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