From 38285e53c340418a1decfe5fd1d958885d23803d Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 28 May 2024 15:45:23 -0700
Subject: [PATCH] working create_batch

---
 litellm/batches/main.py              |  8 +++++++-
 litellm/llms/openai.py               |  2 +-
 litellm/tests/test_openai_batches.py | 28 +++++++++++++++++++---------
 litellm/types/llms/openai.py         |  2 +-
 4 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/litellm/batches/main.py b/litellm/batches/main.py
index 6aa0d8d8846..3bd1d07a475 100644
--- a/litellm/batches/main.py
+++ b/litellm/batches/main.py
@@ -26,6 +26,7 @@ from ..types.llms.openai import (
     CreateFileRequest,
     FileTypes,
     FileObject,
+    Batch,
 )
 
 from typing import Literal, Optional, Dict
@@ -44,6 +45,11 @@ def create_file(
     extra_body: Optional[Dict[str, str]] = None,
     **kwargs,
 ) -> FileObject:
+    """
+    Files are used to upload documents that can be used with features like Assistants, Fine-tuning, and Batch API.
+
+    LiteLLM Equivalent of POST: POST https://api.openai.com/v1/files
+    """
     try:
         optional_params = GenericLiteLLMParams(**kwargs)
         if custom_llm_provider == "openai":
@@ -127,7 +133,7 @@ def create_batch(
     extra_headers: Optional[Dict[str, str]] = None,
     extra_body: Optional[Dict[str, str]] = None,
     **kwargs,
-):
+) -> Batch:
     """
     Creates and executes a batch from an uploaded file of request
 
diff --git a/litellm/llms/openai.py b/litellm/llms/openai.py
index e5f229eb5c9..5c5b837ea6a 100644
--- a/litellm/llms/openai.py
+++ b/litellm/llms/openai.py
@@ -1604,7 +1604,7 @@ class OpenAIBatchesAPI(BaseLLM):
         max_retries: Optional[int],
         organization: Optional[str],
         client: Optional[OpenAI] = None,
-    ):
+    ) -> Batch:
         openai_client: OpenAI = self.get_openai_client(
             api_key=api_key,
             api_base=api_base,
diff --git a/litellm/tests/test_openai_batches.py b/litellm/tests/test_openai_batches.py
index ac282e0ed42..b99991babaa 100644
--- a/litellm/tests/test_openai_batches.py
+++ b/litellm/tests/test_openai_batches.py
@@ -20,6 +20,7 @@ def test_create_batch():
     """
     1. Create File for Batch completion
     2. Create Batch Request
+    3. Retrieve the specific batch
     """
     file_obj = litellm.create_file(
         file=open("openai_batch_completions.jsonl", "rb"),
@@ -33,16 +34,25 @@ def test_create_batch():
         batch_input_file_id is not None
     ), "Failed to create file, expected a non null file_id but got {batch_input_file_id}"
 
-    print("response from creating file=", file_obj)
-    # response = create_batch(
-    #     completion_window="24h",
-    #     endpoint="/v1/chat/completions",
-    #     input_file_id="1",
-    #     custom_llm_provider="openai",
-    #     metadata={"key1": "value1", "key2": "value2"},
-    # )
+    response = litellm.create_batch(
+        completion_window="24h",
+        endpoint="/v1/chat/completions",
+        input_file_id=batch_input_file_id,
+        custom_llm_provider="openai",
+        metadata={"key1": "value1", "key2": "value2"},
+    )
 
-    print("response")
+    print("response from litellm.create_batch=", response)
+
+    assert (
+        response.id is not None
+    ), f"Failed to create batch, expected a non null batch_id but got {response.id}"
+    assert (
+        response.endpoint == "/v1/chat/completions"
+    ), f"Failed to create batch, expected endpoint to be /v1/chat/completions but got {response.endpoint}"
+    assert (
+        response.input_file_id == batch_input_file_id
+    ), f"Failed to create batch, expected input_file_id to be {batch_input_file_id} but got {response.input_file_id}"
     pass
 
 
diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py
index c65fb2d5f78..81a50db6d22 100644
--- a/litellm/types/llms/openai.py
+++ b/litellm/types/llms/openai.py
@@ -19,7 +19,7 @@ from openai.types.beta.threads.run import Run
 from openai.types.beta.assistant import Assistant
 from openai.pagination import SyncCursorPage
 from os import PathLike
-from openai.types import FileObject
+from openai.types import FileObject, Batch
 
 from typing import TypedDict, List, Optional, Tuple, Mapping, IO