From 758ed9e923cd5e794fc2a997c9358b97215f97dc Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 28 May 2024 16:47:27 -0700 Subject: [PATCH] feat - add litellm.acreate_file --- litellm/batches/main.py | 55 +++++++++++++++++++++++++--- litellm/llms/openai.py | 44 ++++++++++++++++++---- litellm/tests/test_openai_batches.py | 54 +++++++++++++++++++++++++++ 3 files changed, 139 insertions(+), 14 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 3963a4e1146..056318c8ddf 100644 --- a/litellm/batches/main.py +++ b/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, diff --git a/litellm/llms/openai.py b/litellm/llms/openai.py index 5c5b837ea6a..05fc5784b61 100644 --- a/litellm/llms/openai.py +++ b/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 diff --git a/litellm/tests/test_openai_batches.py b/litellm/tests/test_openai_batches.py index fc797635b0a..2de417619b2 100644 --- a/litellm/tests/test_openai_batches.py +++ b/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