feat - add litellm.acreate_file

litellm/batches/main.py

@@ -10,11 +10,14 @@ https://platform.openai.com/docs/api-reference/batch
 
 """
 
-from typing import Iterable
 import os
-import litellm
-from openai import OpenAI
+import asyncio
+from functools import partial
+import contextvars
+from typing import Literal, Optional, Dict, Coroutine, Any, Union
 import httpx
+
+import litellm
 from litellm import client
 from litellm.utils import supports_httpx_timeout
 from ..types.router import *
@@ -29,14 +32,51 @@ from ..types.llms.openai import (
     Batch,
 )
 
-from typing import Literal, Optional, Dict
-
 ####### ENVIRONMENT VARIABLES ###################
 openai_batches_instance = OpenAIBatchesAPI()
 openai_files_instance = OpenAIFilesAPI()
 #################################################
 
 
+async def acreate_file(
+    file: FileTypes,
+    purpose: Literal["assistants", "batch", "fine-tune"],
+    custom_llm_provider: Literal["openai"] = "openai",
+    extra_headers: Optional[Dict[str, str]] = None,
+    extra_body: Optional[Dict[str, str]] = None,
+    **kwargs,
+) -> Coroutine[Any, Any, 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
+    """
+    loop = asyncio.get_event_loop()
+    kwargs["acreate_file"] = True
+
+    # Use a partial function to pass your keyword arguments
+    func = partial(
+        create_file,
+        file,
+        purpose,
+        custom_llm_provider,
+        extra_headers,
+        extra_body,
+        **kwargs,
+    )
+
+    # Add the context to the function
+    ctx = contextvars.copy_context()
+    func_with_context = partial(ctx.run, func)
+    init_response = await loop.run_in_executor(None, func_with_context)
+    if asyncio.iscoroutine(init_response):
+        response = await init_response
+    else:
+        response = init_response  # type: ignore
+
+    return response
+
+
 def create_file(
     file: FileTypes,
     purpose: Literal["assistants", "batch", "fine-tune"],
@@ -44,7 +84,7 @@ def create_file(
     extra_headers: Optional[Dict[str, str]] = None,
     extra_body: Optional[Dict[str, str]] = None,
     **kwargs,
-) -> FileObject:
+) -> Union[FileObject | Coroutine[Any, Any, FileObject]]:
     """
     Files are used to upload documents that can be used with features like Assistants, Fine-tuning, and Batch API.
 
@@ -98,7 +138,10 @@ def create_file(
                 extra_body=extra_body,
             )
 
+            _is_async = kwargs.pop("acreate_file", False) is True
+
             response = openai_files_instance.create_file(
+                _is_async=_is_async,
                 api_base=api_base,
                 api_key=api_key,
                 timeout=timeout,

litellm/llms/openai.py

@@ -21,7 +21,7 @@ from litellm.utils import (
     TranscriptionResponse,
     TextCompletionResponse,
 )
-from typing import Callable, Optional
+from typing import Callable, Optional, Coroutine
 import litellm
 from .prompt_templates.factory import prompt_factory, custom_prompt
 from openai import OpenAI, AsyncOpenAI
@@ -1518,42 +1518,70 @@ class OpenAIFilesAPI(BaseLLM):
         timeout: Union[float, httpx.Timeout],
         max_retries: Optional[int],
         organization: Optional[str],
-        client: Optional[OpenAI] = None,
-    ) -> OpenAI:
+        client: Optional[Union[OpenAI, AsyncOpenAI]] = None,
+        _is_async: bool = False,
+    ) -> Optional[Union[OpenAI, AsyncOpenAI]]:
         received_args = locals()
+        openai_client: Optional[Union[OpenAI, AsyncOpenAI]] = None
         if client is None:
             data = {}
             for k, v in received_args.items():
-                if k == "self" or k == "client":
+                if k == "self" or k == "client" or k == "_is_async":
                     pass
                 elif k == "api_base" and v is not None:
                     data["base_url"] = v
                 elif v is not None:
                     data[k] = v
-            openai_client = OpenAI(**data)  # type: ignore
+            if _is_async is True:
+                openai_client = AsyncOpenAI(**data)
+            else:
+                openai_client = OpenAI(**data)  # type: ignore
         else:
             openai_client = client
 
         return openai_client
 
+    async def acreate_file(
+        self,
+        create_file_data: CreateFileRequest,
+        openai_client: AsyncOpenAI,
+    ) -> FileObject:
+        response = await openai_client.files.create(**create_file_data)
+        return response
+
     def create_file(
         self,
+        _is_async: bool,
         create_file_data: CreateFileRequest,
         api_base: str,
         api_key: Optional[str],
         timeout: Union[float, httpx.Timeout],
         max_retries: Optional[int],
         organization: Optional[str],
-        client: Optional[OpenAI] = None,
-    ) -> FileObject:
-        openai_client: OpenAI = self.get_openai_client(
+        client: Optional[Union[OpenAI, AsyncOpenAI]] = None,
+    ) -> Union[FileObject, Coroutine[Any, Any, FileObject]]:
+        openai_client: Optional[Union[OpenAI, AsyncOpenAI]] = self.get_openai_client(
             api_key=api_key,
             api_base=api_base,
             timeout=timeout,
             max_retries=max_retries,
             organization=organization,
             client=client,
+            _is_async=_is_async,
         )
+        if openai_client is None:
+            raise ValueError(
+                "OpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment."
+            )
+
+        if _is_async is True:
+            if not isinstance(openai_client, AsyncOpenAI):
+                raise ValueError(
+                    "OpenAI client is not an instance of AsyncOpenAI. Make sure you passed an AsyncOpenAI client."
+                )
+            return self.acreate_file(  # type: ignore
+                create_file_data=create_file_data, openai_client=openai_client
+            )
         response = openai_client.files.create(**create_file_data)
         return response
 

litellm/tests/test_openai_batches.py

@@ -2,6 +2,7 @@
 ## Unit Tests for OpenAI Batches API
 import sys, os, json
 import traceback
+import asyncio
 from dotenv import load_dotenv
 
 load_dotenv()
@@ -68,6 +69,59 @@ def test_create_batch():
     pass
 
 
+@pytest.mark.asyncio()
+async def test_async_create_batch():
+    """
+    1. Create File for Batch completion
+    2. Create Batch Request
+    3. Retrieve the specific batch
+    """
+    print("Testing async create batch")
+    file_obj = await litellm.acreate_file(
+        file=open("openai_batch_completions.jsonl", "rb"),
+        purpose="batch",
+        custom_llm_provider="openai",
+    )
+    print("Response from creating file=", file_obj)
+
+    batch_input_file_id = file_obj.id
+    assert (
+        batch_input_file_id is not None
+    ), "Failed to create file, expected a non null file_id but got {batch_input_file_id}"
+
+    # create_batch_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 from litellm.create_batch=", create_batch_response)
+
+    # assert (
+    #     create_batch_response.id is not None
+    # ), f"Failed to create batch, expected a non null batch_id but got {create_batch_response.id}"
+    # assert (
+    #     create_batch_response.endpoint == "/v1/chat/completions"
+    # ), f"Failed to create batch, expected endpoint to be /v1/chat/completions but got {create_batch_response.endpoint}"
+    # assert (
+    #     create_batch_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 {create_batch_response.input_file_id}"
+
+    # time.sleep(30)
+
+    # retrieved_batch = litellm.retrieve_batch(
+    #     batch_id=create_batch_response.id, custom_llm_provider="openai"
+    # )
+    # print("retrieved batch=", retrieved_batch)
+    # # just assert that we retrieved a non None batch
+
+    # assert retrieved_batch.id == create_batch_response.id
+
+    pass
+
+
 def test_retrieve_batch():
     pass