diff --git a/docs/my-website/docs/batches.md b/docs/my-website/docs/batches.md index 101d1e505b..898738f632 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,84 @@ 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="" +``` + +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 632702b914..35dced84c9 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/batch-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/__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..99e1707dcf 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() ################################################# @@ -44,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: Literal["openai", "azure"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -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( @@ -183,7 +215,7 @@ def create_batch( async def aretrieve_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, @@ -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( @@ -316,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: Literal["openai", "azure"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -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..e235187801 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, @@ -1966,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 = {} @@ -2814,3 +2820,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/main.py b/litellm/main.py index 16a4f89ed2..8104bfd864 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -4831,7 +4831,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..0e401035a9 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -1,4 +1,4 @@ model_list: - - model_name: "*" + - model_name: "batch-gpt-4o-mini" litellm_params: model: "*" diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 807e02a3a5..f4400b682c 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,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) + response = await litellm.acreate_file(**_create_file_request) # type: ignore ### ALERTING ### asyncio.create_task( @@ -195,6 +203,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 +222,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 +260,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 # type: ignore ) ### ALERTING ### @@ -303,6 +318,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 +337,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 +376,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 # type: ignore ) ### ALERTING ### @@ -412,6 +434,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 +453,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 +491,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 # type: ignore ) ### ALERTING ### @@ -520,6 +549,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 +568,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 +606,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 # type: ignore ) ### 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 c793ffbe3f..3ef5609db3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4884,6 +4884,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)], @@ -4897,6 +4902,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), ): """ @@ -4943,9 +4949,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 # type: ignore ) ### ALERTING ### @@ -5001,6 +5008,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)], @@ -5015,6 +5027,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" ), @@ -5039,9 +5052,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 # type: ignore ) ### ALERTING ### @@ -5098,6 +5112,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)], @@ -5110,6 +5129,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), @@ -5130,9 +5150,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, # type: ignore after=after, limit=limit, ) 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)