(Feat - Batches API) add support for retrieving vertex api batch jobs (#7661)

* add _async_retrieve_batch

* fix aretrieve_batch

* fix _get_batch_id_from_vertex_ai_batch_response

* fix batches docs
This commit is contained in:
Ishaan Jaff 2025-01-09 18:35:03 -08:00 committed by GitHub
parent 2507c275f6
commit 13f364682d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 177 additions and 2 deletions

View file

@ -2480,7 +2480,7 @@ create_batch_response = oai_client.batches.create(
```json
{
"id": "projects/633608382793/locations/us-central1/batchPredictionJobs/986266568679751680",
"id": "3814889423749775360",
"completion_window": "24hrs",
"created_at": 1733392026,
"endpoint": "",
@ -2503,6 +2503,43 @@ create_batch_response = oai_client.batches.create(
}
```
#### 4. Retrieve a batch
```python
retrieved_batch = oai_client.batches.retrieve(
batch_id=create_batch_response.id,
extra_body={"custom_llm_provider": "vertex_ai"}, # tell litellm to use `vertex_ai` for this batch request
)
```
**Expected Response**
```json
{
"id": "3814889423749775360",
"completion_window": "24hrs",
"created_at": 1736500100,
"endpoint": "",
"input_file_id": "gs://example-bucket-1-litellm/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/7b2e47f5-3dd4-436d-920f-f9155bbdc952",
"object": "batch",
"status": "completed",
"cancelled_at": null,
"cancelling_at": null,
"completed_at": null,
"error_file_id": null,
"errors": null,
"expired_at": null,
"expires_at": null,
"failed_at": null,
"finalizing_at": null,
"in_progress_at": null,
"metadata": null,
"output_file_id": "gs://example-bucket-1-litellm/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001",
"request_counts": null
}
```
## **Fine Tuning APIs**