From b9515c3b9641ecd4df294bac04e2863bb02e8b0b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 28 May 2024 12:53:46 -0700 Subject: [PATCH 01/11] feat - add batches types --- litellm/types/llms/openai.py | 52 +++++++++++++++++++++++++++++++++++- 1 file changed, 51 insertions(+), 1 deletion(-) diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 1c60ad6dbb6..60f7a904179 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 @@ -146,3 +146,53 @@ class Thread(BaseModel): object: Literal["thread"] """The object type, which is always `thread`.""" + + +# OpenAI Batches Types +class CreateBatchRequest(BaseModel): + """ + CreateBatchRequest + """ + + completion_window: Literal["24h"] + endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"] + input_file_id: str + metadata: Optional[Dict[str, str]] = None + extra_headers: Optional[Dict[str, str]] = None + extra_body: Optional[Dict[str, str]] = None + timeout: Optional[float] = None + + +class RetrieveBatchRequest(BaseModel): + """ + RetrieveBatchRequest + """ + + batch_id: str + extra_headers: Optional[Dict[str, str]] = None + extra_body: Optional[Dict[str, str]] = None + timeout: Optional[float] = None + + +class CancelBatchRequest(BaseModel): + """ + CancelBatchRequest + """ + + batch_id: str + extra_headers: Optional[Dict[str, str]] = None + extra_body: Optional[Dict[str, str]] = None + timeout: Optional[float] = None + + +class ListBatchRequest(BaseModel): + """ + ListBatchRequest - List your organization's batches + Calls https://api.openai.com/v1/batches + """ + + after: Optional[str] = None + limit: Optional[int] = 20 + extra_headers: Optional[Dict[str, str]] = None + extra_body: Optional[Dict[str, str]] = None + timeout: Optional[float] = None From 0af4c9206fe94d06d8aacd33de63f8fa5963a87c Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 28 May 2024 14:54:58 -0700 Subject: [PATCH 02/11] test - openai batches file --- litellm/tests/openai_batch_completions.jsonl | 2 ++ 1 file changed, 2 insertions(+) create mode 100644 litellm/tests/openai_batch_completions.jsonl 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 From d5dbf084ed2458a7fa6605e889c50fac9fde867e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 28 May 2024 15:35:11 -0700 Subject: [PATCH 03/11] feat - import batches in __init__ --- litellm/__init__.py | 1 + litellm/batches/main.py | 239 +++++++++++++++++++++++++++ litellm/llms/openai.py | 183 ++++++++++++++++++++ litellm/tests/test_openai_batches.py | 58 +++++++ litellm/types/llms/openai.py | 78 ++++++--- 5 files changed, 539 insertions(+), 20 deletions(-) create mode 100644 litellm/batches/main.py create mode 100644 litellm/tests/test_openai_batches.py 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..6aa0d8d8846 --- /dev/null +++ b/litellm/batches/main.py @@ -0,0 +1,239 @@ +""" +Main File for Batches API implementation + +https://platform.openai.com/docs/api-reference/batch + +- create_batch() +- retrieve_batch() +- cancel_batch() +- list_batch() + +""" + +from typing import Iterable +import os +import litellm +from openai import OpenAI +import httpx +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, +) + +from typing import Literal, Optional, Dict + +####### ENVIRONMENT VARIABLES ################### +openai_batches_instance = OpenAIBatchesAPI() +openai_files_instance = OpenAIFilesAPI() +################################################# + + +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, +) -> FileObject: + 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, + ) + + response = openai_files_instance.create_file( + 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 + + +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, +): + """ + 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 + + _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, + ) + 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 retrieve_batch(): + pass + + +def cancel_batch(): + pass + + +def list_batch(): + pass + + +# Async Functions +async def acreate_batch(): + pass + + +async def aretrieve_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..e5f229eb5c9 100644 --- a/litellm/llms/openai.py +++ b/litellm/llms/openai.py @@ -1497,6 +1497,189 @@ 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[OpenAI] = None, + ) -> OpenAI: + received_args = locals() + if client is None: + data = {} + for k, v in received_args.items(): + if k == "self" or k == "client": + 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 + else: + openai_client = client + + return openai_client + + def create_file( + self, + 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( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + organization=organization, + client=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[OpenAI] = None, + ) -> OpenAI: + received_args = locals() + if client is None: + data = {} + for k, v in received_args.items(): + if k == "self" or k == "client": + 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 + else: + openai_client = client + + return openai_client + + def create_batch( + self, + 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[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.create(**create_batch_data) + return response + + def retrieve_batch( + self, + 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: 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.retrieve(**retrieve_batch_data) + return response + + def cancel_batch( + self, + 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: 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.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/test_openai_batches.py b/litellm/tests/test_openai_batches.py new file mode 100644 index 00000000000..ac282e0ed42 --- /dev/null +++ b/litellm/tests/test_openai_batches.py @@ -0,0 +1,58 @@ +# What is this? +## Unit Tests for OpenAI Batches API +import sys, os, json +import traceback +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, +) + + +def test_create_batch(): + """ + 1. Create File for Batch completion + 2. Create Batch Request + """ + file_obj = litellm.create_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}" + + 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"}, + # ) + + print("response") + pass + + +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 60f7a904179..c65fb2d5f78 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -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 -from typing import TypedDict, List, Optional +from typing import TypedDict, List, Optional, Tuple, Mapping, IO + +FileContent = Union[IO[bytes], bytes, PathLike[str]] + +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: @@ -148,8 +163,31 @@ class Thread(BaseModel): """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(BaseModel): +class CreateBatchRequest(TypedDict, total=False): """ CreateBatchRequest """ @@ -157,42 +195,42 @@ class CreateBatchRequest(BaseModel): completion_window: Literal["24h"] endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"] input_file_id: str - metadata: Optional[Dict[str, str]] = None - extra_headers: Optional[Dict[str, str]] = None - extra_body: Optional[Dict[str, str]] = None - timeout: Optional[float] = None + metadata: Optional[Dict[str, str]] + extra_headers: Optional[Dict[str, str]] + extra_body: Optional[Dict[str, str]] + timeout: Optional[float] -class RetrieveBatchRequest(BaseModel): +class RetrieveBatchRequest(TypedDict, total=False): """ RetrieveBatchRequest """ batch_id: str - extra_headers: Optional[Dict[str, str]] = None - extra_body: Optional[Dict[str, str]] = None - timeout: Optional[float] = None + extra_headers: Optional[Dict[str, str]] + extra_body: Optional[Dict[str, str]] + timeout: Optional[float] -class CancelBatchRequest(BaseModel): +class CancelBatchRequest(TypedDict, total=False): """ CancelBatchRequest """ batch_id: str - extra_headers: Optional[Dict[str, str]] = None - extra_body: Optional[Dict[str, str]] = None - timeout: Optional[float] = None + extra_headers: Optional[Dict[str, str]] + extra_body: Optional[Dict[str, str]] + timeout: Optional[float] -class ListBatchRequest(BaseModel): +class ListBatchRequest(TypedDict, total=False): """ ListBatchRequest - List your organization's batches Calls https://api.openai.com/v1/batches """ - after: Optional[str] = None - limit: Optional[int] = 20 - extra_headers: Optional[Dict[str, str]] = None - extra_body: Optional[Dict[str, str]] = None - timeout: Optional[float] = None + after: Union[str, NotGiven] + limit: Union[int, NotGiven] + extra_headers: Optional[Dict[str, str]] + extra_body: Optional[Dict[str, str]] + timeout: Optional[float] From 38285e53c340418a1decfe5fd1d958885d23803d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 28 May 2024 15:45:23 -0700 Subject: [PATCH 04/11] 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 From 938f4703c3da18580222058750d9010cdf769c3e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 28 May 2024 15:55:32 -0700 Subject: [PATCH 05/11] feat- add retrieve_batch() --- litellm/batches/main.py | 87 +++++++++++++++++++++++++++- litellm/tests/test_openai_batches.py | 28 ++++++--- 2 files changed, 105 insertions(+), 10 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 3bd1d07a475..3963a4e1146 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -216,8 +216,91 @@ def create_batch( raise e -def retrieve_batch(): - pass +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, +): + """ + 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, + ) + + response = openai_batches_instance.retrieve_batch( + 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(): diff --git a/litellm/tests/test_openai_batches.py b/litellm/tests/test_openai_batches.py index b99991babaa..fc797635b0a 100644 --- a/litellm/tests/test_openai_batches.py +++ b/litellm/tests/test_openai_batches.py @@ -14,6 +14,7 @@ from litellm import ( create_batch, create_file, ) +import time def test_create_batch(): @@ -34,7 +35,7 @@ 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}" - response = litellm.create_batch( + create_batch_response = litellm.create_batch( completion_window="24h", endpoint="/v1/chat/completions", input_file_id=batch_input_file_id, @@ -42,17 +43,28 @@ def test_create_batch(): metadata={"key1": "value1", "key2": "value2"}, ) - print("response from litellm.create_batch=", response) + print("response from litellm.create_batch=", create_batch_response) assert ( - response.id is not None - ), f"Failed to create batch, expected a non null batch_id but got {response.id}" + 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 ( - response.endpoint == "/v1/chat/completions" - ), f"Failed to create batch, expected endpoint to be /v1/chat/completions but got {response.endpoint}" + 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 ( - 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}" + 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 From 758ed9e923cd5e794fc2a997c9358b97215f97dc Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 28 May 2024 16:47:27 -0700 Subject: [PATCH 06/11] 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 From 1ef7cd923cf11ceb85ad8da3019ae5975285e331 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 28 May 2024 17:03:29 -0700 Subject: [PATCH 07/11] feat - add acreate_batch --- litellm/batches/main.py | 104 ++++++++++++++++++++------- litellm/llms/openai.py | 58 ++++++++++++--- litellm/tests/test_openai_batches.py | 34 ++++----- 3 files changed, 143 insertions(+), 53 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 056318c8ddf..05a6dfd517d 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -51,30 +51,33 @@ async def acreate_file( LiteLLM Equivalent of POST: POST https://api.openai.com/v1/files """ - loop = asyncio.get_event_loop() - kwargs["acreate_file"] = True + 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, - ) + # 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 + # 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 + return response + except Exception as e: + raise e def create_file( @@ -167,6 +170,52 @@ def create_file( 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]: + """ + 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"], @@ -176,7 +225,7 @@ def create_batch( extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, -) -> Batch: +) -> Union[Batch, Coroutine[Any, Any, Batch]]: """ Creates and executes a batch from an uploaded file of request @@ -224,6 +273,8 @@ def create_batch( 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, @@ -240,6 +291,7 @@ def create_batch( create_batch_data=_create_batch_request, timeout=timeout, max_retries=optional_params.max_retries, + _is_async=_is_async, ) else: raise litellm.exceptions.BadRequestError( @@ -320,7 +372,10 @@ def retrieve_batch( 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, @@ -354,11 +409,6 @@ def list_batch(): pass -# Async Functions -async def acreate_batch(): - pass - - async def aretrieve_batch(): pass diff --git a/litellm/llms/openai.py b/litellm/llms/openai.py index 05fc5784b61..fa1f13c70ac 100644 --- a/litellm/llms/openai.py +++ b/litellm/llms/openai.py @@ -1605,47 +1605,76 @@ class OpenAIBatchesAPI(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_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[OpenAI] = None, - ) -> Batch: - openai_client: OpenAI = self.get_openai_client( + 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 def retrieve_batch( self, + _is_async: bool, retrieve_batch_data: RetrieveBatchRequest, api_key: Optional[str], api_base: Optional[str], @@ -1654,19 +1683,25 @@ class OpenAIBatchesAPI(BaseLLM): organization: Optional[str], client: Optional[OpenAI] = None, ): - openai_client: OpenAI = self.get_openai_client( + 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.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], @@ -1675,14 +1710,19 @@ class OpenAIBatchesAPI(BaseLLM): organization: Optional[str], client: Optional[OpenAI] = None, ): - openai_client: OpenAI = self.get_openai_client( + 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 diff --git a/litellm/tests/test_openai_batches.py b/litellm/tests/test_openai_batches.py index 2de417619b2..497662006dc 100644 --- a/litellm/tests/test_openai_batches.py +++ b/litellm/tests/test_openai_batches.py @@ -89,25 +89,25 @@ async def test_async_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}" - # 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"}, - # ) + 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) + 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}" + 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) From 6688215c186256f975c4dccbdb179c237f5b1b9e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 28 May 2024 17:12:41 -0700 Subject: [PATCH 08/11] feat - add aretrieve_batch --- litellm/batches/main.py | 52 ++++++++++++++++++++++++---- litellm/llms/openai.py | 17 +++++++++ litellm/tests/test_openai_batches.py | 16 ++++----- 3 files changed, 69 insertions(+), 16 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 05a6dfd517d..119043e4c1f 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -47,7 +47,7 @@ async def acreate_file( **kwargs, ) -> Coroutine[Any, Any, FileObject]: """ - Files are used to upload documents that can be used with features like Assistants, Fine-tuning, and Batch API. + 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 """ @@ -181,7 +181,7 @@ async def acreate_batch( **kwargs, ) -> Coroutine[Any, Any, Batch]: """ - Creates and executes a batch from an uploaded file of request + Async: Creates and executes a batch from an uploaded file of request LiteLLM Equivalent of POST: https://api.openai.com/v1/batches """ @@ -311,6 +311,48 @@ def create_batch( 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", @@ -318,7 +360,7 @@ def retrieve_batch( extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, -): +) -> Union[Batch, Coroutine[Any, Any, Batch]]: """ Retrieves a batch. @@ -409,10 +451,6 @@ def list_batch(): pass -async def aretrieve_batch(): - pass - - async def acancel_batch(): pass diff --git a/litellm/llms/openai.py b/litellm/llms/openai.py index fa1f13c70ac..43d088f0dbd 100644 --- a/litellm/llms/openai.py +++ b/litellm/llms/openai.py @@ -1672,6 +1672,14 @@ class OpenAIBatchesAPI(BaseLLM): 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, @@ -1696,6 +1704,15 @@ class OpenAIBatchesAPI(BaseLLM): 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 diff --git a/litellm/tests/test_openai_batches.py b/litellm/tests/test_openai_batches.py index 497662006dc..2bf00901281 100644 --- a/litellm/tests/test_openai_batches.py +++ b/litellm/tests/test_openai_batches.py @@ -109,17 +109,15 @@ async def test_async_create_batch(): 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) + await asyncio.sleep(1) - # 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 + 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 - - pass + assert retrieved_batch.id == create_batch_response.id def test_retrieve_batch(): From c6eb004ed1bb0ce73d323c4b12311df5f7d6ea7b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 28 May 2024 17:21:59 -0700 Subject: [PATCH 09/11] fix python 3.8 error --- litellm/types/llms/openai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 81a50db6d22..50ac1335ec1 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -23,7 +23,7 @@ from openai.types import FileObject, Batch from typing import TypedDict, List, Optional, Tuple, Mapping, IO -FileContent = Union[IO[bytes], bytes, PathLike[str]] +FileContent = Union[IO[bytes], bytes, PathLike] FileTypes = Union[ # file (or bytes) From 215f19440c4d08f6b33a4b951aba6c1a8c882866 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 28 May 2024 17:25:08 -0700 Subject: [PATCH 10/11] fix python3.8 error --- litellm/batches/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 119043e4c1f..917ad45e10c 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -87,7 +87,7 @@ def create_file( extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, -) -> Union[FileObject | Coroutine[Any, Any, 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. From 5e9d024c2cf8269c4e70bc44226230f462f1b502 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 28 May 2024 17:29:13 -0700 Subject: [PATCH 11/11] fix - test openai batches --- litellm/tests/test_openai_batches.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_openai_batches.py b/litellm/tests/test_openai_batches.py index 2bf00901281..fc29331a5a5 100644 --- a/litellm/tests/test_openai_batches.py +++ b/litellm/tests/test_openai_batches.py @@ -24,8 +24,12 @@ def test_create_batch(): 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("openai_batch_completions.jsonl", "rb"), + file=open(file_path, "rb"), purpose="batch", custom_llm_provider="openai", ) @@ -77,8 +81,12 @@ async def test_async_create_batch(): 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("openai_batch_completions.jsonl", "rb"), + file=open(file_path, "rb"), purpose="batch", custom_llm_provider="openai", )