Merge pull request #2 from SankalpC10/fix/vertex-batch-output-file-id

added complete filepath in vertex_ai batch response

Author: SankalpC10

litellm/llms/vertex_ai/batches/transformation.py

@@ -55,6 +55,9 @@ def transform_vertex_ai_batch_response_to_openai_batch_response(
             created_at=_convert_vertex_datetime_to_openai_datetime(
                 vertex_datetime=response.get("createTime", "")
             ),
+            completed_at=_convert_vertex_datetime_to_openai_datetime(
+                vertex_datetime=response.get("endTime", "")
+            ),
             endpoint="",
             input_file_id=cls._get_input_file_id_from_vertex_ai_batch_response(
                 response
@@ -124,7 +127,15 @@ def _get_output_file_id_from_vertex_ai_batch_response(
             return output_file_id
 
         output_uri_prefix = gcs_destination.get("outputUriPrefix", "")
-        return output_uri_prefix
+        # return output_uri_prefix
+        output_info = response.get("outputInfo")
+        if output_info is None:
+            return output_file_id
+        output_directory = output_info.get("gcsOutputDirectory")
+        if output_directory is None:
+            return output_file_id
+
+        return output_directory
 
     @classmethod
     def _get_batch_job_status_from_vertex_ai_batch_response(

litellm/types/llms/vertex_ai.py

@@ -511,6 +511,8 @@ class OutputConfig(TypedDict, total=False):
     predictionsFormat: str
     gcsDestination: GcsDestination
 
+class OutputInfo(TypedDict):
+    gcsOutputDirectory: str
 
 class GcsBucketResponse(TypedDict):
     """
@@ -570,8 +572,10 @@ class VertexBatchPredictionResponse(TypedDict, total=False):
     model: str
     inputConfig: InputConfig
     outputConfig: OutputConfig
+    outputInfo: OutputInfo
     state: str
     createTime: str
+    endTime: str
     updateTime: str
     modelVersionId: str