From 80675b22bdc109cc0f3e6f2faae2d4ed5e25be41 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 22 Aug 2024 14:46:51 -0700 Subject: [PATCH 1/7] feat(batches): add azure openai batches endpoint support Closes https://github.com/BerriAI/litellm/issues/5073 --- litellm/__init__.py | 2 +- litellm/batches/main.py | 262 ++++++++++++------ litellm/files/main.py | 223 ++++++++++----- litellm/llms/azure.py | 217 +++++++++++++++ litellm/llms/files_apis/azure.py | 15 +- .../tests/batch_job_results_furniture.jsonl | 2 + .../tests/test_openai_batches_and_files.py | 36 ++- 7 files changed, 584 insertions(+), 173 deletions(-) create mode 100644 litellm/tests/batch_job_results_furniture.jsonl diff --git a/litellm/__init__.py b/litellm/__init__.py index c7648ac07f..850865a36b 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -117,7 +117,7 @@ disable_streaming_logging: bool = False in_memory_llm_clients_cache: dict = {} safe_memory_mode: bool = False ### DEFAULT AZURE API VERSION ### -AZURE_DEFAULT_API_VERSION = "2024-02-01" # this is updated to the latest +AZURE_DEFAULT_API_VERSION = "2024-07-01-preview" # this is updated to the latest ### GUARDRAILS ### llamaguard_model_name: Optional[str] = None openai_moderations_model_name: Optional[str] = None diff --git a/litellm/batches/main.py b/litellm/batches/main.py index a2ebc664ea..de3ddd11c9 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -20,7 +20,8 @@ import httpx import litellm from litellm import client -from litellm.llms.openai import OpenAIBatchesAPI, OpenAIFilesAPI +from litellm.llms.azure import AzureBatchesAPI +from litellm.llms.openai import OpenAIBatchesAPI from litellm.types.llms.openai import ( Batch, CancelBatchRequest, @@ -33,10 +34,11 @@ from litellm.types.llms.openai import ( RetrieveBatchRequest, ) from litellm.types.router import GenericLiteLLMParams -from litellm.utils import supports_httpx_timeout +from litellm.utils import get_secret, supports_httpx_timeout ####### ENVIRONMENT VARIABLES ################### openai_batches_instance = OpenAIBatchesAPI() +azure_batches_instance = AzureBatchesAPI() ################################################# @@ -90,7 +92,7 @@ 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", + custom_llm_provider: Literal["openai", "azure"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -103,6 +105,32 @@ def create_batch( """ try: optional_params = GenericLiteLLMParams(**kwargs) + _is_async = kwargs.pop("acreate_batch", False) is True + ### 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) is 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, + ) + if custom_llm_provider == "openai": # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -125,34 +153,6 @@ def create_batch( 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, @@ -163,6 +163,38 @@ def create_batch( max_retries=optional_params.max_retries, _is_async=_is_async, ) + elif custom_llm_provider == "azure": + api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore + api_version = ( + optional_params.api_version + or litellm.api_version + or get_secret("AZURE_API_VERSION") + ) # type: ignore + + api_key = ( + optional_params.api_key + or litellm.api_key + or litellm.azure_key + or get_secret("AZURE_OPENAI_API_KEY") + or get_secret("AZURE_API_KEY") + ) # type: ignore + + extra_body = optional_params.get("extra_body", {}) + azure_ad_token: Optional[str] = None + if extra_body is not None: + azure_ad_token = extra_body.pop("azure_ad_token", None) + else: + azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore + + response = azure_batches_instance.create_batch( + _is_async=_is_async, + api_base=api_base, + api_key=api_key, + api_version=api_version, + timeout=timeout, + max_retries=optional_params.max_retries, + create_batch_data=_create_batch_request, + ) else: raise litellm.exceptions.BadRequestError( message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format( @@ -225,7 +257,7 @@ async def aretrieve_batch( def retrieve_batch( batch_id: str, - custom_llm_provider: Literal["openai"] = "openai", + custom_llm_provider: Literal["openai", "azure"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -238,6 +270,30 @@ def retrieve_batch( """ try: optional_params = GenericLiteLLMParams(**kwargs) + ### 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) is 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 + if custom_llm_provider == "openai": # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -260,31 +316,6 @@ def retrieve_batch( 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, @@ -295,6 +326,38 @@ def retrieve_batch( timeout=timeout, max_retries=optional_params.max_retries, ) + elif custom_llm_provider == "azure": + api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore + api_version = ( + optional_params.api_version + or litellm.api_version + or get_secret("AZURE_API_VERSION") + ) # type: ignore + + api_key = ( + optional_params.api_key + or litellm.api_key + or litellm.azure_key + or get_secret("AZURE_OPENAI_API_KEY") + or get_secret("AZURE_API_KEY") + ) # type: ignore + + extra_body = optional_params.get("extra_body", {}) + azure_ad_token: Optional[str] = None + if extra_body is not None: + azure_ad_token = extra_body.pop("azure_ad_token", None) + else: + azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore + + response = azure_batches_instance.retrieve_batch( + _is_async=_is_async, + api_base=api_base, + api_key=api_key, + api_version=api_version, + timeout=timeout, + max_retries=optional_params.max_retries, + retrieve_batch_data=_retrieve_batch_request, + ) else: raise litellm.exceptions.BadRequestError( message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format( @@ -357,7 +420,7 @@ async def alist_batches( def list_batches( after: Optional[str] = None, limit: Optional[int] = None, - custom_llm_provider: Literal["openai"] = "openai", + custom_llm_provider: Literal["openai", "azure"] = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -368,7 +431,31 @@ def list_batches( List your organization's batches. """ try: + # set API KEY optional_params = GenericLiteLLMParams(**kwargs) + 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) is 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("alist_batches", False) is True 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 = ( @@ -383,32 +470,6 @@ def list_batches( 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("alist_batches", False) is True response = openai_batches_instance.list_batches( _is_async=_is_async, @@ -420,9 +481,40 @@ def list_batches( timeout=timeout, max_retries=optional_params.max_retries, ) + elif custom_llm_provider == "azure": + api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore + api_version = ( + optional_params.api_version + or litellm.api_version + or get_secret("AZURE_API_VERSION") + ) # type: ignore + + api_key = ( + optional_params.api_key + or litellm.api_key + or litellm.azure_key + or get_secret("AZURE_OPENAI_API_KEY") + or get_secret("AZURE_API_KEY") + ) # type: ignore + + extra_body = optional_params.get("extra_body", {}) + azure_ad_token: Optional[str] = None + if extra_body is not None: + azure_ad_token = extra_body.pop("azure_ad_token", None) + else: + azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore + + response = azure_batches_instance.list_batches( + _is_async=_is_async, + api_base=api_base, + api_key=api_key, + api_version=api_version, + 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( + message="LiteLLM doesn't support {} for 'list_batch'. Only 'openai' is supported.".format( custom_llm_provider ), model="n/a", diff --git a/litellm/files/main.py b/litellm/files/main.py index 49d3553989..1ed1c1e611 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -87,6 +87,24 @@ def file_retrieve( """ try: optional_params = GenericLiteLLMParams(**kwargs) + ### 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) is 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("is_async", False) is True + 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 = ( @@ -108,25 +126,6 @@ def file_retrieve( 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("is_async", False) is True response = openai_files_instance.retrieve_file( file_id=file_id, @@ -137,9 +136,41 @@ def file_retrieve( max_retries=optional_params.max_retries, organization=organization, ) + elif custom_llm_provider == "azure": + api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore + api_version = ( + optional_params.api_version + or litellm.api_version + or get_secret("AZURE_API_VERSION") + ) # type: ignore + + api_key = ( + optional_params.api_key + or litellm.api_key + or litellm.azure_key + or get_secret("AZURE_OPENAI_API_KEY") + or get_secret("AZURE_API_KEY") + ) # type: ignore + + extra_body = optional_params.get("extra_body", {}) + azure_ad_token: Optional[str] = None + if extra_body is not None: + azure_ad_token = extra_body.pop("azure_ad_token", None) + else: + azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore + + response = azure_files_instance.retrieve_file( + _is_async=_is_async, + api_base=api_base, + api_key=api_key, + api_version=api_version, + timeout=timeout, + max_retries=optional_params.max_retries, + file_id=file_id, + ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format( + message="LiteLLM doesn't support {} for 'file_retrieve'. Only 'openai' and 'azure' are supported.".format( custom_llm_provider ), model="n/a", @@ -361,6 +392,23 @@ def file_list( """ try: optional_params = GenericLiteLLMParams(**kwargs) + ### 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("is_async", False) is True 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 = ( @@ -382,25 +430,6 @@ def file_list( 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("is_async", False) is True response = openai_files_instance.list_files( purpose=purpose, @@ -411,9 +440,41 @@ def file_list( max_retries=optional_params.max_retries, organization=organization, ) + elif custom_llm_provider == "azure": + api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore + api_version = ( + optional_params.api_version + or litellm.api_version + or get_secret("AZURE_API_VERSION") + ) # type: ignore + + api_key = ( + optional_params.api_key + or litellm.api_key + or litellm.azure_key + or get_secret("AZURE_OPENAI_API_KEY") + or get_secret("AZURE_API_KEY") + ) # type: ignore + + extra_body = optional_params.get("extra_body", {}) + azure_ad_token: Optional[str] = None + if extra_body is not None: + azure_ad_token = extra_body.pop("azure_ad_token", None) + else: + azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore + + response = azure_files_instance.list_files( + _is_async=_is_async, + api_base=api_base, + api_key=api_key, + api_version=api_version, + timeout=timeout, + max_retries=optional_params.max_retries, + purpose=purpose, + ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'file_list'. Only 'openai' is supported.".format( + message="LiteLLM doesn't support {} for 'file_list'. Only 'openai' and 'azure' are supported.".format( custom_llm_provider ), model="n/a", @@ -645,6 +706,29 @@ def file_content( """ try: optional_params = GenericLiteLLMParams(**kwargs) + ### 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 + + _file_content_request = FileContentRequest( + file_id=file_id, + extra_headers=extra_headers, + extra_body=extra_body, + ) + + _is_async = kwargs.pop("afile_content", False) is True 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 = ( @@ -666,31 +750,6 @@ def file_content( 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 - - _file_content_request = FileContentRequest( - file_id=file_id, - extra_headers=extra_headers, - extra_body=extra_body, - ) - - _is_async = kwargs.pop("afile_content", False) is True response = openai_files_instance.file_content( _is_async=_is_async, @@ -701,9 +760,41 @@ def file_content( max_retries=optional_params.max_retries, organization=organization, ) + elif custom_llm_provider == "azure": + api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore + api_version = ( + optional_params.api_version + or litellm.api_version + or get_secret("AZURE_API_VERSION") + ) # type: ignore + + api_key = ( + optional_params.api_key + or litellm.api_key + or litellm.azure_key + or get_secret("AZURE_OPENAI_API_KEY") + or get_secret("AZURE_API_KEY") + ) # type: ignore + + extra_body = optional_params.get("extra_body", {}) + azure_ad_token: Optional[str] = None + if extra_body is not None: + azure_ad_token = extra_body.pop("azure_ad_token", None) + else: + azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore + + response = azure_files_instance.file_content( + _is_async=_is_async, + api_base=api_base, + api_key=api_key, + api_version=api_version, + timeout=timeout, + max_retries=optional_params.max_retries, + file_content_request=_file_content_request, + ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format( + message="LiteLLM doesn't support {} for 'file_content'. Only 'openai' and 'azure' are supported.".format( custom_llm_provider ), model="n/a", diff --git a/litellm/llms/azure.py b/litellm/llms/azure.py index 00cdb90cfd..a54bef3192 100644 --- a/litellm/llms/azure.py +++ b/litellm/llms/azure.py @@ -47,14 +47,18 @@ from ..types.llms.openai import ( AsyncAssistantEventHandler, AsyncAssistantStreamManager, AsyncCursorPage, + Batch, + CancelBatchRequest, ChatCompletionToolChoiceFunctionParam, ChatCompletionToolChoiceObjectParam, ChatCompletionToolParam, ChatCompletionToolParamFunctionChunk, + CreateBatchRequest, HttpxBinaryResponseContent, MessageData, OpenAICreateThreadParamsMessage, OpenAIMessage, + RetrieveBatchRequest, Run, SyncCursorPage, Thread, @@ -2814,3 +2818,216 @@ class AzureAssistantsAPI(BaseLLM): ) return response + + +class AzureBatchesAPI(BaseLLM): + """ + Azure methods to support for batches + - create_batch() + - retrieve_batch() + - cancel_batch() + - list_batch() + """ + + def __init__(self) -> None: + super().__init__() + + def get_azure_openai_client( + self, + api_key: Optional[str], + api_base: Optional[str], + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + api_version: Optional[str] = None, + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = None, + _is_async: bool = False, + ) -> Optional[Union[AzureOpenAI, AsyncAzureOpenAI]]: + received_args = locals() + openai_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = 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["azure_endpoint"] = v + elif v is not None: + data[k] = v + if "api_version" not in data: + data["api_version"] = litellm.AZURE_DEFAULT_API_VERSION + if _is_async is True: + openai_client = AsyncAzureOpenAI(**data) + else: + openai_client = AzureOpenAI(**data) # type: ignore + else: + openai_client = client + + return openai_client + + async def acreate_batch( + self, + create_batch_data: CreateBatchRequest, + azure_client: AsyncAzureOpenAI, + ) -> Batch: + response = await azure_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], + api_version: Optional[str], + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = None, + ) -> Union[Batch, Coroutine[Any, Any, Batch]]: + azure_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = ( + self.get_azure_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + api_version=api_version, + max_retries=max_retries, + client=client, + _is_async=_is_async, + ) + ) + if azure_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(azure_client, AsyncAzureOpenAI): + 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, azure_client=azure_client + ) + response = azure_client.batches.create(**create_batch_data) + return response + + async def aretrieve_batch( + self, + retrieve_batch_data: RetrieveBatchRequest, + client: AsyncAzureOpenAI, + ) -> Batch: + response = await 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], + api_version: Optional[str], + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + client: Optional[AzureOpenAI] = None, + ): + azure_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = ( + self.get_azure_openai_client( + api_key=api_key, + api_base=api_base, + api_version=api_version, + timeout=timeout, + max_retries=max_retries, + client=client, + _is_async=_is_async, + ) + ) + if azure_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(azure_client, AsyncAzureOpenAI): + 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, client=azure_client + ) + response = azure_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[AzureOpenAI] = None, + ): + azure_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = ( + self.get_azure_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + client=client, + _is_async=_is_async, + ) + ) + if azure_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 = azure_client.batches.cancel(**cancel_batch_data) + return response + + async def alist_batches( + self, + client: AsyncAzureOpenAI, + after: Optional[str] = None, + limit: Optional[int] = None, + ): + response = await client.batches.list(after=after, limit=limit) # type: ignore + return response + + def list_batches( + self, + _is_async: bool, + api_key: Optional[str], + api_base: Optional[str], + api_version: Optional[str], + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + after: Optional[str] = None, + limit: Optional[int] = None, + client: Optional[AzureOpenAI] = None, + ): + azure_client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = ( + self.get_azure_openai_client( + api_key=api_key, + api_base=api_base, + timeout=timeout, + max_retries=max_retries, + api_version=api_version, + client=client, + _is_async=_is_async, + ) + ) + if azure_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(azure_client, AsyncAzureOpenAI): + raise ValueError( + "OpenAI client is not an instance of AsyncOpenAI. Make sure you passed an AsyncOpenAI client." + ) + return self.alist_batches( # type: ignore + client=azure_client, after=after, limit=limit + ) + response = azure_client.batches.list(after=after, limit=limit) # type: ignore + return response diff --git a/litellm/llms/files_apis/azure.py b/litellm/llms/files_apis/azure.py index c4c9ee48af..38d215cc75 100644 --- a/litellm/llms/files_apis/azure.py +++ b/litellm/llms/files_apis/azure.py @@ -121,7 +121,6 @@ class AzureOpenAIFilesAPI(BaseLLM): api_key: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], - organization: Optional[str], api_version: Optional[str] = None, client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = None, ) -> Union[ @@ -134,7 +133,7 @@ class AzureOpenAIFilesAPI(BaseLLM): timeout=timeout, api_version=api_version, max_retries=max_retries, - organization=organization, + organization=None, client=client, _is_async=_is_async, ) @@ -173,7 +172,6 @@ class AzureOpenAIFilesAPI(BaseLLM): api_key: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], - organization: Optional[str], api_version: Optional[str] = None, client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = None, ): @@ -183,7 +181,7 @@ class AzureOpenAIFilesAPI(BaseLLM): api_base=api_base, timeout=timeout, max_retries=max_retries, - organization=organization, + organization=None, api_version=api_version, client=client, _is_async=_is_async, @@ -213,6 +211,9 @@ class AzureOpenAIFilesAPI(BaseLLM): openai_client: AsyncAzureOpenAI, ) -> FileDeleted: response = await openai_client.files.delete(file_id=file_id) + + if not isinstance(response, FileDeleted): # azure returns an empty string + return FileDeleted(id=file_id, deleted=True, object="file") return response def delete_file( @@ -255,6 +256,9 @@ class AzureOpenAIFilesAPI(BaseLLM): ) response = openai_client.files.delete(file_id=file_id) + if not isinstance(response, FileDeleted): # azure returns an empty string + return FileDeleted(id=file_id, deleted=True, object="file") + return response async def alist_files( @@ -275,7 +279,6 @@ class AzureOpenAIFilesAPI(BaseLLM): api_key: Optional[str], timeout: Union[float, httpx.Timeout], max_retries: Optional[int], - organization: Optional[str], purpose: Optional[str] = None, api_version: Optional[str] = None, client: Optional[Union[AzureOpenAI, AsyncAzureOpenAI]] = None, @@ -286,7 +289,7 @@ class AzureOpenAIFilesAPI(BaseLLM): api_base=api_base, timeout=timeout, max_retries=max_retries, - organization=organization, + organization=None, # openai param api_version=api_version, client=client, _is_async=_is_async, diff --git a/litellm/tests/batch_job_results_furniture.jsonl b/litellm/tests/batch_job_results_furniture.jsonl new file mode 100644 index 0000000000..05448952a0 --- /dev/null +++ b/litellm/tests/batch_job_results_furniture.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_and_files.py b/litellm/tests/test_openai_batches_and_files.py index e8bde4d20d..cad5052c2e 100644 --- a/litellm/tests/test_openai_batches_and_files.py +++ b/litellm/tests/test_openai_batches_and_files.py @@ -22,7 +22,8 @@ import litellm from litellm import create_batch, create_file -def test_create_batch(): +@pytest.mark.parametrize("provider", ["openai", "azure"]) +def test_create_batch(provider): """ 1. Create File for Batch completion 2. Create Batch Request @@ -35,7 +36,7 @@ def test_create_batch(): file_obj = litellm.create_file( file=open(file_path, "rb"), purpose="batch", - custom_llm_provider="openai", + custom_llm_provider=provider, ) print("Response from creating file=", file_obj) @@ -44,11 +45,12 @@ 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}" + time.sleep(5) create_batch_response = litellm.create_batch( completion_window="24h", endpoint="/v1/chat/completions", input_file_id=batch_input_file_id, - custom_llm_provider="openai", + custom_llm_provider=provider, metadata={"key1": "value1", "key2": "value2"}, ) @@ -59,13 +61,14 @@ def test_create_batch(): ), 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" + or create_batch_response.endpoint == "/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}" retrieved_batch = litellm.retrieve_batch( - batch_id=create_batch_response.id, custom_llm_provider="openai" + batch_id=create_batch_response.id, custom_llm_provider=provider ) print("retrieved batch=", retrieved_batch) # just assert that we retrieved a non None batch @@ -73,11 +76,11 @@ def test_create_batch(): assert retrieved_batch.id == create_batch_response.id # list all batches - list_batches = litellm.list_batches(custom_llm_provider="openai", limit=2) + list_batches = litellm.list_batches(custom_llm_provider=provider, limit=2) print("list_batches=", list_batches) file_content = litellm.file_content( - file_id=batch_input_file_id, custom_llm_provider="openai" + file_id=batch_input_file_id, custom_llm_provider=provider ) result = file_content.content @@ -90,8 +93,9 @@ def test_create_batch(): pass +@pytest.mark.parametrize("provider", ["openai", "azure"]) @pytest.mark.asyncio() -async def test_async_create_batch(): +async def test_async_create_batch(provider): """ 1. Create File for Batch completion 2. Create Batch Request @@ -105,10 +109,11 @@ async def test_async_create_batch(): file_obj = await litellm.acreate_file( file=open(file_path, "rb"), purpose="batch", - custom_llm_provider="openai", + custom_llm_provider=provider, ) print("Response from creating file=", file_obj) + await asyncio.sleep(5) batch_input_file_id = file_obj.id assert ( batch_input_file_id is not None @@ -118,7 +123,7 @@ async def test_async_create_batch(): completion_window="24h", endpoint="/v1/chat/completions", input_file_id=batch_input_file_id, - custom_llm_provider="openai", + custom_llm_provider=provider, metadata={"key1": "value1", "key2": "value2"}, ) @@ -129,6 +134,7 @@ async def test_async_create_batch(): ), 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" + or create_batch_response.endpoint == "/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 @@ -137,7 +143,7 @@ async def test_async_create_batch(): await asyncio.sleep(1) retrieved_batch = await litellm.aretrieve_batch( - batch_id=create_batch_response.id, custom_llm_provider="openai" + batch_id=create_batch_response.id, custom_llm_provider=provider ) print("retrieved batch=", retrieved_batch) # just assert that we retrieved a non None batch @@ -145,27 +151,27 @@ async def test_async_create_batch(): assert retrieved_batch.id == create_batch_response.id # list all batches - list_batches = await litellm.alist_batches(custom_llm_provider="openai", limit=2) + list_batches = await litellm.alist_batches(custom_llm_provider=provider, limit=2) print("list_batches=", list_batches) # try to get file content for our original file file_content = await litellm.afile_content( - file_id=batch_input_file_id, custom_llm_provider="openai" + file_id=batch_input_file_id, custom_llm_provider=provider ) print("file content = ", file_content) # file obj file_obj = await litellm.afile_retrieve( - file_id=batch_input_file_id, custom_llm_provider="openai" + file_id=batch_input_file_id, custom_llm_provider=provider ) print("file obj = ", file_obj) assert file_obj.id == batch_input_file_id # delete file delete_file_response = await litellm.afile_delete( - file_id=batch_input_file_id, custom_llm_provider="openai" + file_id=batch_input_file_id, custom_llm_provider=provider ) print("delete file response = ", delete_file_response) @@ -173,7 +179,7 @@ async def test_async_create_batch(): assert delete_file_response.id == batch_input_file_id all_files_list = await litellm.afile_list( - custom_llm_provider="openai", + custom_llm_provider=provider, ) print("all_files_list = ", all_files_list) From ada426d652d935a6544d44cc959bd52cd08785bb Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 22 Aug 2024 14:51:14 -0700 Subject: [PATCH 2/7] docs(batches.md): add docs on calling azure batches api --- docs/my-website/docs/batches.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/docs/my-website/docs/batches.md b/docs/my-website/docs/batches.md index 101d1e505b..c73048f8b9 100644 --- a/docs/my-website/docs/batches.md +++ b/docs/my-website/docs/batches.md @@ -5,6 +5,9 @@ import TabItem from '@theme/TabItem'; Covers Batches, Files +Supported Providers: +- Azure OpenAI +- OpenAI ## Quick Start @@ -139,3 +142,12 @@ print("list_batches_response=", list_batches_response) ## [👉 Proxy API Reference](https://litellm-api.up.railway.app/#/batch) + +## Azure Batches API + +Just add the azure env vars to your environment. + +```bash +export AZURE_API_KEY="" +export AZURE_API_BASE="" +``` \ No newline at end of file From 86256634588abc3bd34faa9feac9110a0d626e1b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 22 Aug 2024 15:21:43 -0700 Subject: [PATCH 3/7] feat(proxy_server.py): support azure batch api endpoints --- litellm/batches/main.py | 6 +- litellm/files/main.py | 10 ++-- litellm/proxy/_new_secret_config.yaml | 9 --- .../openai_files_endpoints/files_endpoints.py | 56 ++++++++++++++++--- litellm/proxy/proxy_cli.py | 2 +- litellm/proxy/proxy_server.py | 33 +++++++++-- 6 files changed, 83 insertions(+), 33 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index de3ddd11c9..2da65fd233 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -46,7 +46,7 @@ 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", + custom_llm_provider: str = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -215,7 +215,7 @@ def create_batch( async def aretrieve_batch( batch_id: str, - custom_llm_provider: Literal["openai"] = "openai", + custom_llm_provider: str = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -379,7 +379,7 @@ def retrieve_batch( async def alist_batches( after: Optional[str] = None, limit: Optional[int] = None, - custom_llm_provider: Literal["openai"] = "openai", + custom_llm_provider: str = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, diff --git a/litellm/files/main.py b/litellm/files/main.py index 1ed1c1e611..dc616db29f 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -35,7 +35,7 @@ azure_files_instance = AzureOpenAIFilesAPI() async def afile_retrieve( file_id: str, - custom_llm_provider: Literal["openai", "azure"] = "openai", + custom_llm_provider: str = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -189,7 +189,7 @@ def file_retrieve( # Delete file async def afile_delete( file_id: str, - custom_llm_provider: Literal["openai", "azure"] = "openai", + custom_llm_provider: str = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -339,7 +339,7 @@ def file_delete( # List files async def afile_list( - custom_llm_provider: Literal["openai", "azure"] = "openai", + custom_llm_provider: str = "openai", purpose: Optional[str] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -493,7 +493,7 @@ def file_list( async def acreate_file( file: FileTypes, purpose: Literal["assistants", "batch", "fine-tune"], - custom_llm_provider: Literal["openai", "azure"] = "openai", + custom_llm_provider: str = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -654,7 +654,7 @@ def create_file( async def afile_content( file_id: str, - custom_llm_provider: Literal["openai", "azure"] = "openai", + custom_llm_provider: str = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 2c888a4f30..96a0242a8e 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -2,12 +2,3 @@ model_list: - model_name: "*" litellm_params: model: "*" - -litellm_settings: - success_callback: ["s3"] - cache: true - s3_callback_params: - s3_bucket_name: mytestbucketlitellm # AWS Bucket Name for S3 - s3_region_name: us-west-2 # AWS Region Name for S3 - s3_aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID # us os.environ/ to pass environment variables. This is AWS Access Key ID for S3 - s3_aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY # AWS Secret Access Key for S3 diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 807e02a3a5..cd51434397 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -66,6 +66,11 @@ def get_files_provider_config( return None +@router.post( + "/{provider}/v1/files", + dependencies=[Depends(user_api_key_auth)], + tags=["files"], +) @router.post( "/v1/files", dependencies=[Depends(user_api_key_auth)], @@ -80,6 +85,7 @@ async def create_file( request: Request, fastapi_response: Response, purpose: str = Form(...), + provider: Optional[str] = None, custom_llm_provider: str = Form(default="openai"), file: UploadFile = File(...), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -110,6 +116,8 @@ async def create_file( data: Dict = {} try: + if provider is not None: + custom_llm_provider = provider # Use orjson to parse JSON data, orjson speeds up requests significantly # Read the file content file_content = await file.read() @@ -141,7 +149,9 @@ async def create_file( _create_file_request.update(llm_provider_config) # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch - response = await litellm.acreate_file(**_create_file_request) + response = await litellm.acreate_file( + **_create_file_request, custom_llm_provider=custom_llm_provider + ) ### ALERTING ### asyncio.create_task( @@ -195,6 +205,11 @@ async def create_file( ) +@router.get( + "/{provider}/v1/files/{file_id:path}", + dependencies=[Depends(user_api_key_auth)], + tags=["files"], +) @router.get( "/v1/files/{file_id:path}", dependencies=[Depends(user_api_key_auth)], @@ -209,6 +224,7 @@ async def get_file( request: Request, fastapi_response: Response, file_id: str, + provider: Optional[str] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -246,9 +262,10 @@ async def get_file( proxy_config=proxy_config, ) - # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch + if provider is None: # default to openai + provider = "openai" response = await litellm.afile_retrieve( - custom_llm_provider="openai", file_id=file_id, **data + custom_llm_provider=provider, file_id=file_id, **data ) ### ALERTING ### @@ -303,6 +320,11 @@ async def get_file( ) +@router.delete( + "/{provider}/v1/files/{file_id:path}", + dependencies=[Depends(user_api_key_auth)], + tags=["files"], +) @router.delete( "/v1/files/{file_id:path}", dependencies=[Depends(user_api_key_auth)], @@ -317,6 +339,7 @@ async def delete_file( request: Request, fastapi_response: Response, file_id: str, + provider: Optional[str] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -355,9 +378,10 @@ async def delete_file( proxy_config=proxy_config, ) - # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch + if provider is None: # default to openai + provider = "openai" response = await litellm.afile_delete( - custom_llm_provider="openai", file_id=file_id, **data + custom_llm_provider=provider, file_id=file_id, **data ) ### ALERTING ### @@ -412,6 +436,11 @@ async def delete_file( ) +@router.get( + "/{provider}/v1/files", + dependencies=[Depends(user_api_key_auth)], + tags=["files"], +) @router.get( "/v1/files", dependencies=[Depends(user_api_key_auth)], @@ -426,6 +455,7 @@ async def list_files( request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + provider: Optional[str] = None, purpose: Optional[str] = None, ): """ @@ -463,9 +493,10 @@ async def list_files( proxy_config=proxy_config, ) - # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch + if provider is None: + provider = "openai" response = await litellm.afile_list( - custom_llm_provider="openai", purpose=purpose, **data + custom_llm_provider=provider, purpose=purpose, **data ) ### ALERTING ### @@ -520,6 +551,11 @@ async def list_files( ) +@router.get( + "/{provider}/v1/files/{file_id:path}/content", + dependencies=[Depends(user_api_key_auth)], + tags=["files"], +) @router.get( "/v1/files/{file_id:path}/content", dependencies=[Depends(user_api_key_auth)], @@ -534,6 +570,7 @@ async def get_file_content( request: Request, fastapi_response: Response, file_id: str, + provider: Optional[str] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -571,9 +608,10 @@ async def get_file_content( proxy_config=proxy_config, ) - # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch + if provider is None: + provider = "openai" response = await litellm.afile_content( - custom_llm_provider="openai", file_id=file_id, **data + custom_llm_provider=provider, file_id=file_id, **data ) ### ALERTING ### diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index e3edd1b8c6..cf2638f3c8 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -78,7 +78,7 @@ def is_port_in_use(port): @click.option("--api_base", default=None, help="API base URL.") @click.option( "--api_version", - default="2024-02-01", + default="2024-07-01-preview", help="For azure - pass in the api version.", ) @click.option( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0a9abc09ad..f554c174a2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4877,6 +4877,11 @@ async def run_thread( ###################################################################### +@router.get( + "/{provider}/v1/batches", + dependencies=[Depends(user_api_key_auth)], + tags=["batch"], +) @router.post( "/v1/batches", dependencies=[Depends(user_api_key_auth)], @@ -4890,6 +4895,7 @@ async def run_thread( async def create_batch( request: Request, fastapi_response: Response, + provider: Optional[str] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -4936,9 +4942,10 @@ async def create_batch( _create_batch_data = CreateBatchRequest(**data) - # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch + if provider is None: + provider = "openai" response = await litellm.acreate_batch( - custom_llm_provider="openai", **_create_batch_data + custom_llm_provider=provider, **_create_batch_data ) ### ALERTING ### @@ -4994,6 +5001,11 @@ async def create_batch( ) +@router.get( + "/{provider}/v1/batches/{batch_id:path}", + dependencies=[Depends(user_api_key_auth)], + tags=["batch"], +) @router.get( "/v1/batches/{batch_id:path}", dependencies=[Depends(user_api_key_auth)], @@ -5008,6 +5020,7 @@ async def retrieve_batch( request: Request, fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + provider: Optional[str] = None, batch_id: str = Path( title="Batch ID to retrieve", description="The ID of the batch to retrieve" ), @@ -5032,9 +5045,10 @@ async def retrieve_batch( batch_id=batch_id, ) - # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch + if provider is None: + provider = "openai" response = await litellm.aretrieve_batch( - custom_llm_provider="openai", **_retrieve_batch_request + custom_llm_provider=provider, **_retrieve_batch_request ) ### ALERTING ### @@ -5091,6 +5105,11 @@ async def retrieve_batch( ) +@router.get( + "/{provider}/v1/batches", + dependencies=[Depends(user_api_key_auth)], + tags=["batch"], +) @router.get( "/v1/batches", dependencies=[Depends(user_api_key_auth)], @@ -5103,6 +5122,7 @@ async def retrieve_batch( ) async def list_batches( fastapi_response: Response, + provider: Optional[str] = None, limit: Optional[int] = None, after: Optional[str] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -5123,9 +5143,10 @@ async def list_batches( global proxy_logging_obj verbose_proxy_logger.debug("GET /v1/batches after={} limit={}".format(after, limit)) try: - # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch + if provider is None: + provider = "openai" response = await litellm.alist_batches( - custom_llm_provider="openai", + custom_llm_provider=provider, after=after, limit=limit, ) From 63cd94c32a809c811e9d5de6e7b9fdac2e8e454b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 22 Aug 2024 15:51:59 -0700 Subject: [PATCH 4/7] fix: fix linting errors --- litellm/batches/main.py | 6 +++--- litellm/files/main.py | 10 +++++----- .../proxy/openai_files_endpoints/files_endpoints.py | 10 +++++----- litellm/proxy/proxy_server.py | 6 +++--- 4 files changed, 16 insertions(+), 16 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 2da65fd233..99e1707dcf 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -46,7 +46,7 @@ async def acreate_batch( completion_window: Literal["24h"], endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], input_file_id: str, - custom_llm_provider: str = "openai", + custom_llm_provider: Literal["openai", "azure"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -215,7 +215,7 @@ def create_batch( async def aretrieve_batch( batch_id: str, - custom_llm_provider: str = "openai", + custom_llm_provider: Literal["openai", "azure"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -379,7 +379,7 @@ def retrieve_batch( async def alist_batches( after: Optional[str] = None, limit: Optional[int] = None, - custom_llm_provider: str = "openai", + custom_llm_provider: Literal["openai", "azure"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, diff --git a/litellm/files/main.py b/litellm/files/main.py index dc616db29f..1ed1c1e611 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -35,7 +35,7 @@ azure_files_instance = AzureOpenAIFilesAPI() async def afile_retrieve( file_id: str, - custom_llm_provider: str = "openai", + custom_llm_provider: Literal["openai", "azure"] = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -189,7 +189,7 @@ def file_retrieve( # Delete file async def afile_delete( file_id: str, - custom_llm_provider: str = "openai", + custom_llm_provider: Literal["openai", "azure"] = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -339,7 +339,7 @@ def file_delete( # List files async def afile_list( - custom_llm_provider: str = "openai", + custom_llm_provider: Literal["openai", "azure"] = "openai", purpose: Optional[str] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -493,7 +493,7 @@ def file_list( async def acreate_file( file: FileTypes, purpose: Literal["assistants", "batch", "fine-tune"], - custom_llm_provider: str = "openai", + custom_llm_provider: Literal["openai", "azure"] = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -654,7 +654,7 @@ def create_file( async def afile_content( file_id: str, - custom_llm_provider: str = "openai", + custom_llm_provider: Literal["openai", "azure"] = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index cd51434397..27bce56895 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -150,7 +150,7 @@ async def create_file( # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch response = await litellm.acreate_file( - **_create_file_request, custom_llm_provider=custom_llm_provider + **_create_file_request, custom_llm_provider=custom_llm_provider # type: ignore ) ### ALERTING ### @@ -265,7 +265,7 @@ async def get_file( if provider is None: # default to openai provider = "openai" response = await litellm.afile_retrieve( - custom_llm_provider=provider, file_id=file_id, **data + custom_llm_provider=provider, file_id=file_id, **data # type: ignore ) ### ALERTING ### @@ -381,7 +381,7 @@ async def delete_file( if provider is None: # default to openai provider = "openai" response = await litellm.afile_delete( - custom_llm_provider=provider, file_id=file_id, **data + custom_llm_provider=provider, file_id=file_id, **data # type: ignore ) ### ALERTING ### @@ -496,7 +496,7 @@ async def list_files( if provider is None: provider = "openai" response = await litellm.afile_list( - custom_llm_provider=provider, purpose=purpose, **data + custom_llm_provider=provider, purpose=purpose, **data # type: ignore ) ### ALERTING ### @@ -611,7 +611,7 @@ async def get_file_content( if provider is None: provider = "openai" response = await litellm.afile_content( - custom_llm_provider=provider, file_id=file_id, **data + custom_llm_provider=provider, file_id=file_id, **data # type: ignore ) ### ALERTING ### diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f554c174a2..1d4da51818 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4945,7 +4945,7 @@ async def create_batch( if provider is None: provider = "openai" response = await litellm.acreate_batch( - custom_llm_provider=provider, **_create_batch_data + custom_llm_provider=provider, **_create_batch_data # type: ignore ) ### ALERTING ### @@ -5048,7 +5048,7 @@ async def retrieve_batch( if provider is None: provider = "openai" response = await litellm.aretrieve_batch( - custom_llm_provider=provider, **_retrieve_batch_request + custom_llm_provider=provider, **_retrieve_batch_request # type: ignore ) ### ALERTING ### @@ -5146,7 +5146,7 @@ async def list_batches( if provider is None: provider = "openai" response = await litellm.alist_batches( - custom_llm_provider=provider, + custom_llm_provider=provider, # type: ignore after=after, limit=limit, ) From d7d3eee3496dbc4dcb82197b5ea50eae083d4989 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 22 Aug 2024 16:11:14 -0700 Subject: [PATCH 5/7] feat(azure.py): support health checking azure deployments Fixes https://github.com/BerriAI/litellm/issues/5279 --- docs/my-website/docs/proxy/health.md | 34 +++++++++++++++++++++++++++ litellm/llms/azure.py | 2 ++ litellm/main.py | 2 +- litellm/proxy/_new_secret_config.yaml | 7 ++++-- 4 files changed, 42 insertions(+), 3 deletions(-) diff --git a/docs/my-website/docs/proxy/health.md b/docs/my-website/docs/proxy/health.md index 632702b914..e7ff69aeb3 100644 --- a/docs/my-website/docs/proxy/health.md +++ b/docs/my-website/docs/proxy/health.md @@ -115,6 +115,39 @@ model_list: mode: audio_speech ``` +### Batch Models (Azure Only) + +For Azure models deployed as 'batch' models, set `mode: batch`. + +```yaml +model_list: + - model_name: "batch-gpt-4o-mini" + litellm_params: + model: "azure/gpt-4o-mini" + api_key: os.environ/AZURE_API_KEY + api_base: os.environ/AZURE_API_BASE + model_info: + mode: batch +``` + +Expected Response + + +```bash +{ + "healthy_endpoints": [ + { + "api_base": "https://...", + "model": "azure/gpt-4o-mini", + "x-ms-region": "East US" + } + ], + "unhealthy_endpoints": [], + "healthy_count": 1, + "unhealthy_count": 0 +} +``` + ## Background Health Checks You can enable model health checks being run in the background, to prevent each model from being queried too frequently via `/health`. @@ -244,3 +277,4 @@ curl -X POST 'http://localhost:4000/chat/completions' \ } ' ``` + diff --git a/litellm/llms/azure.py b/litellm/llms/azure.py index a54bef3192..e235187801 100644 --- a/litellm/llms/azure.py +++ b/litellm/llms/azure.py @@ -1970,6 +1970,8 @@ class AzureChatCompletion(BaseLLM): input=prompt, # type: ignore voice="alloy", ) + elif mode == "batch": + completion = await client.batches.with_raw_response.list(limit=1) # type: ignore else: raise Exception("mode not set") response = {} diff --git a/litellm/main.py b/litellm/main.py index 28054537cf..49436f1537 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -4825,7 +4825,7 @@ def speech( async def ahealth_check( model_params: dict, mode: Optional[ - Literal["completion", "embedding", "image_generation", "chat"] + Literal["completion", "embedding", "image_generation", "chat", "batch"] ] = None, prompt: Optional[str] = None, input: Optional[List] = None, diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 96a0242a8e..50a6d993ec 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -1,4 +1,7 @@ model_list: - - model_name: "*" + - model_name: "batch-gpt-4o-mini" litellm_params: - model: "*" + model: "azure/gpt-4o-mini" + api_key: os.environ/AZURE_API_KEY + api_base: os.environ/AZURE_API_BASE + From 7398e944725696254293edfd86f82d0df3c45e95 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 22 Aug 2024 16:42:44 -0700 Subject: [PATCH 6/7] fix(files_endpoints.py): fix multiple args error --- litellm/proxy/openai_files_endpoints/files_endpoints.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 27bce56895..f4400b682c 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -149,9 +149,7 @@ async def create_file( _create_file_request.update(llm_provider_config) # for now use custom_llm_provider=="openai" -> this will change as LiteLLM adds more providers for acreate_batch - response = await litellm.acreate_file( - **_create_file_request, custom_llm_provider=custom_llm_provider # type: ignore - ) + response = await litellm.acreate_file(**_create_file_request) # type: ignore ### ALERTING ### asyncio.create_task( From 56cb94ac5eaf663f2d9033c7ecab3d2ab929f7fd Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 22 Aug 2024 18:51:44 -0700 Subject: [PATCH 7/7] docs(batches.md): add more examples to docs --- docs/my-website/docs/batches.md | 74 +++++++++++++++++++++++++++- docs/my-website/docs/proxy/health.md | 2 +- 2 files changed, 74 insertions(+), 2 deletions(-) diff --git a/docs/my-website/docs/batches.md b/docs/my-website/docs/batches.md index c73048f8b9..898738f632 100644 --- a/docs/my-website/docs/batches.md +++ b/docs/my-website/docs/batches.md @@ -150,4 +150,76 @@ Just add the azure env vars to your environment. ```bash export AZURE_API_KEY="" export AZURE_API_BASE="" -``` \ No newline at end of file +``` + +AND use `/azure/*` for the Batches API calls + +```bash +http://0.0.0.0:4000/azure/v1/batches +``` +### Usage + +**Setup** + +- Add Azure API Keys to your environment + +#### 1. Upload a File + +```bash +curl http://localhost:4000/azure/v1/files \ + -H "Authorization: Bearer sk-1234" \ + -F purpose="batch" \ + -F file="@mydata.jsonl" +``` + +**Example File** + +Note: `model` should be your azure deployment name. + +```json +{"custom_id": "task-0", "method": "POST", "url": "/chat/completions", "body": {"model": "REPLACE-WITH-MODEL-DEPLOYMENT-NAME", "messages": [{"role": "system", "content": "You are an AI assistant that helps people find information."}, {"role": "user", "content": "When was Microsoft founded?"}]}} +{"custom_id": "task-1", "method": "POST", "url": "/chat/completions", "body": {"model": "REPLACE-WITH-MODEL-DEPLOYMENT-NAME", "messages": [{"role": "system", "content": "You are an AI assistant that helps people find information."}, {"role": "user", "content": "When was the first XBOX released?"}]}} +{"custom_id": "task-2", "method": "POST", "url": "/chat/completions", "body": {"model": "REPLACE-WITH-MODEL-DEPLOYMENT-NAME", "messages": [{"role": "system", "content": "You are an AI assistant that helps people find information."}, {"role": "user", "content": "What is Altair Basic?"}]}} +``` + +#### 2. Create a batch + +```bash +curl http://0.0.0.0:4000/azure/v1/batches \ + -H "Authorization: Bearer $LITELLM_API_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "input_file_id": "file-abc123", + "endpoint": "/v1/chat/completions", + "completion_window": "24h" + }' + +``` + +#### 3. Retrieve batch + + +```bash +curl http://0.0.0.0:4000/azure/v1/batches/batch_abc123 \ + -H "Authorization: Bearer $LITELLM_API_KEY" \ + -H "Content-Type: application/json" \ +``` + +#### 4. Cancel batch + +```bash +curl http://0.0.0.0:4000/azure/v1/batches/batch_abc123/cancel \ + -H "Authorization: Bearer $LITELLM_API_KEY" \ + -H "Content-Type: application/json" \ + -X POST +``` + +#### 5. List Batch + +```bash +curl http://0.0.0.0:4000/v1/batches?limit=2 \ + -H "Authorization: Bearer $LITELLM_API_KEY" \ + -H "Content-Type: application/json" +``` + +### [👉 Health Check Azure Batch models](./proxy/health.md#batch-models-azure-only) \ No newline at end of file diff --git a/docs/my-website/docs/proxy/health.md b/docs/my-website/docs/proxy/health.md index e7ff69aeb3..35dced84c9 100644 --- a/docs/my-website/docs/proxy/health.md +++ b/docs/my-website/docs/proxy/health.md @@ -123,7 +123,7 @@ For Azure models deployed as 'batch' models, set `mode: batch`. model_list: - model_name: "batch-gpt-4o-mini" litellm_params: - model: "azure/gpt-4o-mini" + model: "azure/batch-gpt-4o-mini" api_key: os.environ/AZURE_API_KEY api_base: os.environ/AZURE_API_BASE model_info: