working create_batch

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
 

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,

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
 
 

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