diff --git a/litellm/__init__.py b/litellm/__init__.py index 3c78c9b270e..56a2088e7ee 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -797,3 +797,4 @@ from .budget_manager import BudgetManager from .proxy.proxy_cli import run_server from .router import Router from .assistants.main import * +from .batches.main import * diff --git a/litellm/batches/main.py b/litellm/batches/main.py new file mode 100644 index 00000000000..917ad45e10c --- /dev/null +++ b/litellm/batches/main.py @@ -0,0 +1,459 @@ +""" +Main File for Batches API implementation + +https://platform.openai.com/docs/api-reference/batch + +- create_batch() +- retrieve_batch() +- cancel_batch() +- list_batch() + +""" + +import os +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 * +from ..llms.openai import OpenAIBatchesAPI, OpenAIFilesAPI +from ..types.llms.openai import ( + CreateBatchRequest, + RetrieveBatchRequest, + CancelBatchRequest, + CreateFileRequest, + FileTypes, + FileObject, + Batch, +) + +####### 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]: + """ + Async: 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: + 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 + except Exception as e: + raise e + + +def create_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, +) -> 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. + + LiteLLM Equivalent of POST: POST https://api.openai.com/v1/files + """ + try: + optional_params = GenericLiteLLMParams(**kwargs) + if custom_llm_provider == "openai": + # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there + api_base = ( + optional_params.api_base + or litellm.api_base + or os.getenv("OPENAI_API_BASE") + or "https://api.openai.com/v1" + ) + organization = ( + optional_params.organization + or litellm.organization + or os.getenv("OPENAI_ORGANIZATION", None) + or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105 + ) + # set API KEY + api_key = ( + optional_params.api_key + or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there + or litellm.openai_key + or os.getenv("OPENAI_API_KEY") + ) + ### TIMEOUT LOGIC ### + timeout = ( + optional_params.timeout or kwargs.get("request_timeout", 600) or 600 + ) + # set timeout for 10 minutes by default + + if ( + timeout is not None + and isinstance(timeout, httpx.Timeout) + and supports_httpx_timeout(custom_llm_provider) == False + ): + read_timeout = timeout.read or 600 + timeout = read_timeout # default 10 min timeout + elif timeout is not None and not isinstance(timeout, httpx.Timeout): + timeout = float(timeout) # type: ignore + elif timeout is None: + timeout = 600.0 + + _create_file_request = CreateFileRequest( + file=file, + purpose=purpose, + extra_headers=extra_headers, + 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, + max_retries=optional_params.max_retries, + organization=organization, + create_file_data=_create_file_request, + ) + else: + raise litellm.exceptions.BadRequestError( + message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format( + custom_llm_provider + ), + model="n/a", + llm_provider=custom_llm_provider, + response=httpx.Response( + status_code=400, + content="Unsupported provider", + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + ), + ) + return response + except Exception as e: + raise e + + +async def acreate_batch( + completion_window: Literal["24h"], + endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], + input_file_id: str, + custom_llm_provider: Literal["openai"] = "openai", + metadata: Optional[Dict[str, str]] = None, + extra_headers: Optional[Dict[str, str]] = None, + extra_body: Optional[Dict[str, str]] = None, + **kwargs, +) -> Coroutine[Any, Any, Batch]: + """ + Async: Creates and executes a batch from an uploaded file of request + + LiteLLM Equivalent of POST: https://api.openai.com/v1/batches + """ + try: + loop = asyncio.get_event_loop() + kwargs["acreate_batch"] = True + + # Use a partial function to pass your keyword arguments + func = partial( + create_batch, + completion_window, + endpoint, + input_file_id, + custom_llm_provider, + metadata, + 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 + except Exception as e: + raise e + + +def create_batch( + completion_window: Literal["24h"], + endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], + input_file_id: str, + custom_llm_provider: Literal["openai"] = "openai", + metadata: Optional[Dict[str, str]] = None, + extra_headers: Optional[Dict[str, str]] = None, + extra_body: Optional[Dict[str, str]] = None, + **kwargs, +) -> Union[Batch, Coroutine[Any, Any, Batch]]: + """ + Creates and executes a batch from an uploaded file of request + + LiteLLM Equivalent of POST: https://api.openai.com/v1/batches + """ + try: + optional_params = GenericLiteLLMParams(**kwargs) + if custom_llm_provider == "openai": + + # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there + api_base = ( + optional_params.api_base + or litellm.api_base + or os.getenv("OPENAI_API_BASE") + or "https://api.openai.com/v1" + ) + organization = ( + optional_params.organization + or litellm.organization + or os.getenv("OPENAI_ORGANIZATION", None) + or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105 + ) + # set API KEY + api_key = ( + optional_params.api_key + or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there + or litellm.openai_key + or os.getenv("OPENAI_API_KEY") + ) + ### TIMEOUT LOGIC ### + timeout = ( + optional_params.timeout or kwargs.get("request_timeout", 600) or 600 + ) + # set timeout for 10 minutes by default + + if ( + timeout is not None + and isinstance(timeout, httpx.Timeout) + and supports_httpx_timeout(custom_llm_provider) == False + ): + read_timeout = timeout.read or 600 + timeout = read_timeout # default 10 min timeout + elif timeout is not None and not isinstance(timeout, httpx.Timeout): + timeout = float(timeout) # type: ignore + elif timeout is None: + timeout = 600.0 + + _is_async = kwargs.pop("acreate_batch", False) is True + + _create_batch_request = CreateBatchRequest( + completion_window=completion_window, + endpoint=endpoint, + input_file_id=input_file_id, + metadata=metadata, + extra_headers=extra_headers, + extra_body=extra_body, + ) + + response = openai_batches_instance.create_batch( + api_base=api_base, + api_key=api_key, + organization=organization, + create_batch_data=_create_batch_request, + timeout=timeout, + max_retries=optional_params.max_retries, + _is_async=_is_async, + ) + else: + raise litellm.exceptions.BadRequestError( + message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format( + custom_llm_provider + ), + model="n/a", + llm_provider=custom_llm_provider, + response=httpx.Response( + status_code=400, + content="Unsupported provider", + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + ), + ) + return response + except Exception as e: + raise e + + +async def aretrieve_batch( + batch_id: str, + custom_llm_provider: Literal["openai"] = "openai", + metadata: Optional[Dict[str, str]] = None, + extra_headers: Optional[Dict[str, str]] = None, + extra_body: Optional[Dict[str, str]] = None, + **kwargs, +) -> Coroutine[Any, Any, Batch]: + """ + Async: Retrieves a batch. + + LiteLLM Equivalent of GET https://api.openai.com/v1/batches/{batch_id} + """ + try: + loop = asyncio.get_event_loop() + kwargs["aretrieve_batch"] = True + + # Use a partial function to pass your keyword arguments + func = partial( + retrieve_batch, + batch_id, + custom_llm_provider, + metadata, + 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 + except Exception as e: + raise e + + +def retrieve_batch( + batch_id: str, + custom_llm_provider: Literal["openai"] = "openai", + metadata: Optional[Dict[str, str]] = None, + extra_headers: Optional[Dict[str, str]] = None, + extra_body: Optional[Dict[str, str]] = None, + **kwargs, +) -> Union[Batch, Coroutine[Any, Any, Batch]]: + """ + Retrieves a batch. + + LiteLLM Equivalent of GET https://api.openai.com/v1/batches/{batch_id} + """ + try: + optional_params = GenericLiteLLMParams(**kwargs) + if custom_llm_provider == "openai": + + # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there + api_base = ( + optional_params.api_base + or litellm.api_base + or os.getenv("OPENAI_API_BASE") + or "https://api.openai.com/v1" + ) + organization = ( + optional_params.organization + or litellm.organization + or os.getenv("OPENAI_ORGANIZATION", None) + or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105 + ) + # set API KEY + api_key = ( + optional_params.api_key + or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there + or litellm.openai_key + or os.getenv("OPENAI_API_KEY") + ) + ### TIMEOUT LOGIC ### + timeout = ( + optional_params.timeout or kwargs.get("request_timeout", 600) or 600 + ) + # set timeout for 10 minutes by default + + if ( + timeout is not None + and isinstance(timeout, httpx.Timeout) + and supports_httpx_timeout(custom_llm_provider) == False + ): + read_timeout = timeout.read or 600 + timeout = read_timeout # default 10 min timeout + elif timeout is not None and not isinstance(timeout, httpx.Timeout): + timeout = float(timeout) # type: ignore + elif timeout is None: + timeout = 600.0 + + _retrieve_batch_request = RetrieveBatchRequest( + batch_id=batch_id, + extra_headers=extra_headers, + extra_body=extra_body, + ) + + _is_async = kwargs.pop("aretrieve_batch", False) is True + + response = openai_batches_instance.retrieve_batch( + _is_async=_is_async, + retrieve_batch_data=_retrieve_batch_request, + api_base=api_base, + api_key=api_key, + organization=organization, + timeout=timeout, + max_retries=optional_params.max_retries, + ) + else: + raise litellm.exceptions.BadRequestError( + message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format( + custom_llm_provider + ), + model="n/a", + llm_provider=custom_llm_provider, + response=httpx.Response( + status_code=400, + content="Unsupported provider", + request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore + ), + ) + return response + except Exception as e: + raise e + + +def cancel_batch(): + pass + + +def list_batch(): + pass + + +async def acancel_batch(): + pass + + +async def alist_batch(): + pass diff --git a/litellm/llms/openai.py b/litellm/llms/openai.py index 05e6566ffa8..43d088f0dbd 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 @@ -1497,6 +1497,274 @@ class OpenAITextCompletion(BaseLLM): yield transformed_chunk +class OpenAIFilesAPI(BaseLLM): + """ + OpenAI methods to support for batches + - create_file() + - retrieve_file() + - list_files() + - delete_file() + - file_content() + - update_file() + """ + + def __init__(self) -> None: + super().__init__() + + def get_openai_client( + self, + api_key: Optional[str], + api_base: Optional[str], + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + organization: Optional[str], + 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" 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 + 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[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 + + +class OpenAIBatchesAPI(BaseLLM): + """ + OpenAI methods to support for batches + - create_batch() + - retrieve_batch() + - cancel_batch() + - list_batch() + """ + + def __init__(self) -> None: + super().__init__() + + def get_openai_client( + self, + api_key: Optional[str], + api_base: Optional[str], + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + organization: Optional[str], + 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" 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 + 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_batch( + self, + create_batch_data: CreateBatchRequest, + openai_client: AsyncOpenAI, + ) -> Batch: + response = await openai_client.batches.create(**create_batch_data) + return response + + def create_batch( + self, + _is_async: bool, + create_batch_data: CreateBatchRequest, + api_key: Optional[str], + api_base: Optional[str], + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + organization: Optional[str], + client: Optional[Union[OpenAI, AsyncOpenAI]] = None, + ) -> Union[Batch, Coroutine[Any, Any, Batch]]: + 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_batch( # type: ignore + create_batch_data=create_batch_data, openai_client=openai_client + ) + response = openai_client.batches.create(**create_batch_data) + return response + + async def aretrieve_batch( + self, + retrieve_batch_data: RetrieveBatchRequest, + openai_client: AsyncOpenAI, + ) -> Batch: + response = await openai_client.batches.retrieve(**retrieve_batch_data) + return response + + def retrieve_batch( + self, + _is_async: bool, + retrieve_batch_data: RetrieveBatchRequest, + api_key: Optional[str], + api_base: Optional[str], + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + organization: Optional[str], + client: Optional[OpenAI] = None, + ): + 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.aretrieve_batch( # type: ignore + retrieve_batch_data=retrieve_batch_data, openai_client=openai_client + ) + response = openai_client.batches.retrieve(**retrieve_batch_data) + return response + + def cancel_batch( + self, + _is_async: bool, + cancel_batch_data: CancelBatchRequest, + api_key: Optional[str], + api_base: Optional[str], + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + organization: Optional[str], + client: Optional[OpenAI] = None, + ): + 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." + ) + response = openai_client.batches.cancel(**cancel_batch_data) + return response + + # def list_batch( + # self, + # list_batch_data: ListBatchRequest, + # api_key: Optional[str], + # api_base: Optional[str], + # timeout: Union[float, httpx.Timeout], + # max_retries: Optional[int], + # organization: Optional[str], + # client: Optional[OpenAI] = None, + # ): + # openai_client: OpenAI = self.get_openai_client( + # api_key=api_key, + # api_base=api_base, + # timeout=timeout, + # max_retries=max_retries, + # organization=organization, + # client=client, + # ) + # response = openai_client.batches.list(**list_batch_data) + # return response + + class OpenAIAssistantsAPI(BaseLLM): def __init__(self) -> None: super().__init__() diff --git a/litellm/tests/openai_batch_completions.jsonl b/litellm/tests/openai_batch_completions.jsonl new file mode 100644 index 00000000000..05448952a0f --- /dev/null +++ b/litellm/tests/openai_batch_completions.jsonl @@ -0,0 +1,2 @@ +{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo-0125", "messages": [{"role": "system", "content": "You are a helpful assistant."},{"role": "user", "content": "Hello world!"}],"max_tokens": 10}} +{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo-0125", "messages": [{"role": "system", "content": "You are an unhelpful assistant."},{"role": "user", "content": "Hello world!"}],"max_tokens": 10}} \ No newline at end of file diff --git a/litellm/tests/test_openai_batches.py b/litellm/tests/test_openai_batches.py new file mode 100644 index 00000000000..fc29331a5a5 --- /dev/null +++ b/litellm/tests/test_openai_batches.py @@ -0,0 +1,140 @@ +# What is this? +## Unit Tests for OpenAI Batches API +import sys, os, json +import traceback +import asyncio +from dotenv import load_dotenv + +load_dotenv() +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import pytest, logging, asyncio +import litellm +from litellm import ( + create_batch, + create_file, +) +import time + + +def test_create_batch(): + """ + 1. Create File for Batch completion + 2. Create Batch Request + 3. Retrieve the specific batch + """ + file_name = "openai_batch_completions.jsonl" + _current_dir = os.path.dirname(os.path.abspath(__file__)) + file_path = os.path.join(_current_dir, file_name) + + file_obj = litellm.create_file( + file=open(file_path, "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 + + +@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_name = "openai_batch_completions.jsonl" + _current_dir = os.path.dirname(os.path.abspath(__file__)) + file_path = os.path.join(_current_dir, file_name) + file_obj = await litellm.acreate_file( + file=open(file_path, "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 = await litellm.acreate_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}" + + await asyncio.sleep(1) + + retrieved_batch = await litellm.aretrieve_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 + + +def test_retrieve_batch(): + pass + + +def test_cancel_batch(): + pass + + +def test_list_batch(): + pass diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 1c60ad6dbb6..50ac1335ec1 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -6,7 +6,7 @@ from typing import ( Literal, Iterable, ) -from typing_extensions import override, Required +from typing_extensions import override, Required, Dict from pydantic import BaseModel from openai.types.beta.threads.message_content import MessageContent @@ -18,8 +18,23 @@ from openai.types.beta.assistant_tool_param import AssistantToolParam 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, Batch -from typing import TypedDict, List, Optional +from typing import TypedDict, List, Optional, Tuple, Mapping, IO + +FileContent = Union[IO[bytes], bytes, PathLike] + +FileTypes = Union[ + # file (or bytes) + FileContent, + # (filename, file (or bytes)) + Tuple[Optional[str], FileContent], + # (filename, file (or bytes), content_type) + Tuple[Optional[str], FileContent, Optional[str]], + # (filename, file (or bytes), content_type, headers) + Tuple[Optional[str], FileContent, Optional[str], Mapping[str, str]], +] class NotGiven: @@ -146,3 +161,76 @@ class Thread(BaseModel): object: Literal["thread"] """The object type, which is always `thread`.""" + + +# OpenAI Files Types +class CreateFileRequest(TypedDict, total=False): + """ + CreateFileRequest + Used by Assistants API, Batches API, and Fine-Tunes API + + Required Params: + file: FileTypes + purpose: Literal['assistants', 'batch', 'fine-tune'] + + Optional Params: + extra_headers: Optional[Dict[str, str]] + extra_body: Optional[Dict[str, str]] = None + timeout: Optional[float] = None + """ + + file: FileTypes + purpose: Literal["assistants", "batch", "fine-tune"] + extra_headers: Optional[Dict[str, str]] + extra_body: Optional[Dict[str, str]] + timeout: Optional[float] + + +# OpenAI Batches Types +class CreateBatchRequest(TypedDict, total=False): + """ + CreateBatchRequest + """ + + completion_window: Literal["24h"] + endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"] + input_file_id: str + metadata: Optional[Dict[str, str]] + extra_headers: Optional[Dict[str, str]] + extra_body: Optional[Dict[str, str]] + timeout: Optional[float] + + +class RetrieveBatchRequest(TypedDict, total=False): + """ + RetrieveBatchRequest + """ + + batch_id: str + extra_headers: Optional[Dict[str, str]] + extra_body: Optional[Dict[str, str]] + timeout: Optional[float] + + +class CancelBatchRequest(TypedDict, total=False): + """ + CancelBatchRequest + """ + + batch_id: str + extra_headers: Optional[Dict[str, str]] + extra_body: Optional[Dict[str, str]] + timeout: Optional[float] + + +class ListBatchRequest(TypedDict, total=False): + """ + ListBatchRequest - List your organization's batches + Calls https://api.openai.com/v1/batches + """ + + after: Union[str, NotGiven] + limit: Union[int, NotGiven] + extra_headers: Optional[Dict[str, str]] + extra_body: Optional[Dict[str, str]] + timeout: Optional[float]