Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion litellm/llms/vertex_ai/batches/transformation.py
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,8 @@ def _get_output_file_id_from_vertex_ai_batch_response(cls, response: VertexBatch
Gets the output file id from the Vertex AI Batch response
"""

output_file_id: str = response.get("outputInfo", OutputInfo()).get("gcsOutputDirectory", "")
output_info = response.get("outputInfo") or OutputInfo()
Comment thread
greptile-apps[bot] marked this conversation as resolved.
output_file_id: str = output_info.get("gcsOutputDirectory", "")
if output_file_id:
output_file_id = output_file_id.rstrip("/") + "/predictions.jsonl"
if output_file_id and output_file_id != "/predictions.jsonl":
Expand Down
12 changes: 12 additions & 0 deletions tests/test_litellm/llms/vertex_ai/batches/test_transformation.py
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,18 @@ def test_get_output_file_id_empty_output_info_falls_through_to_output_config():
assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl"


def test_get_output_file_id_output_info_explicit_none_falls_through_to_output_config():
resp = {
"outputInfo": None,
"outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg"}},
}
assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl"


def test_get_output_file_id_output_info_explicit_none_and_no_output_config():
assert T._get_output_file_id_from_vertex_ai_batch_response({"outputInfo": None}) == ""


def test_get_output_file_id_no_output_info_and_no_output_config():
assert T._get_output_file_id_from_vertex_ai_batch_response({}) == ""

Expand Down
Loading