From 6df995d0d1bc80e8cdb6500af1e06ec6e8625c1d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 18:27:25 -0700 Subject: [PATCH 01/16] refactor: dead test cleanup --- tests/audio_tests/test_audio_speech.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/audio_tests/test_audio_speech.py b/tests/audio_tests/test_audio_speech.py index 565c11ba1d..abf0df22a6 100644 --- a/tests/audio_tests/test_audio_speech.py +++ b/tests/audio_tests/test_audio_speech.py @@ -286,7 +286,7 @@ async def test_speech_litellm_vertex_async_with_voice_ssml(): def test_audio_speech_cost_calc(): from litellm.integrations.custom_logger import CustomLogger - model = "azure/azure-tts" + model = "azure/tts" api_base = os.getenv("AZURE_TTS_API_BASE") api_key = os.getenv("AZURE_TTS_API_KEY") From 5cd8ca236506b954f97af872ec9075df165aa962 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 18:39:32 -0700 Subject: [PATCH 02/16] refactor: refactor testing --- tests/batches_tests/test_fine_tuning_api.py | 72 --------- tests/llm_translation/test_rerank.py | 42 ----- .../azure_cancel_expected_output.json | 20 +++ .../azure_cancel_raw_response.json | 18 +++ .../azure_cancel_request.json | 3 + .../azure_create_expected_output.json | 20 +++ .../azure_create_raw_response.json | 18 +++ .../azure_create_request.json | 8 + .../azure_list_raw_response.json | 20 +++ .../azure_list_request.json | 4 + .../llms/azure/test_azure_fine_tuning_api.py | 150 ++++++++++++++++++ 11 files changed, 261 insertions(+), 114 deletions(-) create mode 100644 tests/test_litellm/expected_fine_tuning_api/azure_cancel_expected_output.json create mode 100644 tests/test_litellm/expected_fine_tuning_api/azure_cancel_raw_response.json create mode 100644 tests/test_litellm/expected_fine_tuning_api/azure_cancel_request.json create mode 100644 tests/test_litellm/expected_fine_tuning_api/azure_create_expected_output.json create mode 100644 tests/test_litellm/expected_fine_tuning_api/azure_create_raw_response.json create mode 100644 tests/test_litellm/expected_fine_tuning_api/azure_create_request.json create mode 100644 tests/test_litellm/expected_fine_tuning_api/azure_list_raw_response.json create mode 100644 tests/test_litellm/expected_fine_tuning_api/azure_list_request.json create mode 100644 tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py diff --git a/tests/batches_tests/test_fine_tuning_api.py b/tests/batches_tests/test_fine_tuning_api.py index 7e220612ac..034f4e1b49 100644 --- a/tests/batches_tests/test_fine_tuning_api.py +++ b/tests/batches_tests/test_fine_tuning_api.py @@ -123,78 +123,6 @@ async def test_create_fine_tune_jobs_async(): pass -@pytest.mark.asyncio -async def test_azure_create_fine_tune_jobs_async(): - try: - verbose_logger.setLevel(logging.DEBUG) - file_name = "azure_fine_tune.jsonl" - _current_dir = os.path.dirname(os.path.abspath(__file__)) - file_path = os.path.join(_current_dir, file_name) - - file_id = "file-5e4b20ecbd724182b9964f3cd2ab7212" - - create_fine_tuning_response = await litellm.acreate_fine_tuning_job( - model="gpt-35-turbo-1106", - training_file=file_id, - custom_llm_provider="azure", - api_base="https://exampleopenaiendpoint-production.up.railway.app", - ) - - print( - "response from litellm.create_fine_tuning_job=", create_fine_tuning_response - ) - - assert create_fine_tuning_response.id is not None - - # response from Example/mocked endpoint - assert create_fine_tuning_response.model == "davinci-002" - - # list fine tuning jobs - print("listing ft jobs") - ft_jobs = await litellm.alist_fine_tuning_jobs( - limit=2, - custom_llm_provider="azure", - api_base="https://exampleopenaiendpoint-production.up.railway.app", - ) - print("response from litellm.list_fine_tuning_jobs=", ft_jobs) - - # cancel ft job - response = await litellm.acancel_fine_tuning_job( - fine_tuning_job_id=create_fine_tuning_response.id, - custom_llm_provider="azure", - api_key=os.getenv("AZURE_SWEDEN_API_KEY"), - api_base="https://exampleopenaiendpoint-production.up.railway.app", - ) - - print("response from litellm.cancel_fine_tuning_job=", response) - - assert response.status == "cancelled" - assert response.id == create_fine_tuning_response.id - except openai.RateLimitError: - pass - except Exception as e: - if "Job has already completed" in str(e): - pass - else: - pytest.fail(f"Error occurred: {e}") - pass - - -def test_azure_trainingtype_defaults_to_one(): - """ - Azure requires trainingType in extra_body. When omitted, AzureOpenAIFineTuningAPI defaults it to 1. - """ - from litellm.llms.azure.fine_tuning.handler import AzureOpenAIFineTuningAPI - - handler = AzureOpenAIFineTuningAPI() - create_data = {"model": "gpt-4o-mini", "training_file": "file-test"} - - handler._ensure_training_type(create_data) - - assert "extra_body" in create_data - assert create_data["extra_body"]["trainingType"] == 1 - - @pytest.mark.asyncio() async def test_create_vertex_fine_tune_jobs_mocked(): load_vertex_ai_credentials() diff --git a/tests/llm_translation/test_rerank.py b/tests/llm_translation/test_rerank.py index edd5981c16..d784677060 100644 --- a/tests/llm_translation/test_rerank.py +++ b/tests/llm_translation/test_rerank.py @@ -148,48 +148,6 @@ async def test_basic_rerank_together_ai(sync_mode): raise e -@pytest.mark.asyncio() -@pytest.mark.parametrize("sync_mode", [True, False]) -@pytest.mark.skip(reason="Skipping test due to Cohere RBAC issues") -async def test_basic_rerank_azure_ai(sync_mode): - import os - - litellm.set_verbose = True - - if sync_mode is True: - response = litellm.rerank( - model="azure_ai/Cohere-rerank-v3-multilingual-ko", - query="hello", - documents=["hello", "world"], - top_n=3, - api_key=os.getenv("AZURE_AI_COHERE_API_KEY"), - api_base=os.getenv("AZURE_AI_COHERE_API_BASE"), - ) - - print("re rank response: ", response) - - assert response.id is not None - assert response.results is not None - - assert_response_shape(response, custom_llm_provider="together_ai") - else: - response = await litellm.arerank( - model="azure_ai/Cohere-rerank-v3-multilingual-ko", - query="hello", - documents=["hello", "world"], - top_n=3, - api_key=os.getenv("AZURE_AI_COHERE_API_KEY"), - api_base=os.getenv("AZURE_AI_COHERE_API_BASE"), - ) - - print("async re rank response: ", response) - - assert response.id is not None - assert response.results is not None - - assert_response_shape(response, custom_llm_provider="together_ai") - - @pytest.mark.asyncio() @pytest.mark.parametrize("version", ["v1", "v2"]) async def test_rerank_custom_api_base(version): diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_cancel_expected_output.json b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_expected_output.json new file mode 100644 index 0000000000..d0a519fcd1 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_expected_output.json @@ -0,0 +1,20 @@ +{ + "id": "ftjob-azure-create-123", + "object": "fine_tuning.job", + "created_at": 1735689600, + "model": "davinci-002", + "status": "cancelled", + "fine_tuned_model": null, + "training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212", + "hyperparameters": { + "n_epochs": 3, + "batch_size": null, + "learning_rate_multiplier": null + }, + "organization_id": "", + "result_files": [], + "validation_file": null, + "trained_tokens": null, + "estimated_finish": null, + "error": null +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_cancel_raw_response.json b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_raw_response.json new file mode 100644 index 0000000000..0093a04f70 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_raw_response.json @@ -0,0 +1,18 @@ +{ + "id": "ftjob-azure-create-123", + "object": "fine_tuning.job", + "created_at": 1735689600, + "model": "davinci-002", + "status": "canceled", + "fine_tuned_model": null, + "training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212", + "hyperparameters": { + "n_epochs": 3 + }, + "organization_id": null, + "result_files": null, + "validation_file": null, + "trained_tokens": null, + "estimated_finish": null, + "error": null +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_cancel_request.json b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_request.json new file mode 100644 index 0000000000..bdbbdbad07 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_cancel_request.json @@ -0,0 +1,3 @@ +{ + "fine_tuning_job_id": "ftjob-azure-create-123" +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_create_expected_output.json b/tests/test_litellm/expected_fine_tuning_api/azure_create_expected_output.json new file mode 100644 index 0000000000..201ae3047d --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_create_expected_output.json @@ -0,0 +1,20 @@ +{ + "id": "ftjob-azure-create-123", + "object": "fine_tuning.job", + "created_at": 1735689600, + "model": "davinci-002", + "status": "running", + "fine_tuned_model": null, + "training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212", + "hyperparameters": { + "n_epochs": 3, + "batch_size": null, + "learning_rate_multiplier": null + }, + "organization_id": "", + "result_files": [], + "validation_file": null, + "trained_tokens": null, + "estimated_finish": null, + "error": null +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_create_raw_response.json b/tests/test_litellm/expected_fine_tuning_api/azure_create_raw_response.json new file mode 100644 index 0000000000..857b249938 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_create_raw_response.json @@ -0,0 +1,18 @@ +{ + "id": "ftjob-azure-create-123", + "object": "fine_tuning.job", + "created_at": 1735689600, + "model": "davinci-002", + "status": "running", + "fine_tuned_model": null, + "training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212", + "hyperparameters": { + "n_epochs": 3 + }, + "organization_id": null, + "result_files": null, + "validation_file": null, + "trained_tokens": null, + "estimated_finish": null, + "error": null +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_create_request.json b/tests/test_litellm/expected_fine_tuning_api/azure_create_request.json new file mode 100644 index 0000000000..57319b2698 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_create_request.json @@ -0,0 +1,8 @@ +{ + "model": "gpt-35-turbo-1106", + "training_file": "file-5e4b20ecbd724182b9964f3cd2ab7212", + "hyperparameters": {}, + "extra_body": { + "trainingType": 1 + } +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_list_raw_response.json b/tests/test_litellm/expected_fine_tuning_api/azure_list_raw_response.json new file mode 100644 index 0000000000..9986e7d4c5 --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_list_raw_response.json @@ -0,0 +1,20 @@ +{ + "object": "list", + "data": [ + { + "id": "ftjob-azure-create-123", + "object": "fine_tuning.job", + "created_at": 1735689600, + "model": "davinci-002", + "status": "running" + }, + { + "id": "ftjob-azure-prev-000", + "object": "fine_tuning.job", + "created_at": 1735603200, + "model": "davinci-002", + "status": "succeeded" + } + ], + "has_more": false +} diff --git a/tests/test_litellm/expected_fine_tuning_api/azure_list_request.json b/tests/test_litellm/expected_fine_tuning_api/azure_list_request.json new file mode 100644 index 0000000000..6bfbb2ffec --- /dev/null +++ b/tests/test_litellm/expected_fine_tuning_api/azure_list_request.json @@ -0,0 +1,4 @@ +{ + "after": "ftjob-azure-prev-000", + "limit": 2 +} diff --git a/tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py b/tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py new file mode 100644 index 0000000000..8d008d6c07 --- /dev/null +++ b/tests/test_litellm/llms/azure/test_azure_fine_tuning_api.py @@ -0,0 +1,150 @@ +import json +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import pytest +from openai import AsyncAzureOpenAI + +import litellm +from litellm.llms.azure.fine_tuning.handler import AzureOpenAIFineTuningAPI + + +def _expected_dir() -> Path: + return Path(__file__).resolve().parent.parent.parent / "expected_fine_tuning_api" + + +def _load_json(file_name: str) -> dict: + path = _expected_dir() / file_name + assert path.exists(), f"Expected fixture file not found: {path}" + with open(path) as f: + return json.load(f) + + +class _MockSDKResponse: + def __init__(self, payload: dict): + self._payload = payload + + def model_dump(self) -> dict: + return self._payload + + +def _mock_azure_client( + create_payload: dict | None = None, + list_payload: dict | None = None, + cancel_payload: dict | None = None, +): + client = AsyncAzureOpenAI( + api_key="test-key", + api_version="2024-10-21", + azure_endpoint="https://exampleopenaiendpoint-production.up.railway.app", + ) + client.fine_tuning.jobs.create = AsyncMock( + return_value=( + _MockSDKResponse(create_payload) if create_payload is not None else None + ) + ) # type: ignore[method-assign] + client.fine_tuning.jobs.list = AsyncMock( + return_value=list_payload + ) # type: ignore[method-assign] + client.fine_tuning.jobs.cancel = AsyncMock( + return_value=( + _MockSDKResponse(cancel_payload) if cancel_payload is not None else None + ) + ) # type: ignore[method-assign] + return client + + +@pytest.mark.asyncio +async def test_azure_acreate_fine_tuning_job_request_and_output_match_expected_json(): + expected_request = _load_json("azure_create_request.json") + raw_response = _load_json("azure_create_raw_response.json") + expected_output = _load_json("azure_create_expected_output.json") + + mock_client = _mock_azure_client(create_payload=raw_response) + + with patch.object( + AzureOpenAIFineTuningAPI, "get_openai_client", return_value=mock_client + ): + response = await litellm.acreate_fine_tuning_job( + model="gpt-35-turbo-1106", + training_file="file-5e4b20ecbd724182b9964f3cd2ab7212", + custom_llm_provider="azure", + api_base="https://exampleopenaiendpoint-production.up.railway.app", + api_key="test-key", + api_version="2024-10-21", + ) + + request_kwargs = mock_client.fine_tuning.jobs.create.call_args.kwargs + assert request_kwargs == expected_request + + response_dict = response.model_dump(exclude={"_hidden_params"}) + for key, expected_value in expected_output.items(): + assert key in response_dict, f"Missing key in response: {key}" + assert response_dict[key] == expected_value + + assert response.id is not None + assert response.model == "davinci-002" + + +@pytest.mark.asyncio +async def test_azure_alist_fine_tuning_jobs_request_matches_expected_json(): + expected_request = _load_json("azure_list_request.json") + raw_list_response = _load_json("azure_list_raw_response.json") + + mock_client = _mock_azure_client(list_payload=raw_list_response) + + with patch.object( + AzureOpenAIFineTuningAPI, "get_openai_client", return_value=mock_client + ): + response = await litellm.alist_fine_tuning_jobs( + after=expected_request["after"], + limit=expected_request["limit"], + custom_llm_provider="azure", + api_base="https://exampleopenaiendpoint-production.up.railway.app", + api_key="test-key", + api_version="2024-10-21", + ) + + request_kwargs = mock_client.fine_tuning.jobs.list.call_args.kwargs + assert request_kwargs == expected_request + assert response == raw_list_response + + +@pytest.mark.asyncio +async def test_azure_acancel_fine_tuning_job_request_and_output_match_expected_json(): + expected_request = _load_json("azure_cancel_request.json") + raw_response = _load_json("azure_cancel_raw_response.json") + expected_output = _load_json("azure_cancel_expected_output.json") + + mock_client = _mock_azure_client(cancel_payload=raw_response) + + with patch.object( + AzureOpenAIFineTuningAPI, "get_openai_client", return_value=mock_client + ): + response = await litellm.acancel_fine_tuning_job( + fine_tuning_job_id=expected_request["fine_tuning_job_id"], + custom_llm_provider="azure", + api_base="https://exampleopenaiendpoint-production.up.railway.app", + api_key="test-key", + api_version="2024-10-21", + ) + + request_kwargs = mock_client.fine_tuning.jobs.cancel.call_args.kwargs + assert request_kwargs == expected_request + + response_dict = response.model_dump(exclude={"_hidden_params"}) + for key, expected_value in expected_output.items(): + assert key in response_dict, f"Missing key in response: {key}" + assert response_dict[key] == expected_value + + assert response.status == "cancelled" + + +def test_azure_trainingtype_defaults_to_one(): + handler = AzureOpenAIFineTuningAPI() + create_data = {"model": "gpt-4o-mini", "training_file": "file-test"} + + handler._ensure_training_type(create_data) + + assert "extra_body" in create_data + assert create_data["extra_body"]["trainingType"] == 1 From 564ad3195b347121a9c898474f9f30ae2df49a7a Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 18:42:55 -0700 Subject: [PATCH 03/16] test: update testing --- tests/local_testing/test_embedding.py | 132 ++++++++++++++------------ 1 file changed, 72 insertions(+), 60 deletions(-) diff --git a/tests/local_testing/test_embedding.py b/tests/local_testing/test_embedding.py index 48d535fe45..d0383f7002 100644 --- a/tests/local_testing/test_embedding.py +++ b/tests/local_testing/test_embedding.py @@ -108,7 +108,11 @@ def test_openai_embedding_3(): "model, api_base, api_key", [ # ("azure/text-embedding-ada-002", None, None), - ("together_ai/BAAI/bge-base-en-v1.5", None, None), # Updated to current Together AI embedding model + ( + "together_ai/BAAI/bge-base-en-v1.5", + None, + None, + ), # Updated to current Together AI embedding model ], ) @pytest.mark.parametrize("sync_mode", [True, False]) @@ -193,9 +197,9 @@ def _azure_ai_image_mock_response(*args, **kwargs): "model, api_base, api_key", [ ( - "azure_ai/Cohere-embed-v3-multilingual-jzu", - "https://Cohere-embed-v3-multilingual-jzu.eastus2.models.ai.azure.com", - os.getenv("AZURE_AI_COHERE_API_KEY_2"), + "azure_ai/Cohere-embed-v3-multilingual-2", + os.getenv("AZURE_AI_COHERE_API_BASE"), + os.getenv("AZURE_AI_COHERE_API_KEY"), ) ], ) @@ -353,11 +357,11 @@ def test_openai_azure_embedding_optional_arg(): ) mock_client.assert_called_once_with( - model="test", - input=["test"], - extra_body={"azure_ad_token": "test"}, - timeout=600, - extra_headers={"X-Stainless-Raw-Response": "true"} + model="test", + input=["test"], + extra_body={"azure_ad_token": "test"}, + timeout=600, + extra_headers={"X-Stainless-Raw-Response": "true"}, ) # Verify azure_ad_token is passed in extra_body, not as a direct parameter assert "azure_ad_token" not in mock_client.call_args.kwargs @@ -811,7 +815,7 @@ def test_watsonx_embeddings(monkeypatch): monkeypatch.setenv("WATSONX_PROJECT_ID", "mock-project-id") client = HTTPHandler() - + # Track the actual request made captured_request = {} @@ -820,7 +824,7 @@ def test_watsonx_embeddings(monkeypatch): captured_request["url"] = url captured_request["headers"] = kwargs.get("headers", {}) captured_request["data"] = kwargs.get("data") - + mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = {"Content-Type": "application/json"} @@ -843,10 +847,12 @@ def test_watsonx_embeddings(monkeypatch): print(f"response: {response}") assert isinstance(response.usage, litellm.Usage) - + # Verify the request was made correctly assert "Authorization" in captured_request["headers"] - assert captured_request["headers"]["Authorization"] == "Bearer mock-watsonx-token" + assert ( + captured_request["headers"]["Authorization"] == "Bearer mock-watsonx-token" + ) assert "us-south.ml.cloud.ibm.com" in captured_request["url"] except litellm.RateLimitError as e: pass @@ -1256,9 +1262,7 @@ def test_jina_ai_img_embeddings(input_data, expected_payload_input): # Call the function we want to test try: - litellm.embedding( - model="jina_ai/jina-embeddings-v4", input=input_data - ) + litellm.embedding(model="jina_ai/jina-embeddings-v4", input=input_data) except Exception as e: pytest.fail( f"litellm.embedding call failed with an unexpected exception: {e}" @@ -1285,105 +1289,113 @@ def test_jina_ai_img_embeddings(input_data, expected_payload_input): def test_encoding_format_none_not_omitted_from_openai_sdk(): """ Test that encoding_format=None is explicitly sent to OpenAI SDK. - + This test verifies that when encoding_format is not provided by the user, liteLLM explicitly sets it to None rather than omitting it. This prevents the OpenAI SDK from adding its default value of 'base64'. - + Without this fix: - OpenAI SDK adds encoding_format='base64' as default when parameter is missing - This causes issues with providers that don't support encoding_format (like Gemini) - + With this fix: - encoding_format=None is explicitly passed - OpenAI SDK respects the explicit None and doesn't add defaults """ - with patch("litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client") as mock_get_client: + with patch( + "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" + ) as mock_get_client: # Create a mock client instance mock_client_instance = MagicMock() mock_get_client.return_value = mock_client_instance - + # Mock the embeddings.with_raw_response.create method mock_response = MagicMock() mock_response.parse.return_value = MagicMock( model_dump=lambda: { - 'data': [{'embedding': [0.1, 0.2, 0.3], 'index': 0}], - 'model': 'text-embedding-ada-002', - 'object': 'list', - 'usage': {'prompt_tokens': 1, 'total_tokens': 1} + "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], + "model": "text-embedding-ada-002", + "object": "list", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, } ) mock_response.headers = {} - - mock_client_instance.embeddings.with_raw_response.create.return_value = mock_response - + + mock_client_instance.embeddings.with_raw_response.create.return_value = ( + mock_response + ) + # Call the embedding function without encoding_format response = embedding( model="text-embedding-ada-002", input="Hello world", ) - + # Get the call arguments to verify what was sent to OpenAI SDK call_args = mock_client_instance.embeddings.with_raw_response.create.call_args - assert call_args is not None, "OpenAI SDK embeddings.create should have been called" - + assert ( + call_args is not None + ), "OpenAI SDK embeddings.create should have been called" + call_kwargs = call_args[1] # Get kwargs - + # The key assertion: encoding_format should be in the request with value None # This prevents OpenAI SDK from adding its default 'base64' value - assert 'encoding_format' in call_kwargs, ( + assert "encoding_format" in call_kwargs, ( "encoding_format should be explicitly passed to OpenAI SDK " "(even if None) to prevent SDK from adding default value" ) - assert call_kwargs['encoding_format'] is None, ( - "encoding_format should be None when not provided by user" - ) - + assert ( + call_kwargs["encoding_format"] is None + ), "encoding_format should be None when not provided by user" + print("✅ PASS: encoding_format=None is correctly passed to OpenAI SDK") def test_encoding_format_explicit_value_preserved(): """ Test that explicitly provided encoding_format values are preserved. - - When user provides encoding_format='float' or 'base64', it should be + + When user provides encoding_format='float' or 'base64', it should be sent as-is to the OpenAI SDK. """ - with patch("litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client") as mock_get_client: + with patch( + "litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client" + ) as mock_get_client: # Create a mock client instance mock_client_instance = MagicMock() mock_get_client.return_value = mock_client_instance - + # Mock the embeddings.with_raw_response.create method mock_response = MagicMock() mock_response.parse.return_value = MagicMock( model_dump=lambda: { - 'data': [{'embedding': [0.1, 0.2, 0.3], 'index': 0}], - 'model': 'text-embedding-ada-002', - 'object': 'list', - 'usage': {'prompt_tokens': 1, 'total_tokens': 1} + "data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}], + "model": "text-embedding-ada-002", + "object": "list", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, } ) mock_response.headers = {} - - mock_client_instance.embeddings.with_raw_response.create.return_value = mock_response - + + mock_client_instance.embeddings.with_raw_response.create.return_value = ( + mock_response + ) + # Test with explicit encoding_format='float' response = embedding( - model="text-embedding-ada-002", - input="Hello world", - encoding_format="float" + model="text-embedding-ada-002", input="Hello world", encoding_format="float" ) - + # Verify the encoding_format was passed correctly call_args = mock_client_instance.embeddings.with_raw_response.create.call_args call_kwargs = call_args[1] - - assert 'encoding_format' in call_kwargs, ( - "encoding_format should be in the request" - ) - assert call_kwargs['encoding_format'] == 'float', ( - "encoding_format should be 'float' when explicitly provided" - ) - + + assert ( + "encoding_format" in call_kwargs + ), "encoding_format should be in the request" + assert ( + call_kwargs["encoding_format"] == "float" + ), "encoding_format should be 'float' when explicitly provided" + print("✅ PASS: encoding_format='float' is correctly preserved") From bee16072484ebe321f07efe0145adabb485a4eda Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 18:57:27 -0700 Subject: [PATCH 04/16] test: update test apis --- tests/local_testing/test_embedding.py | 4 ++-- tests/vector_store_tests/test_azure_ai_vector_store.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/local_testing/test_embedding.py b/tests/local_testing/test_embedding.py index d0383f7002..9a63c6c83c 100644 --- a/tests/local_testing/test_embedding.py +++ b/tests/local_testing/test_embedding.py @@ -198,8 +198,8 @@ def _azure_ai_image_mock_response(*args, **kwargs): [ ( "azure_ai/Cohere-embed-v3-multilingual-2", - os.getenv("AZURE_AI_COHERE_API_BASE"), - os.getenv("AZURE_AI_COHERE_API_KEY"), + os.getenv("AZURE_AI_API_BASE"), + os.getenv("AZURE_AI_API_KEY"), ) ], ) diff --git a/tests/vector_store_tests/test_azure_ai_vector_store.py b/tests/vector_store_tests/test_azure_ai_vector_store.py index 52eb6635a9..58e45f259a 100644 --- a/tests/vector_store_tests/test_azure_ai_vector_store.py +++ b/tests/vector_store_tests/test_azure_ai_vector_store.py @@ -17,10 +17,10 @@ async def test_basic_search_vector_store(sync_mode): "vector_store_id": "my-vector-index", "custom_llm_provider": "azure_ai", "azure_search_service_name": "azure-kb-search", - "litellm_embedding_model": "azure/text-embedding-3-large", + "litellm_embedding_model": "azure_ai/text-embedding-3-large", "litellm_embedding_config": { - "api_base": os.getenv("AZURE_AI_SEARCH_EMBEDDING_API_BASE"), - "api_key": os.getenv("AZURE_AI_SEARCH_EMBEDDING_API_KEY"), + "api_base": os.getenv("AZURE_AI_API_BASE"), + "api_key": os.getenv("AZURE_AI_API_KEY"), }, "api_key": os.getenv("AZURE_SEARCH_API_KEY"), } From a41ba7bb6a454b01ac5dc6f846e5ca0652b208d0 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 19:01:55 -0700 Subject: [PATCH 05/16] test: update test --- .../test_azure_anthropic_structured_output.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/pass_through_unit_tests/messages_api_structured_output/test_azure_anthropic_structured_output.py b/tests/pass_through_unit_tests/messages_api_structured_output/test_azure_anthropic_structured_output.py index da46016b35..b2470bf6b6 100644 --- a/tests/pass_through_unit_tests/messages_api_structured_output/test_azure_anthropic_structured_output.py +++ b/tests/pass_through_unit_tests/messages_api_structured_output/test_azure_anthropic_structured_output.py @@ -30,7 +30,7 @@ class TestAzureAnthropicStructuredOutput(BaseAnthropicMessagesStructuredOutputTe return "azure_ai/claude-opus-4-5" def get_api_base(self) -> Optional[str]: - return "https://krish-mh44t553-eastus2.services.ai.azure.com/" + return "https://krris-mnb3t0vd-swedencentral.services.ai.azure.com" def get_api_key(self) -> Optional[str]: - return os.environ.get("AZURE_ANTHROPIC_API_KEY") \ No newline at end of file + return os.environ.get("AZURE_ANTHROPIC_API_KEY") From bc829d51f2650609a0f4079e2412e4cd5ceaa042 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 19:17:38 -0700 Subject: [PATCH 06/16] test: test --- .../test_prometheus_unit_tests.py | 6 +- .../proxy/auth/test_user_api_key_auth.py | 2 +- tests/image_gen_tests/test_image_edits.py | 253 +++++++------ .../image_gen_tests/test_image_generation.py | 41 ++- .../litellm_utils_tests/test_health_check.py | 27 +- .../test_azure_responses_api.py | 65 ++-- tests/llm_translation/test_azure_ai.py | 107 ++++-- tests/llm_translation/test_azure_o_series.py | 40 ++- tests/llm_translation/test_azure_openai.py | 76 ++-- .../test_router_llm_translation_tests.py | 16 +- .../example_config_yaml/azure_config.yaml | 4 +- tests/local_testing/test_acooldowns_router.py | 2 +- tests/local_testing/test_assistants.py | 7 +- tests/local_testing/test_azure_openai.py | 10 +- tests/local_testing/test_azure_perf.py | 8 +- tests/local_testing/test_caching.py | 153 ++++++-- tests/local_testing/test_caching_ssl.py | 4 +- tests/local_testing/test_class.py | 8 +- tests/local_testing/test_completion.py | 128 ++----- tests/local_testing/test_config.py | 21 +- .../test_configs/test_bad_config.yaml | 8 +- ...st_cloudflare_azure_with_cache_config.yaml | 2 +- .../test_configs/test_config_no_auth.yaml | 14 +- .../test_configs/test_custom_logger.yaml | 8 +- tests/local_testing/test_embedding.py | 12 +- tests/local_testing/test_exceptions.py | 102 +++--- tests/local_testing/test_loadtest_router.py | 4 +- .../test_prompt_injection_detection.py | 4 +- .../test_router_budget_limiter.py | 5 +- tests/local_testing/test_router_caching.py | 4 +- .../local_testing/test_router_client_init.py | 2 +- .../test_router_cooldown_handlers.py | 4 +- tests/local_testing/test_router_debug_logs.py | 4 +- tests/local_testing/test_router_fallbacks.py | 78 ++-- tests/local_testing/test_router_init.py | 28 +- tests/local_testing/test_router_timeout.py | 4 +- tests/local_testing/test_router_utils.py | 29 +- tests/local_testing/test_timeout.py | 8 +- .../local_testing/test_tpm_rpm_routing_v2.py | 16 +- tests/logging_callback_tests/test_alerting.py | 25 +- .../test_custom_callback_router.py | 48 +-- tests/old_proxy_tests/tests/load_test_q.py | 4 +- .../test_unit_test_passthrough_router.py | 2 +- .../example_config_yaml/azure_config.yaml | 4 +- .../test_configs/test_bad_config.yaml | 8 +- ...st_cloudflare_azure_with_cache_config.yaml | 2 +- .../test_configs/test_config_no_auth.yaml | 14 +- .../test_proxy_pass_user_config.py | 9 +- tests/proxy_unit_tests/test_proxy_server.py | 185 +++++----- .../llms/azure/test_azure_common_utils.py | 4 +- ...e_anthropic_count_tokens_transformation.py | 3 +- .../test_openai_responses_transformation.py | 26 +- .../proxy/auth/test_user_api_key_auth.py | 12 +- .../test_files_endpoint.py | 339 ++++++++++++------ tests/test_litellm/test_router.py | 84 +++-- tests/test_litellm/test_setup_wizard.py | 16 +- 56 files changed, 1192 insertions(+), 907 deletions(-) diff --git a/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py b/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py index ab37a0a84c..76a5778347 100644 --- a/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py +++ b/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py @@ -169,9 +169,9 @@ async def test_prometheus_metric_tracking(): "model_name": "gpt-3.5-turbo", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_key": os.getenv("AZURE_AI_API_KEY"), + "api_version": os.getenv("AZURE_AI_API_VERSION"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "model_info": {"id": "azure-model-id"}, }, diff --git a/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py b/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py index a45df5df00..0a2930aebb 100644 --- a/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py +++ b/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py @@ -72,7 +72,7 @@ async def test_enterprise_custom_auth_returns_string(): auth_obj = await _user_api_key_auth_builder( request=request, api_key="my-custom-key", - azure_api_key_header="", + AZURE_AI_API_KEY_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, azure_apim_header=None, diff --git a/tests/image_gen_tests/test_image_edits.py b/tests/image_gen_tests/test_image_edits.py index 393b4cb67a..8fe2e25a99 100644 --- a/tests/image_gen_tests/test_image_edits.py +++ b/tests/image_gen_tests/test_image_edits.py @@ -23,10 +23,11 @@ from litellm.types.utils import StandardLoggingPayload # Configure pytest marks to avoid warnings pytestmark = pytest.mark.asyncio + class TestCustomLogger(CustomLogger): def __init__(self): self.standard_logging_payload: Optional[StandardLoggingPayload] = None - + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): self.standard_logging_payload = kwargs.get("standard_logging_object", None) pass @@ -80,12 +81,12 @@ class BaseLLMImageEditTest(ABC): result = self.image_edit_function(**call_args) else: result = await self.async_image_edit_function(**call_args) - + print("result from image edit", result) # Validate the response meets expected schema ImageResponse.model_validate(result) - + if isinstance(result, ImageResponse) and result.data: image_base64 = result.data[0].b64_json if image_base64: @@ -97,6 +98,7 @@ class BaseLLMImageEditTest(ABC): except litellm.ContentPolicyViolationError as e: pass + # Get the current directory of the file being run pwd = os.path.dirname(os.path.realpath(__file__)) @@ -107,6 +109,7 @@ TEST_IMAGES = [ SINGLE_TEST_IMAGE = open(os.path.join(pwd, "ishaan_github.png"), "rb") + def get_test_images_as_bytesio(): """Helper function to get test images as BytesIO objects""" bytesio_images = [] @@ -129,6 +132,7 @@ class TestOpenAIImageEditGPTImage1(BaseLLMImageEditTest): "image": TEST_IMAGES, } + class TestOpenAIImageEditDallE2(BaseLLMImageEditTest): """ Concrete implementation of BaseLLMImageEditTest for OpenAI DALL-E-2 image edits. @@ -155,7 +159,7 @@ class TestAzureAIFlux2ImageEdit(BaseLLMImageEditTest): "model": "azure_ai/flux.2-pro", "image": SINGLE_TEST_IMAGE, "api_base": "https://litellm-ci-cd-prod.services.ai.azure.com", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": "preview", } @@ -187,7 +191,7 @@ async def test_openai_image_edit_litellm_router(): # Validate the response meets expected schema ImageResponse.model_validate(result) - + if isinstance(result, ImageResponse) and result.data: image_base64 = result.data[0].b64_json if image_base64: @@ -199,17 +203,19 @@ async def test_openai_image_edit_litellm_router(): except litellm.ContentPolicyViolationError as e: pass + @pytest.mark.flaky(retries=3, delay=2) @pytest.mark.asyncio async def test_openai_image_edit_with_bytesio(): """Test image editing using BytesIO objects instead of file readers""" from litellm import image_edit, aimage_edit + litellm._turn_on_debug() try: prompt = """ Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO. """ - + # Get images as BytesIO objects bytesio_images = get_test_images_as_bytesio() @@ -222,7 +228,7 @@ async def test_openai_image_edit_with_bytesio(): # Validate the response meets expected schema ImageResponse.model_validate(result) - + if isinstance(result, ImageResponse) and result.data: image_base64 = result.data[0].b64_json if image_base64: @@ -239,7 +245,7 @@ async def test_openai_image_edit_with_bytesio(): async def test_azure_image_edit_litellm_sdk(): """Test Azure image edit with mocked httpx request to validate request body and URL""" from litellm import image_edit, aimage_edit - + # Mock response for Azure image edit mock_response = { "created": 1589478378, @@ -247,7 +253,7 @@ async def test_azure_image_edit_litellm_sdk(): { "b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==" } - ] + ], } class MockResponse: @@ -267,16 +273,16 @@ async def test_azure_image_edit_litellm_sdk(): mock_post.return_value = MockResponse(mock_response, 200) litellm._turn_on_debug() - + prompt = """ Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO. """ - + # Set up test environment variables test_api_base = "https://ai-api-gw-uae-north.openai.azure.com" test_api_key = "test-api-key" test_api_version = "2025-04-01-preview" - + result = await aimage_edit( prompt=prompt, model="azure/gpt-image-1", @@ -285,41 +291,54 @@ async def test_azure_image_edit_litellm_sdk(): api_version=test_api_version, image=TEST_IMAGES, ) - + # Verify the request was made correctly mock_post.assert_called_once() - + # Check the URL call_args = mock_post.call_args expected_url = f"{test_api_base}/openai/deployments/gpt-image-1/images/edits?api-version={test_api_version}" - actual_url = call_args.args[0] if call_args.args else call_args.kwargs.get('url') + actual_url = ( + call_args.args[0] if call_args.args else call_args.kwargs.get("url") + ) print(f"Expected URL: {expected_url}") print(f"Actual URL: {actual_url}") - assert actual_url == expected_url, f"URL mismatch. Expected: {expected_url}, Got: {actual_url}" - + assert ( + actual_url == expected_url + ), f"URL mismatch. Expected: {expected_url}, Got: {actual_url}" + # Check the request body - if 'data' in call_args.kwargs: + if "data" in call_args.kwargs: # For multipart form data, check the data parameter - form_data = call_args.kwargs['data'] - print("Form data keys:", list(form_data.keys()) if hasattr(form_data, 'keys') else "Not a dict") - + form_data = call_args.kwargs["data"] + print( + "Form data keys:", + list(form_data.keys()) if hasattr(form_data, "keys") else "Not a dict", + ) + # Validate that model and prompt are in the form data - assert 'model' in form_data, "model should be in form data" - assert 'prompt' in form_data, "prompt should be in form data" - assert form_data['model'] == 'gpt-image-1', f"Expected model 'gpt-image-1', got {form_data['model']}" - assert prompt.strip() in form_data['prompt'], f"Expected prompt to contain '{prompt.strip()}'" - + assert "model" in form_data, "model should be in form data" + assert "prompt" in form_data, "prompt should be in form data" + assert ( + form_data["model"] == "gpt-image-1" + ), f"Expected model 'gpt-image-1', got {form_data['model']}" + assert ( + prompt.strip() in form_data["prompt"] + ), f"Expected prompt to contain '{prompt.strip()}'" + # Check headers - headers = call_args.kwargs.get('headers', {}) + headers = call_args.kwargs.get("headers", {}) print("Request headers:", headers) - assert 'Authorization' in headers, "Authorization header should be present" - assert headers['Authorization'].startswith('Bearer '), "Authorization should be Bearer token" - + assert "Authorization" in headers, "Authorization header should be present" + assert headers["Authorization"].startswith( + "Bearer " + ), "Authorization should be Bearer token" + print("result from image edit", result) # Validate the response meets expected schema ImageResponse.model_validate(result) - + if isinstance(result, ImageResponse) and result.data: image_base64 = result.data[0].b64_json if image_base64: @@ -330,15 +349,15 @@ async def test_azure_image_edit_litellm_sdk(): f.write(image_bytes) - @pytest.mark.asyncio async def test_openai_image_edit_cost_tracking(): """Test OpenAI image edit cost tracking with custom logger""" from litellm import image_edit, aimage_edit + test_custom_logger = TestCustomLogger() litellm.logging_callback_manager._reset_all_callbacks() litellm.callbacks = [test_custom_logger] - + # Mock response for Azure image edit with usage data for cost tracking mock_response = { "created": 1589478378, @@ -350,12 +369,9 @@ async def test_openai_image_edit_cost_tracking(): "usage": { "total_tokens": 1100, "input_tokens": 100, - "input_tokens_details": { - "image_tokens": 50, - "text_tokens": 50 - }, - "output_tokens": 1000 - } + "input_tokens_details": {"image_tokens": 50, "text_tokens": 50}, + "output_tokens": 1000, + }, } class MockResponse: @@ -375,26 +391,25 @@ async def test_openai_image_edit_cost_tracking(): mock_post.return_value = MockResponse(mock_response, 200) litellm._turn_on_debug() - + prompt = """ Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO. """ - + # Set up test environment variables - + result = await aimage_edit( prompt=prompt, model="openai/gpt-image-1", image=TEST_IMAGES, ) - + # Verify the request was made correctly mock_post.assert_called_once() - # Validate the response meets expected schema ImageResponse.model_validate(result) - + if isinstance(result, ImageResponse) and result.data: image_base64 = result.data[0].b64_json if image_base64: @@ -403,30 +418,36 @@ async def test_openai_image_edit_cost_tracking(): # Save the image to a file with open("test_image_edit.png", "wb") as f: f.write(image_bytes) - await asyncio.sleep(5) - print("standard logging payload", json.dumps(test_custom_logger.standard_logging_payload, indent=4, default=str)) + print( + "standard logging payload", + json.dumps( + test_custom_logger.standard_logging_payload, indent=4, default=str + ), + ) # check model assert test_custom_logger.standard_logging_payload["model"] == "gpt-image-1" - assert test_custom_logger.standard_logging_payload["custom_llm_provider"] == "openai" + assert ( + test_custom_logger.standard_logging_payload["custom_llm_provider"] + == "openai" + ) # check response_cost assert test_custom_logger.standard_logging_payload["response_cost"] is not None assert test_custom_logger.standard_logging_payload["response_cost"] > 0 - - @pytest.mark.asyncio async def test_azure_image_edit_cost_tracking(): """Test Azure image edit cost tracking with custom logger""" from litellm import image_edit, aimage_edit + test_custom_logger = TestCustomLogger() litellm.logging_callback_manager._reset_all_callbacks() litellm.callbacks = [test_custom_logger] - + # Mock response for Azure image edit with usage data for cost tracking mock_response = { "created": 1589478378, @@ -438,12 +459,9 @@ async def test_azure_image_edit_cost_tracking(): "usage": { "total_tokens": 1100, "input_tokens": 100, - "input_tokens_details": { - "image_tokens": 50, - "text_tokens": 50 - }, - "output_tokens": 1000 - } + "input_tokens_details": {"image_tokens": 50, "text_tokens": 50}, + "output_tokens": 1000, + }, } class MockResponse: @@ -463,27 +481,26 @@ async def test_azure_image_edit_cost_tracking(): mock_post.return_value = MockResponse(mock_response, 200) litellm._turn_on_debug() - + prompt = """ Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO. """ - + # Set up test environment variables - + result = await aimage_edit( prompt=prompt, model="azure/CUSTOM_AZURE_DEPLOYMENT_NAME", base_model="azure/gpt-image-1", image=TEST_IMAGES, ) - + # Verify the request was made correctly mock_post.assert_called_once() - # Validate the response meets expected schema ImageResponse.model_validate(result) - + if isinstance(result, ImageResponse) and result.data: image_base64 = result.data[0].b64_json if image_base64: @@ -492,14 +509,24 @@ async def test_azure_image_edit_cost_tracking(): # Save the image to a file with open("test_image_edit.png", "wb") as f: f.write(image_bytes) - await asyncio.sleep(5) - print("standard logging payload", json.dumps(test_custom_logger.standard_logging_payload, indent=4, default=str)) + print( + "standard logging payload", + json.dumps( + test_custom_logger.standard_logging_payload, indent=4, default=str + ), + ) # check model - assert test_custom_logger.standard_logging_payload["model"] == "CUSTOM_AZURE_DEPLOYMENT_NAME" - assert test_custom_logger.standard_logging_payload["custom_llm_provider"] == "azure" + assert ( + test_custom_logger.standard_logging_payload["model"] + == "CUSTOM_AZURE_DEPLOYMENT_NAME" + ) + assert ( + test_custom_logger.standard_logging_payload["custom_llm_provider"] + == "azure" + ) # check response_cost assert test_custom_logger.standard_logging_payload["response_cost"] is not None @@ -511,6 +538,7 @@ async def test_azure_image_edit_cost_tracking(): async def test_recraft_image_edit_api(): from litellm import aimage_edit import requests + litellm._turn_on_debug() global TEST_IMAGES try: @@ -526,10 +554,10 @@ async def test_recraft_image_edit_api(): # Validate the response meets expected schema ImageResponse.model_validate(result) - + if isinstance(result, ImageResponse) and result.data: image_url = result.data[0].url - + # download the image image_bytes = requests.get(image_url).content with open("test_image_edit.png", "wb") as f: @@ -545,51 +573,55 @@ def test_recraft_image_edit_config(): from litellm.llms.recraft.image_edit.transformation import RecraftImageEditConfig from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.router import GenericLiteLLMParams - + config = RecraftImageEditConfig() - + # Test supported OpenAI params supported_params = config.get_supported_openai_params("recraftv3") expected_params = ["n", "response_format", "style"] assert supported_params == expected_params - + # Test parameter mapping (reuses OpenAI logic with filtering) - image_edit_params = ImageEditOptionalRequestParams({ - "n": 2, - "response_format": "b64_json", - "style": "realistic_image", - "size": "1024x1024", # Should be dropped - "quality": "high" # Should be dropped - }) - - mapped_params = config.map_openai_params(image_edit_params, "recraftv3", drop_params=True) - + image_edit_params = ImageEditOptionalRequestParams( + { + "n": 2, + "response_format": "b64_json", + "style": "realistic_image", + "size": "1024x1024", # Should be dropped + "quality": "high", # Should be dropped + } + ) + + mapped_params = config.map_openai_params( + image_edit_params, "recraftv3", drop_params=True + ) + # Should only contain supported params assert mapped_params["n"] == 2 assert mapped_params["response_format"] == "b64_json" assert mapped_params["style"] == "realistic_image" assert "size" not in mapped_params # Should be dropped assert "quality" not in mapped_params # Should be dropped - + # Test request transformation (reuses OpenAI file handling) mock_image = b"fake_image_data" prompt = "winter landscape" litellm_params = GenericLiteLLMParams(api_key="test_key") - + data, files = config.transform_image_edit_request( model="recraftv3", prompt=prompt, image=mock_image, image_edit_optional_request_params={"strength": 0.7, "n": 1}, litellm_params=litellm_params, - headers={} + headers={}, ) - + # Check data structure (like OpenAI but with Recraft additions) assert data["prompt"] == prompt assert data["strength"] == 0.7 # Recraft-specific parameter assert data["model"] == "recraftv3" - + # Check file structure (reuses OpenAI logic) assert len(files) == 1 assert files[0][0] == "image" # Field name (not image[] like OpenAI) @@ -603,11 +635,12 @@ def test_recraft_image_edit_config(): async def test_multiple_vs_single_image_edit(sync_mode): """Test that both single and multiple image editing work correctly""" from litellm import image_edit, aimage_edit + litellm._turn_on_debug() - + try: prompt = "Add a soft blue tint to the image(s)" - + # Test single image if sync_mode: single_result = image_edit( @@ -621,10 +654,10 @@ async def test_multiple_vs_single_image_edit(sync_mode): model="gpt-image-1", image=SINGLE_TEST_IMAGE, ) - + print("Single image result:", single_result) ImageResponse.model_validate(single_result) - + # Test multiple images if sync_mode: multiple_result = image_edit( @@ -638,10 +671,10 @@ async def test_multiple_vs_single_image_edit(sync_mode): model="gpt-image-1", image=TEST_IMAGES, ) - + print("Multiple images result:", multiple_result) ImageResponse.model_validate(multiple_result) - + # Both should return valid responses assert single_result is not None assert multiple_result is not None @@ -649,7 +682,7 @@ async def test_multiple_vs_single_image_edit(sync_mode): assert multiple_result.data is not None assert len(single_result.data) > 0 assert len(multiple_result.data) > 0 - + except litellm.ContentPolicyViolationError as e: pytest.skip(f"Content policy violation: {e}") @@ -659,36 +692,37 @@ async def test_multiple_vs_single_image_edit(sync_mode): async def test_multiple_image_edit_with_different_formats(): """Test multiple images editing with different file formats and types""" from litellm import aimage_edit + litellm._turn_on_debug() - + try: prompt = "Create a cohesive artistic style across all images" - + # Test with mixed BytesIO and file objects mixed_images = [ SINGLE_TEST_IMAGE, # File object - get_test_images_as_bytesio()[1] # BytesIO object + get_test_images_as_bytesio()[1], # BytesIO object ] - + result = await aimage_edit( prompt=prompt, model="gpt-image-1", image=mixed_images, ) - + print("Mixed format images result:", result) ImageResponse.model_validate(result) - + assert result is not None assert result.data is not None assert len(result.data) > 0 - + # Save result if available if result.data and result.data[0].b64_json: image_bytes = base64.b64decode(result.data[0].b64_json) with open("test_multiple_image_edit_mixed.png", "wb") as f: f.write(image_bytes) - + except litellm.ContentPolicyViolationError as e: pytest.skip(f"Content policy violation: {e}") @@ -698,7 +732,7 @@ async def test_multiple_image_edit_with_different_formats(): async def test_image_edit_array_handling(): """Test that the image parameter correctly handles both single items and arrays""" from litellm import aimage_edit - + # Mock response mock_response = { "created": 1589478378, @@ -706,7 +740,7 @@ async def test_image_edit_array_handling(): { "b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==" } - ] + ], } class MockResponse: @@ -723,29 +757,26 @@ async def test_image_edit_array_handling(): new_callable=AsyncMock, ) as mock_post: mock_post.return_value = MockResponse(mock_response, 200) - + prompt = "Test prompt" - + # Test 1: Single image (should be converted to list internally) result1 = await aimage_edit( prompt=prompt, model="gpt-image-1", image=SINGLE_TEST_IMAGE, ) - + # Test 2: Multiple images (already a list) result2 = await aimage_edit( prompt=prompt, model="gpt-image-1", image=TEST_IMAGES, ) - # Both valid calls should succeed ImageResponse.model_validate(result1) ImageResponse.model_validate(result2) - + # Verify that both calls were made to the API assert mock_post.call_count == 2 - - diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index 3b4abeeb82..aca5c41de3 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -121,6 +121,7 @@ class TestVertexImageGeneration(BaseImageGenTest): class TestVertexAIGeminiImageGeneration(BaseImageGenTest): """Test Gemini image generation models (Nano Banana)""" + def get_base_image_generation_call_args(self) -> dict: # comment this when running locally load_vertex_ai_credentials() @@ -212,7 +213,9 @@ class TestAimlImageGeneration(BaseImageGenTest): custom_logger = TestCustomLogger() litellm.logging_callback_manager._reset_all_callbacks() litellm.callbacks = [custom_logger] - base_image_generation_call_args = self.get_base_image_generation_call_args() + base_image_generation_call_args = ( + self.get_base_image_generation_call_args() + ) litellm.set_verbose = True # Pass dummy api_key so validate_environment passes; HTTP is mocked response = await litellm.aimage_generation( @@ -229,7 +232,9 @@ class TestAimlImageGeneration(BaseImageGenTest): # print("response_cost", response._hidden_params["response_cost"]) logged_standard_logging_payload = custom_logger.standard_logging_payload - print("logged_standard_logging_payload", logged_standard_logging_payload) + print( + "logged_standard_logging_payload", logged_standard_logging_payload + ) assert logged_standard_logging_payload is not None assert logged_standard_logging_payload["response_cost"] is not None assert logged_standard_logging_payload["response_cost"] > 0 @@ -244,7 +249,9 @@ class TestAimlImageGeneration(BaseImageGenTest): response_dict["usage"] = dict(response_dict["usage"]) print("response usage=", response_dict.get("usage")) - assert response.data is not None # type guard for iteration (base fails here if None) + assert ( + response.data is not None + ) # type guard for iteration (base fails here if None) for d in response.data: assert isinstance(d, Image) print("data in response.data", d) @@ -266,25 +273,27 @@ class TestGoogleImageGen(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: return {"model": "gemini/imagen-4.0-generate-001"} + @pytest.mark.skip(reason="Runwayml image generation API only tested locally") class TestRunwaymlImageGeneration(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: return {"model": "runwayml/gen4_image"} -class TestAzureOpenAIDalle3(BaseImageGenTest): - def get_base_image_generation_call_args(self) -> dict: - return { - "model": "azure/dall-e-3", - "api_version": "2024-02-01", - "api_base": os.getenv("AZURE_API_BASE"), - "api_key": os.getenv("AZURE_API_KEY"), - "metadata": { - "model_info": { - "base_model": "azure/dall-e-3", - } - }, - } +## AZURE AI DALL-E 3 is deprecated and new deployments cannot be made +# class TestAzureOpenAIDalle3(BaseImageGenTest): +# def get_base_image_generation_call_args(self) -> dict: +# return { +# "model": "azure/dall-e-3", +# "api_version": "2024-02-01", +# "api_base": os.getenv("AZURE_AI_API_BASE"), +# "api_key": os.getenv("AZURE_AI_API_KEY"), +# "metadata": { +# "model_info": { +# "base_model": "azure/dall-e-3", +# } +# }, +# } @pytest.mark.skip(reason="model EOL") diff --git a/tests/litellm_utils_tests/test_health_check.py b/tests/litellm_utils_tests/test_health_check.py index b048590d51..f87f2005d3 100644 --- a/tests/litellm_utils_tests/test_health_check.py +++ b/tests/litellm_utils_tests/test_health_check.py @@ -21,9 +21,9 @@ async def test_azure_health_check(): model_params={ "model": "azure/gpt-4.1-mini", "messages": [{"role": "user", "content": "Hey, how's it going?"}], - "api_key": os.getenv("AZURE_API_KEY"), - "api_base": os.getenv("AZURE_API_BASE"), - "api_version": os.getenv("AZURE_API_VERSION"), + "api_key": os.getenv("AZURE_AI_API_KEY"), + "api_base": os.getenv("AZURE_AI_API_BASE"), + "api_version": os.getenv("AZURE_AI_API_VERSION"), } ) print(f"response: {response}") @@ -51,9 +51,9 @@ async def test_azure_embedding_health_check(): response = await litellm.ahealth_check( model_params={ "model": "azure/text-embedding-ada-002", - "api_key": os.getenv("AZURE_API_KEY"), - "api_base": os.getenv("AZURE_API_BASE"), - "api_version": os.getenv("AZURE_API_VERSION"), + "api_key": os.getenv("AZURE_AI_API_KEY"), + "api_base": os.getenv("AZURE_AI_API_BASE"), + "api_version": os.getenv("AZURE_AI_API_VERSION"), }, input=["test for litellm"], mode="embedding", @@ -83,7 +83,9 @@ async def test_openai_img_gen_health_check(): # asyncio.run(test_openai_img_gen_health_check()) -@pytest.mark.skip(reason="Azure DALL-E 3 model deployment is deprecated (410 ModelDeprecated)") +@pytest.mark.skip( + reason="Azure DALL-E 3 model deployment is deprecated (410 ModelDeprecated)" +) @pytest.mark.asyncio async def test_azure_img_gen_health_check(): """ @@ -98,8 +100,8 @@ async def test_azure_img_gen_health_check(): response = await litellm.ahealth_check( model_params={ "model": "azure/dall-e-3", - "api_base": os.getenv("AZURE_API_BASE"), - "api_key": os.getenv("AZURE_API_KEY"), + "api_base": os.getenv("AZURE_AI_API_BASE"), + "api_key": os.getenv("AZURE_AI_API_KEY"), }, mode="image_generation", prompt="cute baby sea otter", @@ -500,7 +502,9 @@ async def test_perform_health_check_filters_by_model_id(): async def mock_perform_health_check(m_list, details=True, **kwargs): captured_list.append(m_list) - return [{"model": "gpt-4", "api_key": m_list[0]["litellm_params"]["api_key"]}], [] + return [ + {"model": "gpt-4", "api_key": m_list[0]["litellm_params"]["api_key"]} + ], [] with patch( "litellm.proxy.health_check._perform_health_check", @@ -657,7 +661,8 @@ async def test_health_check_creates_only_bounded_initial_tasks(): return real_create_task(coro) with patch("litellm.ahealth_check", side_effect=mock_health_check), patch( - "litellm.proxy.health_check.asyncio.create_task", side_effect=tracked_create_task + "litellm.proxy.health_check.asyncio.create_task", + side_effect=tracked_create_task, ): perform_task = real_create_task( _perform_health_check(model_list, max_concurrency=2) diff --git a/tests/llm_responses_api_testing/test_azure_responses_api.py b/tests/llm_responses_api_testing/test_azure_responses_api.py index 5e876ce084..fed9e9e11f 100644 --- a/tests/llm_responses_api_testing/test_azure_responses_api.py +++ b/tests/llm_responses_api_testing/test_azure_responses_api.py @@ -25,11 +25,11 @@ class TestAzureResponsesAPITest(BaseResponsesAPITest): return { "model": "azure/gpt-4.1-mini", "truncation": "auto", - "api_base": os.getenv("AZURE_API_BASE"), - "api_key": os.getenv("AZURE_API_KEY"), + "api_base": os.getenv("AZURE_AI_API_BASE"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": "2025-03-01-preview", } - + def get_advanced_model_for_shell_tool(self) -> Optional[str]: """If specified, overrides the model used by test_responses_api_shell_tool_streaming_sees_shell_output (e.g. openai/gpt-5.2 for shell support).""" return "azure/gpt-5-mini" @@ -45,8 +45,8 @@ async def test_azure_responses_api_preview_api_version(): model="azure/gpt-5-mini", truncation="auto", api_version="preview", - api_base=os.getenv("AZURE_API_BASE"), - api_key=os.getenv("AZURE_API_KEY"), + api_base=os.getenv("AZURE_AI_API_BASE"), + api_key=os.getenv("AZURE_AI_API_KEY"), input="Hello, can you tell me a short joke?", ) @@ -108,7 +108,9 @@ async def test_azure_responses_api_status_error(): "role": "assistant", "type": "message", "status": "completed", - "content": [{"type": "output_text", "text": "Here's an interesting fact."}], + "content": [ + {"type": "output_text", "text": "Here's an interesting fact."} + ], } ], } @@ -124,7 +126,7 @@ async def test_azure_responses_api_status_error(): captured_request_body = json.loads(kwargs["data"]) import httpx - + # Create a proper httpx Response object response_content = json.dumps(mock_response_data).encode("utf-8") response = httpx.Response( @@ -149,18 +151,17 @@ async def test_azure_responses_api_status_error(): ) # Verify that 'status' field is not present in any of the input messages - print("Final request body:", json.dumps(captured_request_body, indent=4, default=str)) + print( + "Final request body:", json.dumps(captured_request_body, indent=4, default=str) + ) assert "input" in captured_request_body, "Request body should contain 'input' field" - + expected_input = [ - { - "content": "tell me an interesting fact", - "role": "user" - }, + {"content": "tell me an interesting fact", "role": "user"}, { "id": "rs_0ab687487834d9df0068e462a1b2d88197aabbc832c9ba5316", "summary": [], - "type": "reasoning" + "type": "reasoning", }, { "id": "msg_0ab687487834d9df0068e462a1df188197b74b1eef05102c18", @@ -169,18 +170,15 @@ async def test_azure_responses_api_status_error(): "annotations": [], "text": "very good morning", "type": "output_text", - "logprobs": [] + "logprobs": [], } ], "role": "assistant", - "type": "message" + "type": "message", }, - { - "role": "user", - "content": "tell me another" - } + {"role": "user", "content": "tell me another"}, ] - + assert captured_request_body["input"] == expected_input, ( f"Request body input should match expected format without 'status' field.\n" f"Expected: {json.dumps(expected_input, indent=2)}\n" @@ -193,9 +191,9 @@ async def test_azure_responses_api_headers_with_llm_provider_prefix(): """ Test that Azure-specific headers like 'x-request-id' and 'apim-request-id' are properly forwarded with 'llm_provider-' prefix in response._hidden_params["headers"]. - + Issue: https://github.com/BerriAI/litellm/issues/16538 - + The fix ensures that processed headers (with llm_provider- prefix) are stored in response._hidden_params["headers"] instead of additional_headers, making them accessible via completion.headers in the same way as the completion API. @@ -253,12 +251,12 @@ async def test_azure_responses_api_headers_with_llm_provider_prefix(): # Check that the response has the expected headers structure assert hasattr(response, "_hidden_params"), "Response should have _hidden_params" - assert "additional_headers" in response._hidden_params, ( - "Response _hidden_params should contain 'additional_headers' with the LLM provider headers" - ) + assert ( + "additional_headers" in response._hidden_params + ), "Response _hidden_params should contain 'additional_headers' with the LLM provider headers" headers = response._hidden_params["additional_headers"] - + # Verify that Azure-specific headers are present with llm_provider- prefix assert "llm_provider-x-request-id" in headers, ( f"Response should contain 'llm_provider-x-request-id' header. " @@ -268,12 +266,17 @@ async def test_azure_responses_api_headers_with_llm_provider_prefix(): f"Response should contain 'llm_provider-apim-request-id' header. " f"Headers: {list(headers.keys())}" ) - + # Verify the header values match - assert headers["llm_provider-x-request-id"] == "12086715-aca3-4006-a29f-2f1e1d552043" - assert headers["llm_provider-apim-request-id"] == "25664b0d-cf4b-4e10-8d27-c7272e7efd49" + assert ( + headers["llm_provider-x-request-id"] == "12086715-aca3-4006-a29f-2f1e1d552043" + ) + assert ( + headers["llm_provider-apim-request-id"] + == "25664b0d-cf4b-4e10-8d27-c7272e7efd49" + ) assert headers["llm_provider-x-ms-region"] == "Sweden Central" - + # Also verify openai-compatible headers are included assert "x-ratelimit-limit-tokens" in headers assert "x-ratelimit-remaining-tokens" in headers diff --git a/tests/llm_translation/test_azure_ai.py b/tests/llm_translation/test_azure_ai.py index 972ba34a17..823fad7648 100644 --- a/tests/llm_translation/test_azure_ai.py +++ b/tests/llm_translation/test_azure_ai.py @@ -283,8 +283,8 @@ async def test_azure_ai_request_format(): litellm._turn_on_debug() # Set up the test parameters - api_key = os.getenv("AZURE_API_KEY") - api_base = os.getenv("AZURE_API_BASE") + api_key = os.getenv("AZURE_AI_API_KEY") + api_base = os.getenv("AZURE_AI_API_BASE") model = "azure_ai/gpt-4.1-mini" messages = [ {"role": "user", "content": "hi"}, @@ -310,17 +310,17 @@ async def test_azure_gpt5_reasoning(model): messages=[{"role": "user", "content": "What is the capital of France?"}], reasoning_effort="minimal", max_tokens=10, - api_base=os.getenv("AZURE_API_BASE"), - api_key=os.getenv("AZURE_API_KEY"), + api_base=os.getenv("AZURE_AI_API_BASE"), + api_key=os.getenv("AZURE_AI_API_KEY"), ) print("response: ", response) assert response.choices[0].message.content is not None - def test_completion_azure(): try: from litellm import completion_cost + litellm.set_verbose = False ## Test azure call response = completion( @@ -331,7 +331,7 @@ def test_completion_azure(): "content": "Hello, how are you?", } ], - api_key="os.environ/AZURE_API_KEY", + api_key="os.environ/AZURE_AI_API_KEY", ) print(f"response: {response}") print(f"response hidden params: {response._hidden_params}") @@ -358,7 +358,7 @@ def test_completion_azure_ai_gpt_4o_with_flexible_api_base(api_base): response = completion( model="azure_ai/gpt-4.1-mini", api_base=api_base, - api_key=os.getenv("AZURE_API_KEY"), + api_key=os.getenv("AZURE_AI_API_KEY"), messages=[{"role": "user", "content": "What is the meaning of life?"}], ) @@ -374,13 +374,15 @@ async def test_azure_ai_model_router(): """ Test Azure AI model router non-streaming response cost tracking. Verifies that the flat cost of $0.14 per M input tokens is applied. - + Tests the pattern: azure_ai/model_router/ Where deployment-name is the Azure deployment (e.g., "azure-model-router"). The model_router prefix is stripped before sending to Azure API. """ - from litellm.llms.azure_ai.cost_calculator import calculate_azure_model_router_flat_cost - + from litellm.llms.azure_ai.cost_calculator import ( + calculate_azure_model_router_flat_cost, + ) + litellm._turn_on_debug() response = await litellm.acompletion( model="azure_ai/model_router/azure-model-router", @@ -394,23 +396,22 @@ async def test_azure_ai_model_router(): tracked_cost = response._hidden_params["response_cost"] assert tracked_cost > 0 print("Tracked cost: ", tracked_cost) - + # Verify flat cost is included using the helper function usage = response.usage if usage and usage.prompt_tokens: expected_flat_cost = calculate_azure_model_router_flat_cost( - model="model_router/azure-model-router", - prompt_tokens=usage.prompt_tokens + model="model_router/azure-model-router", prompt_tokens=usage.prompt_tokens ) print(f"Prompt tokens: {usage.prompt_tokens}") print(f"Expected flat cost: ${expected_flat_cost:.9f}") print(f"Total tracked cost: ${tracked_cost:.9f}") - + # Total cost should be at least the flat cost - assert tracked_cost >= expected_flat_cost, ( - f"Cost ${tracked_cost:.9f} should be >= flat cost ${expected_flat_cost:.9f}" - ) - + assert ( + tracked_cost >= expected_flat_cost + ), f"Cost ${tracked_cost:.9f} should be >= flat cost ${expected_flat_cost:.9f}" + # Verify the flat cost is non-zero assert expected_flat_cost > 0, "Flat cost should be greater than 0" @@ -445,15 +446,20 @@ async def test_azure_ai_model_router_streaming_model_in_chunk(): # The model should NOT be azure-model-router (the request model) # It should be the actual model from the response (e.g., gpt-4.1-nano, gpt-5-nano, etc.) for model in chunks_with_model: - assert model != "azure-model-router", f"Chunk model should be actual model, not request model. Got: {model}" + assert ( + model != "azure-model-router" + ), f"Chunk model should be actual model, not request model. Got: {model}" # The actual model should be a real model name like gpt-4.1-nano, gpt-5-nano, etc. print(f"Verified chunk has actual model: {model}") -class AzureModelRouterStreamingCallback(litellm.integrations.custom_logger.CustomLogger): +class AzureModelRouterStreamingCallback( + litellm.integrations.custom_logger.CustomLogger +): """ Custom callback to capture streaming cost tracking for Azure Model Router. """ + def __init__(self): self.standard_logging_payload = None self.response_cost = None @@ -466,17 +472,21 @@ class AzureModelRouterStreamingCallback(litellm.integrations.custom_logger.Custo self.async_success_called = True self.standard_logging_payload = kwargs.get("standard_logging_object") self.complete_streaming_response = kwargs.get("complete_streaming_response") - + if self.standard_logging_payload: self.response_cost = self.standard_logging_payload.get("response_cost") - print(f"standard_logging_payload model: {self.standard_logging_payload.get('model')}") + print( + f"standard_logging_payload model: {self.standard_logging_payload.get('model')}" + ) print(f"standard_logging_payload response_cost: {self.response_cost}") - + if self.complete_streaming_response: - print(f"complete_streaming_response model: {self.complete_streaming_response.model}") - print(f"complete_streaming_response usage: {self.complete_streaming_response.usage}") - - + print( + f"complete_streaming_response model: {self.complete_streaming_response.model}" + ) + print( + f"complete_streaming_response usage: {self.complete_streaming_response.usage}" + ) @pytest.mark.asyncio @@ -504,10 +514,16 @@ async def test_azure_ai_model_router_streaming_cost_with_stream_options(): full_response = "" chunks_with_model = [] async for chunk in response: - print(f"Chunk: model={chunk.model}, choices={len(chunk.choices) if chunk.choices else 0}") + print( + f"Chunk: model={chunk.model}, choices={len(chunk.choices) if chunk.choices else 0}" + ) if chunk.model: chunks_with_model.append(chunk.model) - if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content: + if ( + chunk.choices + and chunk.choices[0].delta + and chunk.choices[0].delta.content + ): full_response += chunk.choices[0].delta.content print(f"Full streamed response: {full_response}") @@ -515,27 +531,42 @@ async def test_azure_ai_model_router_streaming_cost_with_stream_options(): # Give async logging time to complete import asyncio + await asyncio.sleep(1) # Verify callback was called - assert test_callback.async_success_called is True, "async_log_success_event was not called" - assert test_callback.standard_logging_payload is not None, "standard_logging_payload is None" + assert ( + test_callback.async_success_called is True + ), "async_log_success_event was not called" + assert ( + test_callback.standard_logging_payload is not None + ), "standard_logging_payload is None" # Check response cost print(f"Final response_cost: {test_callback.response_cost}") - + # The first chunk may have the request model (azure-model-router) because it's created # before the API response is received. Subsequent chunks should have the actual model. # At least some chunks should have the actual model (not azure-model-router) - actual_model_chunks = [m for m in chunks_with_model if m != "azure-model-router"] - assert len(actual_model_chunks) > 0, "No chunks had the actual model from the API response" + actual_model_chunks = [ + m for m in chunks_with_model if m != "azure-model-router" + ] + assert ( + len(actual_model_chunks) > 0 + ), "No chunks had the actual model from the API response" print(f"Chunks with actual model: {actual_model_chunks}") # Verify response cost is tracked - this is the main goal of this test - assert test_callback.response_cost is not None, "response_cost is None with stream_options" - assert test_callback.response_cost > 0, f"response_cost should be > 0, got {test_callback.response_cost}" - print(f"Streaming cost tracking with stream_options passed. Cost: {test_callback.response_cost}") + assert ( + test_callback.response_cost is not None + ), "response_cost is None with stream_options" + assert ( + test_callback.response_cost > 0 + ), f"response_cost should be > 0, got {test_callback.response_cost}" + print( + f"Streaming cost tracking with stream_options passed. Cost: {test_callback.response_cost}" + ) finally: litellm.logging_callback_manager._reset_all_callbacks() - litellm.callbacks = [] \ No newline at end of file + litellm.callbacks = [] diff --git a/tests/llm_translation/test_azure_o_series.py b/tests/llm_translation/test_azure_o_series.py index 925453e68c..181d3b677d 100644 --- a/tests/llm_translation/test_azure_o_series.py +++ b/tests/llm_translation/test_azure_o_series.py @@ -24,9 +24,9 @@ class TestAzureOpenAIO3Mini(BaseOSeriesModelsTest, BaseLLMChatTest): litellm.in_memory_llm_clients_cache.flush_cache() return { "model": "azure/o3-mini", - "api_key": os.getenv("AZURE_API_KEY"), - "api_base": os.getenv("AZURE_API_BASE"), - "api_version": "2024-12-01-preview" + "api_key": os.getenv("AZURE_AI_API_KEY"), + "api_base": os.getenv("AZURE_AI_API_BASE"), + "api_version": "2024-12-01-preview", } def get_client(self): @@ -187,13 +187,31 @@ async def test_azure_o1_series_response_format_extra_params(): litellm.set_verbose = True client = AsyncAzureOpenAI( - api_key="fake-api-key", - base_url="https://openai-prod-test.openai.azure.com/openai/deployments/o1/chat/completions?api-version=2025-01-01-preview", - api_version="2025-01-01-preview" + api_key="fake-api-key", + base_url="https://openai-prod-test.openai.azure.com/openai/deployments/o1/chat/completions?api-version=2025-01-01-preview", + api_version="2025-01-01-preview", ) - tools = [{'type': 'function', 'function': {'name': 'get_current_time', 'description': 'Get the current time in a given location.', 'parameters': {'type': 'object', 'properties': {'location': {'type': 'string', 'description': 'The city name, e.g. San Francisco'}}, 'required': ['location']}}}] - response_format = {'type': 'json_object'} + tools = [ + { + "type": "function", + "function": { + "name": "get_current_time", + "description": "Get the current time in a given location.", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city name, e.g. San Francisco", + } + }, + "required": ["location"], + }, + }, + } + ] + response_format = {"type": "json_object"} tool_choice = "auto" with patch.object( client.chat.completions.with_raw_response, "create" @@ -208,7 +226,7 @@ async def test_azure_o1_series_response_format_extra_params(): messages=[{"role": "user", "content": "Hello! return a json object"}], tools=tools, response_format=response_format, - tool_choice=tool_choice + tool_choice=tool_choice, ) except Exception as e: print(f"Error: {e}") @@ -220,7 +238,3 @@ async def test_azure_o1_series_response_format_extra_params(): assert request_body["tools"] == tools assert request_body["response_format"] == response_format assert request_body["tool_choice"] == tool_choice - - - - diff --git a/tests/llm_translation/test_azure_openai.py b/tests/llm_translation/test_azure_openai.py index 1da380b57a..f3dc954020 100644 --- a/tests/llm_translation/test_azure_openai.py +++ b/tests/llm_translation/test_azure_openai.py @@ -208,8 +208,8 @@ class TestAzureEmbedding(BaseLLMEmbeddingTest): def get_base_embedding_call_args(self) -> dict: return { "model": "azure/text-embedding-ada-002", - "api_key": os.getenv("AZURE_API_KEY"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_key": os.getenv("AZURE_AI_API_KEY"), + "api_base": os.getenv("AZURE_AI_API_BASE"), } def get_custom_llm_provider(self) -> litellm.LlmProviders: @@ -618,8 +618,8 @@ def test_azure_safety_result(): response = completion( model="azure/gpt-4.1-mini", - api_key=os.getenv("AZURE_API_KEY"), - api_base=os.getenv("AZURE_API_BASE"), + api_key=os.getenv("AZURE_AI_API_KEY"), + api_base=os.getenv("AZURE_AI_API_BASE"), api_version="2024-12-01-preview", messages=[{"role": "user", "content": "Hello world"}], ) @@ -671,6 +671,8 @@ def test_completion_azure_deployment_id(): ) # Add any assertions here to check the response print(response) + + def test_azure_with_content_safety_error(): """ Verify user can access innererror from the Azure OpenAI exception @@ -679,55 +681,55 @@ def test_azure_with_content_safety_error(): from litellm.exceptions import ContentPolicyViolationError from litellm.litellm_core_utils.exception_mapping_utils import exception_type from unittest.mock import MagicMock - - mock_exception = Exception("The response was filtered due to the prompt triggering Azure OpenAI's content management policy") + + mock_exception = Exception( + "The response was filtered due to the prompt triggering Azure OpenAI's content management policy" + ) mock_exception.body = { "innererror": { "code": "ResponsibleAIPolicyViolation", "content_filter_result": { - "hate": { - "filtered": False, - "severity": "safe" - }, - "jailbreak": { - "filtered": False, - "detected": False - }, - "self_harm": { - "filtered": False, - "severity": "safe" - }, - "sexual": { - "filtered": False, - "severity": "safe" - }, - "violence": { - "filtered": True, - "severity": "high" - } - } + "hate": {"filtered": False, "severity": "safe"}, + "jailbreak": {"filtered": False, "detected": False}, + "self_harm": {"filtered": False, "severity": "safe"}, + "sexual": {"filtered": False, "severity": "safe"}, + "violence": {"filtered": True, "severity": "high"}, + }, } } - + mock_response = MagicMock() mock_response.status_code = 400 mock_exception.response = mock_response - + with pytest.raises(ContentPolicyViolationError) as exc_info: exception_type( model="azure/gpt-4o-new-test", original_exception=mock_exception, - custom_llm_provider="azure" + custom_llm_provider="azure", ) - + e = exc_info.value print("got exception=", e) assert e.provider_specific_fields is not None print("got provider_specific_fields=", e.provider_specific_fields) assert e.provider_specific_fields.get("innererror") is not None - assert e.provider_specific_fields["innererror"]["code"] == "ResponsibleAIPolicyViolation" - assert e.provider_specific_fields["innererror"]["content_filter_result"]["violence"]["filtered"] is True - assert e.provider_specific_fields["innererror"]["content_filter_result"]["violence"]["severity"] == "high" + assert ( + e.provider_specific_fields["innererror"]["code"] + == "ResponsibleAIPolicyViolation" + ) + assert ( + e.provider_specific_fields["innererror"]["content_filter_result"]["violence"][ + "filtered" + ] + is True + ) + assert ( + e.provider_specific_fields["innererror"]["content_filter_result"]["violence"][ + "severity" + ] + == "high" + ) def test_azure_openai_with_prompt_cache_key(): @@ -737,9 +739,9 @@ def test_azure_openai_with_prompt_cache_key(): litellm._turn_on_debug() response = litellm.completion( model="azure/gpt-4.1-mini", - api_key=os.getenv("AZURE_API_KEY"), - api_base=os.getenv("AZURE_API_BASE"), + api_key=os.getenv("AZURE_AI_API_KEY"), + api_base=os.getenv("AZURE_AI_API_BASE"), api_version="2024-12-01-preview", messages=[{"role": "user", "content": "What is the weather in San Francisco?"}], prompt_cache_key="test_streaming_azure_openai", - ) \ No newline at end of file + ) diff --git a/tests/llm_translation/test_router_llm_translation_tests.py b/tests/llm_translation/test_router_llm_translation_tests.py index 61446ce613..26456ab0a3 100644 --- a/tests/llm_translation/test_router_llm_translation_tests.py +++ b/tests/llm_translation/test_router_llm_translation_tests.py @@ -66,8 +66,8 @@ def test_router_azure_acompletion(): print("Router Test Azure - Acompletion, Acompletion with stream") # remove api key from env to repro how proxy passes key to router - old_api_key = os.environ["AZURE_API_KEY"] - os.environ.pop("AZURE_API_KEY", None) + old_api_key = os.environ["AZURE_AI_API_KEY"] + os.environ.pop("AZURE_AI_API_KEY", None) model_list = [ { @@ -75,8 +75,8 @@ def test_router_azure_acompletion(): "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", "api_key": old_api_key, - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_version": os.getenv("AZURE_AI_API_VERSION"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "rpm": 1800, }, @@ -85,8 +85,8 @@ def test_router_azure_acompletion(): "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", "api_key": old_api_key, - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_version": os.getenv("AZURE_AI_API_VERSION"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "rpm": 1800, }, @@ -126,9 +126,9 @@ def test_router_azure_acompletion(): asyncio.run(test2()) print("\n Passed Streaming") - os.environ["AZURE_API_KEY"] = old_api_key + os.environ["AZURE_AI_API_KEY"] = old_api_key router.reset() except Exception as e: - os.environ["AZURE_API_KEY"] = old_api_key + os.environ["AZURE_AI_API_KEY"] = old_api_key print(f"FAILED TEST") pytest.fail(f"Got unexpected exception on router! - {e}") diff --git a/tests/local_testing/example_config_yaml/azure_config.yaml b/tests/local_testing/example_config_yaml/azure_config.yaml index 0a015aefde..05ba0c9bf5 100644 --- a/tests/local_testing/example_config_yaml/azure_config.yaml +++ b/tests/local_testing/example_config_yaml/azure_config.yaml @@ -4,12 +4,12 @@ model_list: model: azure/gpt-4.1-mini api_base: https://openai-gpt-4-test-v-1.openai.azure.com/ api_version: "2023-05-15" - api_key: os.environ/AZURE_API_KEY + api_key: os.environ/AZURE_AI_API_KEY tpm: 20_000 - model_name: gpt-4-team2 litellm_params: model: azure/gpt-4 - api_key: os.environ/AZURE_API_KEY + api_key: os.environ/AZURE_AI_API_KEY api_base: https://openai-gpt-4-test-v-2.openai.azure.com/ tpm: 100_000 diff --git a/tests/local_testing/test_acooldowns_router.py b/tests/local_testing/test_acooldowns_router.py index ff99210298..18dc26bda9 100644 --- a/tests/local_testing/test_acooldowns_router.py +++ b/tests/local_testing/test_acooldowns_router.py @@ -31,7 +31,7 @@ def _make_model_list(): "model": "azure/gpt-4.1-mini", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, diff --git a/tests/local_testing/test_assistants.py b/tests/local_testing/test_assistants.py index df0df6b79c..93ebba3c34 100644 --- a/tests/local_testing/test_assistants.py +++ b/tests/local_testing/test_assistants.py @@ -37,10 +37,11 @@ V0 Scope: - Run Thread -> `/v1/threads/{thread_id}/run` """ + def _add_azure_related_dynamic_params(data: dict) -> dict: data["api_version"] = "2024-02-15-preview" - data["api_base"] = os.getenv("AZURE_API_BASE") - data["api_key"] = os.getenv("AZURE_API_KEY") + data["api_base"] = os.getenv("AZURE_AI_API_BASE") + data["api_key"] = os.getenv("AZURE_AI_API_KEY") return data @@ -236,8 +237,6 @@ async def test_aarun_thread_litellm(sync_mode, provider, is_streaming): """ import openai - - try: get_assistants_data = { "custom_llm_provider": provider, diff --git a/tests/local_testing/test_azure_openai.py b/tests/local_testing/test_azure_openai.py index e95c1b6fcc..1b99140b6e 100644 --- a/tests/local_testing/test_azure_openai.py +++ b/tests/local_testing/test_azure_openai.py @@ -40,11 +40,13 @@ async def test_aaaaazure_tenant_id_auth(respx_mock: MockRouter): PROD Test """ - litellm.disable_aiohttp_transport = True # since this uses respx, we need to set use_aiohttp_transport to False - + litellm.disable_aiohttp_transport = ( + True # since this uses respx, we need to set use_aiohttp_transport to False + ) + # Clear the HTTP client cache to ensure respx mocking works # This is critical because respx only intercepts clients created AFTER mocking is active - if hasattr(litellm, 'in_memory_llm_clients_cache'): + if hasattr(litellm, "in_memory_llm_clients_cache"): litellm.in_memory_llm_clients_cache.flush_cache() router = Router( @@ -53,7 +55,7 @@ async def test_aaaaazure_tenant_id_auth(respx_mock: MockRouter): "model_name": "gpt-3.5-turbo", "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), "tenant_id": os.getenv("AZURE_TENANT_ID"), "client_id": os.getenv("AZURE_CLIENT_ID"), "client_secret": os.getenv("AZURE_CLIENT_SECRET"), diff --git a/tests/local_testing/test_azure_perf.py b/tests/local_testing/test_azure_perf.py index 1e2d5cc4f7..57d56a24a1 100644 --- a/tests/local_testing/test_azure_perf.py +++ b/tests/local_testing/test_azure_perf.py @@ -9,8 +9,8 @@ # from openai import AsyncAzureOpenAI # client = AsyncAzureOpenAI( -# api_key=os.getenv("AZURE_API_KEY"), -# azure_endpoint=os.getenv("AZURE_API_BASE"), # type: ignore +# api_key=os.getenv("AZURE_AI_API_KEY"), +# azure_endpoint=os.getenv("AZURE_AI_API_BASE"), # type: ignore # api_version=os.getenv("AZURE_API_VERSION"), # ) @@ -19,8 +19,8 @@ # "model_name": "azure-test", # "litellm_params": { # "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_API_KEY"), -# "api_base": os.getenv("AZURE_API_BASE"), +# "api_key": os.getenv("AZURE_AI_API_KEY"), +# "api_base": os.getenv("AZURE_AI_API_BASE"), # "api_version": os.getenv("AZURE_API_VERSION"), # }, # } diff --git a/tests/local_testing/test_caching.py b/tests/local_testing/test_caching.py index 01004e4bfa..0b06b7195b 100644 --- a/tests/local_testing/test_caching.py +++ b/tests/local_testing/test_caching.py @@ -147,7 +147,12 @@ def test_caching_dynamic_args(): # test in memory cache port=_redis_port_env, password=_redis_password_env, ) - response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test") + response1 = completion( + model="gpt-3.5-turbo", + messages=messages, + caching=True, + mock_response="Hello world from cache test", + ) response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True) print(f"response1: {response1}") print(f"response2: {response2}") @@ -173,7 +178,12 @@ def test_caching_v2(): # test in memory cache try: litellm.set_verbose = True litellm.cache = Cache() - response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test") + response1 = completion( + model="gpt-3.5-turbo", + messages=messages, + caching=True, + mock_response="Hello world from cache test", + ) response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True) print(f"response1: {response1}") print(f"response2: {response2}") @@ -200,9 +210,18 @@ def test_caching_with_ttl(): litellm.set_verbose = True litellm.cache = Cache() response1 = completion( - model="gpt-3.5-turbo", messages=messages, caching=True, ttl=0, mock_response="Hello world from cache test 1" + model="gpt-3.5-turbo", + messages=messages, + caching=True, + ttl=0, + mock_response="Hello world from cache test 1", + ) + response2 = completion( + model="gpt-3.5-turbo", + messages=messages, + caching=True, + mock_response="Hello world from cache test 2", ) - response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test 2") print(f"response1: {response1}") print(f"response2: {response2}") litellm.cache = None # disable cache @@ -221,8 +240,18 @@ def test_caching_with_default_ttl(): try: litellm.set_verbose = True litellm.cache = Cache(ttl=0) - response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test") - response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test") + response1 = completion( + model="gpt-3.5-turbo", + messages=messages, + caching=True, + mock_response="Hello world from cache test", + ) + response2 = completion( + model="gpt-3.5-turbo", + messages=messages, + caching=True, + mock_response="Hello world from cache test", + ) print(f"response1: {response1}") print(f"response2: {response2}") litellm.cache = None # disable cache @@ -247,10 +276,16 @@ async def test_caching_with_cache_controls(sync_flag): if sync_flag: ## TTL = 0 response1 = completion( - model="gpt-3.5-turbo", messages=messages, cache={"ttl": 0}, mock_response="Hello world" + model="gpt-3.5-turbo", + messages=messages, + cache={"ttl": 0}, + mock_response="Hello world", ) response2 = completion( - model="gpt-3.5-turbo", messages=messages, cache={"s-maxage": 10}, mock_response="Hello world" + model="gpt-3.5-turbo", + messages=messages, + cache={"s-maxage": 10}, + mock_response="Hello world", ) assert response2["id"] != response1["id"] @@ -322,9 +357,19 @@ def test_caching_with_models_v2(): litellm.cache = Cache() print("test2 for caching") litellm.set_verbose = True - response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True, mock_response="Hello world from cache test") + response1 = completion( + model="gpt-3.5-turbo", + messages=messages, + caching=True, + mock_response="Hello world from cache test", + ) response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True) - response3 = completion(model="gpt-4.1-nano", messages=messages, caching=True, mock_response="Different model response") + response3 = completion( + model="gpt-4.1-nano", + messages=messages, + caching=True, + mock_response="Different model response", + ) print(f"response1: {response1}") print(f"response2: {response2}") print(f"response3: {response3}") @@ -423,7 +468,10 @@ def test_embedding_caching(): text_to_embed = [embedding_large_text] start_time = time.time() embedding1 = embedding( - model="text-embedding-ada-002", input=text_to_embed, caching=True, mock_response="0.1,0.2,0.3,0.4,0.5" + model="text-embedding-ada-002", + input=text_to_embed, + caching=True, + mock_response="0.1,0.2,0.3,0.4,0.5", ) end_time = time.time() print(f"Embedding 1 response time: {end_time - start_time} seconds") @@ -459,12 +507,18 @@ async def test_embedding_caching_individual_items_and_then_list(): "world", ] embedding1 = await aembedding( - model="text-embedding-ada-002", input=text_to_embed[0], caching=True, mock_response="0.1,0.2,0.3,0.4,0.5" + model="text-embedding-ada-002", + input=text_to_embed[0], + caching=True, + mock_response="0.1,0.2,0.3,0.4,0.5", ) initial_prompt_tokens = embedding1.usage.prompt_tokens await asyncio.sleep(1) embedding2 = await aembedding( - model="text-embedding-ada-002", input=text_to_embed[1], caching=True, mock_response="0.6,0.7,0.8,0.9,1.0" + model="text-embedding-ada-002", + input=text_to_embed[1], + caching=True, + mock_response="0.6,0.7,0.8,0.9,1.0", ) await asyncio.sleep(1) embedding3 = await aembedding( @@ -480,7 +534,10 @@ async def test_embedding_caching_individual_items_and_then_list(): additional_text = "this is a new text" text_to_embed.append(additional_text) embedding4 = await aembedding( - model="text-embedding-ada-002", input=text_to_embed, caching=True, mock_response="0.1,0.2,0.3,0.4,0.5" + model="text-embedding-ada-002", + input=text_to_embed, + caching=True, + mock_response="0.1,0.2,0.3,0.4,0.5", ) assert embedding4.usage.prompt_tokens > embedding3.usage.prompt_tokens @@ -490,7 +547,10 @@ async def test_embedding_caching_individual_items(): litellm.cache = Cache() text_to_embed = "hello" embedding1 = await aembedding( - model="text-embedding-ada-002", input=text_to_embed, caching=True, mock_response="0.1,0.2,0.3,0.4,0.5" + model="text-embedding-ada-002", + input=text_to_embed, + caching=True, + mock_response="0.1,0.2,0.3,0.4,0.5", ) await asyncio.sleep(1) @@ -512,13 +572,13 @@ def test_embedding_caching_azure(): litellm.cache = Cache() text_to_embed = [embedding_large_text] - api_key = os.environ["AZURE_API_KEY"] - api_base = os.environ["AZURE_API_BASE"] + api_key = os.environ["AZURE_AI_API_KEY"] + api_base = os.environ["AZURE_AI_API_BASE"] api_version = os.environ["AZURE_API_VERSION"] os.environ["AZURE_API_VERSION"] = "" - os.environ["AZURE_API_BASE"] = "" - os.environ["AZURE_API_KEY"] = "" + os.environ["AZURE_AI_API_BASE"] = "" + os.environ["AZURE_AI_API_KEY"] = "" start_time = time.time() print("AZURE CONFIGS") @@ -560,8 +620,8 @@ def test_embedding_caching_azure(): pytest.fail("Error occurred: Embedding caching failed") os.environ["AZURE_API_VERSION"] = api_version - os.environ["AZURE_API_BASE"] = api_base - os.environ["AZURE_API_KEY"] = api_key + os.environ["AZURE_AI_API_BASE"] = api_base + os.environ["AZURE_AI_API_KEY"] = api_key # test_embedding_caching_azure() @@ -851,9 +911,18 @@ def test_redis_cache_completion(): model="gpt-3.5-turbo", messages=messages, caching=True, max_tokens=20 ) response3 = completion( - model="gpt-3.5-turbo", messages=messages, caching=True, temperature=0.5, mock_response="Different params response" + model="gpt-3.5-turbo", + messages=messages, + caching=True, + temperature=0.5, + mock_response="Different params response", + ) + response4 = completion( + model="gpt-4o-mini", + messages=messages, + caching=True, + mock_response="Different model response", ) - response4 = completion(model="gpt-4o-mini", messages=messages, caching=True, mock_response="Different model response") print("\nresponse 1", response1) print("\nresponse 2", response2) @@ -1127,7 +1196,11 @@ async def test_redis_cache_atext_completion(): print("test for caching, atext_completion") response1 = await litellm.atext_completion( - model="gpt-3.5-turbo-instruct", prompt=prompt, max_tokens=40, temperature=1, mock_response="Hello world from cache test" + model="gpt-3.5-turbo-instruct", + prompt=prompt, + max_tokens=40, + temperature=1, + mock_response="Hello world from cache test", ) await asyncio.sleep(0.5) @@ -1458,11 +1531,17 @@ def test_cache_override(): # test embedding response1 = embedding( - model="text-embedding-ada-002", input=["hello who are you"], caching=False, mock_response="0.1,0.2,0.3,0.4,0.5" + model="text-embedding-ada-002", + input=["hello who are you"], + caching=False, + mock_response="0.1,0.2,0.3,0.4,0.5", ) response2 = embedding( - model="text-embedding-ada-002", input=["hello who are you"], caching=False, mock_response="0.6,0.7,0.8,0.9,1.0" + model="text-embedding-ada-002", + input=["hello who are you"], + caching=False, + mock_response="0.6,0.7,0.8,0.9,1.0", ) # When caching=False, responses should have different IDs @@ -2787,7 +2866,7 @@ def test_caching_thinking_args_hit(): # test in memory cache async def test_cache_key_in_hidden_params_acompletion(): """ Test that cache_key is present in _hidden_params on cache hits for acompletion. - + Validates fix for missing x-litellm-cache-key header on proxy cache hits. """ litellm.cache = Cache( @@ -2796,10 +2875,10 @@ async def test_cache_key_in_hidden_params_acompletion(): port=os.environ["REDIS_PORT"], password=os.environ["REDIS_PASSWORD"], ) - + unique_content = f"test cache key hidden params {uuid.uuid4()}" messages = [{"role": "user", "content": unique_content}] - + # First call - cache miss response1 = await litellm.acompletion( model="gpt-3.5-turbo", @@ -2807,12 +2886,12 @@ async def test_cache_key_in_hidden_params_acompletion(): mock_response="test response", caching=True, ) - + print(f"Response 1 _hidden_params: {response1._hidden_params}") assert response1._hidden_params.get("cache_hit") is not True - + await asyncio.sleep(0.5) - + # Second call - cache hit response2 = await litellm.acompletion( model="gpt-3.5-turbo", @@ -2820,17 +2899,17 @@ async def test_cache_key_in_hidden_params_acompletion(): mock_response="test response", caching=True, ) - + print(f"Response 2 _hidden_params: {response2._hidden_params}") - + # Verify cache hit occurred assert response2._hidden_params.get("cache_hit") is True - + # Verify cache_key is present in _hidden_params assert "cache_key" in response2._hidden_params assert response2._hidden_params["cache_key"] is not None - + # Verify both responses have same ID (cache hit) assert response1.id == response2.id - + litellm.cache = None diff --git a/tests/local_testing/test_caching_ssl.py b/tests/local_testing/test_caching_ssl.py index 523976f123..a8a38e747a 100644 --- a/tests/local_testing/test_caching_ssl.py +++ b/tests/local_testing/test_caching_ssl.py @@ -59,9 +59,9 @@ def test_caching_router(): "model_name": "gpt-3.5-turbo", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, diff --git a/tests/local_testing/test_class.py b/tests/local_testing/test_class.py index e02f59a294..b4b4f85a9d 100644 --- a/tests/local_testing/test_class.py +++ b/tests/local_testing/test_class.py @@ -56,9 +56,9 @@ # # "model_name": "gpt-3.5-turbo", # openai model name # # "litellm_params": { # params for litellm completion/embedding call # # "model": "azure/gpt-4.1-mini", -# # "api_key": os.getenv("AZURE_API_KEY"), +# # "api_key": os.getenv("AZURE_AI_API_KEY"), # # "api_version": os.getenv("AZURE_API_VERSION"), -# # "api_base": os.getenv("AZURE_API_BASE"), +# # "api_base": os.getenv("AZURE_AI_API_BASE"), # # }, # # } # # ] @@ -94,9 +94,9 @@ # # "model_name": "gpt-3.5-turbo", # openai model name # # "litellm_params": { # params for litellm completion/embedding call # # "model": "azure/gpt-4.1-mini", -# # "api_key": os.getenv("AZURE_API_KEY"), +# # "api_key": os.getenv("AZURE_AI_API_KEY"), # # "api_version": os.getenv("AZURE_API_VERSION"), -# # "api_base": os.getenv("AZURE_API_BASE"), +# # "api_base": os.getenv("AZURE_AI_API_BASE"), # # }, # # } # # ], diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index e6f5cd8651..af888e23cf 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -132,7 +132,6 @@ def test_null_role_response(): assert response.choices[0].message.role == "assistant" - def predibase_mock_post(url, data=None, json=None, headers=None, timeout=None): mock_response = MagicMock() mock_response.status_code = 200 @@ -286,7 +285,9 @@ def test_completion_claude_3_empty_response(): }, ] try: - response = litellm.completion(model="claude-sonnet-4-5-20250929", messages=messages) + response = litellm.completion( + model="claude-sonnet-4-5-20250929", messages=messages + ) print(response) except litellm.InternalServerError as e: pytest.skip(f"InternalServerError - {str(e)}") @@ -849,8 +850,8 @@ def test_completion_mistral_azure(): litellm.set_verbose = True response = completion( model="mistral/Mistral-large-nmefg", - api_key=os.environ["MISTRAL_AZURE_API_KEY"], - api_base=os.environ["MISTRAL_AZURE_API_BASE"], + api_key=os.environ["MISTRAL_AZURE_AI_API_KEY"], + api_base=os.environ["MISTRAL_AZURE_AI_API_BASE"], max_tokens=5, messages=[ { @@ -996,59 +997,6 @@ def test_completion_gpt4_vision(): pytest.fail(f"Error occurred: {e}") -# test_completion_gpt4_vision() - - -def test_completion_azure_gpt4_vision(): - # azure/gpt-4, vision takes 5-seconds to respond - try: - litellm.set_verbose = True - response = completion( - model="azure/gpt-4-vision", - timeout=5, - messages=[ - { - "role": "user", - "content": [ - {"type": "text", "text": "Whats in this image?"}, - { - "type": "image_url", - "image_url": { - "url": "https://avatars.githubusercontent.com/u/29436595?v=4" - }, - }, - ], - } - ], - base_url="https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions", - api_key=os.getenv("AZURE_VISION_API_KEY"), - enhancements={"ocr": {"enabled": True}, "grounding": {"enabled": True}}, - dataSources=[ - { - "type": "AzureComputerVision", - "parameters": { - "endpoint": "https://gpt-4-vision-enhancement.cognitiveservices.azure.com/", - "key": os.environ["AZURE_VISION_ENHANCE_KEY"], - }, - } - ], - ) - print(response) - except openai.APIError as e: - pass - except openai.APITimeoutError: - print("got a timeout error") - pass - except openai.RateLimitError as e: - print("got a rate liimt error", e) - pass - except openai.APIStatusError as e: - print("got an api status error", e) - pass - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - # test_completion_azure_gpt4_vision() @@ -1751,7 +1699,6 @@ def test_completion_openai_pydantic(model, api_version): pytest.fail(f"Error occurred: {e}") - def test_completion_text_openai(): try: # litellm.set_verbose =True @@ -2341,9 +2288,9 @@ def test_completion_azure_extra_headers(): response = completion( model="azure/gpt-4.1-mini", messages=messages, - api_base=os.getenv("AZURE_API_BASE"), + api_base=os.getenv("AZURE_AI_API_BASE"), api_version="2023-07-01-preview", - api_key=os.getenv("AZURE_API_KEY"), + api_key=os.getenv("AZURE_AI_API_KEY"), extra_headers={ "Authorization": "my-bad-key", "Ocp-Apim-Subscription-Key": "hello-world-testing", @@ -2379,8 +2326,8 @@ def test_completion_azure_ad_token(): litellm.set_verbose = True - old_key = os.environ["AZURE_API_KEY"] - os.environ.pop("AZURE_API_KEY", None) + old_key = os.environ["AZURE_AI_API_KEY"] + os.environ.pop("AZURE_AI_API_KEY", None) http_client = Client() @@ -2396,7 +2343,7 @@ def test_completion_azure_ad_token(): except Exception as e: pass finally: - os.environ["AZURE_API_KEY"] = old_key + os.environ["AZURE_AI_API_KEY"] = old_key mock_client.assert_called_once() request = mock_client.call_args[0][0] @@ -2412,8 +2359,8 @@ def test_completion_azure_key_completion_arg(): # DO NOT REMOVE THIS TEST. No MATTER WHAT Happens! # If you want to remove it, speak to Ishaan! # Ishaan will be very disappointed if this test is removed -> this is a standard way to pass api_key + the router + proxy use this - old_key = os.environ["AZURE_API_KEY"] - os.environ.pop("AZURE_API_KEY", None) + old_key = os.environ["AZURE_AI_API_KEY"] + os.environ.pop("AZURE_AI_API_KEY", None) try: print("azure gpt-3.5 test\n\n") litellm.set_verbose = True @@ -2430,9 +2377,9 @@ def test_completion_azure_key_completion_arg(): print("Hidden Params", response._hidden_params) assert response._hidden_params["custom_llm_provider"] == "azure" - os.environ["AZURE_API_KEY"] = old_key + os.environ["AZURE_AI_API_KEY"] = old_key except Exception as e: - os.environ["AZURE_API_KEY"] = old_key + os.environ["AZURE_AI_API_KEY"] = old_key pytest.fail(f"Error occurred: {e}") @@ -2443,8 +2390,8 @@ async def test_re_use_azure_async_client(): import openai client = openai.AsyncAzureOpenAI( - azure_endpoint=os.environ["AZURE_API_BASE"], - api_key=os.environ["AZURE_API_KEY"], + azure_endpoint=os.environ["AZURE_AI_API_BASE"], + api_key=os.environ["AZURE_AI_API_KEY"], api_version="2023-07-01-preview", ) ## Test azure call @@ -2525,13 +2472,13 @@ def test_completion_azure2(): try: print("azure gpt-3.5 test\n\n") litellm.set_verbose = False - api_base = os.environ["AZURE_API_BASE"] - api_key = os.environ["AZURE_API_KEY"] + api_base = os.environ["AZURE_AI_API_BASE"] + api_key = os.environ["AZURE_AI_API_KEY"] api_version = os.environ["AZURE_API_VERSION"] - os.environ["AZURE_API_BASE"] = "" + os.environ["AZURE_AI_API_BASE"] = "" os.environ["AZURE_API_VERSION"] = "" - os.environ["AZURE_API_KEY"] = "" + os.environ["AZURE_AI_API_KEY"] = "" ## Test azure call response = completion( @@ -2546,9 +2493,9 @@ def test_completion_azure2(): # Add any assertions here to check the response print(response) - os.environ["AZURE_API_BASE"] = api_base + os.environ["AZURE_AI_API_BASE"] = api_base os.environ["AZURE_API_VERSION"] = api_version - os.environ["AZURE_API_KEY"] = api_key + os.environ["AZURE_AI_API_KEY"] = api_key except Exception as e: pytest.fail(f"Error occurred: {e}") @@ -2562,13 +2509,13 @@ def test_completion_azure3(): try: print("azure gpt-3.5 test\n\n") litellm.set_verbose = True - litellm.api_base = os.environ["AZURE_API_BASE"] - litellm.api_key = os.environ["AZURE_API_KEY"] + litellm.api_base = os.environ["AZURE_AI_API_BASE"] + litellm.api_key = os.environ["AZURE_AI_API_KEY"] litellm.api_version = os.environ["AZURE_API_VERSION"] - os.environ["AZURE_API_BASE"] = "" + os.environ["AZURE_AI_API_BASE"] = "" os.environ["AZURE_API_VERSION"] = "" - os.environ["AZURE_API_KEY"] = "" + os.environ["AZURE_AI_API_KEY"] = "" ## Test azure call response = completion( @@ -2580,9 +2527,9 @@ def test_completion_azure3(): # Add any assertions here to check the response print(response) - os.environ["AZURE_API_BASE"] = litellm.api_base + os.environ["AZURE_AI_API_BASE"] = litellm.api_base os.environ["AZURE_API_VERSION"] = litellm.api_version - os.environ["AZURE_API_KEY"] = litellm.api_key + os.environ["AZURE_AI_API_KEY"] = litellm.api_key except Exception as e: pytest.fail(f"Error occurred: {e}") @@ -2594,7 +2541,7 @@ def test_completion_azure3(): # new azure test for using litellm. vars, # use the following vars in this test and make an azure_api_call # litellm.api_type = self.azure_api_type -# litellm.api_base = self.azure_api_base +# litellm.api_base = self.AZURE_AI_API_BASE # litellm.api_version = self.azure_api_version # litellm.api_key = self.api_key def test_completion_azure_with_litellm_key(): @@ -2604,14 +2551,14 @@ def test_completion_azure_with_litellm_key(): #### set litellm vars litellm.api_type = "azure" - litellm.api_base = os.environ["AZURE_API_BASE"] + litellm.api_base = os.environ["AZURE_AI_API_BASE"] litellm.api_version = os.environ["AZURE_API_VERSION"] - litellm.api_key = os.environ["AZURE_API_KEY"] + litellm.api_key = os.environ["AZURE_AI_API_KEY"] ######### UNSET ENV VARs for this ################ - os.environ["AZURE_API_BASE"] = "" + os.environ["AZURE_AI_API_BASE"] = "" os.environ["AZURE_API_VERSION"] = "" - os.environ["AZURE_API_KEY"] = "" + os.environ["AZURE_AI_API_KEY"] = "" ######### UNSET OpenAI vars for this ############## openai.api_type = "" @@ -2627,9 +2574,9 @@ def test_completion_azure_with_litellm_key(): print(response) ######### RESET ENV VARs for this ################ - os.environ["AZURE_API_BASE"] = litellm.api_base + os.environ["AZURE_AI_API_BASE"] = litellm.api_base os.environ["AZURE_API_VERSION"] = litellm.api_version - os.environ["AZURE_API_KEY"] = litellm.api_key + os.environ["AZURE_AI_API_KEY"] = litellm.api_key ######### UNSET litellm vars litellm.api_type = None @@ -3081,7 +3028,6 @@ async def test_completion_bedrock_httpx_models(sync_mode, model): pytest.fail(f"An error occurred - {str(e)}") - # test_completion_bedrock_titan() @@ -3256,7 +3202,6 @@ def test_completion_anyscale_api(): pytest.fail(f"Error occurred: {e}") - @pytest.mark.skip(reason="anyscale stopped serving public api endpoints") def test_completion_anyscale_2(): try: @@ -3871,9 +3816,6 @@ async def test_dynamic_azure_params(stream, sync_mode): raise e - - - @pytest.mark.parametrize( "model", ["gpt-4o", "azure/gpt-4.1-mini"], diff --git a/tests/local_testing/test_config.py b/tests/local_testing/test_config.py index e74f92ca7a..2c5d04d381 100644 --- a/tests/local_testing/test_config.py +++ b/tests/local_testing/test_config.py @@ -47,8 +47,8 @@ async def test_delete_deployment(): litellm_params = LiteLLM_Params( model="azure/gpt-4.1-mini", - api_key=os.getenv("AZURE_API_KEY"), - api_base=os.getenv("AZURE_API_BASE"), + api_key=os.getenv("AZURE_AI_API_KEY"), + api_base=os.getenv("AZURE_AI_API_BASE"), api_version=os.getenv("AZURE_API_VERSION"), ) encrypted_litellm_params = litellm_params.dict(exclude_none=True) @@ -131,8 +131,8 @@ async def test_add_existing_deployment(): litellm_params = LiteLLM_Params( model="gpt-3.5-turbo", - api_key=os.getenv("AZURE_API_KEY"), - api_base=os.getenv("AZURE_API_BASE"), + api_key=os.getenv("AZURE_AI_API_KEY"), + api_base=os.getenv("AZURE_AI_API_BASE"), api_version=os.getenv("AZURE_API_VERSION"), ) deployment = Deployment(model_name="gpt-3.5-turbo", litellm_params=litellm_params) @@ -186,8 +186,8 @@ async def test_db_error_new_model_check(): litellm_params = LiteLLM_Params( model="gpt-3.5-turbo", - api_key=os.getenv("AZURE_API_KEY"), - api_base=os.getenv("AZURE_API_BASE"), + api_key=os.getenv("AZURE_AI_API_KEY"), + api_base=os.getenv("AZURE_AI_API_BASE"), api_version=os.getenv("AZURE_API_VERSION"), ) deployment = Deployment(model_name="gpt-3.5-turbo", litellm_params=litellm_params) @@ -233,8 +233,8 @@ async def test_db_error_new_model_check(): litellm_params = LiteLLM_Params( model="azure/gpt-4.1-mini", - api_key=os.getenv("AZURE_API_KEY"), - api_base=os.getenv("AZURE_API_BASE"), + api_key=os.getenv("AZURE_AI_API_KEY"), + api_base=os.getenv("AZURE_AI_API_BASE"), api_version=os.getenv("AZURE_API_VERSION"), ) @@ -251,8 +251,8 @@ def _create_model_list(flag_value: Literal[0, 1], master_key: str): new_litellm_params = LiteLLM_Params( model="azure/gpt-4.1-mini-3", - api_key=os.getenv("AZURE_API_KEY"), - api_base=os.getenv("AZURE_API_BASE"), + api_key=os.getenv("AZURE_AI_API_KEY"), + api_base=os.getenv("AZURE_AI_API_BASE"), api_version=os.getenv("AZURE_API_VERSION"), ) @@ -421,4 +421,3 @@ def test_litellm_proxy_responses_api_config(): assert ( config.custom_llm_provider == LlmProviders.LITELLM_PROXY ), "custom_llm_provider should be LITELLM_PROXY" - diff --git a/tests/local_testing/test_configs/test_bad_config.yaml b/tests/local_testing/test_configs/test_bad_config.yaml index 4a70886a93..4bc4c7cc54 100644 --- a/tests/local_testing/test_configs/test_bad_config.yaml +++ b/tests/local_testing/test_configs/test_bad_config.yaml @@ -6,16 +6,16 @@ model_list: - model_name: working-azure-gpt-3.5-turbo litellm_params: model: azure/gpt-4.1-mini - api_base: os.environ/AZURE_API_BASE - api_key: os.environ/AZURE_API_KEY + api_base: os.environ/AZURE_AI_API_BASE + api_key: os.environ/AZURE_AI_API_KEY - model_name: azure-gpt-3.5-turbo litellm_params: model: azure/gpt-4.1-mini - api_base: os.environ/AZURE_API_BASE + api_base: os.environ/AZURE_AI_API_BASE api_key: bad-key - model_name: azure-embedding litellm_params: model: azure/text-embedding-ada-002 - api_base: os.environ/AZURE_API_BASE + api_base: os.environ/AZURE_AI_API_BASE api_key: bad-key \ No newline at end of file diff --git a/tests/local_testing/test_configs/test_cloudflare_azure_with_cache_config.yaml b/tests/local_testing/test_configs/test_cloudflare_azure_with_cache_config.yaml index 9902835618..24240008fe 100644 --- a/tests/local_testing/test_configs/test_cloudflare_azure_with_cache_config.yaml +++ b/tests/local_testing/test_configs/test_cloudflare_azure_with_cache_config.yaml @@ -3,7 +3,7 @@ model_list: litellm_params: model: azure/gpt-4.1-mini api_base: https://gateway.ai.cloudflare.com/v1/0399b10e77ac6668c80404a5ff49eb37/litellm-test/azure-openai/openai-gpt-4-test-v-1 - api_key: os.environ/AZURE_API_KEY + api_key: os.environ/AZURE_AI_API_KEY api_version: 2023-07-01-preview litellm_settings: diff --git a/tests/local_testing/test_configs/test_config_no_auth.yaml b/tests/local_testing/test_configs/test_config_no_auth.yaml index cdc447a5ee..f489621704 100644 --- a/tests/local_testing/test_configs/test_config_no_auth.yaml +++ b/tests/local_testing/test_configs/test_config_no_auth.yaml @@ -11,7 +11,7 @@ model_list: model_name: azure-model - litellm_params: api_base: https://gateway.ai.cloudflare.com/v1/0399b10e77ac6668c80404a5ff49eb37/litellm-test/azure-openai/openai-gpt-4-test-v-1 - api_key: os.environ/AZURE_API_KEY + api_key: os.environ/AZURE_AI_API_KEY model: azure/gpt-4.1-mini model_name: azure-cloudflare-model - litellm_params: @@ -49,8 +49,8 @@ model_list: id: 79fc75bf-8e1b-47d5-8d24-9365a854af03 model_name: test_openai_models - litellm_params: - api_base: os.environ/AZURE_API_BASE - api_key: os.environ/AZURE_API_KEY + api_base: os.environ/AZURE_AI_API_BASE + api_key: os.environ/AZURE_AI_API_KEY api_version: 2023-07-01-preview model: azure/text-embedding-ada-002 model_info: @@ -94,16 +94,16 @@ model_list: mode: image_generation model_name: dall-e-3 - litellm_params: - api_base: os.environ/AZURE_API_BASE - api_key: os.environ/AZURE_API_KEY + api_base: os.environ/AZURE_AI_API_BASE + api_key: os.environ/AZURE_AI_API_KEY api_version: 2023-06-01-preview model: azure/ model_info: mode: image_generation model_name: dall-e-2 - litellm_params: - api_base: os.environ/AZURE_API_BASE - api_key: os.environ/AZURE_API_KEY + api_base: os.environ/AZURE_AI_API_BASE + api_key: os.environ/AZURE_AI_API_KEY api_version: 2023-07-01-preview model: azure/text-embedding-ada-002 model_info: diff --git a/tests/local_testing/test_configs/test_custom_logger.yaml b/tests/local_testing/test_configs/test_custom_logger.yaml index 22bbfe42be..464d66bd78 100644 --- a/tests/local_testing/test_configs/test_custom_logger.yaml +++ b/tests/local_testing/test_configs/test_custom_logger.yaml @@ -2,8 +2,8 @@ model_list: - model_name: Azure OpenAI GPT-4 Canada litellm_params: model: azure/gpt-4.1-mini - api_base: os.environ/AZURE_API_BASE - api_key: os.environ/AZURE_API_KEY + api_base: os.environ/AZURE_AI_API_BASE + api_key: os.environ/AZURE_AI_API_KEY api_version: "2023-07-01-preview" model_info: mode: chat @@ -12,8 +12,8 @@ model_list: - model_name: azure-embedding-model litellm_params: model: azure/text-embedding-ada-002 - api_base: os.environ/AZURE_API_BASE - api_key: os.environ/AZURE_API_KEY + api_base: os.environ/AZURE_AI_API_BASE + api_key: os.environ/AZURE_AI_API_KEY api_version: "2023-07-01-preview" model_info: mode: embedding diff --git a/tests/local_testing/test_embedding.py b/tests/local_testing/test_embedding.py index 9a63c6c83c..2c265d444a 100644 --- a/tests/local_testing/test_embedding.py +++ b/tests/local_testing/test_embedding.py @@ -296,13 +296,13 @@ def test_openai_embedding_timeouts(): def test_openai_azure_embedding(): try: - api_key = os.environ["AZURE_API_KEY"] - api_base = os.environ["AZURE_API_BASE"] + api_key = os.environ["AZURE_AI_API_KEY"] + api_base = os.environ["AZURE_AI_API_BASE"] api_version = os.environ["AZURE_API_VERSION"] os.environ["AZURE_API_VERSION"] = "" - os.environ["AZURE_API_BASE"] = "" - os.environ["AZURE_API_KEY"] = "" + os.environ["AZURE_AI_API_BASE"] = "" + os.environ["AZURE_AI_API_KEY"] = "" response = embedding( model="azure/text-embedding-ada-002", @@ -314,8 +314,8 @@ def test_openai_azure_embedding(): print(response) os.environ["AZURE_API_VERSION"] = api_version - os.environ["AZURE_API_BASE"] = api_base - os.environ["AZURE_API_KEY"] = api_key + os.environ["AZURE_AI_API_BASE"] = api_base + os.environ["AZURE_AI_API_KEY"] = api_key except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/tests/local_testing/test_exceptions.py b/tests/local_testing/test_exceptions.py index 2c950d7906..a0ae71c4f1 100644 --- a/tests/local_testing/test_exceptions.py +++ b/tests/local_testing/test_exceptions.py @@ -162,8 +162,8 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th temporary_secret_key = os.environ["AWS_SECRET_ACCESS_KEY"] os.environ["AWS_SECRET_ACCESS_KEY"] = "bad-key" elif model == "azure/gpt-4.1-mini": - temporary_key = os.environ["AZURE_API_KEY"] - os.environ["AZURE_API_KEY"] = "bad-key" + temporary_key = os.environ["AZURE_AI_API_KEY"] + os.environ["AZURE_AI_API_KEY"] = "bad-key" elif model == "claude-3-5-haiku-20241022": temporary_key = os.environ["ANTHROPIC_API_KEY"] os.environ["ANTHROPIC_API_KEY"] = "bad-key" @@ -175,9 +175,7 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th os.environ["AI21_API_KEY"] = "bad-key" elif "togethercomputer" in model: temporary_key = os.environ["TOGETHERAI_API_KEY"] - os.environ["TOGETHERAI_API_KEY"] = ( - "sk-test-togetherai-key-808" - ) + os.environ["TOGETHERAI_API_KEY"] = "sk-test-togetherai-key-808" elif model in litellm.openrouter_models: temporary_key = os.environ["OPENROUTER_API_KEY"] os.environ["OPENROUTER_API_KEY"] = "bad-key" @@ -212,7 +210,7 @@ def invalid_auth(model): # set the model key to an invalid key, depending on th if model == "gpt-3.5-turbo": os.environ["OPENAI_API_KEY"] = temporary_key elif model == "chatgpt-test": - os.environ["AZURE_API_KEY"] = temporary_key + os.environ["AZURE_AI_API_KEY"] = temporary_key azure = True elif model == "claude-3-5-haiku-20241022": os.environ["ANTHROPIC_API_KEY"] = temporary_key @@ -259,17 +257,17 @@ def test_completion_azure_exception(): print("azure gpt-3.5 test\n\n") litellm.set_verbose = True ## Test azure call - old_azure_key = os.environ["AZURE_API_KEY"] - os.environ["AZURE_API_KEY"] = "good morning" + old_azure_key = os.environ["AZURE_AI_API_KEY"] + os.environ["AZURE_AI_API_KEY"] = "good morning" response = completion( model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hello"}], ) - os.environ["AZURE_API_KEY"] = old_azure_key + os.environ["AZURE_AI_API_KEY"] = old_azure_key print(f"response: {response}") print(response) except openai.AuthenticationError as e: - os.environ["AZURE_API_KEY"] = old_azure_key + os.environ["AZURE_AI_API_KEY"] = old_azure_key print("good job got the correct error for azure when key not set") except Exception as e: pytest.fail(f"Error occurred: {e}") @@ -303,8 +301,8 @@ async def asynctest_completion_azure_exception(): print("azure gpt-3.5 test\n\n") litellm.set_verbose = True ## Test azure call - old_azure_key = os.environ["AZURE_API_KEY"] - os.environ["AZURE_API_KEY"] = "good morning" + old_azure_key = os.environ["AZURE_AI_API_KEY"] + os.environ["AZURE_AI_API_KEY"] = "good morning" response = await litellm.acompletion( model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hello"}], @@ -312,7 +310,7 @@ async def asynctest_completion_azure_exception(): print(f"response: {response}") print(response) except openai.AuthenticationError as e: - os.environ["AZURE_API_KEY"] = old_azure_key + os.environ["AZURE_AI_API_KEY"] = old_azure_key print("good job got the correct error for azure when key not set") print(e) except Exception as e: @@ -495,6 +493,7 @@ def test_completion_bedrock_invalid_role_exception(): == "litellm.BadRequestError: Invalid Message passed in {'role': 'very-bad-role', 'content': 'hello'}" ) + @pytest.mark.skip(reason="OpenAI exception changed to a generic error") def test_content_policy_exceptionimage_generation_openai(): try: @@ -773,7 +772,15 @@ def test_litellm_predibase_exception(): @pytest.mark.parametrize( - "provider", ["predibase", "vertex_ai_beta", "anthropic", "databricks", "watsonx", "fireworks_ai"] + "provider", + [ + "predibase", + "vertex_ai_beta", + "anthropic", + "databricks", + "watsonx", + "fireworks_ai", + ], ) def test_exception_mapping(provider): """ @@ -826,14 +833,14 @@ def test_fireworks_ai_exception_mapping(): 2. Text-based rate limit detection (the main issue fixed) 3. Generic 400 errors that should NOT be rate limits 4. ExceptionCheckers utility function - + Related to: https://github.com/BerriAI/litellm/pull/11455 Based on Fireworks AI documentation: https://docs.fireworks.ai/tools-sdks/python-client/api-reference """ import litellm from litellm.llms.fireworks_ai.common_utils import FireworksAIException from litellm.litellm_core_utils.exception_mapping_utils import ExceptionCheckers - + # Test scenarios covering all important cases test_scenarios = [ { @@ -855,57 +862,63 @@ def test_fireworks_ai_exception_mapping(): "expected_exception": litellm.BadRequestError, }, ] - + # Test each scenario for scenario in test_scenarios: mock_exception = FireworksAIException( - status_code=scenario["status_code"], - message=scenario["message"], - headers={} + status_code=scenario["status_code"], message=scenario["message"], headers={} ) - + try: response = litellm.completion( model="fireworks_ai/llama-v3p1-70b-instruct", messages=[{"role": "user", "content": "Hello"}], mock_response=mock_exception, ) - pytest.fail(f"Expected {scenario['expected_exception'].__name__} to be raised") + pytest.fail( + f"Expected {scenario['expected_exception'].__name__} to be raised" + ) except scenario["expected_exception"] as e: if scenario["expected_exception"] == litellm.RateLimitError: assert "rate limit" in str(e).lower() or "429" in str(e) except Exception as e: - pytest.fail(f"Expected {scenario['expected_exception'].__name__} but got {type(e).__name__}: {e}") - + pytest.fail( + f"Expected {scenario['expected_exception'].__name__} but got {type(e).__name__}: {e}" + ) + # Test ExceptionCheckers.is_error_str_rate_limit() method directly - + # Test cases that should return True (rate limit detected) rate_limit_strings = [ "429 rate limit exceeded", - "Rate limit exceeded, please try again later", + "Rate limit exceeded, please try again later", "RATE LIMIT ERROR", "Error 429: rate limit", '{"error":{"type":"invalid_request_error","message":"rate limit exceeded, please try again later"}}', "HTTP 429 Too Many Requests", ] - + for error_str in rate_limit_strings: - assert ExceptionCheckers.is_error_str_rate_limit(error_str), f"Should detect rate limit in: {error_str}" - + assert ExceptionCheckers.is_error_str_rate_limit( + error_str + ), f"Should detect rate limit in: {error_str}" + # Test cases that should return False (not rate limit) non_rate_limit_strings = [ "400 Bad Request", - "Authentication failed", + "Authentication failed", "Invalid model specified", "Context window exceeded", "Internal server error", "", "Some other error message", ] - + for error_str in non_rate_limit_strings: - assert not ExceptionCheckers.is_error_str_rate_limit(error_str), f"Should NOT detect rate limit in: {error_str}" - + assert not ExceptionCheckers.is_error_str_rate_limit( + error_str + ), f"Should NOT detect rate limit in: {error_str}" + # Test edge cases assert not ExceptionCheckers.is_error_str_rate_limit(None) # type: ignore assert not ExceptionCheckers.is_error_str_rate_limit(42) # type: ignore @@ -1142,6 +1155,7 @@ def test_openai_gateway_timeout_error(): """ openai_client = OpenAI() mapped_target = openai_client.chat.completions.with_raw_response # type: ignore + def _return_exception(*args, **kwargs): import datetime @@ -1175,13 +1189,17 @@ def test_openai_gateway_timeout_error(): setattr(exception, k, v) raise exception - try: + try: with patch.object( mapped_target, "create", side_effect=_return_exception, ): - litellm.completion(model="openai/gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello world"}], client=openai_client) + litellm.completion( + model="openai/gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello world"}], + client=openai_client, + ) pytest.fail("Expected to raise Timeout") except litellm.Timeout as e: assert e.status_code == 504 @@ -1350,7 +1368,7 @@ def test_context_window_exceeded_error_from_litellm_proxy(): def test_bad_request_error_with_response_without_request(): """ Test that BadRequestError handles Response objects without a request attribute. - + This simulates a real scenario where a Response is created without a request (e.g., in tests or when manually creating error responses), and we need to ensure it doesn't raise RuntimeError when the exception is created. @@ -1362,8 +1380,7 @@ def test_bad_request_error_with_response_without_request(): # Create a Response without a request (simulates the scenario that was failing) response_without_request = Response(status_code=400, text="Bad Request") - - + # Test that extract_and_raise_litellm_exception can handle this args = { "response": response_without_request, @@ -1371,17 +1388,17 @@ def test_bad_request_error_with_response_without_request(): "model": "gpt-3.5-turbo", "custom_llm_provider": "openai", } - + # This should raise BadRequestError without RuntimeError with pytest.raises(litellm.BadRequestError) as exc_info: extract_and_raise_litellm_exception(**args) - + # Verify the exception was created successfully error = exc_info.value assert error is not None assert error.model == "gpt-3.5-turbo" assert error.llm_provider == "openai" - + # Verify the exception has a response (should be minimal error response) assert error.response is not None # The response should have a request (minimal error response has one) @@ -1420,6 +1437,3 @@ async def test_exception_bubbling_up(sync_mode, stream_mode, model): assert exc_info.value.code == "invalid_value" assert exc_info.value.param is not None assert exc_info.value.type == "invalid_request_error" - - - diff --git a/tests/local_testing/test_loadtest_router.py b/tests/local_testing/test_loadtest_router.py index 3f6e4af4fb..3d1062f0d2 100644 --- a/tests/local_testing/test_loadtest_router.py +++ b/tests/local_testing/test_loadtest_router.py @@ -39,8 +39,8 @@ # "model_name": "gpt-3.5-turbo", # "litellm_params": { # "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_API_KEY"), -# "api_base": os.getenv("AZURE_API_BASE"), +# "api_key": os.getenv("AZURE_AI_API_KEY"), +# "api_base": os.getenv("AZURE_AI_API_BASE"), # "api_version": os.getenv("AZURE_API_VERSION"), # }, # }, diff --git a/tests/local_testing/test_prompt_injection_detection.py b/tests/local_testing/test_prompt_injection_detection.py index c4cc4cde32..b1a9aff158 100644 --- a/tests/local_testing/test_prompt_injection_detection.py +++ b/tests/local_testing/test_prompt_injection_detection.py @@ -108,9 +108,9 @@ async def test_prompt_injection_llm_eval(): "model_name": "gpt-3.5-turbo", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, diff --git a/tests/local_testing/test_router_budget_limiter.py b/tests/local_testing/test_router_budget_limiter.py index 05ce7c1f53..1a36e9de8f 100644 --- a/tests/local_testing/test_router_budget_limiter.py +++ b/tests/local_testing/test_router_budget_limiter.py @@ -75,9 +75,9 @@ async def test_provider_budgets_e2e_test(): "model_name": "gpt-3.5-turbo", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "model_info": {"id": "azure-model-id"}, }, @@ -609,6 +609,7 @@ async def test_deployment_budgets_e2e_test_expect_to_fail(): assert "Exceeded budget for deployment" in str(exc_info.value) + @pytest.mark.flaky(retries=6, delay=2) @pytest.mark.asyncio async def test_tag_budgets_e2e_test_expect_to_fail(): diff --git a/tests/local_testing/test_router_caching.py b/tests/local_testing/test_router_caching.py index 6fc220bf72..cb223b661b 100644 --- a/tests/local_testing/test_router_caching.py +++ b/tests/local_testing/test_router_caching.py @@ -268,8 +268,8 @@ async def test_acompletion_caching_on_router_caching_groups(): "model_name": "azure-gpt-3.5-turbo", "litellm_params": { "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_key": os.getenv("AZURE_AI_API_KEY"), + "api_base": os.getenv("AZURE_AI_API_BASE"), "api_version": os.getenv("AZURE_API_VERSION"), }, "tpm": 100000, diff --git a/tests/local_testing/test_router_client_init.py b/tests/local_testing/test_router_client_init.py index f2541601c3..ea808f96f3 100644 --- a/tests/local_testing/test_router_client_init.py +++ b/tests/local_testing/test_router_client_init.py @@ -85,7 +85,7 @@ def test_router_init_azure_service_principal_with_secret_with_environment_variab To allow for local testing without real credentials, first must mock Azure SDK authentication functions and environment variables. """ - monkeypatch.delenv("AZURE_API_KEY", raising=False) + monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) litellm.enable_azure_ad_token_refresh = True # mock the token provider function mocked_func_generating_token = MagicMock(return_value="test_token") diff --git a/tests/local_testing/test_router_cooldown_handlers.py b/tests/local_testing/test_router_cooldown_handlers.py index 7be8289abf..fdc89fc04e 100644 --- a/tests/local_testing/test_router_cooldown_handlers.py +++ b/tests/local_testing/test_router_cooldown_handlers.py @@ -45,9 +45,9 @@ async def test_cooldown_badrequest_error(): "model_name": "gpt-3.5-turbo", "litellm_params": { "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, } ], diff --git a/tests/local_testing/test_router_debug_logs.py b/tests/local_testing/test_router_debug_logs.py index 1004e7747e..0b4771b526 100644 --- a/tests/local_testing/test_router_debug_logs.py +++ b/tests/local_testing/test_router_debug_logs.py @@ -34,9 +34,9 @@ def test_async_fallbacks(caplog): "model_name": "azure/gpt-3.5-turbo", "litellm_params": { "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), "mock_response": "Hello world", }, "tpm": 240000, diff --git a/tests/local_testing/test_router_fallbacks.py b/tests/local_testing/test_router_fallbacks.py index c586fa8c93..383ad10457 100644 --- a/tests/local_testing/test_router_fallbacks.py +++ b/tests/local_testing/test_router_fallbacks.py @@ -70,7 +70,7 @@ def test_sync_fallbacks(): "model": "azure/gpt-4.1-mini", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -79,9 +79,9 @@ def test_sync_fallbacks(): "model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -92,7 +92,7 @@ def test_sync_fallbacks(): "model": "azure/chatgpt-functioncalling", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -132,7 +132,9 @@ def test_sync_fallbacks(): response = router.completion(**kwargs) print(f"response: {response}") time.sleep(0.05) # allow a delay as success_callbacks are on a separate thread - assert customHandler.previous_models == 3 # 1 init call + 2 retries (fallback not counted as previous) + assert ( + customHandler.previous_models == 3 + ) # 1 init call + 2 retries (fallback not counted as previous) print("Passed ! Test router_fallbacks: test_sync_fallbacks()") router.reset() @@ -153,7 +155,7 @@ async def test_async_fallbacks(): "model": "azure/gpt-4.1-mini", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -162,9 +164,9 @@ async def test_async_fallbacks(): "model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -175,7 +177,7 @@ async def test_async_fallbacks(): "model": "azure/chatgpt-functioncalling", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -220,7 +222,9 @@ async def test_async_fallbacks(): await asyncio.sleep( 0.05 ) # allow a delay as success_callbacks are on a separate thread - assert customHandler.previous_models == 3 # 1 init call + 2 retries (fallback not counted as previous) + assert ( + customHandler.previous_models == 3 + ) # 1 init call + 2 retries (fallback not counted as previous) router.reset() except litellm.Timeout as e: pass @@ -242,7 +246,7 @@ def test_sync_fallbacks_embeddings(): "model": "azure/text-embedding-ada-002", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -292,7 +296,7 @@ async def test_async_fallbacks_embeddings(): "model": "azure/text-embedding-ada-002", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -348,7 +352,7 @@ def test_dynamic_fallbacks_sync(): "model": "azure/gpt-4.1-mini", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -357,9 +361,9 @@ def test_dynamic_fallbacks_sync(): "model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -370,7 +374,7 @@ def test_dynamic_fallbacks_sync(): "model": "azure/chatgpt-functioncalling", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -403,7 +407,9 @@ def test_dynamic_fallbacks_sync(): response = router.completion(**kwargs) print(f"response: {response}") time.sleep(0.05) # allow a delay as success_callbacks are on a separate thread - assert customHandler.previous_models >= 3 # 1 init call, retries, 1 fallback (count varies with cooldown timing) + assert ( + customHandler.previous_models >= 3 + ) # 1 init call, retries, 1 fallback (count varies with cooldown timing) router.reset() except Exception as e: pytest.fail(f"An exception occurred - {e}") @@ -425,7 +431,7 @@ async def test_dynamic_fallbacks_async(): "model": "azure/gpt-4.1-mini", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -434,9 +440,9 @@ async def test_dynamic_fallbacks_async(): "model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -447,7 +453,7 @@ async def test_dynamic_fallbacks_async(): "model": "azure/chatgpt-functioncalling", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -489,7 +495,9 @@ async def test_dynamic_fallbacks_async(): await asyncio.sleep( 0.05 ) # allow a delay as success_callbacks are on a separate thread - assert customHandler.previous_models >= 3 # 1 init call, retries, 1 fallback (count varies with cooldown timing) + assert ( + customHandler.previous_models >= 3 + ) # 1 init call, retries, 1 fallback (count varies with cooldown timing) router.reset() except Exception as e: pytest.fail(f"An exception occurred - {e}") @@ -562,7 +570,7 @@ def test_sync_fallbacks_streaming(): "model": "azure/gpt-4.1-mini", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -571,9 +579,9 @@ def test_sync_fallbacks_streaming(): "model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -584,7 +592,7 @@ def test_sync_fallbacks_streaming(): "model": "azure/chatgpt-functioncalling", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -643,7 +651,7 @@ async def test_async_fallbacks_max_retries_per_request(): "model": "azure/gpt-4.1-mini", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -652,9 +660,9 @@ async def test_async_fallbacks_max_retries_per_request(): "model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -665,7 +673,7 @@ async def test_async_fallbacks_max_retries_per_request(): "model": "azure/chatgpt-functioncalling", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -750,9 +758,9 @@ def test_ausage_based_routing_fallbacks(): def get_azure_params(deployment_name: str): params = { "model": f"azure/{deployment_name}", - "api_key": os.environ["AZURE_API_KEY"], + "api_key": os.environ["AZURE_AI_API_KEY"], "api_version": os.environ["AZURE_API_VERSION"], - "api_base": os.environ["AZURE_API_BASE"], + "api_base": os.environ["AZURE_AI_API_BASE"], } return params @@ -855,7 +863,7 @@ def test_custom_cooldown_times(): "model": "azure/gpt-4.1-mini", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 24000000, }, @@ -863,9 +871,9 @@ def test_custom_cooldown_times(): "model_name": "gpt-3.5-turbo", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 1, }, diff --git a/tests/local_testing/test_router_init.py b/tests/local_testing/test_router_init.py index e232ff105d..6f3453cbdd 100644 --- a/tests/local_testing/test_router_init.py +++ b/tests/local_testing/test_router_init.py @@ -41,9 +41,9 @@ # "model_name": "gpt-3.5-turbo", # "litellm_params": { # "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_API_KEY"), +# "api_key": os.getenv("AZURE_AI_API_KEY"), # "api_version": os.getenv("AZURE_API_VERSION"), -# "api_base": os.getenv("AZURE_API_BASE"), +# "api_base": os.getenv("AZURE_AI_API_BASE"), # "timeout": 0.01, # "stream_timeout": 0.000_001, # "max_retries": 7, @@ -97,9 +97,9 @@ # "model_name": "gpt-3.5-turbo", # "litellm_params": { # "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_API_KEY"), +# "api_key": os.getenv("AZURE_AI_API_KEY"), # "api_version": os.getenv("AZURE_API_VERSION"), -# "api_base": os.getenv("AZURE_API_BASE"), +# "api_base": os.getenv("AZURE_AI_API_BASE"), # }, # }, # ] @@ -135,7 +135,7 @@ # "model_name": "azure-cloudflare", # "litellm_params": { # "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_API_KEY"), +# "api_key": os.getenv("AZURE_AI_API_KEY"), # "api_version": os.getenv("AZURE_API_VERSION"), # "api_base": "https://gateway.ai.cloudflare.com/v1/0399b10e77ac6668c80404a5ff49eb37/litellm-test/azure-openai/openai-gpt-4-test-v-1", # }, @@ -202,9 +202,9 @@ # "model_name": "gpt-3.5-turbo", # "litellm_params": { # "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_API_KEY"), +# "api_key": os.getenv("AZURE_AI_API_KEY"), # "api_version": os.getenv("AZURE_API_VERSION"), -# "api_base": os.getenv("AZURE_API_BASE"), +# "api_base": os.getenv("AZURE_AI_API_BASE"), # "timeout": 0.000001, # "stream_timeout": 0.000_001, # }, @@ -255,9 +255,9 @@ # "model_name": "gpt-3.5-turbo", # "litellm_params": { # "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_API_KEY"), +# "api_key": os.getenv("AZURE_AI_API_KEY"), # "api_version": os.getenv("AZURE_API_VERSION"), -# "api_base": os.getenv("AZURE_API_BASE"), +# "api_base": os.getenv("AZURE_AI_API_BASE"), # "timeout": 200, # regular calls will not timeout, stream calls will # "stream_timeout": 10, # }, @@ -349,7 +349,7 @@ # "model_name": "gpt-4-vision-enhancements", # "litellm_params": { # "model": "azure/gpt-4-vision", -# "api_key": os.getenv("AZURE_API_KEY"), +# "api_key": os.getenv("AZURE_AI_API_KEY"), # "base_url": "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/", # "dataSources": [ # { @@ -616,9 +616,9 @@ # "model_name": "gpt-3.5-turbo", # "litellm_params": { # "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_API_KEY"), +# "api_key": os.getenv("AZURE_AI_API_KEY"), # "api_version": os.getenv("AZURE_API_VERSION"), -# "api_base": os.getenv("AZURE_API_BASE"), +# "api_base": os.getenv("AZURE_AI_API_BASE"), # "timeout": 0.01, # "stream_timeout": 0.000_001, # "max_retries": 7, @@ -661,9 +661,9 @@ # "model_name": "gpt-3.5-turbo", # "litellm_params": { # "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_API_KEY"), +# "api_key": os.getenv("AZURE_AI_API_KEY"), # "api_version": os.getenv("AZURE_API_VERSION"), -# "api_base": os.getenv("AZURE_API_BASE"), +# "api_base": os.getenv("AZURE_AI_API_BASE"), # "timeout": 0.01, # "stream_timeout": 0.000_001, # "max_retries": 7, diff --git a/tests/local_testing/test_router_timeout.py b/tests/local_testing/test_router_timeout.py index 1d09f1f1e0..943a5413d6 100644 --- a/tests/local_testing/test_router_timeout.py +++ b/tests/local_testing/test_router_timeout.py @@ -31,8 +31,8 @@ def test_router_timeouts(): "model_name": "openai-gpt-4", "litellm_params": { "model": "azure/gpt-4.1-mini", - "api_key": "os.environ/AZURE_API_KEY", - "api_base": "os.environ/AZURE_API_BASE", + "api_key": "os.environ/AZURE_AI_API_KEY", + "api_base": "os.environ/AZURE_AI_API_BASE", "api_version": "os.environ/AZURE_API_VERSION", }, "tpm": 80000, diff --git a/tests/local_testing/test_router_utils.py b/tests/local_testing/test_router_utils.py index 4f7f53cef0..f2fd2fdf55 100644 --- a/tests/local_testing/test_router_utils.py +++ b/tests/local_testing/test_router_utils.py @@ -35,7 +35,7 @@ def test_returned_settings(): "model": "azure/gpt-4.1-mini", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -99,7 +99,7 @@ def test_update_kwargs_before_fallbacks_unit_test(): "model": "azure/gpt-4.1-mini", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, } ], @@ -136,7 +136,7 @@ async def test_update_kwargs_before_fallbacks(call_type): "model": "azure/gpt-4.1-mini", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, } ], @@ -266,6 +266,7 @@ async def test_call_router_callbacks_on_success(): ) assert increment["increment_value"] == 1 + @pytest.mark.serial @pytest.mark.asyncio async def test_call_router_callbacks_on_failure(): @@ -486,7 +487,9 @@ def test_router_get_deployment_credentials_with_provider(): ) # Test getting credentials by model_id - credentials = router.get_deployment_credentials_with_provider(model_id="openai-deployment-1") + credentials = router.get_deployment_credentials_with_provider( + model_id="openai-deployment-1" + ) assert credentials is not None assert credentials["api_key"] == "sk-test-123" assert credentials["custom_llm_provider"] == "openai" @@ -499,14 +502,16 @@ def test_router_get_deployment_credentials_with_provider(): assert credentials2["custom_llm_provider"] == "anthropic" # Test with non-existent model - credentials3 = router.get_deployment_credentials_with_provider(model_id="non-existent") + credentials3 = router.get_deployment_credentials_with_provider( + model_id="non-existent" + ) assert credentials3 is None def test_router_get_deployment_credentials_with_provider_wildcard(): """ Test that get_deployment_credentials_with_provider handles wildcard patterns. - + When a model like openai/gpt-4o is requested and the config has openai/*, the method should resolve the wildcard pattern and return credentials. """ @@ -533,20 +538,26 @@ def test_router_get_deployment_credentials_with_provider_wildcard(): ) # Test wildcard pattern matching for OpenAI - credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-4o") + credentials = router.get_deployment_credentials_with_provider( + model_id="openai/gpt-4o" + ) assert credentials is not None assert credentials["api_key"] == "sk-wildcard-123" assert credentials["custom_llm_provider"] == "openai" assert credentials["api_base"] == "https://api.openai.com/v1" # Test wildcard pattern matching for Anthropic - credentials2 = router.get_deployment_credentials_with_provider(model_id="anthropic/claude-3-opus") + credentials2 = router.get_deployment_credentials_with_provider( + model_id="anthropic/claude-3-opus" + ) assert credentials2 is not None assert credentials2["api_key"] == "sk-ant-wildcard-456" assert credentials2["custom_llm_provider"] == "anthropic" # Test with non-matching model - credentials3 = router.get_deployment_credentials_with_provider(model_id="vertex_ai/gemini-pro") + credentials3 = router.get_deployment_credentials_with_provider( + model_id="vertex_ai/gemini-pro" + ) assert credentials3 is None diff --git a/tests/local_testing/test_timeout.py b/tests/local_testing/test_timeout.py index 4128a595d7..e2ec9ba9a0 100644 --- a/tests/local_testing/test_timeout.py +++ b/tests/local_testing/test_timeout.py @@ -111,8 +111,8 @@ def test_hanging_request_azure(): "model_name": "azure-gpt", "litellm_params": { "model": "azure/gpt-4.1-mini", - "api_base": os.environ["AZURE_API_BASE"], - "api_key": os.environ["AZURE_API_KEY"], + "api_base": os.environ["AZURE_AI_API_BASE"], + "api_key": os.environ["AZURE_AI_API_KEY"], }, }, { @@ -175,8 +175,8 @@ def test_hanging_request_openai(): "model_name": "azure-gpt", "litellm_params": { "model": "azure/gpt-4.1-mini", - "api_base": os.environ["AZURE_API_BASE"], - "api_key": os.environ["AZURE_API_KEY"], + "api_base": os.environ["AZURE_AI_API_BASE"], + "api_key": os.environ["AZURE_AI_API_KEY"], }, }, { diff --git a/tests/local_testing/test_tpm_rpm_routing_v2.py b/tests/local_testing/test_tpm_rpm_routing_v2.py index a3218bc987..c7449ef6e2 100644 --- a/tests/local_testing/test_tpm_rpm_routing_v2.py +++ b/tests/local_testing/test_tpm_rpm_routing_v2.py @@ -39,9 +39,7 @@ from create_mock_standard_logging_payload import create_standard_logging_payload def test_tpm_rpm_updated(): test_cache = DualCache() - lowest_tpm_logger = LowestTPMLoggingHandler( - router_cache=test_cache - ) + lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache) model_group = "gpt-3.5-turbo" deployment_id = "1234" deployment = "azure/gpt-4.1-mini" @@ -108,9 +106,7 @@ def test_get_available_deployments(): "model_info": {"id": "5678"}, }, ] - lowest_tpm_logger = LowestTPMLoggingHandler( - router_cache=test_cache - ) + lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache) model_group = "gpt-3.5-turbo" ## DEPLOYMENT 1 ## total_tokens = 50 @@ -669,9 +665,7 @@ def test_return_potential_deployments(): """ test_cache = DualCache() - lowest_tpm_logger = LowestTPMLoggingHandler( - router_cache=test_cache - ) + lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache) args: Dict = { "healthy_deployments": [ @@ -731,8 +725,8 @@ async def test_tpm_rpm_routing_model_name_checks(): "model_name": "gpt-3.5-turbo", "litellm_params": { "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_key": os.getenv("AZURE_AI_API_KEY"), + "api_base": os.getenv("AZURE_AI_API_BASE"), "mock_response": "Hey, how's it going?", }, } diff --git a/tests/logging_callback_tests/test_alerting.py b/tests/logging_callback_tests/test_alerting.py index f77cbcb4bb..2278b6ff93 100644 --- a/tests/logging_callback_tests/test_alerting.py +++ b/tests/logging_callback_tests/test_alerting.py @@ -641,7 +641,7 @@ async def test_outage_alerting_called( "model_name": model, "litellm_params": { "model": model, - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_base": api_base, "vertex_location": vertex_location, "vertex_project": vertex_project, @@ -749,7 +749,7 @@ async def test_region_outage_alerting_called( "model_name": model, "litellm_params": { "model": model, - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_base": api_base, "vertex_location": vertex_location, "vertex_project": vertex_project, @@ -760,7 +760,7 @@ async def test_region_outage_alerting_called( "model_name": model, "litellm_params": { "model": model, - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_base": api_base, "vertex_location": vertex_location, "vertex_project": "vertex_project-2", @@ -868,7 +868,9 @@ async def test_langfuse_trace_id(): returned_trace_id = trace_url.split("/")[-1] - assert returned_trace_id == litellm_logging_obj._get_trace_id(service_name="langfuse") + assert returned_trace_id == litellm_logging_obj._get_trace_id( + service_name="langfuse" + ) @pytest.mark.asyncio @@ -1007,7 +1009,7 @@ async def test_soft_budget_alerts(): # Verify alert message contains correct percentage alert_message = mock_send_alert.call_args[1]["message"] - + print("GOT MESSAGE\n\n", alert_message) expected_message = ( @@ -1077,10 +1079,10 @@ key_no_max_budget_info = CallInfo( async def test_soft_budget_alerts_webhook(entity_info): """ Tests that soft budget alerts are triggered for different entity types. - + Tests: - Key with max budget - - Team + - Team - User - Key without max budget """ @@ -1097,7 +1099,7 @@ async def test_soft_budget_alerts_webhook(entity_info): # Verify the webhook event call_args = mock_send_alert.call_args[1] logged_webhook_event: WebhookEvent = call_args["user_info"] - + # Validate the webhook event has all expected fields assert logged_webhook_event.spend == entity_info.spend assert logged_webhook_event.soft_budget == entity_info.soft_budget @@ -1106,10 +1108,3 @@ async def test_soft_budget_alerts_webhook(entity_info): assert logged_webhook_event.user_email == entity_info.user_email assert logged_webhook_event.key_alias == entity_info.key_alias assert logged_webhook_event.event_group == entity_info.event_group - - - - - - - \ No newline at end of file diff --git a/tests/logging_callback_tests/test_custom_callback_router.py b/tests/logging_callback_tests/test_custom_callback_router.py index f6c7f2fa02..63d8b14f48 100644 --- a/tests/logging_callback_tests/test_custom_callback_router.py +++ b/tests/logging_callback_tests/test_custom_callback_router.py @@ -267,7 +267,10 @@ class CompletionCustomHandler( try: print("CompletionCustomHandler.async_log_success_event, kwargs: ", kwargs) self.states.append("async_success") - print("############### CompletionCustomHandler async success, kwargs: ", kwargs) + print( + "############### CompletionCustomHandler async success, kwargs: ", + kwargs, + ) ## START TIME assert isinstance(start_time, datetime) ## END TIME @@ -396,9 +399,9 @@ async def test_async_chat_azure(): "model_name": "gpt-4.1-nano", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "model_info": {"base_model": "azure/gpt-4.1-mini"}, "tpm": 240000, @@ -443,7 +446,7 @@ async def test_async_chat_azure(): "model": "azure/gpt-4o-new-test", "api_key": "my-bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -483,9 +486,9 @@ async def test_async_embedding_azure(): "model_name": "azure-embedding-model", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/text-embedding-ada-002", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -506,7 +509,7 @@ async def test_async_embedding_azure(): "model": "azure/text-embedding-ada-002", "api_key": "my-bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -549,7 +552,7 @@ async def test_async_chat_azure_with_fallbacks(): "model": "azure/gpt-4.1-mini", "api_key": "my-bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -608,9 +611,9 @@ async def test_async_completion_azure_caching(): "model_name": "gpt-4.1-nano", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -664,23 +667,23 @@ async def test_async_completion_azure_caching_streaming(): ) litellm.callbacks = [customHandler_caching] unique_time = uuid.uuid4() - + # Use Router instead of direct litellm.acompletion to get router-specific metadata model_list = [ { "model_name": "gpt-4.1-nano", "litellm_params": { "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, }, ] router = Router(model_list=model_list) - + response1 = await router.acompletion( model="gpt-4.1-nano", messages=[ @@ -725,12 +728,16 @@ async def test_async_embedding_azure_caching(): port=os.environ["REDIS_PORT"], password=os.environ["REDIS_PASSWORD"], ) - router = Router(model_list=[{ - "model_name": "text-embedding-ada-002", - "litellm_params": { - "model": "openai/text-embedding-ada-002", - }, - }]) + router = Router( + model_list=[ + { + "model_name": "text-embedding-ada-002", + "litellm_params": { + "model": "openai/text-embedding-ada-002", + }, + } + ] + ) litellm.callbacks = [customHandler_caching] unique_time = time.time() response1 = await router.aembedding( @@ -818,4 +825,3 @@ async def test_rate_limit_error_callback(): assert "original_model_group" in mock_client.call_args.kwargs assert mock_client.call_args.kwargs["original_model_group"] == "my-test-gpt" - diff --git a/tests/old_proxy_tests/tests/load_test_q.py b/tests/old_proxy_tests/tests/load_test_q.py index a8f2c0a322..89137c306a 100644 --- a/tests/old_proxy_tests/tests/load_test_q.py +++ b/tests/old_proxy_tests/tests/load_test_q.py @@ -26,7 +26,7 @@ config = { "model_name": "gpt-3.5-turbo", "litellm_params": { "model": "azure/gpt-4.1-mini", - "api_key": os.environ["AZURE_API_KEY"], + "api_key": os.environ["AZURE_AI_API_KEY"], "api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/", "api_version": "2023-07-01-preview", }, @@ -34,7 +34,7 @@ config = { ] } print("STARTING LOAD TEST Q") -print(os.environ["AZURE_API_KEY"]) +print(os.environ["AZURE_AI_API_KEY"]) response = requests.post( url=f"{base_url}/key/generate", diff --git a/tests/pass_through_unit_tests/test_unit_test_passthrough_router.py b/tests/pass_through_unit_tests/test_unit_test_passthrough_router.py index 8e016b68d0..3f62f55791 100644 --- a/tests/pass_through_unit_tests/test_unit_test_passthrough_router.py +++ b/tests/pass_through_unit_tests/test_unit_test_passthrough_router.py @@ -102,7 +102,7 @@ class TestPassthroughEndpointRouter(unittest.TestCase): mock_get_secret.return_value = "env_azure_key" result = self.router.get_credentials("azure", None) self.assertEqual(result, "env_azure_key") - mock_get_secret.assert_called_once_with("AZURE_API_KEY") + mock_get_secret.assert_called_once_with("AZURE_AI_API_KEY") def test_default_env_variable_method(self): """ diff --git a/tests/proxy_unit_tests/example_config_yaml/azure_config.yaml b/tests/proxy_unit_tests/example_config_yaml/azure_config.yaml index 0a015aefde..05ba0c9bf5 100644 --- a/tests/proxy_unit_tests/example_config_yaml/azure_config.yaml +++ b/tests/proxy_unit_tests/example_config_yaml/azure_config.yaml @@ -4,12 +4,12 @@ model_list: model: azure/gpt-4.1-mini api_base: https://openai-gpt-4-test-v-1.openai.azure.com/ api_version: "2023-05-15" - api_key: os.environ/AZURE_API_KEY + api_key: os.environ/AZURE_AI_API_KEY tpm: 20_000 - model_name: gpt-4-team2 litellm_params: model: azure/gpt-4 - api_key: os.environ/AZURE_API_KEY + api_key: os.environ/AZURE_AI_API_KEY api_base: https://openai-gpt-4-test-v-2.openai.azure.com/ tpm: 100_000 diff --git a/tests/proxy_unit_tests/test_configs/test_bad_config.yaml b/tests/proxy_unit_tests/test_configs/test_bad_config.yaml index 4a70886a93..4bc4c7cc54 100644 --- a/tests/proxy_unit_tests/test_configs/test_bad_config.yaml +++ b/tests/proxy_unit_tests/test_configs/test_bad_config.yaml @@ -6,16 +6,16 @@ model_list: - model_name: working-azure-gpt-3.5-turbo litellm_params: model: azure/gpt-4.1-mini - api_base: os.environ/AZURE_API_BASE - api_key: os.environ/AZURE_API_KEY + api_base: os.environ/AZURE_AI_API_BASE + api_key: os.environ/AZURE_AI_API_KEY - model_name: azure-gpt-3.5-turbo litellm_params: model: azure/gpt-4.1-mini - api_base: os.environ/AZURE_API_BASE + api_base: os.environ/AZURE_AI_API_BASE api_key: bad-key - model_name: azure-embedding litellm_params: model: azure/text-embedding-ada-002 - api_base: os.environ/AZURE_API_BASE + api_base: os.environ/AZURE_AI_API_BASE api_key: bad-key \ No newline at end of file diff --git a/tests/proxy_unit_tests/test_configs/test_cloudflare_azure_with_cache_config.yaml b/tests/proxy_unit_tests/test_configs/test_cloudflare_azure_with_cache_config.yaml index 9902835618..24240008fe 100644 --- a/tests/proxy_unit_tests/test_configs/test_cloudflare_azure_with_cache_config.yaml +++ b/tests/proxy_unit_tests/test_configs/test_cloudflare_azure_with_cache_config.yaml @@ -3,7 +3,7 @@ model_list: litellm_params: model: azure/gpt-4.1-mini api_base: https://gateway.ai.cloudflare.com/v1/0399b10e77ac6668c80404a5ff49eb37/litellm-test/azure-openai/openai-gpt-4-test-v-1 - api_key: os.environ/AZURE_API_KEY + api_key: os.environ/AZURE_AI_API_KEY api_version: 2023-07-01-preview litellm_settings: diff --git a/tests/proxy_unit_tests/test_configs/test_config_no_auth.yaml b/tests/proxy_unit_tests/test_configs/test_config_no_auth.yaml index cdc447a5ee..f489621704 100644 --- a/tests/proxy_unit_tests/test_configs/test_config_no_auth.yaml +++ b/tests/proxy_unit_tests/test_configs/test_config_no_auth.yaml @@ -11,7 +11,7 @@ model_list: model_name: azure-model - litellm_params: api_base: https://gateway.ai.cloudflare.com/v1/0399b10e77ac6668c80404a5ff49eb37/litellm-test/azure-openai/openai-gpt-4-test-v-1 - api_key: os.environ/AZURE_API_KEY + api_key: os.environ/AZURE_AI_API_KEY model: azure/gpt-4.1-mini model_name: azure-cloudflare-model - litellm_params: @@ -49,8 +49,8 @@ model_list: id: 79fc75bf-8e1b-47d5-8d24-9365a854af03 model_name: test_openai_models - litellm_params: - api_base: os.environ/AZURE_API_BASE - api_key: os.environ/AZURE_API_KEY + api_base: os.environ/AZURE_AI_API_BASE + api_key: os.environ/AZURE_AI_API_KEY api_version: 2023-07-01-preview model: azure/text-embedding-ada-002 model_info: @@ -94,16 +94,16 @@ model_list: mode: image_generation model_name: dall-e-3 - litellm_params: - api_base: os.environ/AZURE_API_BASE - api_key: os.environ/AZURE_API_KEY + api_base: os.environ/AZURE_AI_API_BASE + api_key: os.environ/AZURE_AI_API_KEY api_version: 2023-06-01-preview model: azure/ model_info: mode: image_generation model_name: dall-e-2 - litellm_params: - api_base: os.environ/AZURE_API_BASE - api_key: os.environ/AZURE_API_KEY + api_base: os.environ/AZURE_AI_API_BASE + api_key: os.environ/AZURE_AI_API_KEY api_version: 2023-07-01-preview model: azure/text-embedding-ada-002 model_info: diff --git a/tests/proxy_unit_tests/test_proxy_pass_user_config.py b/tests/proxy_unit_tests/test_proxy_pass_user_config.py index 4828014e33..6beb86eca7 100644 --- a/tests/proxy_unit_tests/test_proxy_pass_user_config.py +++ b/tests/proxy_unit_tests/test_proxy_pass_user_config.py @@ -54,8 +54,9 @@ def client_no_auth(): @pytest.mark.skipif( - os.environ.get("AZURE_API_KEY") is None or os.environ.get("OPENAI_API_KEY") is None, - reason="AZURE_API_KEY or OPENAI_API_KEY not set - skipping integration test" + os.environ.get("AZURE_AI_API_KEY") is None + or os.environ.get("OPENAI_API_KEY") is None, + reason="AZURE_AI_API_KEY or OPENAI_API_KEY not set - skipping integration test", ) def test_chat_completion(client_no_auth): global headers @@ -69,9 +70,9 @@ def test_chat_completion(client_no_auth): model_name="user-azure-instance", litellm_params=CompletionRequest( model="azure/gpt-4.1-mini", - api_key=os.getenv("AZURE_API_KEY"), + api_key=os.getenv("AZURE_AI_API_KEY"), api_version=os.getenv("AZURE_API_VERSION"), - api_base=os.getenv("AZURE_API_BASE"), + api_base=os.getenv("AZURE_AI_API_BASE"), timeout=10, ), tpm=240000, diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 61a2f3055a..59b9297d1b 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -119,7 +119,7 @@ def fake_env_vars(monkeypatch): # Set some fake environment variables monkeypatch.setenv("OPENAI_API_KEY", "fake_openai_api_key") monkeypatch.setenv("OPENAI_API_BASE", "http://fake-openai-api-base") - monkeypatch.setenv("AZURE_API_BASE", "http://fake-azure-api-base") + monkeypatch.setenv("AZURE_AI_API_BASE", "http://fake-azure-api-base") monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake_azure_openai_api_key") monkeypatch.setenv("AZURE_SWEDEN_API_BASE", "http://fake-azure-sweden-api-base") monkeypatch.setenv("REDIS_HOST", "localhost") @@ -178,7 +178,7 @@ def test_chat_completion(mock_acompletion, client_no_auth): def test_chat_completion_malformed_messages_returns_400(client_no_auth): """ Test that malformed messages (strings instead of dicts) return 400 instead of 500. - + This test verifies that when a client sends messages as raw strings instead of {role, content} objects, LiteLLM returns a 400 invalid_request_error instead of a 500 Internal Server Error. @@ -188,33 +188,41 @@ def test_chat_completion_malformed_messages_returns_400(client_no_auth): # Test data with malformed messages (string instead of dict) test_data = { "model": "gpt-3.5-turbo", - "messages": ["hi how are you"], # Invalid: should be [{"role": "user", "content": "hi how are you"}] + "messages": [ + "hi how are you" + ], # Invalid: should be [{"role": "user", "content": "hi how are you"}] } print("testing proxy server with malformed messages") - response = client_no_auth.post("/v1/chat/completions", json=test_data, headers=headers) - + response = client_no_auth.post( + "/v1/chat/completions", json=test_data, headers=headers + ) + print(f"response status: {response.status_code}") print(f"response text: {response.text}") - + # Should return 400, not 500 - assert response.status_code == 400, f"Expected 400, got {response.status_code}. Response: {response.text}" - + assert ( + response.status_code == 400 + ), f"Expected 400, got {response.status_code}. Response: {response.text}" + # Verify error format result = response.json() assert "error" in result, "Response should contain 'error' key" error = result["error"] - + # Verify error type and message - assert error.get("type") == "invalid_request_error" or error.get("type") is None, \ - f"Expected invalid_request_error or None, got {error.get('type')}" - assert error.get("code") == "400" or error.get("code") == 400, \ - f"Expected code 400, got {error.get('code')}" - + assert ( + error.get("type") == "invalid_request_error" or error.get("type") is None + ), f"Expected invalid_request_error or None, got {error.get('type')}" + assert ( + error.get("code") == "400" or error.get("code") == 400 + ), f"Expected code 400, got {error.get('code')}" + # Error message should indicate invalid request format error_message = error.get("message", "") assert len(error_message) > 0, "Error message should not be empty" - + except Exception as e: pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}") @@ -342,7 +350,7 @@ def test_chat_completion_forward_llm_provider_auth_headers( """ Test that LLM provider auth headers (x-api-key, x-goog-api-key) are forwarded when forward_llm_provider_auth_headers=True. - + This allows clients to send their own LLM provider API keys through the proxy. """ try: @@ -351,7 +359,7 @@ def test_chat_completion_forward_llm_provider_auth_headers( gs["forward_client_headers_to_llm_api"] = True gs["forward_llm_provider_auth_headers"] = forward_llm_auth_headers setattr(litellm.proxy.proxy_server, "general_settings", gs) - + # Test data test_data = { "model": "gpt-3.5-turbo", @@ -360,7 +368,7 @@ def test_chat_completion_forward_llm_provider_auth_headers( ], "max_tokens": 10, } - + # Headers including LLM provider auth request_headers = { "Authorization": "Bearer sk-proxy-auth-123", # Proxy auth (should be stripped) @@ -368,17 +376,17 @@ def test_chat_completion_forward_llm_provider_auth_headers( "x-goog-api-key": "google-api-key-123", # Google API key "X-Custom-Header": "custom-value", # Custom header (should be forwarded) } - + # Make request response = client_no_auth.post( "/v1/chat/completions", json=test_data, headers=request_headers ) - + assert response.status_code == 200 - + # Check forwarded headers forwarded_headers = mock_acompletion.call_args.kwargs.get("headers", {}) - + if forward_llm_auth_headers: # LLM provider auth headers should be forwarded assert "x-api-key" in forwarded_headers @@ -389,19 +397,23 @@ def test_chat_completion_forward_llm_provider_auth_headers( # LLM provider auth headers should be stripped assert "x-api-key" not in forwarded_headers assert "x-goog-api-key" not in forwarded_headers - + # Custom headers should always be forwarded (when forward_client_headers_to_llm_api=True) assert "x-custom-header" in forwarded_headers assert forwarded_headers["x-custom-header"] == "custom-value" - + # Proxy Authorization should never be forwarded assert "authorization" not in forwarded_headers - - print(f"✓ Test passed with forward_llm_provider_auth_headers={forward_llm_auth_headers}") + + print( + f"✓ Test passed with forward_llm_provider_auth_headers={forward_llm_auth_headers}" + ) print(f" Forwarded headers: {list(forwarded_headers.keys())}") - + except Exception as e: - pytest.fail(f"Test failed with forward_llm_auth_headers={forward_llm_auth_headers}: {str(e)}") + pytest.fail( + f"Test failed with forward_llm_auth_headers={forward_llm_auth_headers}: {str(e)}" + ) finally: # Clean up gs = getattr(litellm.proxy.proxy_server, "general_settings") @@ -2406,9 +2418,7 @@ async def test_run_background_health_check_reflects_llm_model_list(monkeypatch): test_model_list_2 = [{"model_name": "model-b"}] called_model_lists = [] - async def fake_perform_health_check( - model_list, details, max_concurrency=None - ): + async def fake_perform_health_check(model_list, details, max_concurrency=None): called_model_lists.append(copy.deepcopy(model_list)) return (["healthy"], ["unhealthy"]) @@ -2452,13 +2462,14 @@ async def test_background_health_check_skip_disabled_models(monkeypatch): test_model_list = [ {"model_name": "model-a"}, - {"model_name": "model-b", "model_info": {"disable_background_health_check": True}}, + { + "model_name": "model-b", + "model_info": {"disable_background_health_check": True}, + }, ] called_model_lists = [] - async def fake_perform_health_check( - model_list, details, max_concurrency=None - ): + async def fake_perform_health_check(model_list, details, max_concurrency=None): called_model_lists.append(copy.deepcopy(model_list)) return (["healthy"], []) @@ -2500,15 +2511,15 @@ def test_get_timeout_from_request(): @pytest.mark.parametrize( "ui_exists, ui_has_content", [ - (True, True), # UI path exists and has content + (True, True), # UI path exists and has content (True, False), # UI path exists but is empty - (False, False), # UI path doesn't exist + (False, False), # UI path doesn't exist ], ) def test_non_root_ui_path_logic(monkeypatch, tmp_path, ui_exists, ui_has_content): """ Test the non-root Docker UI path detection logic. - + Tests that when LITELLM_NON_ROOT is set to "true": - If UI path exists and has content, it should be used - If UI path doesn't exist or is empty, proper error logging occurs @@ -2516,44 +2527,54 @@ def test_non_root_ui_path_logic(monkeypatch, tmp_path, ui_exists, ui_has_content import tempfile import shutil from unittest.mock import MagicMock - + # Create a temporary directory to act as /tmp/litellm_ui test_ui_path = tmp_path / "litellm_ui" - + if ui_exists: test_ui_path.mkdir(parents=True, exist_ok=True) if ui_has_content: # Create some dummy files to simulate built UI (test_ui_path / "index.html").write_text("") (test_ui_path / "app.js").write_text("console.log('test');") - + # Mock the environment variable and os.path operations monkeypatch.setenv("LITELLM_NON_ROOT", "true") - + # Create a mock logger to capture log messages mock_logger = MagicMock() - + # We need to reimport or reload the relevant code section # Since this is module-level code, we'll test the logic directly ui_path = None non_root_ui_path = str(test_ui_path) - + # Simulate the logic from proxy_server.py lines 909-920 if os.getenv("LITELLM_NON_ROOT", "").lower() == "true": if os.path.exists(non_root_ui_path) and os.listdir(non_root_ui_path): - mock_logger.info(f"Using pre-built UI for non-root Docker: {non_root_ui_path}") - mock_logger.info(f"UI files found: {len(os.listdir(non_root_ui_path))} items") + mock_logger.info( + f"Using pre-built UI for non-root Docker: {non_root_ui_path}" + ) + mock_logger.info( + f"UI files found: {len(os.listdir(non_root_ui_path))} items" + ) ui_path = non_root_ui_path else: - mock_logger.error(f"UI not found at {non_root_ui_path}. UI will not be available.") - mock_logger.error(f"Path exists: {os.path.exists(non_root_ui_path)}, Has content: {os.path.exists(non_root_ui_path) and bool(os.listdir(non_root_ui_path))}") - + mock_logger.error( + f"UI not found at {non_root_ui_path}. UI will not be available." + ) + mock_logger.error( + f"Path exists: {os.path.exists(non_root_ui_path)}, Has content: {os.path.exists(non_root_ui_path) and bool(os.listdir(non_root_ui_path))}" + ) + # Verify behavior based on test parameters if ui_exists and ui_has_content: # UI should be found and used assert ui_path == non_root_ui_path assert mock_logger.info.call_count == 2 - mock_logger.info.assert_any_call(f"Using pre-built UI for non-root Docker: {non_root_ui_path}") + mock_logger.info.assert_any_call( + f"Using pre-built UI for non-root Docker: {non_root_ui_path}" + ) # Verify the second info call mentions the number of items info_calls = [call[0][0] for call in mock_logger.info.call_args_list] assert any("UI files found:" in call and "items" in call for call in info_calls) @@ -2562,7 +2583,9 @@ def test_non_root_ui_path_logic(monkeypatch, tmp_path, ui_exists, ui_has_content # UI should not be found, error should be logged assert ui_path is None assert mock_logger.error.call_count == 2 - mock_logger.error.assert_any_call(f"UI not found at {non_root_ui_path}. UI will not be available.") + mock_logger.error.assert_any_call( + f"UI not found at {non_root_ui_path}. UI will not be available." + ) # Verify the second error call has path existence info error_calls = [call[0][0] for call in mock_logger.error.call_args_list] assert any("Path exists:" in call for call in error_calls) @@ -2574,17 +2597,17 @@ async def test_get_config_callbacks_with_all_types(client_no_auth): """ Test that /get/config/callbacks returns all three callback types: - success_callback with type="success" - - failure_callback with type="failure" + - failure_callback with type="failure" - callbacks (success_and_failure) with type="success_and_failure" """ from litellm.proxy.proxy_server import ProxyConfig - + # Create a mock config with all three callback types mock_config_data = { "litellm_settings": { "success_callback": ["langfuse", "braintrust"], "failure_callback": ["sentry"], - "callbacks": ["otel", "langsmith"] + "callbacks": ["otel", "langsmith"], }, "environment_variables": { "LANGFUSE_PUBLIC_KEY": "test-public-key", @@ -2595,51 +2618,53 @@ async def test_get_config_callbacks_with_all_types(client_no_auth): "OTEL_ENDPOINT": "http://localhost:4317", "LANGSMITH_API_KEY": "test-langsmith-key", }, - "general_settings": {} + "general_settings": {}, } - + proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config") - + with patch.object( proxy_config, "get_config", new=AsyncMock(return_value=mock_config_data) ): response = client_no_auth.get("/get/config/callbacks") - + assert response.status_code == 200 result = response.json() - + # Verify response structure assert "status" in result assert result["status"] == "success" assert "callbacks" in result - + callbacks = result["callbacks"] - + # Verify we have all 5 callbacks (2 success + 1 failure + 2 success_and_failure) assert len(callbacks) == 5 - + # Group callbacks by type success_callbacks = [cb for cb in callbacks if cb.get("type") == "success"] failure_callbacks = [cb for cb in callbacks if cb.get("type") == "failure"] - success_and_failure_callbacks = [cb for cb in callbacks if cb.get("type") == "success_and_failure"] - + success_and_failure_callbacks = [ + cb for cb in callbacks if cb.get("type") == "success_and_failure" + ] + # Verify all callbacks have required fields for callback in callbacks: assert "name" in callback assert "variables" in callback assert "type" in callback assert callback["type"] in ["success", "failure", "success_and_failure"] - + # Verify success callbacks assert len(success_callbacks) == 2 success_names = [cb["name"] for cb in success_callbacks] assert "langfuse" in success_names assert "braintrust" in success_names - + # Verify failure callbacks assert len(failure_callbacks) == 1 assert failure_callbacks[0]["name"] == "sentry" - + # Verify success_and_failure callbacks assert len(success_and_failure_callbacks) == 2 success_and_failure_names = [cb["name"] for cb in success_and_failure_callbacks] @@ -2654,13 +2679,13 @@ async def test_get_config_callbacks_environment_variables(client_no_auth): for each callback type. Values are returned as-is from the config (no decryption). """ from litellm.proxy.proxy_server import ProxyConfig - + # Create a mock config with callbacks and their env vars mock_config_data = { "litellm_settings": { "success_callback": ["langfuse"], "failure_callback": [], - "callbacks": ["otel"] + "callbacks": ["otel"], }, "environment_variables": { "LANGFUSE_PUBLIC_KEY": "test-public-key", @@ -2670,21 +2695,21 @@ async def test_get_config_callbacks_environment_variables(client_no_auth): "OTEL_ENDPOINT": "http://localhost:4317", "OTEL_HEADERS": "key=value", }, - "general_settings": {} + "general_settings": {}, } - + proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config") - + with patch.object( proxy_config, "get_config", new=AsyncMock(return_value=mock_config_data) ): response = client_no_auth.get("/get/config/callbacks") - + assert response.status_code == 200 result = response.json() - + callbacks = result["callbacks"] - + # Find langfuse callback (success type) langfuse_callback = next( (cb for cb in callbacks if cb["name"] == "langfuse"), None @@ -2692,7 +2717,7 @@ async def test_get_config_callbacks_environment_variables(client_no_auth): assert langfuse_callback is not None assert langfuse_callback["type"] == "success" assert "variables" in langfuse_callback - + # Verify langfuse env vars are present (values returned as-is, no decryption) langfuse_vars = langfuse_callback["variables"] assert "LANGFUSE_PUBLIC_KEY" in langfuse_vars @@ -2701,15 +2726,13 @@ async def test_get_config_callbacks_environment_variables(client_no_auth): assert langfuse_vars["LANGFUSE_SECRET_KEY"] == "test-secret-key" assert "LANGFUSE_HOST" in langfuse_vars assert langfuse_vars["LANGFUSE_HOST"] == "https://cloud.langfuse.com" - + # Find otel callback (success_and_failure type) - otel_callback = next( - (cb for cb in callbacks if cb["name"] == "otel"), None - ) + otel_callback = next((cb for cb in callbacks if cb["name"] == "otel"), None) assert otel_callback is not None assert otel_callback["type"] == "success_and_failure" assert "variables" in otel_callback - + # Verify otel env vars are present otel_vars = otel_callback["variables"] assert "OTEL_EXPORTER" in otel_vars diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index 8c6096a9fc..9ed801b360 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -460,7 +460,7 @@ async def test_ensure_initialize_azure_sdk_client_always_used(call_type): "api_key": "test-api-key", "api_version": os.getenv("AZURE_API_VERSION", "2023-05-15"), "api_base": os.getenv( - "AZURE_API_BASE", "https://test.openai.azure.com" + "AZURE_AI_API_BASE", "https://test.openai.azure.com" ), }, } @@ -670,7 +670,7 @@ async def test_ensure_initialize_azure_sdk_client_always_used_azure_text(call_ty "api_key": "test-api-key", "api_version": os.getenv("AZURE_API_VERSION", "2023-05-15"), "api_base": os.getenv( - "AZURE_API_BASE", "https://test.openai.azure.com" + "AZURE_AI_API_BASE", "https://test.openai.azure.com" ), }, } diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py index 7880683168..d66798a572 100644 --- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py @@ -3,6 +3,7 @@ Tests for Azure AI Anthropic CountTokens transformation. Verifies that the CountTokens API uses the correct authentication headers. """ + import os import sys @@ -40,7 +41,7 @@ class TestAzureAIAnthropicCountTokensConfig: assert headers["anthropic-version"] == "2023-06-01" assert "anthropic-beta" in headers - def test_get_required_headers_includes_azure_api_key(self): + def test_get_required_headers_includes_AZURE_AI_API_KEY(self): """ Test that get_required_headers includes Azure api-key header. diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py index d3214a8801..802868aa7a 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_transformation.py @@ -336,7 +336,9 @@ class TestOpenAIResponsesAPIConfig: ) assert isinstance(result, ImageGenerationPartialImageEvent) - assert result.type == ResponsesAPIStreamEvents.IMAGE_GENERATION_PARTIAL_IMAGE + assert ( + result.type == ResponsesAPIStreamEvents.IMAGE_GENERATION_PARTIAL_IMAGE + ) assert result.partial_image_index == idx assert result.b64_json == chunk["b64_json"] @@ -689,9 +691,7 @@ class TestTransformListInputItemsRequest: def test_openai_transform_compact_response_api_request_query_params_preserved(self): """Test compact URL construction preserves query params and appends path.""" # Setup - azure_style_api_base = ( - "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview" - ) + azure_style_api_base = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview" # Execute url, data = self.openai_config.transform_compact_response_api_request( @@ -731,12 +731,12 @@ class TestTransformListInputItemsRequest: def test_azure_transform_list_input_items_request_minimal(self): """Test Azure implementation with minimal parameters""" # Setup - azure_api_base = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview" + AZURE_AI_API_BASE = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview" # Execute url, params = self.azure_config.transform_list_input_items_request( response_id=self.response_id, - api_base=azure_api_base, + api_base=AZURE_AI_API_BASE, litellm_params=self.litellm_params, headers=self.headers, ) @@ -749,12 +749,12 @@ class TestTransformListInputItemsRequest: def test_azure_transform_list_input_items_request_url_construction(self): """Test Azure implementation URL construction with response_id in path""" # Setup - azure_api_base = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview" + AZURE_AI_API_BASE = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview" # Execute url, params = self.azure_config.transform_list_input_items_request( response_id=self.response_id, - api_base=azure_api_base, + api_base=AZURE_AI_API_BASE, litellm_params=self.litellm_params, headers=self.headers, ) @@ -768,12 +768,12 @@ class TestTransformListInputItemsRequest: def test_azure_transform_list_input_items_request_with_all_params(self): """Test Azure implementation with all optional parameters""" # Setup - azure_api_base = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview" + AZURE_AI_API_BASE = "https://test.openai.azure.com/openai/responses?api-version=2024-05-01-preview" # Execute url, params = self.azure_config.transform_list_input_items_request( response_id=self.response_id, - api_base=azure_api_base, + api_base=AZURE_AI_API_BASE, litellm_params=self.litellm_params, headers=self.headers, after="cursor_after_123", @@ -1128,9 +1128,9 @@ class TestPhaseParameter: phase = getattr(output_item, "phase", None) expected = "commentary" if idx == 0 else "final_answer" - assert phase == expected, ( - f"output[{idx}] phase={phase!r}, expected {expected!r}" - ) + assert ( + phase == expected + ), f"output[{idx}] phase={phase!r}, expected {expected!r}" def test_streaming_output_item_done_preserves_phase(self): """OutputItemDoneEvent must preserve phase on its item.""" diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 6e000f9c0e..ca70b4bfa8 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -28,7 +28,7 @@ def test_get_api_key(): assert get_api_key( custom_litellm_key_header=None, api_key=bearer_token, - azure_api_key_header=None, + AZURE_AI_API_KEY_header=None, anthropic_api_key_header=None, google_ai_studio_api_key_header=None, azure_apim_header=None, @@ -59,7 +59,7 @@ def test_get_api_key_with_custom_litellm_key_header( assert get_api_key( custom_litellm_key_header=custom_litellm_key_header, api_key=None, - azure_api_key_header=None, + AZURE_AI_API_KEY_header=None, anthropic_api_key_header=None, google_ai_studio_api_key_header=None, azure_apim_header=None, @@ -334,7 +334,6 @@ async def test_proxy_admin_expired_key_from_cache(): "litellm.proxy.auth.user_api_key_auth._delete_cache_key_object", new_callable=AsyncMock, ) as mock_delete_cache: - mock_get_key_object.return_value = expired_token # Set attributes on proxy_server module (these are imported inside _user_api_key_auth_builder) @@ -372,7 +371,7 @@ async def test_proxy_admin_expired_key_from_cache(): await _user_api_key_auth_builder( request=request, api_key=f"Bearer {api_key}", # Add Bearer prefix - azure_api_key_header="", + AZURE_AI_API_KEY_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, azure_apim_header=None, @@ -609,7 +608,6 @@ class TestJWTOAuth2Coexistence: "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", new_callable=AsyncMock, ) as mock_jwt_auth: - litellm.proxy.proxy_server.jwt_handler.update_environment( prisma_client=None, user_api_key_cache=DualCache(), @@ -674,7 +672,6 @@ class TestJWTOAuth2Coexistence: new_callable=AsyncMock, return_value=mock_jwt_result, ) as mock_jwt_auth: - litellm.proxy.proxy_server.jwt_handler.update_environment( prisma_client=None, user_api_key_cache=DualCache(), @@ -726,7 +723,6 @@ class TestJWTOAuth2Coexistence: new_callable=AsyncMock, return_value=mock_oauth2_response, ) as mock_oauth2: - result = await user_api_key_auth( request=mock_request, api_key=f"Bearer {jwt_like_token}", @@ -846,7 +842,7 @@ async def test_user_api_key_auth_builder_no_blocking_calls(): await _user_api_key_auth_builder( request=request, api_key=f"Bearer {api_key}", - azure_api_key_header="", + AZURE_AI_API_KEY_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, azure_apim_header=None, diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index ed5b8cbd81..c6a03cf4ec 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -34,8 +34,8 @@ def llm_router() -> Router: "model_name": "azure-gpt-3-5-turbo", "litellm_params": { "model": "azure/chatgpt-v-2", - "api_key": "azure_api_key", - "api_base": "azure_api_base", + "api_key": "AZURE_AI_API_KEY", + "api_base": "AZURE_AI_API_BASE", "api_version": "azure_api_version", }, "model_info": { @@ -106,16 +106,18 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: Asserts 'create_file' is called with the correct arguments """ import litellm + import litellm.proxy.proxy_server as ps from litellm import Router from litellm.proxy._types import LitellmUserRoles - import litellm.proxy.proxy_server as ps from litellm.proxy.utils import ProxyLogging monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) # Mock create_file as an async function - mock_create_file = mocker.patch("litellm.files.main.create_file", new=mocker.AsyncMock()) + mock_create_file = mocker.patch( + "litellm.files.main.create_file", new=mocker.AsyncMock() + ) proxy_logging_obj = ProxyLogging( user_api_key_cache=DualCache(default_in_memory_ttl=1) @@ -127,7 +129,14 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: from litellm.llms.base_llm.files.transformation import BaseFileEndpoints class DummyManagedFiles(BaseFileEndpoints): - async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict): + async def acreate_file( + self, + llm_router, + create_file_request, + target_model_names_list, + litellm_parent_otel_span, + user_api_key_dict, + ): # Handle both dict and object forms of create_file_request if isinstance(create_file_request, dict): file_data = create_file_request.get("file") @@ -135,12 +144,12 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: else: file_data = create_file_request.file purpose_data = create_file_request.purpose - + # Call the mocked litellm.files.main.create_file to ensure asserts work await litellm.files.main.create_file( custom_llm_provider="azure", model="azure/chatgpt-v-2", - api_key="azure_api_key", + api_key="AZURE_AI_API_KEY", file=file_data[1], purpose=purpose_data, ) @@ -153,6 +162,7 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: ) # Return a dummy response object as needed by the test from litellm.types.llms.openai import OpenAIFileObject + return OpenAIFileObject( id="dummy-id", object="file", @@ -162,17 +172,21 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: purpose=purpose_data, status="uploaded", ) - + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): raise NotImplementedError("Not implemented for test") - + async def afile_list(self, purpose, litellm_parent_otel_span): raise NotImplementedError("Not implemented for test") - - async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_delete( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - - async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_content( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") # Manually add the hook to the proxy_hook_mapping @@ -214,7 +228,7 @@ def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: if ( kwargs.get("custom_llm_provider") == "azure" and kwargs.get("model") == "azure/chatgpt-v-2" - and kwargs.get("api_key") == "azure_api_key" + and kwargs.get("api_key") == "AZURE_AI_API_KEY" ): azure_call_found = True break @@ -245,8 +259,8 @@ def test_target_storage_invokes_storage_backend( """ Ensure target_storage is parsed and invokes the storage backend service. """ - from litellm.proxy._types import LitellmUserRoles import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) @@ -304,8 +318,8 @@ def test_target_storage_with_target_models( """ Ensure target_storage and target_model_names are parsed and passed through. """ - from litellm.proxy._types import LitellmUserRoles import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) @@ -611,7 +625,9 @@ def test_create_file_for_each_model( assert openai_call_found, "OpenAI call not found with expected parameters" -def test_create_file_with_expires_after(mocker: MockerFixture, monkeypatch, llm_router: Router): +def test_create_file_with_expires_after( + mocker: MockerFixture, monkeypatch, llm_router: Router +): """ Test that expires_after is properly parsed and passed through when creating a file """ @@ -624,18 +640,25 @@ def test_create_file_with_expires_after(mocker: MockerFixture, monkeypatch, llm_ proxy_logging_obj._add_proxy_hooks(llm_router) class DummyManagedFiles(BaseFileEndpoints): - async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict): + async def acreate_file( + self, + llm_router, + create_file_request, + target_model_names_list, + litellm_parent_otel_span, + user_api_key_dict, + ): # Verify expires_after is in the request if isinstance(create_file_request, dict): expires_after = create_file_request.get("expires_after") else: expires_after = getattr(create_file_request, "expires_after", None) - + # Verify expires_after was passed correctly assert expires_after is not None, "expires_after should be in the request" assert expires_after["anchor"] == "created_at" assert expires_after["seconds"] == 2592000 - + # Return a dummy response return OpenAIFileObject( id="file-abc123", @@ -646,17 +669,21 @@ def test_create_file_with_expires_after(mocker: MockerFixture, monkeypatch, llm_ purpose="fine-tune", status="uploaded", ) - + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): raise NotImplementedError("Not implemented for test") - + async def afile_list(self, purpose, litellm_parent_otel_span): raise NotImplementedError("Not implemented for test") - - async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_delete( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - - async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_content( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") proxy_logging_obj.proxy_hook_mapping["managed_files"] = DummyManagedFiles() @@ -688,7 +715,9 @@ def test_create_file_with_expires_after(mocker: MockerFixture, monkeypatch, llm_ assert result["purpose"] == "fine-tune" -def test_create_file_with_expires_after_missing_anchor(mocker: MockerFixture, monkeypatch, llm_router: Router): +def test_create_file_with_expires_after_missing_anchor( + mocker: MockerFixture, monkeypatch, llm_router: Router +): """ Test that an error is returned when expires_after[anchor] is missing """ @@ -717,10 +746,15 @@ def test_create_file_with_expires_after_missing_anchor(mocker: MockerFixture, mo assert response.status_code == 400 error_detail = response.json() - assert "expires_after" in error_detail["error"]["message"].lower() or "both" in error_detail["error"]["message"].lower() + assert ( + "expires_after" in error_detail["error"]["message"].lower() + or "both" in error_detail["error"]["message"].lower() + ) -def test_create_file_with_expires_after_missing_seconds(mocker: MockerFixture, monkeypatch, llm_router: Router): +def test_create_file_with_expires_after_missing_seconds( + mocker: MockerFixture, monkeypatch, llm_router: Router +): """ Test that an error is returned when expires_after[seconds] is missing """ @@ -749,10 +783,15 @@ def test_create_file_with_expires_after_missing_seconds(mocker: MockerFixture, m assert response.status_code == 400 error_detail = response.json() - assert "expires_after" in error_detail["error"]["message"].lower() or "both" in error_detail["error"]["message"].lower() + assert ( + "expires_after" in error_detail["error"]["message"].lower() + or "both" in error_detail["error"]["message"].lower() + ) -def test_create_file_with_expires_after_valid_values(mocker: MockerFixture, monkeypatch, llm_router: Router): +def test_create_file_with_expires_after_valid_values( + mocker: MockerFixture, monkeypatch, llm_router: Router +): """ Test that expires_after works with valid anchor and seconds values """ @@ -765,18 +804,25 @@ def test_create_file_with_expires_after_valid_values(mocker: MockerFixture, monk proxy_logging_obj._add_proxy_hooks(llm_router) class DummyManagedFiles(BaseFileEndpoints): - async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict): + async def acreate_file( + self, + llm_router, + create_file_request, + target_model_names_list, + litellm_parent_otel_span, + user_api_key_dict, + ): # Verify expires_after is in the request if isinstance(create_file_request, dict): expires_after = create_file_request.get("expires_after") else: expires_after = getattr(create_file_request, "expires_after", None) - + # Verify expires_after was passed correctly assert expires_after is not None, "expires_after should be in the request" assert expires_after["anchor"] == "created_at" assert expires_after["seconds"] == 3600 - + return OpenAIFileObject( id="file-abc123", object="file", @@ -786,17 +832,21 @@ def test_create_file_with_expires_after_valid_values(mocker: MockerFixture, monk purpose="fine-tune", status="uploaded", ) - + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): raise NotImplementedError("Not implemented for test") - + async def afile_list(self, purpose, litellm_parent_otel_span): raise NotImplementedError("Not implemented for test") - - async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_delete( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - - async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_content( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") proxy_logging_obj.proxy_hook_mapping["managed_files"] = DummyManagedFiles() @@ -827,7 +877,9 @@ def test_create_file_with_expires_after_valid_values(mocker: MockerFixture, monk assert result["purpose"] == "fine-tune" -def test_create_file_without_expires_after(mocker: MockerFixture, monkeypatch, llm_router: Router): +def test_create_file_without_expires_after( + mocker: MockerFixture, monkeypatch, llm_router: Router +): """ Test that file creation works normally without expires_after """ @@ -840,16 +892,25 @@ def test_create_file_without_expires_after(mocker: MockerFixture, monkeypatch, l proxy_logging_obj._add_proxy_hooks(llm_router) class DummyManagedFiles(BaseFileEndpoints): - async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict): + async def acreate_file( + self, + llm_router, + create_file_request, + target_model_names_list, + litellm_parent_otel_span, + user_api_key_dict, + ): # Verify expires_after is None when not provided if isinstance(create_file_request, dict): expires_after = create_file_request.get("expires_after") else: expires_after = getattr(create_file_request, "expires_after", None) - + # expires_after should be None when not provided - assert expires_after is None, "expires_after should be None when not provided" - + assert ( + expires_after is None + ), "expires_after should be None when not provided" + return OpenAIFileObject( id="file-abc123", object="file", @@ -859,17 +920,21 @@ def test_create_file_without_expires_after(mocker: MockerFixture, monkeypatch, l purpose="fine-tune", status="uploaded", ) - + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): raise NotImplementedError("Not implemented for test") - + async def afile_list(self, purpose, litellm_parent_otel_span): raise NotImplementedError("Not implemented for test") - - async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_delete( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - - async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_content( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") proxy_logging_obj.proxy_hook_mapping["managed_files"] = DummyManagedFiles() @@ -898,11 +963,13 @@ def test_create_file_without_expires_after(mocker: MockerFixture, monkeypatch, l assert result["purpose"] == "fine-tune" -def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, llm_router: Router): +def test_managed_files_with_loadbalancing( + mocker: MockerFixture, monkeypatch, llm_router: Router +): """ Test that managed files work with loadbalancing when both target_model_names and enable_loadbalancing_on_batch_endpoints are enabled. - + This ensures that the priority order is correct: - managed files should take precedence over deprecated loadbalancing - managed files internally use llm_router.acreate_file() which provides loadbalancing @@ -912,28 +979,34 @@ def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, ll # Enable loadbalancing on batch endpoints monkeypatch.setattr("litellm.enable_loadbalancing_on_batch_endpoints", True) - + proxy_logging_obj = ProxyLogging( user_api_key_cache=DualCache(default_in_memory_ttl=1) ) proxy_logging_obj._add_proxy_hooks(llm_router) - + # Track calls to verify loadbalancing through router router_acreate_file_calls = [] - + class ManagedFilesWithLoadbalancing(BaseFileEndpoints): - async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict): + async def acreate_file( + self, + llm_router, + create_file_request, + target_model_names_list, + litellm_parent_otel_span, + user_api_key_dict, + ): # Verify we receive the target model names - assert len(target_model_names_list) > 0, "Should have target_model_names_list" - + assert ( + len(target_model_names_list) > 0 + ), "Should have target_model_names_list" + # Simulate what managed files does - call llm_router.acreate_file for each model # This is where loadbalancing happens internally for model in target_model_names_list: - router_acreate_file_calls.append({ - "model": model, - "via_router": True - }) - + router_acreate_file_calls.append({"model": model, "via_router": True}) + # Return a managed file ID (base64 encoded) return OpenAIFileObject( id="litellm_managed_file_abc123", @@ -944,23 +1017,29 @@ def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, ll purpose="batch", status="uploaded", ) - + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): raise NotImplementedError("Not implemented for test") - + async def afile_list(self, purpose, litellm_parent_otel_span): raise NotImplementedError("Not implemented for test") - - async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_delete( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - - async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_content( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - + import litellm.proxy.proxy_server as ps from litellm.proxy._types import LitellmUserRoles - proxy_logging_obj.proxy_hook_mapping["managed_files"] = ManagedFilesWithLoadbalancing() + proxy_logging_obj.proxy_hook_mapping[ + "managed_files" + ] = ManagedFilesWithLoadbalancing() monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) monkeypatch.setattr( "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj @@ -971,12 +1050,12 @@ def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, ll app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( api_key="test-key", user_role=LitellmUserRoles.PROXY_ADMIN ) - + try: # Create batch file content test_file_content = b'{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hello"}]}}' test_file = ("batch_data.jsonl", test_file_content, "application/jsonl") - + # Make request with both target_model_names AND enable_loadbalancing_on_batch_endpoints response = client.post( "/v1/files", @@ -987,7 +1066,7 @@ def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, ll }, headers={"Authorization": "Bearer test-key"}, ) - + # Verify success assert response.status_code == 200, response.text finally: @@ -995,13 +1074,17 @@ def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, ll result = response.json() assert result["id"] == "litellm_managed_file_abc123" assert result["purpose"] == "batch" - + # Verify that managed files was called (via router for loadbalancing) # This proves that managed files took precedence over deprecated loadbalancing - assert len(router_acreate_file_calls) == 2, "Should have called router for both models" + assert ( + len(router_acreate_file_calls) == 2 + ), "Should have called router for both models" assert router_acreate_file_calls[0]["model"] == "azure-gpt-3-5-turbo" assert router_acreate_file_calls[1]["model"] == "gpt-3.5-turbo" - assert all(call["via_router"] for call in router_acreate_file_calls), "All calls should go through router" + assert all( + call["via_router"] for call in router_acreate_file_calls + ), "All calls should go through router" def test_create_file_with_nested_litellm_metadata( @@ -1009,22 +1092,29 @@ def test_create_file_with_nested_litellm_metadata( ): """ Test that nested litellm_metadata is correctly parsed from form data in bracket notation. - + Regression test for: litellm_metadata[spend_logs_metadata][owner] format should be correctly parsed into nested dictionary structure. """ from litellm.llms.base_llm.files.transformation import BaseFileEndpoints from litellm.types.llms.openai import OpenAIFileObject - + proxy_logging_obj = ProxyLogging( user_api_key_cache=DualCache(default_in_memory_ttl=1) ) proxy_logging_obj._add_proxy_hooks(llm_router) - + captured_litellm_metadata = {} - + class DummyManagedFiles(BaseFileEndpoints): - async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict): + async def acreate_file( + self, + llm_router, + create_file_request, + target_model_names_list, + litellm_parent_otel_span, + user_api_key_dict, + ): # Capture litellm_metadata for verification if isinstance(create_file_request, dict): captured_litellm_metadata.update( @@ -1034,7 +1124,7 @@ def test_create_file_with_nested_litellm_metadata( captured_litellm_metadata.update( getattr(create_file_request, "litellm_metadata", {}) ) - + return OpenAIFileObject( id="file-test-123", object="file", @@ -1044,28 +1134,32 @@ def test_create_file_with_nested_litellm_metadata( purpose="fine-tune", status="uploaded", ) - + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): raise NotImplementedError("Not implemented for test") - + async def afile_list(self, purpose, litellm_parent_otel_span): raise NotImplementedError("Not implemented for test") - - async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_delete( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - - async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_content( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - + proxy_logging_obj.proxy_hook_mapping["managed_files"] = DummyManagedFiles() monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) monkeypatch.setattr( "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj ) - + test_file_content = b'{"prompt": "Hello", "completion": "Hi"}' test_file = ("test.jsonl", test_file_content, "application/jsonl") - + # Test with nested litellm_metadata in bracket notation response = client.post( "/v1/files", @@ -1080,12 +1174,12 @@ def test_create_file_with_nested_litellm_metadata( }, headers={"Authorization": "Bearer test-key"}, ) - + # Verify success assert response.status_code == 200 result = response.json() assert result["id"] == "file-test-123" - + # Verify nested metadata was correctly parsed assert "spend_logs_metadata" in captured_litellm_metadata assert captured_litellm_metadata["spend_logs_metadata"]["owner"] == "john_doe" @@ -1099,26 +1193,33 @@ def test_create_file_with_deep_nested_litellm_metadata( ): """ Test that deeply nested litellm_metadata is correctly parsed from form data. - + Regression test for: litellm_metadata[a][b][c] format should be correctly parsed. """ + import litellm.proxy.proxy_server as ps from litellm.llms.base_llm.files.transformation import BaseFileEndpoints from litellm.proxy._types import LitellmUserRoles - import litellm.proxy.proxy_server as ps from litellm.types.llms.openai import OpenAIFileObject - + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) - + proxy_logging_obj = ProxyLogging( user_api_key_cache=DualCache(default_in_memory_ttl=1) ) proxy_logging_obj._add_proxy_hooks(llm_router) - + captured_litellm_metadata = {} - + class DummyManagedFiles(BaseFileEndpoints): - async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict): + async def acreate_file( + self, + llm_router, + create_file_request, + target_model_names_list, + litellm_parent_otel_span, + user_api_key_dict, + ): if isinstance(create_file_request, dict): captured_litellm_metadata.update( create_file_request.get("litellm_metadata", {}) @@ -1127,7 +1228,7 @@ def test_create_file_with_deep_nested_litellm_metadata( captured_litellm_metadata.update( getattr(create_file_request, "litellm_metadata", {}) ) - + return OpenAIFileObject( id="file-test-456", object="file", @@ -1137,33 +1238,37 @@ def test_create_file_with_deep_nested_litellm_metadata( purpose="batch", status="uploaded", ) - + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): raise NotImplementedError("Not implemented for test") - + async def afile_list(self, purpose, litellm_parent_otel_span): raise NotImplementedError("Not implemented for test") - - async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_delete( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - - async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + + async def afile_content( + self, file_id, litellm_parent_otel_span, llm_router, **data + ): raise NotImplementedError("Not implemented for test") - + proxy_logging_obj.proxy_hook_mapping["managed_files"] = DummyManagedFiles() monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) monkeypatch.setattr( "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj ) - + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user" ) - + try: test_file_content = b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo"}}' test_file = ("nested.jsonl", test_file_content, "application/jsonl") - + # Test with deeply nested metadata response = client.post( "/v1/files", @@ -1177,12 +1282,12 @@ def test_create_file_with_deep_nested_litellm_metadata( }, headers={"Authorization": "Bearer test-key"}, ) - + # Verify success assert response.status_code == 200, response.text result = response.json() assert result["id"] == "file-test-456" - + # Verify deeply nested metadata was correctly parsed assert "config" in captured_litellm_metadata assert "database" in captured_litellm_metadata["config"] @@ -1356,7 +1461,9 @@ def test_file_team_injects_when_caller_sends_nothing( # --------------------------------------------------------------------------- -def _post_file_raw(monkeypatch, llm_router: Router, team_metadata: dict, form_data: dict): +def _post_file_raw( + monkeypatch, llm_router: Router, team_metadata: dict, form_data: dict +): """POST /v1/files and return the raw response (no status assertion).""" from litellm.proxy.auth.user_api_key_auth import user_api_key_auth diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index f4a7c4f499..a02464e946 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -24,9 +24,9 @@ def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): "model_name": "gpt-3.5-turbo", "litellm_params": { "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, } ], @@ -646,8 +646,15 @@ def test_arouter_responses_api_bridge(): mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = {"content-type": "application/json"} - mock_response.json.return_value = {"id": "resp_test", "object": "response", "status": "completed", "output": []} - mock_response.text = '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' + mock_response.json.return_value = { + "id": "resp_test", + "object": "response", + "status": "completed", + "output": [], + } + mock_response.text = ( + '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' + ) with patch.object(client, "post", return_value=mock_response) as mock_post: try: @@ -2147,7 +2154,10 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint() ) assert credentials is not None - assert credentials["aws_bedrock_runtime_endpoint"] == "https://bedrock-runtime.us-east-1.amazonaws.com" + assert ( + credentials["aws_bedrock_runtime_endpoint"] + == "https://bedrock-runtime.us-east-1.amazonaws.com" + ) assert credentials["aws_access_key_id"] == "test-access-key" assert credentials["aws_secret_access_key"] == "test-secret-key" assert credentials["aws_region_name"] == "us-east-1" @@ -2169,11 +2179,11 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): credential_values={ "api_key": "resolved-api-key", "api_base": "https://resolved.openai.azure.com", - "api_version": "2024-02-01" - } + "api_version": "2024-02-01", + }, ) ] - + router = litellm.Router( model_list=[ { @@ -2197,7 +2207,7 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): assert credentials["custom_llm_provider"] == "azure" # Ensure credential name is removed after resolution assert "litellm_credential_name" not in credentials - + # Cleanup litellm.credential_list = [] @@ -2302,7 +2312,10 @@ async def test_aguardrail_helper(): # Mock the original function async def mock_original_function(**kwargs): - return {"result": "success", "selected_guardrail": kwargs.get("selected_guardrail")} + return { + "result": "success", + "selected_guardrail": kwargs.get("selected_guardrail"), + } result = await router._aguardrail_helper( model="content-filter", @@ -2336,7 +2349,10 @@ async def test_aguardrail(): # Mock the original function async def mock_original_function(**kwargs): - return {"result": "success", "selected_guardrail": kwargs.get("selected_guardrail")} + return { + "result": "success", + "selected_guardrail": kwargs.get("selected_guardrail"), + } result = await router.aguardrail( guardrail_name="content-filter", @@ -2346,6 +2362,7 @@ async def test_aguardrail(): assert result["result"] == "success" assert result["selected_guardrail"]["id"] == "guardrail-1" + @pytest.mark.asyncio async def test_anthropic_messages_call_type_is_cached(): """ @@ -2417,36 +2434,33 @@ async def test_anthropic_messages_call_type_is_cached(): additional_headers=None, ), ) - + cache = DualCache() deployment_check = PromptCachingDeploymentCheck(cache=cache) prompt_cache = PromptCachingCache(cache=cache) - + # Create messages with enough tokens to pass the caching threshold test_messages = [ { - "role": "user", + "role": "user", "content": [ { - "type": "text", + "type": "text", "text": "test long message here" * 1024, - "cache_control": { - "type": "ephemeral", - "ttl": "5m" - } + "cache_control": {"type": "ephemeral", "ttl": "5m"}, } - ] + ], } ] test_model_id = "test-model-id-123" - + # Create a payload with anthropic_messages call type payload = create_standard_logging_payload() payload["call_type"] = CallTypes.anthropic_messages.value payload["messages"] = test_messages payload["model"] = "anthropic/claude-3-5-sonnet-20240620" payload["model_id"] = test_model_id - + # Log the success event (should cache the model_id) await deployment_check.async_log_success_event( kwargs={"standard_logging_object": payload}, @@ -2454,19 +2468,23 @@ async def test_anthropic_messages_call_type_is_cached(): start_time=1234567890.0, end_time=1234567891.0, ) - + # Small delay to ensure cache write completes await asyncio.sleep(0.1) - + # Verify that the model_id was actually cached cached_result = await prompt_cache.async_get_model_id( messages=test_messages, tools=None, ) - + # This assertion will FAIL if anthropic_messages is filtered out - assert cached_result is not None, "Model ID should be cached for anthropic_messages call type" - assert cached_result["model_id"] == test_model_id, f"Expected {test_model_id}, got {cached_result['model_id']}" + assert ( + cached_result is not None + ), "Model ID should be cached for anthropic_messages call type" + assert ( + cached_result["model_id"] == test_model_id + ), f"Expected {test_model_id}, got {cached_result['model_id']}" def test_update_kwargs_with_deployment_propagates_model_tags(): @@ -2682,9 +2700,7 @@ def test_credential_name_injected_as_tag(): ) kwargs: dict = {"metadata": {"tags": ["A.101"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="xai-model" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="xai-model") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert "Credential: xAI" in kwargs["metadata"]["tags"] @@ -2709,9 +2725,7 @@ def test_credential_name_not_duplicated_in_tags(): ) kwargs: dict = {"metadata": {"tags": ["Credential: xAI", "A.101"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="xai-model" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="xai-model") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert kwargs["metadata"]["tags"].count("Credential: xAI") == 1 @@ -2733,9 +2747,7 @@ def test_credential_name_not_injected_when_absent(): ) kwargs: dict = {"metadata": {"tags": ["A.101"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-model" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-model") router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) assert kwargs["metadata"]["tags"] == ["A.101"] diff --git a/tests/test_litellm/test_setup_wizard.py b/tests/test_litellm/test_setup_wizard.py index e10bd893e3..c96d6d7ed6 100644 --- a/tests/test_litellm/test_setup_wizard.py +++ b/tests/test_litellm/test_setup_wizard.py @@ -58,7 +58,7 @@ _ANTHROPIC = { _AZURE = { "id": "azure", "name": "Azure OpenAI", - "env_key": "AZURE_API_KEY", + "env_key": "AZURE_AI_API_KEY", "models": [], "test_model": None, "needs_api_base": True, @@ -123,8 +123,8 @@ def test_build_config_master_key_quoted(): def test_build_config_does_not_mutate_env_vars(): """_build_config must not modify the caller's env_vars dict.""" env_vars = { - "AZURE_API_KEY": "az-key", - "_LITELLM_AZURE_API_BASE_AZURE": "https://my.azure.com", + "AZURE_AI_API_KEY": "az-key", + "_LITELLM_AZURE_AI_API_BASE_AZURE": "https://my.azure.com", "_LITELLM_AZURE_DEPLOYMENT_AZURE": "my-deployment", } original_keys = set(env_vars.keys()) @@ -134,8 +134,8 @@ def test_build_config_does_not_mutate_env_vars(): def test_build_config_azure_uses_deployment_name(): env_vars = { - "AZURE_API_KEY": "az-key", - "_LITELLM_AZURE_API_BASE_AZURE": "https://my.azure.com", + "AZURE_AI_API_KEY": "az-key", + "_LITELLM_AZURE_AI_API_BASE_AZURE": "https://my.azure.com", "_LITELLM_AZURE_DEPLOYMENT_AZURE": "my-gpt4o", } config = SetupWizard._build_config([_AZURE], env_vars, "sk-master") @@ -147,7 +147,7 @@ def test_build_config_azure_uses_deployment_name(): def test_build_config_azure_no_deployment_skipped(): """Azure without a deployment name should emit nothing (not fallback to gpt-4o).""" - env_vars = {"AZURE_API_KEY": "az-key"} # no deployment sentinel + env_vars = {"AZURE_AI_API_KEY": "az-key"} # no deployment sentinel config = SetupWizard._build_config([_AZURE], env_vars, "sk-master") # No azure model entry should be emitted when deployment name is absent assert "model: azure/" not in config @@ -157,7 +157,7 @@ def test_build_config_no_display_name_collision_openai_and_azure(): """OpenAI gpt-4o and azure gpt-4o should get distinct model_name values.""" env_vars = { "OPENAI_API_KEY": "sk-openai", - "AZURE_API_KEY": "az-key", + "AZURE_AI_API_KEY": "az-key", "_LITELLM_AZURE_DEPLOYMENT_AZURE": "gpt-4o", } config = SetupWizard._build_config([_OPENAI, _AZURE], env_vars, "sk-master") @@ -182,7 +182,7 @@ def test_build_config_internal_sentinel_keys_excluded(): """_LITELLM_ prefixed sentinel keys must not appear in environment_variables.""" env_vars = { "OPENAI_API_KEY": "sk-real", - "_LITELLM_AZURE_API_BASE_AZURE": "https://x.azure.com", + "_LITELLM_AZURE_AI_API_BASE_AZURE": "https://x.azure.com", } config = SetupWizard._build_config([_OPENAI], env_vars, "sk-master") assert "_LITELLM_" not in config From cdcab8a243e4d48639ab628b161a0351131e7768 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 19:39:11 -0700 Subject: [PATCH 07/16] refactor: cleanup deprecated models --- .../test_openai_batches_and_files.py | 46 +++--- .../litellm_utils_tests/test_health_check.py | 29 ---- tests/llm_translation/test_azure_ai.py | 31 +--- tests/local_testing/test_router.py | 142 ++++++------------ 4 files changed, 76 insertions(+), 172 deletions(-) diff --git a/tests/batches_tests/test_openai_batches_and_files.py b/tests/batches_tests/test_openai_batches_and_files.py index e1165812e2..2ece740bf5 100644 --- a/tests/batches_tests/test_openai_batches_and_files.py +++ b/tests/batches_tests/test_openai_batches_and_files.py @@ -234,9 +234,9 @@ def cleanup_azure_ft_models(): import requests client = AzureOpenAI( - api_key=os.getenv("AZURE_FT_API_KEY"), - azure_endpoint=os.getenv("AZURE_FT_API_BASE"), - api_version=os.getenv("AZURE_API_VERSION"), + api_key=os.getenv("AZURE_AI_API_KEY"), + azure_endpoint=os.getenv("AZURE_AI_API_BASE"), + api_version=os.getenv("AZURE_AI_API_VERSION"), ) _list_ft_jobs = client.fine_tuning.jobs.list() @@ -577,7 +577,10 @@ async def test_vertex_list_batches(monkeypatch): monkeypatch.setattr( "litellm.llms.vertex_ai.batches.handler.VertexAIBatchPrediction._ensure_access_token", - lambda self, credentials, project_id, custom_llm_provider: ("mock-token", "litellm-test-project"), + lambda self, credentials, project_id, custom_llm_provider: ( + "mock-token", + "litellm-test-project", + ), ) with patch( @@ -648,7 +651,7 @@ async def test_vertex_async_create_batch_logs_error_body_on_http_error(): async def test_delete_batch_output_file(): """ Test that deleting a batch output file works correctly. - + This test verifies the fix for: - When a batch is retrieved and has an output_file_id, the file object is properly stored - The output file can be deleted without validation errors @@ -656,11 +659,11 @@ async def test_delete_batch_output_file(): """ litellm._turn_on_debug() print("Testing delete batch output file") - + file_name = "openai_batch_completions.jsonl" _current_dir = os.path.dirname(os.path.abspath(__file__)) file_path = os.path.join(_current_dir, file_name) - + # Create file for batch file_obj = await litellm.acreate_file( file=open(file_path, "rb"), @@ -669,7 +672,7 @@ async def test_delete_batch_output_file(): ) print("Response from creating file=", file_obj) batch_input_file_id = file_obj.id - + # Create batch create_batch_response = await litellm.acreate_batch( completion_window="24h", @@ -678,36 +681,37 @@ async def test_delete_batch_output_file(): custom_llm_provider="openai", ) print("Batch created with ID=", create_batch_response.id) - + # Retrieve batch to get output_file_id 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="openai" ) print("Retrieved batch=", retrieved_batch) - + # If batch has completed and has output file, test deleting it if retrieved_batch.output_file_id: print(f"Testing deletion of output file: {retrieved_batch.output_file_id}") - + # This is the key test - deleting the output file should work # without validation errors (file_object should not be None) delete_output_file_response = await litellm.afile_delete( - file_id=retrieved_batch.output_file_id, - custom_llm_provider="openai" + file_id=retrieved_batch.output_file_id, custom_llm_provider="openai" ) - + print("Delete output file response=", delete_output_file_response) assert delete_output_file_response.id == retrieved_batch.output_file_id - assert delete_output_file_response.deleted is True or hasattr(delete_output_file_response, 'id') + assert delete_output_file_response.deleted is True or hasattr( + delete_output_file_response, "id" + ) print("✓ Successfully deleted batch output file") else: - print("⚠ Batch has not completed yet or no output file available, skipping output file deletion test") - + print( + "⚠ Batch has not completed yet or no output file available, skipping output file deletion test" + ) + # Clean up - delete the input file delete_input_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="openai" ) print("Delete input file response=", delete_input_file_response) assert delete_input_file_response.id == batch_input_file_id diff --git a/tests/litellm_utils_tests/test_health_check.py b/tests/litellm_utils_tests/test_health_check.py index f87f2005d3..d014b3f72a 100644 --- a/tests/litellm_utils_tests/test_health_check.py +++ b/tests/litellm_utils_tests/test_health_check.py @@ -246,35 +246,6 @@ async def test_audio_transcription_health_check(): print(response) -@pytest.mark.asyncio -@pytest.mark.parametrize( - "model", ["azure/gpt-4o-realtime-preview", "openai/gpt-4o-realtime-preview"] -) -async def test_async_realtime_health_check(model, mocker): - """ - Test Health Check with Valid models passes - - """ - mock_websocket = AsyncMock() - mock_connect = AsyncMock().__aenter__.return_value = mock_websocket - mocker.patch("websockets.connect", return_value=mock_connect) - - litellm.set_verbose = True - model_params = { - "model": model, - } - if model == "azure/gpt-4o-realtime-preview": - model_params["api_base"] = os.getenv("AZURE_REALTIME_API_BASE") - model_params["api_key"] = os.getenv("AZURE_REALTIME_API_KEY") - model_params["api_version"] = os.getenv("AZURE_REALTIME_API_VERSION") - response = await litellm.ahealth_check( - model_params=model_params, - mode="realtime", - ) - print(response) - assert response == {} - - def test_update_litellm_params_for_health_check(): """ Test if _update_litellm_params_for_health_check correctly: diff --git a/tests/llm_translation/test_azure_ai.py b/tests/llm_translation/test_azure_ai.py index 823fad7648..b6f53e4e95 100644 --- a/tests/llm_translation/test_azure_ai.py +++ b/tests/llm_translation/test_azure_ai.py @@ -188,35 +188,6 @@ def test_azure_ai_services_with_api_version(): ) -@pytest.mark.skip(reason="Skipping due to cohere ssl issues") -def test_completion_azure_ai_command_r(): - try: - import os - - litellm.set_verbose = True - - os.environ["AZURE_AI_API_BASE"] = os.getenv("AZURE_COHERE_API_BASE", "") - os.environ["AZURE_AI_API_KEY"] = os.getenv("AZURE_COHERE_API_KEY", "") - - response = completion( - model="azure_ai/command-r-plus", - messages=[ - { - "role": "user", - "content": [ - {"type": "text", "text": "What is the meaning of life?"} - ], - } - ], - ) # type: ignore - - assert "azure_ai" in response.model - except litellm.Timeout as e: - pass - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - def test_azure_deepseek_reasoning_content(): import json @@ -387,7 +358,7 @@ async def test_azure_ai_model_router(): response = await litellm.acompletion( model="azure_ai/model_router/azure-model-router", messages=[{"role": "user", "content": "hi who is this"}], - api_base="https://ishaa-mh6uutut-swedencentral.cognitiveservices.azure.com/openai/v1/", + api_base=os.getenv("AZURE_MODEL_ROUTER_API_BASE"), api_key=os.getenv("AZURE_MODEL_ROUTER_API_KEY"), ) print("response: ", response) diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py index e68d271a11..f7885fb8a0 100644 --- a/tests/local_testing/test_router.py +++ b/tests/local_testing/test_router.py @@ -126,7 +126,9 @@ async def test_router_provider_wildcard_routing(): print("response 3 = ", response3) response4 = await router.acompletion( - model=os.environ.get("CI_CD_DEFAULT_ANTHROPIC_MODEL", "claude-haiku-4-5-20251001"), + model=os.environ.get( + "CI_CD_DEFAULT_ANTHROPIC_MODEL", "claude-haiku-4-5-20251001" + ), messages=[{"role": "user", "content": "hello"}], ) @@ -356,51 +358,6 @@ async def test_router_retries(sync_mode): print(response.choices[0].message) -@pytest.mark.parametrize( - "mistral_api_base", - [ - "os.environ/AZURE_MISTRAL_API_BASE", - "https://Mistral-large-nmefg-serverless.eastus2.inference.ai.azure.com/v1/", - "https://Mistral-large-nmefg-serverless.eastus2.inference.ai.azure.com/v1", - "https://Mistral-large-nmefg-serverless.eastus2.inference.ai.azure.com/", - "https://Mistral-large-nmefg-serverless.eastus2.inference.ai.azure.com", - ], -) -@pytest.mark.skip( - reason="Router no longer creates clients, this is delegated to the provider integration." -) -def test_router_azure_ai_studio_init(mistral_api_base): - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "azure/mistral-large-latest", - "api_key": "os.environ/AZURE_MISTRAL_API_KEY", - "api_base": mistral_api_base, - }, - "model_info": {"id": 1234}, - } - ] - ) - - # model_client = router._get_client( - # deployment={"model_info": {"id": 1234}}, client_type="sync_client", kwargs={} - # ) - # url = getattr(model_client, "_base_url") - # uri_reference = str(getattr(url, "_uri_reference")) - - # print(f"uri_reference: {uri_reference}") - - # assert "/v1/" in uri_reference - # assert uri_reference.count("v1") == 1 - response = router.completion( - model="azure/mistral-large-latest", - messages=[{"role": "user", "content": "Hey, how's it going?"}], - ) - assert response is not None - - def test_exception_raising(): # this tests if the router raises an exception when invalid params are set # in this test both deployments have bad keys - Keep this test. It validates if the router raises the most recent exception @@ -409,8 +366,8 @@ def test_exception_raising(): try: print("testing if router raises an exception") - old_api_key = os.environ["AZURE_API_KEY"] - os.environ["AZURE_API_KEY"] = "" + old_api_key = os.environ["AZURE_AI_API_KEY"] + os.environ["AZURE_AI_API_KEY"] = "" model_list = [ { "model_name": "gpt-3.5-turbo", # openai model name @@ -418,7 +375,7 @@ def test_exception_raising(): "model": "azure/gpt-4.1-mini", "api_key": "bad-key", "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -446,16 +403,16 @@ def test_exception_raising(): model="gpt-3.5-turbo", messages=[{"role": "user", "content": "hello this request will fail"}], ) - os.environ["AZURE_API_KEY"] = old_api_key + os.environ["AZURE_AI_API_KEY"] = old_api_key pytest.fail(f"Should have raised an Auth Error") except openai.AuthenticationError: print( "Test Passed: Caught an OPENAI AUTH Error, Good job. This is what we needed!" ) - os.environ["AZURE_API_KEY"] = old_api_key + os.environ["AZURE_AI_API_KEY"] = old_api_key router.reset() except Exception as e: - os.environ["AZURE_API_KEY"] = old_api_key + os.environ["AZURE_AI_API_KEY"] = old_api_key print("Got unexpected exception on router!", e) @@ -530,7 +487,7 @@ def test_call_one_endpoint(): # this test makes a completion calls azure/gpt-4.1-mini, it should work try: print("Testing calling a specific deployment") - old_api_key = os.environ["AZURE_API_KEY"] + old_api_key = os.environ["AZURE_AI_API_KEY"] model_list = [ { @@ -539,7 +496,7 @@ def test_call_one_endpoint(): "model": "azure/gpt-4.1-mini", "api_key": old_api_key, "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), }, "tpm": 240000, "rpm": 1800, @@ -548,8 +505,8 @@ def test_call_one_endpoint(): "model_name": "text-embedding-ada-002", "litellm_params": { "model": "azure/text-embedding-ada-002", - "api_key": os.environ["AZURE_API_KEY"], - "api_base": os.environ["AZURE_API_BASE"], + "api_key": os.environ["AZURE_AI_API_KEY"], + "api_base": os.environ["AZURE_AI_API_BASE"], }, "tpm": 100000, "rpm": 10000, @@ -562,7 +519,7 @@ def test_call_one_endpoint(): set_verbose=True, num_retries=1, ) # type: ignore - old_api_base = os.environ.pop("AZURE_API_BASE", None) + old_api_base = os.environ.pop("AZURE_AI_API_BASE", None) async def call_azure_completion(): response = await router.acompletion( @@ -584,8 +541,8 @@ def test_call_one_endpoint(): asyncio.run(call_azure_completion()) asyncio.run(call_azure_embedding()) - os.environ["AZURE_API_BASE"] = old_api_base - os.environ["AZURE_API_KEY"] = old_api_key + os.environ["AZURE_AI_API_BASE"] = old_api_base + os.environ["AZURE_AI_API_KEY"] = old_api_key except Exception as e: print(f"FAILED TEST") pytest.fail(f"Got unexpected exception on router! - {e}") @@ -594,7 +551,6 @@ def test_call_one_endpoint(): # test_call_one_endpoint() - @pytest.mark.asyncio @pytest.mark.parametrize("sync_mode", [True, False]) async def test_async_router_context_window_fallback(sync_mode): @@ -708,9 +664,9 @@ def test_router_context_window_check_pre_call_check_in_group_custom_model_info() "model_name": "gpt-3.5-turbo", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), "base_model": "azure/gpt-35-turbo", "mock_response": "Hello world 1!", }, @@ -762,9 +718,9 @@ def test_router_context_window_check_pre_call_check(): "model_name": "gpt-3.5-turbo", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), "base_model": "azure/gpt-35-turbo", "mock_response": "Hello world 1!", }, @@ -816,9 +772,9 @@ def test_router_context_window_check_pre_call_check_out_group(): "model_name": "gpt-3.5-turbo-small", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), "base_model": "azure/gpt-35-turbo", }, }, @@ -896,9 +852,9 @@ def test_router_region_pre_call_check(allowed_model_region): "model_name": "gpt-3.5-turbo", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_API_KEY"), + "api_key": os.getenv("AZURE_AI_API_KEY"), "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_API_BASE"), + "api_base": os.getenv("AZURE_AI_API_BASE"), "base_model": "azure/gpt-35-turbo", "region_name": allowed_model_region, }, @@ -1173,8 +1129,8 @@ def test_azure_embedding_on_router(): "model_name": "text-embedding-ada-002", "litellm_params": { "model": "azure/text-embedding-ada-002", - "api_key": os.environ["AZURE_API_KEY"], - "api_base": os.environ["AZURE_API_BASE"], + "api_key": os.environ["AZURE_AI_API_KEY"], + "api_base": os.environ["AZURE_AI_API_BASE"], }, "tpm": 100000, "rpm": 10000, @@ -1381,8 +1337,8 @@ def test_reading_keys_os_environ(): "model_name": "gpt-3.5-turbo", "litellm_params": { "model": "gpt-3.5-turbo", - "api_key": "os.environ/AZURE_API_KEY", - "api_base": "os.environ/AZURE_API_BASE", + "api_key": "os.environ/AZURE_AI_API_KEY", + "api_base": "os.environ/AZURE_AI_API_BASE", "api_version": "os.environ/AZURE_API_VERSION", "timeout": "os.environ/AZURE_TIMEOUT", "stream_timeout": "os.environ/AZURE_STREAM_TIMEOUT", @@ -1394,11 +1350,11 @@ def test_reading_keys_os_environ(): router = Router(model_list=model_list) for model in router.model_list: assert ( - model["litellm_params"]["api_key"] == os.environ["AZURE_API_KEY"] - ), f"{model['litellm_params']['api_key']} vs {os.environ['AZURE_API_KEY']}" + model["litellm_params"]["api_key"] == os.environ["AZURE_AI_API_KEY"] + ), f"{model['litellm_params']['api_key']} vs {os.environ['AZURE_AI_API_KEY']}" assert ( - model["litellm_params"]["api_base"] == os.environ["AZURE_API_BASE"] - ), f"{model['litellm_params']['api_base']} vs {os.environ['AZURE_API_BASE']}" + model["litellm_params"]["api_base"] == os.environ["AZURE_AI_API_BASE"] + ), f"{model['litellm_params']['api_base']} vs {os.environ['AZURE_AI_API_BASE']}" assert ( model["litellm_params"]["api_version"] == os.environ["AZURE_API_VERSION"] @@ -1415,8 +1371,8 @@ def test_reading_keys_os_environ(): print("passed testing of reading keys from os.environ") model_id = model["model_info"]["id"] async_client: openai.AsyncAzureOpenAI = router.cache.get_cache(f"{model_id}_async_client") # type: ignore - assert async_client.api_key == os.environ["AZURE_API_KEY"] - assert async_client.base_url == os.environ["AZURE_API_BASE"] + assert async_client.api_key == os.environ["AZURE_AI_API_KEY"] + assert async_client.base_url == os.environ["AZURE_AI_API_BASE"] assert async_client.max_retries == int( os.environ["AZURE_MAX_RETRIES"] ), f"{async_client.max_retries} vs {os.environ['AZURE_MAX_RETRIES']}" @@ -1428,8 +1384,8 @@ def test_reading_keys_os_environ(): print("\n Testing async streaming client") stream_async_client: openai.AsyncAzureOpenAI = router.cache.get_cache(f"{model_id}_stream_async_client") # type: ignore - assert stream_async_client.api_key == os.environ["AZURE_API_KEY"] - assert stream_async_client.base_url == os.environ["AZURE_API_BASE"] + assert stream_async_client.api_key == os.environ["AZURE_AI_API_KEY"] + assert stream_async_client.base_url == os.environ["AZURE_AI_API_BASE"] assert stream_async_client.max_retries == int( os.environ["AZURE_MAX_RETRIES"] ), f"{stream_async_client.max_retries} vs {os.environ['AZURE_MAX_RETRIES']}" @@ -1440,8 +1396,8 @@ def test_reading_keys_os_environ(): print("\n Testing sync client") client: openai.AzureOpenAI = router.cache.get_cache(f"{model_id}_client") # type: ignore - assert client.api_key == os.environ["AZURE_API_KEY"] - assert client.base_url == os.environ["AZURE_API_BASE"] + assert client.api_key == os.environ["AZURE_AI_API_KEY"] + assert client.base_url == os.environ["AZURE_AI_API_BASE"] assert client.max_retries == int( os.environ["AZURE_MAX_RETRIES"] ), f"{client.max_retries} vs {os.environ['AZURE_MAX_RETRIES']}" @@ -1452,8 +1408,8 @@ def test_reading_keys_os_environ(): print("\n Testing sync stream client") stream_client: openai.AzureOpenAI = router.cache.get_cache(f"{model_id}_stream_client") # type: ignore - assert stream_client.api_key == os.environ["AZURE_API_KEY"] - assert stream_client.base_url == os.environ["AZURE_API_BASE"] + assert stream_client.api_key == os.environ["AZURE_AI_API_KEY"] + assert stream_client.base_url == os.environ["AZURE_AI_API_BASE"] assert stream_client.max_retries == int( os.environ["AZURE_MAX_RETRIES"] ), f"{stream_client.max_retries} vs {os.environ['AZURE_MAX_RETRIES']}" @@ -1503,7 +1459,7 @@ def test_reading_openai_keys_os_environ(): for model in router.model_list: assert ( model["litellm_params"]["api_key"] == os.environ["OPENAI_API_KEY"] - ), f"{model['litellm_params']['api_key']} vs {os.environ['AZURE_API_KEY']}" + ), f"{model['litellm_params']['api_key']} vs {os.environ['AZURE_AI_API_KEY']}" assert float(model["litellm_params"]["timeout"]) == float( os.environ["AZURE_TIMEOUT"] ), f"{model['litellm_params']['timeout']} vs {os.environ['AZURE_TIMEOUT']}" @@ -1574,7 +1530,9 @@ def test_router_anthropic_key_dynamic(): { "model_name": "anthropic-claude", "litellm_params": { - "model": os.environ.get("CI_CD_DEFAULT_ANTHROPIC_MODEL", "claude-haiku-4-5-20251001"), + "model": os.environ.get( + "CI_CD_DEFAULT_ANTHROPIC_MODEL", "claude-haiku-4-5-20251001" + ), "api_key": anthropic_api_key, }, } @@ -2273,8 +2231,8 @@ async def test_router_batch_endpoints(provider): "model_name": "my-custom-name", "litellm_params": { "model": "azure/gpt-4o-mini", - "api_base": os.getenv("AZURE_API_BASE"), - "api_key": os.getenv("AZURE_API_KEY"), + "api_base": os.getenv("AZURE_AI_API_BASE"), + "api_key": os.getenv("AZURE_AI_API_KEY"), }, }, ] @@ -2452,8 +2410,8 @@ def test_is_team_specific_model(): # "model_name": "gpt-3.5-turbo", # "litellm_params": { # "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_API_KEY"), -# "api_base": os.getenv("AZURE_API_BASE"), +# "api_key": os.getenv("AZURE_AI_API_KEY"), +# "api_base": os.getenv("AZURE_AI_API_BASE"), # "tpm": 100000, # "rpm": 100000, # }, @@ -2462,8 +2420,8 @@ def test_is_team_specific_model(): # "model_name": "gpt-3.5-turbo", # "litellm_params": { # "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_API_KEY"), -# "api_base": os.getenv("AZURE_API_BASE"), +# "api_key": os.getenv("AZURE_AI_API_KEY"), +# "api_base": os.getenv("AZURE_AI_API_BASE"), # "tpm": 500, # "rpm": 500, # }, From ce26dd4101ff6f0636e129bced8a669c355e6e2e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 19:42:14 -0700 Subject: [PATCH 08/16] test: update tests --- tests/local_testing/test_router_client_init.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/local_testing/test_router_client_init.py b/tests/local_testing/test_router_client_init.py index ea808f96f3..6cf8b4e59b 100644 --- a/tests/local_testing/test_router_client_init.py +++ b/tests/local_testing/test_router_client_init.py @@ -174,9 +174,9 @@ async def test_audio_speech_router(): { "model_name": "tts", "litellm_params": { - "model": "azure/azure-tts", - "api_base": os.getenv("AZURE_SWEDEN_API_BASE"), - "api_key": os.getenv("AZURE_SWEDEN_API_KEY"), + "model": "azure/tts", + "api_base": os.getenv("AZURE_TTS_API_BASE"), + "api_key": os.getenv("AZURE_TTS_API_KEY"), }, }, ] From 1fb677702dc0ad4dea16639a2a54095463f0c7a3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 20:19:05 -0700 Subject: [PATCH 09/16] test: update to new vertex ai keys --- tests/image_gen_tests/vertex_key.json | 8 ++++---- tests/litellm_utils_tests/vertex_key.json | 8 ++++---- tests/load_tests/vertex_key.json | 8 ++++---- tests/local_testing/vertex_key.json | 8 ++++---- tests/ocr_tests/vertex_key.json | 8 ++++---- tests/pass_through_tests/vertex_key.json | 8 ++++---- tests/proxy_unit_tests/vertex_key.json | 8 ++++---- tests/unified_google_tests/vertex_key.json | 8 ++++---- 8 files changed, 32 insertions(+), 32 deletions(-) diff --git a/tests/image_gen_tests/vertex_key.json b/tests/image_gen_tests/vertex_key.json index 45ca6acc01..800969fb30 100644 --- a/tests/image_gen_tests/vertex_key.json +++ b/tests/image_gen_tests/vertex_key.json @@ -1,13 +1,13 @@ { "type": "service_account", - "project_id": "pathrise-convert-1606954137718", + "project_id": "litellm-ci-cd", "private_key_id": "", "private_key": "", - "client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com", - "client_id": "109577393201924326488", + "client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com", + "client_id": "116563532503305622785", "auth_uri": "https://accounts.google.com/o/oauth2/auth", "token_uri": "https://oauth2.googleapis.com/token", "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", - "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com", + "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com", "universe_domain": "googleapis.com" } diff --git a/tests/litellm_utils_tests/vertex_key.json b/tests/litellm_utils_tests/vertex_key.json index 45ca6acc01..800969fb30 100644 --- a/tests/litellm_utils_tests/vertex_key.json +++ b/tests/litellm_utils_tests/vertex_key.json @@ -1,13 +1,13 @@ { "type": "service_account", - "project_id": "pathrise-convert-1606954137718", + "project_id": "litellm-ci-cd", "private_key_id": "", "private_key": "", - "client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com", - "client_id": "109577393201924326488", + "client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com", + "client_id": "116563532503305622785", "auth_uri": "https://accounts.google.com/o/oauth2/auth", "token_uri": "https://oauth2.googleapis.com/token", "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", - "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com", + "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com", "universe_domain": "googleapis.com" } diff --git a/tests/load_tests/vertex_key.json b/tests/load_tests/vertex_key.json index 45ca6acc01..800969fb30 100644 --- a/tests/load_tests/vertex_key.json +++ b/tests/load_tests/vertex_key.json @@ -1,13 +1,13 @@ { "type": "service_account", - "project_id": "pathrise-convert-1606954137718", + "project_id": "litellm-ci-cd", "private_key_id": "", "private_key": "", - "client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com", - "client_id": "109577393201924326488", + "client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com", + "client_id": "116563532503305622785", "auth_uri": "https://accounts.google.com/o/oauth2/auth", "token_uri": "https://oauth2.googleapis.com/token", "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", - "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com", + "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com", "universe_domain": "googleapis.com" } diff --git a/tests/local_testing/vertex_key.json b/tests/local_testing/vertex_key.json index 45ca6acc01..800969fb30 100644 --- a/tests/local_testing/vertex_key.json +++ b/tests/local_testing/vertex_key.json @@ -1,13 +1,13 @@ { "type": "service_account", - "project_id": "pathrise-convert-1606954137718", + "project_id": "litellm-ci-cd", "private_key_id": "", "private_key": "", - "client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com", - "client_id": "109577393201924326488", + "client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com", + "client_id": "116563532503305622785", "auth_uri": "https://accounts.google.com/o/oauth2/auth", "token_uri": "https://oauth2.googleapis.com/token", "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", - "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com", + "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com", "universe_domain": "googleapis.com" } diff --git a/tests/ocr_tests/vertex_key.json b/tests/ocr_tests/vertex_key.json index 45ca6acc01..800969fb30 100644 --- a/tests/ocr_tests/vertex_key.json +++ b/tests/ocr_tests/vertex_key.json @@ -1,13 +1,13 @@ { "type": "service_account", - "project_id": "pathrise-convert-1606954137718", + "project_id": "litellm-ci-cd", "private_key_id": "", "private_key": "", - "client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com", - "client_id": "109577393201924326488", + "client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com", + "client_id": "116563532503305622785", "auth_uri": "https://accounts.google.com/o/oauth2/auth", "token_uri": "https://oauth2.googleapis.com/token", "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", - "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com", + "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com", "universe_domain": "googleapis.com" } diff --git a/tests/pass_through_tests/vertex_key.json b/tests/pass_through_tests/vertex_key.json index 45ca6acc01..800969fb30 100644 --- a/tests/pass_through_tests/vertex_key.json +++ b/tests/pass_through_tests/vertex_key.json @@ -1,13 +1,13 @@ { "type": "service_account", - "project_id": "pathrise-convert-1606954137718", + "project_id": "litellm-ci-cd", "private_key_id": "", "private_key": "", - "client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com", - "client_id": "109577393201924326488", + "client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com", + "client_id": "116563532503305622785", "auth_uri": "https://accounts.google.com/o/oauth2/auth", "token_uri": "https://oauth2.googleapis.com/token", "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", - "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com", + "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com", "universe_domain": "googleapis.com" } diff --git a/tests/proxy_unit_tests/vertex_key.json b/tests/proxy_unit_tests/vertex_key.json index 45ca6acc01..800969fb30 100644 --- a/tests/proxy_unit_tests/vertex_key.json +++ b/tests/proxy_unit_tests/vertex_key.json @@ -1,13 +1,13 @@ { "type": "service_account", - "project_id": "pathrise-convert-1606954137718", + "project_id": "litellm-ci-cd", "private_key_id": "", "private_key": "", - "client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com", - "client_id": "109577393201924326488", + "client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com", + "client_id": "116563532503305622785", "auth_uri": "https://accounts.google.com/o/oauth2/auth", "token_uri": "https://oauth2.googleapis.com/token", "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", - "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com", + "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com", "universe_domain": "googleapis.com" } diff --git a/tests/unified_google_tests/vertex_key.json b/tests/unified_google_tests/vertex_key.json index 45ca6acc01..800969fb30 100644 --- a/tests/unified_google_tests/vertex_key.json +++ b/tests/unified_google_tests/vertex_key.json @@ -1,13 +1,13 @@ { "type": "service_account", - "project_id": "pathrise-convert-1606954137718", + "project_id": "litellm-ci-cd", "private_key_id": "", "private_key": "", - "client_email": "ci-cd-723@pathrise-convert-1606954137718.iam.gserviceaccount.com", - "client_id": "109577393201924326488", + "client_email": "test-litellm-ci-cd@litellm-ci-cd.iam.gserviceaccount.com", + "client_id": "116563532503305622785", "auth_uri": "https://accounts.google.com/o/oauth2/auth", "token_uri": "https://oauth2.googleapis.com/token", "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", - "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/ci-cd-723%40pathrise-convert-1606954137718.iam.gserviceaccount.com", + "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-litellm-ci-cd%40litellm-ci-cd.iam.gserviceaccount.com", "universe_domain": "googleapis.com" } From 0fef88d67c63b8a045d1085ad997b632d392cdf4 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 20:23:44 -0700 Subject: [PATCH 10/16] test: remove dead tests --- .../test_anthropic_completion.py | 6 ++++ tests/logging_callback_tests/test_alerting.py | 34 ------------------- 2 files changed, 6 insertions(+), 34 deletions(-) diff --git a/tests/llm_translation/test_anthropic_completion.py b/tests/llm_translation/test_anthropic_completion.py index 3c1162cd01..b6a98e4b2c 100644 --- a/tests/llm_translation/test_anthropic_completion.py +++ b/tests/llm_translation/test_anthropic_completion.py @@ -1350,6 +1350,9 @@ def test_anthropic_text_editor(): @pytest.mark.parametrize("spec", ["anthropic", "openai"]) +@pytest.mark.skipif( + os.getenv("ZAPIER_CI_CD_MCP_TOKEN") is None, reason="ZAPIER_CI_CD_MCP_TOKEN not set" +) def test_anthropic_mcp_server_tool_use(spec: str): litellm._turn_on_debug() @@ -1391,6 +1394,9 @@ def test_anthropic_mcp_server_tool_use(spec: str): @pytest.mark.parametrize( "model", ["openai/gpt-4.1", "anthropic/claude-sonnet-4-20250514"] ) +@pytest.mark.skipif( + os.getenv("ZAPIER_CI_CD_MCP_TOKEN") is None, reason="ZAPIER_CI_CD_MCP_TOKEN not set" +) def test_anthropic_mcp_server_responses_api(model: str): from litellm import responses diff --git a/tests/logging_callback_tests/test_alerting.py b/tests/logging_callback_tests/test_alerting.py index 2278b6ff93..86588bbd14 100644 --- a/tests/logging_callback_tests/test_alerting.py +++ b/tests/logging_callback_tests/test_alerting.py @@ -788,40 +788,6 @@ async def test_region_outage_alerting_called( mock_send_alert.assert_not_called() -@pytest.mark.asyncio -@pytest.mark.skip(reason="test only needs to run locally ") -async def test_alerting(): - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - "api_key": "bad_key", - }, - } - ], - debug_level="DEBUG", - set_verbose=True, - alerting_config=AlertingConfig( - alerting_threshold=10, # threshold for slow / hanging llm responses (in seconds). Defaults to 300 seconds - webhook_url=os.getenv( - "SLACK_WEBHOOK_URL" - ), # webhook you want to send alerts to - ), - ) - try: - await router.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hey, how's it going?"}], - ) - - except Exception: - pass - finally: - await asyncio.sleep(3) - - @pytest.mark.asyncio async def test_langfuse_trace_id(): """ From fe379fd7386c2b30794f4bd7e89ae3566654862f Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 20:27:39 -0700 Subject: [PATCH 11/16] test: rename env var --- tests/local_testing/test_arize_ai.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/local_testing/test_arize_ai.py b/tests/local_testing/test_arize_ai.py index 3b497d638a..138858cee0 100644 --- a/tests/local_testing/test_arize_ai.py +++ b/tests/local_testing/test_arize_ai.py @@ -48,8 +48,8 @@ async def test_async_dynamic_arize_config(): messages=[{"role": "user", "content": "hi test from arize dynamic config"}], temperature=0.1, user="OTEL_USER", - arize_api_key=os.getenv("ARIZE_SPACE_2_API_KEY"), - arize_space_key=os.getenv("ARIZE_SPACE_2_KEY"), + arize_api_key=os.getenv("ARIZE_SPACE_API_KEY"), + arize_space_key=os.getenv("ARIZE_SPACE_KEY"), ) await asyncio.sleep(2) From a92b31a63613a64413d628ba2b3d08effb137505 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 20:36:08 -0700 Subject: [PATCH 12/16] test: remove dead tests --- .../test_bedrock_completion.py | 75 +++--------- .../test_clarifai_completion.py | 109 ------------------ tests/local_testing/test_completion.py | 45 -------- tests/local_testing/test_streaming.py | 88 ++------------ 4 files changed, 24 insertions(+), 293 deletions(-) delete mode 100644 tests/llm_translation/test_clarifai_completion.py diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index b71e4e5187..f57c77db54 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -217,60 +217,6 @@ def test_completion_bedrock_claude_external_client_auth(): # test_completion_bedrock_claude_external_client_auth() -@pytest.mark.skip(reason="Expired token, need to renew") -def test_completion_bedrock_claude_sts_client_auth(): - print("\ncalling bedrock claude external client auth") - import os - - aws_access_key_id = os.environ["AWS_TEMP_ACCESS_KEY_ID"] - aws_secret_access_key = os.environ["AWS_TEMP_SECRET_ACCESS_KEY"] - aws_region_name = os.environ["AWS_REGION_NAME"] - aws_role_name = os.environ["AWS_TEMP_ROLE_NAME"] - - try: - import boto3 - - litellm.set_verbose = True - - response = completion( - model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", - messages=messages, - max_tokens=10, - temperature=0.1, - aws_region_name=aws_region_name, - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - aws_role_name=aws_role_name, - aws_session_name="my-test-session", - ) - - response = embedding( - model="cohere.embed-multilingual-v3", - input=["hello world"], - aws_region_name="us-east-1", - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - aws_role_name=aws_role_name, - aws_session_name="my-test-session", - ) - - response = completion( - model="gpt-3.5-turbo", - messages=messages, - aws_region_name="us-east-1", - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - aws_role_name=aws_role_name, - aws_session_name="my-test-session", - ) - # Add any assertions here to check the response - print(response) - except RateLimitError: - pass - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - @pytest.fixture() def bedrock_session_token_creds(): print("\ncalling oidc auto to get aws_session_token credentials") @@ -3413,7 +3359,8 @@ def test_bedrock_openai_imported_model(): print(f"URL: {url}") assert "bedrock-runtime.us-east-1.amazonaws.com" in url assert ( - "arn:aws:bedrock:us-east-1:117159858402:imported-model%2Fm4gc1mrfuddy" in url + "arn:aws:bedrock:us-east-1:117159858402:imported-model%2Fm4gc1mrfuddy" + in url ) assert "/invoke" in url @@ -3850,10 +3797,12 @@ def test_bedrock_openai_error_handling(): assert exc_info.value.status_code == 422 print("✓ Error handling works correctly") + # ============================================================================ # Nova Grounding (web_search_options) Unit Tests (Mocked) # ============================================================================ + def test_bedrock_nova_grounding_web_search_options_non_streaming(): """ Unit test for Nova grounding using web_search_options parameter (non-streaming). @@ -3907,7 +3856,9 @@ def test_bedrock_nova_grounding_web_search_options_non_streaming(): break assert system_tool_found, "systemTool with nova_grounding should be present" - print(f"✓ web_search_options correctly transformed to systemTool (non-streaming)") + print( + f"✓ web_search_options correctly transformed to systemTool (non-streaming)" + ) def test_bedrock_nova_grounding_with_function_tools(): @@ -3987,7 +3938,9 @@ def test_bedrock_nova_grounding_with_function_tools(): assert tool["systemTool"]["name"] == "nova_grounding" system_tool_found = True - assert function_tool_found, "Function tool (get_stock_price) should be present" + assert ( + function_tool_found + ), "Function tool (get_stock_price) should be present" assert system_tool_found, "systemTool (nova_grounding) should be present" print(f"✓ Both function tools and web_search_options correctly combined") @@ -4092,10 +4045,12 @@ def test_bedrock_nova_grounding_request_transformation(): mock_post.return_value = MagicMock( status_code=200, json=lambda: { - "output": {"message": {"role": "assistant", "content": [{"text": "Test"}]}}, + "output": { + "message": {"role": "assistant", "content": [{"text": "Test"}]} + }, "stopReason": "end_turn", - "usage": {"inputTokens": 10, "outputTokens": 5} - } + "usage": {"inputTokens": 10, "outputTokens": 5}, + }, ) try: diff --git a/tests/llm_translation/test_clarifai_completion.py b/tests/llm_translation/test_clarifai_completion.py deleted file mode 100644 index 5080413f2e..0000000000 --- a/tests/llm_translation/test_clarifai_completion.py +++ /dev/null @@ -1,109 +0,0 @@ -import sys, os -import traceback -from dotenv import load_dotenv -import asyncio, logging - -load_dotenv() -import os, io - -sys.path.insert( - 0, os.path.abspath("../..") -) # Adds the parent directory to the system path -import pytest -import litellm -from litellm import ( - embedding, - completion, - acompletion, - acreate, - completion_cost, - Timeout, - ModelResponse, -) -from litellm import RateLimitError - -# litellm.num_retries = 3 -litellm.cache = None -litellm.success_callback = [] -user_message = "Write a short poem about the sky" -messages = [{"content": user_message, "role": "user"}] - - -@pytest.fixture(autouse=True) -def reset_callbacks(): - print("\npytest fixture - resetting callbacks") - litellm.success_callback = [] - litellm._async_success_callback = [] - litellm.failure_callback = [] - litellm.callbacks = [] - - -@pytest.mark.skip(reason="Account rate limited.") -def test_completion_clarifai_claude_2_1(): - print("calling clarifai claude completion") - import os - - clarifai_pat = os.environ["CLARIFAI_API_KEY"] - - try: - response = completion( - model="clarifai/anthropic.completion.claude-2_1", - num_retries=3, - messages=messages, - max_tokens=10, - temperature=0.1, - ) - print(response) - - except RateLimitError: - pass - - except Exception as e: - pytest.fail(f"Error occured: {e}") - - -@pytest.mark.skip(reason="Account rate limited") -def test_completion_clarifai_mistral_large(): - try: - litellm.set_verbose = True - response: ModelResponse = completion( - model="clarifai/mistralai.completion.mistral-small", - messages=messages, - num_retries=3, - max_tokens=10, - temperature=0.78, - ) - # Add any assertions here to check the response - assert len(response.choices) > 0 - assert len(response.choices[0].message.content) > 0 - except RateLimitError: - pass - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - -@pytest.mark.skip(reason="Account rate limited") -@pytest.mark.asyncio -def test_async_completion_clarifai(): - import asyncio - - litellm.set_verbose = True - - async def test_get_response(): - user_message = "Hello, how are you?" - messages = [{"content": user_message, "role": "user"}] - try: - response = await acompletion( - model="clarifai/openai.chat-completion.GPT-4", - messages=messages, - num_retries=3, - timeout=10, - api_key=os.getenv("CLARIFAI_API_KEY"), - ) - print(f"response: {response}") - except litellm.Timeout as e: - pass - except Exception as e: - pytest.fail(f"An exception occurred: {e}") - - asyncio.run(test_get_response()) diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index af888e23cf..d5da0719e0 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -176,34 +176,6 @@ def predibase_mock_post(url, data=None, json=None, headers=None, timeout=None): return mock_response -# @pytest.mark.skip(reason="local-only test") -@pytest.mark.asyncio -async def test_completion_predibase(): - try: - litellm.set_verbose = True - - # with patch("requests.post", side_effect=predibase_mock_post): - response = await litellm.acompletion( - model="predibase/llama-3-8b-instruct", - tenant_id="c4768f95", - api_key=os.getenv("PREDIBASE_API_KEY"), - messages=[{"role": "user", "content": "who are u?"}], - max_tokens=10, - timeout=5, - ) - - print(response) - except litellm.Timeout as e: - print("got a timeout error from predibase") - pass - except litellm.ServiceUnavailableError as e: - pass - except litellm.InternalServerError: - pass - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - # test_completion_predibase() @@ -3238,23 +3210,6 @@ def test_mistral_anyscale_stream(): print(chunk["choices"][0]["delta"].get("content", ""), end="") -# test_completion_anyscale_2() -# def test_completion_with_fallbacks_multiple_keys(): -# print(f"backup key 1: {os.getenv('BACKUP_OPENAI_API_KEY_1')}") -# print(f"backup key 2: {os.getenv('BACKUP_OPENAI_API_KEY_2')}") -# backup_keys = [{"api_key": os.getenv("BACKUP_OPENAI_API_KEY_1")}, {"api_key": os.getenv("BACKUP_OPENAI_API_KEY_2")}] -# try: -# api_key = "bad-key" -# response = completion( -# model="gpt-3.5-turbo", messages=messages, force_timeout=120, fallbacks=backup_keys, api_key=api_key -# ) -# # Add any assertions here to check the response -# print(response) -# except Exception as e: -# error_str = traceback.format_exc() -# pytest.fail(f"Error occurred: {error_str}") - - # test_completion_with_fallbacks_multiple_keys() def test_petals(): try: diff --git a/tests/local_testing/test_streaming.py b/tests/local_testing/test_streaming.py index 56f0e5fe82..97d6875f16 100644 --- a/tests/local_testing/test_streaming.py +++ b/tests/local_testing/test_streaming.py @@ -552,7 +552,6 @@ async def test_completion_predibase_streaming(sync_mode): pytest.fail(f"Error occurred: {e}") - def test_completion_azure_function_calling_stream(): try: litellm.set_verbose = False @@ -1655,80 +1654,9 @@ def test_sagemaker_weird_response(): # test_sagemaker_weird_response() -@pytest.mark.skip(reason="Move to being a mock endpoint") -@pytest.mark.asyncio -async def test_sagemaker_streaming_async(): - try: - messages = [{"role": "user", "content": "Hey, how's it going?"}] - litellm.set_verbose = True - response = await litellm.acompletion( - model="sagemaker/jumpstart-dft-hf-llm-mistral-7b-ins-20240329-150233", - model_id="huggingface-llm-mistral-7b-instruct-20240329-150233", - messages=messages, - temperature=0.2, - max_tokens=80, - aws_region_name=os.getenv("AWS_REGION_NAME_2"), - aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID_2"), - aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY_2"), - stream=True, - ) - # Add any assertions here to check the response - print(response) - complete_response = "" - has_finish_reason = False - # Add any assertions here to check the response - idx = 0 - async for chunk in response: - # print - chunk, finished = streaming_format_tests(idx, chunk) - has_finish_reason = finished - complete_response += chunk - if finished: - break - idx += 1 - if has_finish_reason is False: - raise Exception("finish reason not set for last chunk") - if complete_response.strip() == "": - raise Exception("Empty response received") - print(f"completion_response: {complete_response}") - except Exception as e: - pytest.fail(f"An exception occurred - {str(e)}") - - # asyncio.run(test_sagemaker_streaming_async()) -@pytest.mark.skip(reason="costly sagemaker deployment. Move to mock implementation") -def test_completion_sagemaker_stream(): - try: - response = completion( - model="sagemaker/jumpstart-dft-hf-llm-mistral-7b-ins-20240329-150233", - model_id="huggingface-llm-mistral-7b-instruct-20240329-150233", - messages=messages, - temperature=0.2, - max_tokens=80, - aws_region_name=os.getenv("AWS_REGION_NAME_2"), - aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID_2"), - aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY_2"), - stream=True, - ) - complete_response = "" - has_finish_reason = False - # Add any assertions here to check the response - for idx, chunk in enumerate(response): - chunk, finished = streaming_format_tests(idx, chunk) - has_finish_reason = finished - if finished: - break - complete_response += chunk - if has_finish_reason is False: - raise Exception("finish reason not set for last chunk") - if complete_response.strip() == "": - raise Exception("Empty response received") - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - @pytest.mark.skip(reason="Account deleted by IBM.") @pytest.mark.asyncio async def test_completion_watsonx_stream(): @@ -2725,8 +2653,8 @@ def test_azure_streaming_and_function_calling(): tool_choice="auto", messages=messages, stream=True, - api_base=os.getenv("AZURE_API_BASE"), - api_key=os.getenv("AZURE_API_KEY"), + api_base=os.getenv("AZURE_AI_API_BASE"), + api_key=os.getenv("AZURE_AI_API_KEY"), api_version="2024-02-15-preview", ) # Add any assertions here to check the response @@ -2796,8 +2724,8 @@ async def test_azure_astreaming_and_function_calling(): tool_choice="auto", messages=messages, stream=True, - api_base=os.getenv("AZURE_API_BASE"), - api_key=os.getenv("AZURE_API_KEY"), + api_base=os.getenv("AZURE_AI_API_BASE"), + api_key=os.getenv("AZURE_AI_API_KEY"), api_version="2024-02-15-preview", caching=True, ) @@ -2827,8 +2755,8 @@ async def test_azure_astreaming_and_function_calling(): tool_choice="auto", messages=messages, stream=True, - api_base=os.getenv("AZURE_API_BASE"), - api_key=os.getenv("AZURE_API_KEY"), + api_base=os.getenv("AZURE_AI_API_BASE"), + api_key=os.getenv("AZURE_AI_API_KEY"), api_version="2024-02-15-preview", caching=True, ) @@ -3109,7 +3037,9 @@ def test_unit_test_custom_stream_wrapper_repeating_chunk( print(f"expected_chunk_fail: {expected_chunk_fail}") if (loop_amount > litellm.REPEATED_STREAMING_CHUNK_LIMIT) and expected_chunk_fail: - with pytest.raises((litellm.InternalServerError, litellm.exceptions.MidStreamFallbackError)): + with pytest.raises( + (litellm.InternalServerError, litellm.exceptions.MidStreamFallbackError) + ): for chunk in response: continue else: From 25f2baad717e8ff8d291794558e69952bc6b6bef Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 20:49:02 -0700 Subject: [PATCH 13/16] test: cleanup dead tests --- tests/llm_translation/test_snowflake.py | 176 ++--- tests/local_testing/test_router_init.py | 704 ------------------ .../test_amazing_s3_logs.py | 122 +-- tests/test_team_logging.py | 134 ---- 4 files changed, 86 insertions(+), 1050 deletions(-) delete mode 100644 tests/local_testing/test_router_init.py diff --git a/tests/llm_translation/test_snowflake.py b/tests/llm_translation/test_snowflake.py index 83aa5635f4..63878c8fc6 100644 --- a/tests/llm_translation/test_snowflake.py +++ b/tests/llm_translation/test_snowflake.py @@ -46,45 +46,51 @@ def mock_snowflake_streaming_response_chunks() -> List[str]: Mock streaming response chunks for Snowflake. """ return [ - json.dumps({ - "id": "chatcmpl-snowflake-stream-123", - "object": "chat.completion.chunk", - "created": 1700000000, - "model": "mistral-7b", - "choices": [ - { - "index": 0, - "delta": {"role": "assistant", "content": "The"}, - "finish_reason": None, - } - ], - }), - json.dumps({ - "id": "chatcmpl-snowflake-stream-123", - "object": "chat.completion.chunk", - "created": 1700000000, - "model": "mistral-7b", - "choices": [ - { - "index": 0, - "delta": {"content": " sky"}, - "finish_reason": None, - } - ], - }), - json.dumps({ - "id": "chatcmpl-snowflake-stream-123", - "object": "chat.completion.chunk", - "created": 1700000000, - "model": "mistral-7b", - "choices": [ - { - "index": 0, - "delta": {"content": " is blue"}, - "finish_reason": "stop", - } - ], - }), + json.dumps( + { + "id": "chatcmpl-snowflake-stream-123", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "mistral-7b", + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "The"}, + "finish_reason": None, + } + ], + } + ), + json.dumps( + { + "id": "chatcmpl-snowflake-stream-123", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "mistral-7b", + "choices": [ + { + "index": 0, + "delta": {"content": " sky"}, + "finish_reason": None, + } + ], + } + ), + json.dumps( + { + "id": "chatcmpl-snowflake-stream-123", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "mistral-7b", + "choices": [ + { + "index": 0, + "delta": {"content": " is blue"}, + "finish_reason": "stop", + } + ], + } + ), ] @@ -120,6 +126,7 @@ def test_chat_completion_snowflake(sync_mode): async_handler = AsyncHTTPHandler() with patch.object(AsyncHTTPHandler, "post", return_value=mock_response): import asyncio + response = asyncio.run( acompletion( model="snowflake/mistral-7b", @@ -148,16 +155,16 @@ def test_chat_completion_snowflake_stream(sync_mode): if sync_mode: sync_handler = HTTPHandler() mock_chunks = mock_snowflake_streaming_response_chunks() - + def mock_iter_lines(): for chunk in mock_chunks: for line in [f"data: {chunk}", "data: [DONE]"]: yield line - + mock_response = MagicMock() mock_response.iter_lines.side_effect = mock_iter_lines mock_response.status_code = 200 - + with patch.object(HTTPHandler, "post", return_value=mock_response): response = completion( model="snowflake/mistral-7b", @@ -167,28 +174,28 @@ def test_chat_completion_snowflake_stream(sync_mode): api_base="https://exampleopenaiendpoint-production.up.railway.app/v1/chat/completions", client=sync_handler, ) - + chunks_received = [] for chunk in response: chunks_received.append(chunk) - + assert len(chunks_received) > 0 else: async_handler = AsyncHTTPHandler() mock_chunks = mock_snowflake_streaming_response_chunks() - + async def mock_iter_lines(): for chunk in mock_chunks: for line in [f"data: {chunk}", "data: [DONE]"]: yield line - + mock_response = MagicMock() mock_response.iter_lines.side_effect = mock_iter_lines mock_response.status_code = 200 - + with patch.object(AsyncHTTPHandler, "post", return_value=mock_response): import asyncio - + async def test_async_stream(): response = await acompletion( model="snowflake/mistral-7b", @@ -198,78 +205,11 @@ def test_chat_completion_snowflake_stream(sync_mode): api_base="https://exampleopenaiendpoint-production.up.railway.app/v1/chat/completions", client=async_handler, ) - + chunks_received = [] async for chunk in response: chunks_received.append(chunk) - + assert len(chunks_received) > 0 - + asyncio.run(test_async_stream()) - - -@pytest.mark.skip(reason="Requires Snowflake credentials - run manually when needed") -def test_snowflake_tool_calling_responses_api(): - """ - Test Snowflake tool calling with Responses API. - Requires SNOWFLAKE_JWT and SNOWFLAKE_ACCOUNT_ID environment variables. - """ - import litellm - - # Skip if credentials not available - if not os.getenv("SNOWFLAKE_JWT") or not os.getenv("SNOWFLAKE_ACCOUNT_ID"): - pytest.skip("Snowflake credentials not available") - - litellm.drop_params = False # We now support tools! - - tools = [ - { - "type": "function", - "name": "get_weather", - "description": "Get the current weather in a given location", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "The city and state, e.g. San Francisco, CA", - } - }, - "required": ["location"], - }, - } - ] - - try: - # Test with tool_choice to force tool use - response = responses( - model="snowflake/claude-3-5-sonnet", - input="What's the weather in Paris?", - tools=tools, - tool_choice={"type": "function", "function": {"name": "get_weather"}}, - max_output_tokens=200, - ) - - assert response is not None - assert hasattr(response, "output") - assert len(response.output) > 0 - - # Verify tool call was made - tool_call_found = False - for item in response.output: - if hasattr(item, "type") and item.type == "function_call": - tool_call_found = True - assert item.name == "get_weather" - assert hasattr(item, "arguments") - print(f"✅ Tool call detected: {item.name}({item.arguments})") - break - - assert tool_call_found, "Expected tool call but none was found" - - except APIConnectionError as e: - if "JWT token is invalid" in str(e): - pytest.skip("Invalid Snowflake JWT token") - elif "Application failed to respond" in str(e) or "502" in str(e): - pytest.skip(f"Snowflake API unavailable: {e}") - else: - raise diff --git a/tests/local_testing/test_router_init.py b/tests/local_testing/test_router_init.py deleted file mode 100644 index 6f3453cbdd..0000000000 --- a/tests/local_testing/test_router_init.py +++ /dev/null @@ -1,704 +0,0 @@ -# # this tests if the router is initialized correctly -# import asyncio -# import os -# import sys -# import time -# import traceback - -# import pytest - -# sys.path.insert( -# 0, os.path.abspath("../..") -# ) # Adds the parent directory to the system path -# from collections import defaultdict -# from concurrent.futures import ThreadPoolExecutor - -# from dotenv import load_dotenv - -# import litellm -# from litellm import Router - -# load_dotenv() - -# # every time we load the router we should have 4 clients: -# # Async -# # Sync -# # Async + Stream -# # Sync + Stream - - -# def test_init_clients(): -# litellm.set_verbose = True -# import logging - -# from litellm._logging import verbose_router_logger - -# verbose_router_logger.setLevel(logging.DEBUG) -# try: -# print("testing init 4 clients with diff timeouts") -# model_list = [ -# { -# "model_name": "gpt-3.5-turbo", -# "litellm_params": { -# "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_AI_API_KEY"), -# "api_version": os.getenv("AZURE_API_VERSION"), -# "api_base": os.getenv("AZURE_AI_API_BASE"), -# "timeout": 0.01, -# "stream_timeout": 0.000_001, -# "max_retries": 7, -# }, -# }, -# ] -# router = Router(model_list=model_list, set_verbose=True) -# for elem in router.model_list: -# model_id = elem["model_info"]["id"] -# assert router.cache.get_cache(f"{model_id}_client") is not None -# assert router.cache.get_cache(f"{model_id}_async_client") is not None -# assert router.cache.get_cache(f"{model_id}_stream_client") is not None -# assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None - -# # check if timeout for stream/non stream clients is set correctly -# async_client = router.cache.get_cache(f"{model_id}_async_client") -# stream_async_client = router.cache.get_cache( -# f"{model_id}_stream_async_client" -# ) - -# assert async_client.timeout == 0.01 -# assert stream_async_client.timeout == 0.000_001 -# print(vars(async_client)) -# print() -# print(async_client._base_url) -# assert ( -# async_client._base_url -# == "https://openai-gpt-4-test-v-1.openai.azure.com/openai/" -# ) -# assert ( -# stream_async_client._base_url -# == "https://openai-gpt-4-test-v-1.openai.azure.com/openai/" -# ) - -# print("PASSED !") - -# except Exception as e: -# traceback.print_exc() -# pytest.fail(f"Error occurred: {e}") - - -# # test_init_clients() - - -# def test_init_clients_basic(): -# litellm.set_verbose = True -# try: -# print("Test basic client init") -# model_list = [ -# { -# "model_name": "gpt-3.5-turbo", -# "litellm_params": { -# "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_AI_API_KEY"), -# "api_version": os.getenv("AZURE_API_VERSION"), -# "api_base": os.getenv("AZURE_AI_API_BASE"), -# }, -# }, -# ] -# router = Router(model_list=model_list) -# for elem in router.model_list: -# model_id = elem["model_info"]["id"] -# assert router.cache.get_cache(f"{model_id}_client") is not None -# assert router.cache.get_cache(f"{model_id}_async_client") is not None -# assert router.cache.get_cache(f"{model_id}_stream_client") is not None -# assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None -# print("PASSED !") - -# # see if we can init clients without timeout or max retries set -# except Exception as e: -# traceback.print_exc() -# pytest.fail(f"Error occurred: {e}") - - -# # test_init_clients_basic() - - -# def test_init_clients_basic_azure_cloudflare(): -# # init azure + cloudflare -# # init OpenAI gpt-3.5 -# # init OpenAI text-embedding -# # init OpenAI comptaible - Mistral/mistral-medium -# # init OpenAI compatible - xinference/bge -# litellm.set_verbose = True -# try: -# print("Test basic client init") -# model_list = [ -# { -# "model_name": "azure-cloudflare", -# "litellm_params": { -# "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_AI_API_KEY"), -# "api_version": os.getenv("AZURE_API_VERSION"), -# "api_base": "https://gateway.ai.cloudflare.com/v1/0399b10e77ac6668c80404a5ff49eb37/litellm-test/azure-openai/openai-gpt-4-test-v-1", -# }, -# }, -# { -# "model_name": "gpt-openai", -# "litellm_params": { -# "model": "gpt-3.5-turbo", -# "api_key": os.getenv("OPENAI_API_KEY"), -# }, -# }, -# { -# "model_name": "text-embedding-ada-002", -# "litellm_params": { -# "model": "text-embedding-ada-002", -# "api_key": os.getenv("OPENAI_API_KEY"), -# }, -# }, -# { -# "model_name": "mistral", -# "litellm_params": { -# "model": "mistral/mistral-tiny", -# "api_key": os.getenv("MISTRAL_API_KEY"), -# }, -# }, -# { -# "model_name": "bge-base-en", -# "litellm_params": { -# "model": "xinference/bge-base-en", -# "api_base": "http://127.0.0.1:9997/v1", -# "api_key": os.getenv("OPENAI_API_KEY"), -# }, -# }, -# ] -# router = Router(model_list=model_list) -# for elem in router.model_list: -# model_id = elem["model_info"]["id"] -# assert router.cache.get_cache(f"{model_id}_client") is not None -# assert router.cache.get_cache(f"{model_id}_async_client") is not None -# assert router.cache.get_cache(f"{model_id}_stream_client") is not None -# assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None -# print("PASSED !") - -# # see if we can init clients without timeout or max retries set -# except Exception as e: -# traceback.print_exc() -# pytest.fail(f"Error occurred: {e}") - - -# # test_init_clients_basic_azure_cloudflare() - - -# def test_timeouts_router(): -# """ -# Test the timeouts of the router with multiple clients. This HASas to raise a timeout error -# """ -# import openai - -# litellm.set_verbose = True -# try: -# print("testing init 4 clients with diff timeouts") -# model_list = [ -# { -# "model_name": "gpt-3.5-turbo", -# "litellm_params": { -# "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_AI_API_KEY"), -# "api_version": os.getenv("AZURE_API_VERSION"), -# "api_base": os.getenv("AZURE_AI_API_BASE"), -# "timeout": 0.000001, -# "stream_timeout": 0.000_001, -# }, -# }, -# ] -# router = Router(model_list=model_list, num_retries=0) - -# print("PASSED !") - -# async def test(): -# try: -# await router.acompletion( -# model="gpt-3.5-turbo", -# messages=[ -# {"role": "user", "content": "hello, write a 20 pg essay"} -# ], -# ) -# except Exception as e: -# raise e - -# asyncio.run(test()) -# except openai.APITimeoutError as e: -# print( -# "Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e -# ) -# print(type(e)) -# pass -# except Exception as e: -# pytest.fail( -# f"Did not raise error `openai.APITimeoutError`. Instead raised error type: {type(e)}, Error: {e}" -# ) - - -# # test_timeouts_router() - - -# def test_stream_timeouts_router(): -# """ -# Test the stream timeouts router. See if it selected the correct client with stream timeout -# """ -# import openai - -# litellm.set_verbose = True -# try: -# print("testing init 4 clients with diff timeouts") -# model_list = [ -# { -# "model_name": "gpt-3.5-turbo", -# "litellm_params": { -# "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_AI_API_KEY"), -# "api_version": os.getenv("AZURE_API_VERSION"), -# "api_base": os.getenv("AZURE_AI_API_BASE"), -# "timeout": 200, # regular calls will not timeout, stream calls will -# "stream_timeout": 10, -# }, -# }, -# ] -# router = Router(model_list=model_list) - -# print("PASSED !") -# data = { -# "model": "gpt-3.5-turbo", -# "messages": [{"role": "user", "content": "hello, write a 20 pg essay"}], -# "stream": True, -# } -# selected_client = router._get_client( -# deployment=router.model_list[0], -# kwargs=data, -# client_type=None, -# ) -# print("Select client timeout", selected_client.timeout) -# assert selected_client.timeout == 10 - -# # make actual call -# response = router.completion(**data) - -# for chunk in response: -# print(f"chunk: {chunk}") -# except openai.APITimeoutError as e: -# print( -# "Passed: Raised correct exception. Got openai.APITimeoutError\nGood Job", e -# ) -# print(type(e)) -# pass -# except Exception as e: -# pytest.fail( -# f"Did not raise error `openai.APITimeoutError`. Instead raised error type: {type(e)}, Error: {e}" -# ) - - -# # test_stream_timeouts_router() - - -# def test_xinference_embedding(): -# # [Test Init Xinference] this tests if we init xinference on the router correctly -# # [Test Exception Mapping] tests that xinference is an openai comptiable provider -# print("Testing init xinference") -# print( -# "this tests if we create an OpenAI client for Xinference, with the correct API BASE" -# ) - -# model_list = [ -# { -# "model_name": "xinference", -# "litellm_params": { -# "model": "xinference/bge-base-en", -# "api_base": "os.environ/XINFERENCE_API_BASE", -# }, -# } -# ] - -# router = Router(model_list=model_list) - -# print(router.model_list) -# print(router.model_list[0]) - -# assert ( -# router.model_list[0]["litellm_params"]["api_base"] == "http://0.0.0.0:9997" -# ) # set in env - -# openai_client = router._get_client( -# deployment=router.model_list[0], -# kwargs={"input": ["hello"], "model": "xinference"}, -# ) - -# assert openai_client._base_url == "http://0.0.0.0:9997" -# assert "xinference" in litellm.openai_compatible_providers -# print("passed") - - -# # test_xinference_embedding() - - -# def test_router_init_gpt_4_vision_enhancements(): -# try: -# # tests base_url set when any base_url with /openai/deployments passed to router -# print("Testing Azure GPT_Vision enhancements") - -# model_list = [ -# { -# "model_name": "gpt-4-vision-enhancements", -# "litellm_params": { -# "model": "azure/gpt-4-vision", -# "api_key": os.getenv("AZURE_AI_API_KEY"), -# "base_url": "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/", -# "dataSources": [ -# { -# "type": "AzureComputerVision", -# "parameters": { -# "endpoint": "os.environ/AZURE_VISION_ENHANCE_ENDPOINT", -# "key": "os.environ/AZURE_VISION_ENHANCE_KEY", -# }, -# } -# ], -# }, -# } -# ] - -# router = Router(model_list=model_list) - -# print(router.model_list) -# print(router.model_list[0]) - -# assert ( -# router.model_list[0]["litellm_params"]["base_url"] -# == "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/" -# ) # set in env - -# assert ( -# router.model_list[0]["litellm_params"]["dataSources"][0]["parameters"][ -# "endpoint" -# ] -# == os.environ["AZURE_VISION_ENHANCE_ENDPOINT"] -# ) - -# assert ( -# router.model_list[0]["litellm_params"]["dataSources"][0]["parameters"][ -# "key" -# ] -# == os.environ["AZURE_VISION_ENHANCE_KEY"] -# ) - -# azure_client = router._get_client( -# deployment=router.model_list[0], -# kwargs={"stream": True, "model": "gpt-4-vision-enhancements"}, -# client_type="async", -# ) - -# assert ( -# azure_client._base_url -# == "https://gpt-4-vision-resource.openai.azure.com/openai/deployments/gpt-4-vision/extensions/" -# ) -# print("passed") -# except Exception as e: -# pytest.fail(f"Error occurred: {e}") - - -# @pytest.mark.parametrize("sync_mode", [True, False]) -# @pytest.mark.asyncio -# async def test_openai_with_organization(sync_mode): -# try: -# print("Testing OpenAI with organization") -# model_list = [ -# { -# "model_name": "openai-bad-org", -# "litellm_params": { -# "model": "gpt-3.5-turbo", -# "organization": "org-ikDc4ex8NB", -# }, -# }, -# { -# "model_name": "openai-good-org", -# "litellm_params": {"model": "gpt-3.5-turbo"}, -# }, -# ] - -# router = Router(model_list=model_list) - -# print(router.model_list) -# print(router.model_list[0]) - -# if sync_mode: -# openai_client = router._get_client( -# deployment=router.model_list[0], -# kwargs={"input": ["hello"], "model": "openai-bad-org"}, -# ) -# print(vars(openai_client)) - -# assert openai_client.organization == "org-ikDc4ex8NB" - -# # bad org raises error - -# try: -# response = router.completion( -# model="openai-bad-org", -# messages=[{"role": "user", "content": "this is a test"}], -# ) -# pytest.fail( -# "Request should have failed - This organization does not exist" -# ) -# except Exception as e: -# print("Got exception: " + str(e)) -# assert "header should match organization for API key" in str( -# e -# ) or "No such organization" in str(e) - -# # good org works -# response = router.completion( -# model="openai-good-org", -# messages=[{"role": "user", "content": "this is a test"}], -# max_tokens=5, -# ) -# else: -# openai_client = router._get_client( -# deployment=router.model_list[0], -# kwargs={"input": ["hello"], "model": "openai-bad-org"}, -# client_type="async", -# ) -# print(vars(openai_client)) - -# assert openai_client.organization == "org-ikDc4ex8NB" - -# # bad org raises error - -# try: -# response = await router.acompletion( -# model="openai-bad-org", -# messages=[{"role": "user", "content": "this is a test"}], -# ) -# pytest.fail( -# "Request should have failed - This organization does not exist" -# ) -# except Exception as e: -# print("Got exception: " + str(e)) -# assert "header should match organization for API key" in str( -# e -# ) or "No such organization" in str(e) - -# # good org works -# response = await router.acompletion( -# model="openai-good-org", -# messages=[{"role": "user", "content": "this is a test"}], -# max_tokens=5, -# ) - -# except Exception as e: -# pytest.fail(f"Error occurred: {e}") - - -# def test_init_clients_azure_command_r_plus(): -# # This tests that the router uses the OpenAI client for Azure/Command-R+ -# # For azure/command-r-plus we need to use openai.OpenAI because of how the Azure provider requires requests being sent -# litellm.set_verbose = True -# import logging - -# from litellm._logging import verbose_router_logger - -# verbose_router_logger.setLevel(logging.DEBUG) -# try: -# print("testing init 4 clients with diff timeouts") -# model_list = [ -# { -# "model_name": "gpt-3.5-turbo", -# "litellm_params": { -# "model": "azure/command-r-plus", -# "api_key": os.getenv("AZURE_COHERE_API_KEY"), -# "api_base": os.getenv("AZURE_COHERE_API_BASE"), -# "timeout": 0.01, -# "stream_timeout": 0.000_001, -# "max_retries": 7, -# }, -# }, -# ] -# router = Router(model_list=model_list, set_verbose=True) -# for elem in router.model_list: -# model_id = elem["model_info"]["id"] -# async_client = router.cache.get_cache(f"{model_id}_async_client") -# stream_async_client = router.cache.get_cache( -# f"{model_id}_stream_async_client" -# ) -# # Assert the Async Clients used are OpenAI clients and not Azure -# # For using Azure/Command-R-Plus and Azure/Mistral the clients NEED to be OpenAI clients used -# # this is weirdness introduced on Azure's side - -# assert "openai.AsyncOpenAI" in str(async_client) -# assert "openai.AsyncOpenAI" in str(stream_async_client) -# print("PASSED !") - -# except Exception as e: -# traceback.print_exc() -# pytest.fail(f"Error occurred: {e}") - - -# @pytest.mark.asyncio -# async def test_aaaaatext_completion_with_organization(): -# try: -# print("Testing Text OpenAI with organization") -# model_list = [ -# { -# "model_name": "openai-bad-org", -# "litellm_params": { -# "model": "text-completion-openai/gpt-3.5-turbo-instruct", -# "api_key": os.getenv("OPENAI_API_KEY", None), -# "organization": "org-ikDc4ex8NB", -# }, -# }, -# { -# "model_name": "openai-good-org", -# "litellm_params": { -# "model": "text-completion-openai/gpt-3.5-turbo-instruct", -# "api_key": os.getenv("OPENAI_API_KEY", None), -# "organization": os.getenv("OPENAI_ORGANIZATION", None), -# }, -# }, -# ] - -# router = Router(model_list=model_list) - -# print(router.model_list) -# print(router.model_list[0]) - -# openai_client = router._get_client( -# deployment=router.model_list[0], -# kwargs={"input": ["hello"], "model": "openai-bad-org"}, -# ) -# print(vars(openai_client)) - -# assert openai_client.organization == "org-ikDc4ex8NB" - -# # bad org raises error - -# try: -# response = await router.atext_completion( -# model="openai-bad-org", -# prompt="this is a test", -# ) -# pytest.fail("Request should have failed - This organization does not exist") -# except Exception as e: -# print("Got exception: " + str(e)) -# assert "header should match organization for API key" in str( -# e -# ) or "No such organization" in str(e) - -# # good org works -# response = await router.atext_completion( -# model="openai-good-org", -# prompt="this is a test", -# max_tokens=5, -# ) -# print("working response: ", response) - -# except Exception as e: -# pytest.fail(f"Error occurred: {e}") - - -# def test_init_clients_async_mode(): -# litellm.set_verbose = True -# import logging - -# from litellm._logging import verbose_router_logger -# from litellm.types.router import RouterGeneralSettings - -# verbose_router_logger.setLevel(logging.DEBUG) -# try: -# print("testing init 4 clients with diff timeouts") -# model_list = [ -# { -# "model_name": "gpt-3.5-turbo", -# "litellm_params": { -# "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_AI_API_KEY"), -# "api_version": os.getenv("AZURE_API_VERSION"), -# "api_base": os.getenv("AZURE_AI_API_BASE"), -# "timeout": 0.01, -# "stream_timeout": 0.000_001, -# "max_retries": 7, -# }, -# }, -# ] -# router = Router( -# model_list=model_list, -# set_verbose=True, -# router_general_settings=RouterGeneralSettings(async_only_mode=True), -# ) -# for elem in router.model_list: -# model_id = elem["model_info"]["id"] - -# # sync clients not initialized in async_only_mode=True -# assert router.cache.get_cache(f"{model_id}_client") is None -# assert router.cache.get_cache(f"{model_id}_stream_client") is None - -# # only async clients initialized in async_only_mode=True -# assert router.cache.get_cache(f"{model_id}_async_client") is not None -# assert router.cache.get_cache(f"{model_id}_stream_async_client") is not None -# except Exception as e: -# pytest.fail(f"Error occurred: {e}") - - -# @pytest.mark.parametrize( -# "environment,expected_models", -# [ -# ("development", ["gpt-3.5-turbo"]), -# ("production", ["gpt-4", "gpt-3.5-turbo", "gpt-4o"]), -# ], -# ) -# def test_init_router_with_supported_environments(environment, expected_models): -# """ -# Tests that the correct models are setup on router when LITELLM_ENVIRONMENT is set -# """ -# os.environ["LITELLM_ENVIRONMENT"] = environment -# model_list = [ -# { -# "model_name": "gpt-3.5-turbo", -# "litellm_params": { -# "model": "azure/gpt-4.1-mini", -# "api_key": os.getenv("AZURE_AI_API_KEY"), -# "api_version": os.getenv("AZURE_API_VERSION"), -# "api_base": os.getenv("AZURE_AI_API_BASE"), -# "timeout": 0.01, -# "stream_timeout": 0.000_001, -# "max_retries": 7, -# }, -# "model_info": {"supported_environments": ["development", "production"]}, -# }, -# { -# "model_name": "gpt-4", -# "litellm_params": { -# "model": "openai/gpt-4", -# "api_key": os.getenv("OPENAI_API_KEY"), -# "timeout": 0.01, -# "stream_timeout": 0.000_001, -# "max_retries": 7, -# }, -# "model_info": {"supported_environments": ["production"]}, -# }, -# { -# "model_name": "gpt-4o", -# "litellm_params": { -# "model": "openai/gpt-4o", -# "api_key": os.getenv("OPENAI_API_KEY"), -# "timeout": 0.01, -# "stream_timeout": 0.000_001, -# "max_retries": 7, -# }, -# "model_info": {"supported_environments": ["production"]}, -# }, -# ] -# router = Router(model_list=model_list, set_verbose=True) -# _model_list = router.get_model_names() - -# print("model_list: ", _model_list) -# print("expected_models: ", expected_models) - -# assert set(_model_list) == set(expected_models) - -# os.environ.pop("LITELLM_ENVIRONMENT") diff --git a/tests/logging_callback_tests/test_amazing_s3_logs.py b/tests/logging_callback_tests/test_amazing_s3_logs.py index f074026926..c1e6884278 100644 --- a/tests/logging_callback_tests/test_amazing_s3_logs.py +++ b/tests/logging_callback_tests/test_amazing_s3_logs.py @@ -74,15 +74,13 @@ async def test_basic_s3_logging(sync_mode, streaming): s3.delete_object(Bucket="load-testing-oct", Key=key) - @pytest.mark.asyncio -@pytest.mark.parametrize( - "streaming", [(True)] -) +@pytest.mark.parametrize("streaming", [True]) @pytest.mark.flaky(retries=3, delay=1) async def test_basic_s3_v2_logging(streaming): from blockbuster import BlockBuster from litellm.integrations.s3_v2 import S3Logger + s3_v2_logger = S3Logger(s3_flush_interval=1) litellm.callbacks = [s3_v2_logger] blockbuster = BlockBuster() @@ -120,7 +118,7 @@ async def test_basic_s3_v2_logging(streaming): print(f"all_s3_keys: {all_s3_keys}") - #assert that atlest one key has response.id in it + # assert that atlest one key has response.id in it assert any(response_id in key for key in all_s3_keys) s3 = boto3.client("s3") # delete all objects @@ -134,22 +132,22 @@ async def test_basic_s3_v2_logging_failure(): """Test that S3 v2 logger makes httpx PUT request when logging failures""" from unittest.mock import AsyncMock, MagicMock, patch from litellm.integrations.s3_v2 import S3Logger - + # Create S3 logger with short flush interval s3_v2_logger = S3Logger(s3_flush_interval=1) - + # Mock the httpx client to capture the PUT request mock_response = MagicMock() mock_response.status_code = 200 mock_response.raise_for_status = MagicMock() - + s3_v2_logger.async_httpx_client = AsyncMock() s3_v2_logger.async_httpx_client.put.return_value = mock_response - + # Track the upload method calls original_upload = s3_v2_logger.async_upload_data_to_s3 upload_called = False - + async def mock_upload(batch_logging_element): nonlocal upload_called upload_called = True @@ -157,12 +155,12 @@ async def test_basic_s3_v2_logging_failure(): url = f"https://test-bucket.s3.us-west-2.amazonaws.com/{batch_logging_element.s3_object_key}" headers = {"Content-Type": "application/json"} data = '{"model": "gpt-4o-mini"}' - + # Make the actual httpx call we want to test await s3_v2_logger.async_httpx_client.put(url=url, headers=headers, data=data) - + s3_v2_logger.async_upload_data_to_s3 = mock_upload - + # Configure S3 callback params litellm.callbacks = [s3_v2_logger] litellm.s3_callback_params = { @@ -172,7 +170,7 @@ async def test_basic_s3_v2_logging_failure(): "s3_region_name": "us-west-2", } litellm.set_verbose = True - + # Trigger a failure by using invalid API key try: response = await litellm.acompletion( @@ -182,33 +180,33 @@ async def test_basic_s3_v2_logging_failure(): ) except Exception as e: print(f"Expected error: {e}") - + # Wait for logger to process the failure await asyncio.sleep(5) - + # Verify that our mock upload was called assert upload_called, "S3 upload method was not called" print("✓ S3 upload method was called") - + # Verify that httpx PUT was called s3_v2_logger.async_httpx_client.put.assert_called() - + # Get the call arguments to verify the S3 URL call_args = s3_v2_logger.async_httpx_client.put.call_args assert call_args is not None - url = call_args[1]['url'] if 'url' in call_args[1] else call_args[0][0] - + url = call_args[1]["url"] if "url" in call_args[1] else call_args[0][0] + # Verify the URL contains expected S3 endpoint assert "test-bucket.s3.us-west-2.amazonaws.com" in url print(f"✓ S3 PUT request made to: {url}") - + # Verify headers include expected content type - headers = call_args[1]['headers'] - assert headers['Content-Type'] == 'application/json' + headers = call_args[1]["headers"] + assert headers["Content-Type"] == "application/json" print("✓ S3 request headers are correct") - + # Verify JSON data was included - data = call_args[1]['data'] + data = call_args[1]["data"] assert data is not None assert '"model": "gpt-4o-mini"' in data print("✓ S3 request data contains expected log payload") @@ -411,83 +409,19 @@ async def make_async_calls(): return total_time -@pytest.mark.skip(reason="flaky test on ci/cd") -def test_s3_logging_r2(): - # all s3 requests need to be in one test function - # since we are modifying stdout, and pytests runs tests in parallel - # on circle ci - we only test litellm.acompletion() - try: - # redirect stdout to log_file - # litellm.cache = litellm.Cache( - # type="s3", s3_bucket_name="litellm-r2-bucket", s3_region_name="us-west-2" - # ) - litellm.set_verbose = True - from litellm._logging import verbose_logger - import logging - - verbose_logger.setLevel(level=logging.DEBUG) - - litellm.success_callback = ["s3"] - litellm.s3_callback_params = { - "s3_bucket_name": "litellm-r2-bucket", - "s3_aws_secret_access_key": "os.environ/R2_S3_ACCESS_KEY", - "s3_aws_access_key_id": "os.environ/R2_S3_ACCESS_ID", - "s3_endpoint_url": "os.environ/R2_S3_URL", - "s3_region_name": "os.environ/R2_S3_REGION_NAME", - } - print("Testing async s3 logging") - - expected_keys = [] - - import time - - curr_time = str(time.time()) - - async def _test(): - return await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": f"This is a test {curr_time}"}], - max_tokens=10, - temperature=0.7, - user="ishaan-2", - ) - - response = asyncio.run(_test()) - print(f"response: {response}") - expected_keys.append(response.id) - - import boto3 - - s3 = boto3.client( - "s3", - endpoint_url=os.getenv("R2_S3_URL"), - region_name=os.getenv("R2_S3_REGION_NAME"), - aws_access_key_id=os.getenv("R2_S3_ACCESS_ID"), - aws_secret_access_key=os.getenv("R2_S3_ACCESS_KEY"), - ) - - bucket_name = "litellm-r2-bucket" - # List objects in the bucket - response = s3.list_objects(Bucket=bucket_name) - - except Exception as e: - pytest.fail(f"An exception occurred - {e}") - finally: - # post, close log file and verify - # Reset stdout to the original value - print("Passed! Testing async s3 logging") - from litellm.integrations.s3_v2 import S3Logger + class TestS3Logger(S3Logger): def __init__(self, *args, **kwargs): self.recorded_requests = {} self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None super().__init__(*args, **kwargs) - + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): self.recorded_requests[response_obj["id"]] = start_time print("recorded request", self.recorded_requests) self.logged_standard_logging_payload = kwargs["standard_logging_object"] - return await super().async_log_success_event(kwargs, response_obj, start_time, end_time) - + return await super().async_log_success_event( + kwargs, response_obj, start_time, end_time + ) diff --git a/tests/test_team_logging.py b/tests/test_team_logging.py index 913b1e1949..9e89d945ed 100644 --- a/tests/test_team_logging.py +++ b/tests/test_team_logging.py @@ -59,137 +59,3 @@ async def chat_completion(session, key, model="azure-gpt-3.5", request_metadata= if status != 200: raise Exception(f"Request did not return a 200 status code: {status}") - - -@pytest.mark.skip(reason="flaky test - covered by simpler unit testing.") -@pytest.mark.asyncio -@pytest.mark.flaky(retries=12, delay=2) -async def test_aaateam_logging(): - """ - -> Team 1 logs to project 1 - -> Create Key - -> Make chat/completions call - -> Fetch logs from langfuse - """ - try: - async with aiohttp.ClientSession() as session: - - key = await generate_key( - session, models=["fake-openai-endpoint"], team_id="team-1" - ) # team-1 logs to project 1 - - from litellm._uuid import uuid - - _trace_id = f"trace-{uuid.uuid4()}" - _request_metadata = { - "trace_id": _trace_id, - } - - await chat_completion( - session, - key["key"], - model="fake-openai-endpoint", - request_metadata=_request_metadata, - ) - - # Test - if the logs were sent to the correct team on langfuse - import langfuse - - print(f"langfuse_public_key: {os.getenv('LANGFUSE_PROJECT1_PUBLIC')}") - print(f"langfuse_secret_key: {os.getenv('LANGFUSE_HOST')}") - langfuse_client = langfuse.Langfuse( - public_key=os.getenv("LANGFUSE_PROJECT1_PUBLIC"), - secret_key=os.getenv("LANGFUSE_PROJECT1_SECRET"), - host="https://us.cloud.langfuse.com", - ) - - await asyncio.sleep(30) - - print(f"searching for trace_id={_trace_id} on langfuse") - - generations = langfuse_client.get_generations(trace_id=_trace_id).data - print(generations) - assert len(generations) == 1 - except Exception as e: - pytest.fail(f"Unexpected error: {str(e)}") - - -@pytest.mark.skip(reason="todo fix langfuse credential error") -@pytest.mark.asyncio -async def test_team_2logging(): - """ - -> Team 1 logs to project 2 - -> Create Key - -> Make chat/completions call - -> Fetch logs from langfuse - """ - langfuse_public_key = os.getenv("LANGFUSE_PROJECT2_PUBLIC") - - print(f"langfuse_public_key: {langfuse_public_key}") - langfuse_secret_key = os.getenv("LANGFUSE_PROJECT2_SECRET") - print(f"langfuse_secret_key: {langfuse_secret_key}") - langfuse_host = "https://us.cloud.langfuse.com" - - try: - assert langfuse_public_key is not None - assert langfuse_secret_key is not None - except Exception as e: - # skip test if langfuse credentials are not set - return - - try: - async with aiohttp.ClientSession() as session: - - key = await generate_key( - session, models=["fake-openai-endpoint"], team_id="team-2" - ) # team-1 logs to project 1 - - from litellm._uuid import uuid - - _trace_id = f"trace-{uuid.uuid4()}" - _request_metadata = { - "trace_id": _trace_id, - } - - await chat_completion( - session, - key["key"], - model="fake-openai-endpoint", - request_metadata=_request_metadata, - ) - - # Test - if the logs were sent to the correct team on langfuse - import langfuse - - langfuse_client = langfuse.Langfuse( - public_key=langfuse_public_key, - secret_key=langfuse_secret_key, - host=langfuse_host, - ) - - await asyncio.sleep(30) - - print(f"searching for trace_id={_trace_id} on langfuse") - - generations = langfuse_client.get_generations(trace_id=_trace_id).data - print("Team 2 generations", generations) - - # team-2 should have 1 generation with this trace id - assert len(generations) == 1 - - # team-1 should have 0 generations with this trace id - langfuse_client_1 = langfuse.Langfuse( - public_key=os.getenv("LANGFUSE_PROJECT1_PUBLIC"), - secret_key=os.getenv("LANGFUSE_PROJECT1_SECRET"), - host="https://us.cloud.langfuse.com", - ) - - generations_team_1 = langfuse_client_1.get_generations( - trace_id=_trace_id - ).data - print("Team 1 generations", generations_team_1) - - assert len(generations_team_1) == 0 - - except Exception as e: - pytest.fail("Team 2 logging failed: " + str(e)) From 9e070143fbe9f90d27a160257abc4a32a2a96411 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 21:13:16 -0700 Subject: [PATCH 14/16] test: update key names --- .../test_openai_batches_and_files.py | 4 +- tests/local_testing/test_alangfuse.py | 263 ------------------ tests/local_testing/test_embedding.py | 2 +- tests/local_testing/test_gcs_bucket.py | 51 +++- 4 files changed, 46 insertions(+), 274 deletions(-) diff --git a/tests/batches_tests/test_openai_batches_and_files.py b/tests/batches_tests/test_openai_batches_and_files.py index 2ece740bf5..e4e1180eb1 100644 --- a/tests/batches_tests/test_openai_batches_and_files.py +++ b/tests/batches_tests/test_openai_batches_and_files.py @@ -75,8 +75,8 @@ def load_vertex_ai_credentials(): service_account_key_data = {} # Update the service_account_key_data with environment variables - private_key_id = os.environ.get("GCS_PRIVATE_KEY_ID", "") - private_key = os.environ.get("GCS_PRIVATE_KEY", "") + private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "") + private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "") private_key = private_key.replace("\\n", "\n") service_account_key_data["private_key_id"] = private_key_id service_account_key_data["private_key"] = private_key diff --git a/tests/local_testing/test_alangfuse.py b/tests/local_testing/test_alangfuse.py index 306c7749f1..7c2ec7e9f6 100644 --- a/tests/local_testing/test_alangfuse.py +++ b/tests/local_testing/test_alangfuse.py @@ -599,233 +599,6 @@ def test_langfuse_logging_function_calling(): # test_langfuse_logging_function_calling() -@pytest.mark.skip(reason="Need to address this on main") -def test_aaalangfuse_existing_trace_id(): - """ - When existing trace id is passed, don't set trace params -> prevents overwriting the trace - - Pass 1 logging object with a trace - - Pass 2nd logging object with the trace id - - Assert no changes to the trace - """ - # Test - if the logs were sent to the correct team on langfuse - import datetime - - import litellm - from litellm.integrations.langfuse.langfuse import LangFuseLogger - - langfuse_Logger = LangFuseLogger( - langfuse_public_key=os.getenv("LANGFUSE_PROJECT2_PUBLIC"), - langfuse_secret=os.getenv("LANGFUSE_PROJECT2_SECRET"), - ) - litellm.success_callback = ["langfuse"] - - # langfuse_args = {'kwargs': { 'start_time': 'end_time': datetime.datetime(2024, 5, 1, 7, 31, 29, 903685), 'user_id': None, 'print_verbose': , 'level': 'DEFAULT', 'status_message': None} - response_obj = litellm.ModelResponse( - id="chatcmpl-9K5HUAbVRqFrMZKXL0WoC295xhguY", - choices=[ - litellm.Choices( - finish_reason="stop", - index=0, - message=litellm.Message( - content="I'm sorry, I am an AI assistant and do not have real-time information. I recommend checking a reliable weather website or app for the most up-to-date weather information in Boston.", - role="assistant", - ), - ) - ], - created=1714573888, - model="gpt-3.5-turbo-0125", - object="chat.completion", - system_fingerprint="fp_3b956da36b", - usage=litellm.Usage(completion_tokens=37, prompt_tokens=14, total_tokens=51), - ) - - ### NEW TRACE ### - message = [{"role": "user", "content": "what's the weather in boston"}] - langfuse_args = { - "response_obj": response_obj, - "kwargs": { - "model": "gpt-3.5-turbo", - "litellm_params": { - "acompletion": False, - "api_key": None, - "force_timeout": 600, - "logger_fn": None, - "verbose": False, - "custom_llm_provider": "openai", - "api_base": "https://api.openai.com/v1/", - "litellm_call_id": None, - "model_alias_map": {}, - "completion_call_id": None, - "metadata": None, - "model_info": None, - "proxy_server_request": None, - "preset_cache_key": None, - "no-log": False, - "stream_response": {}, - }, - "messages": message, - "optional_params": {"temperature": 0.1, "extra_body": {}}, - "start_time": "2024-05-01 07:31:27.986164", - "stream": False, - "user": None, - "call_type": "completion", - "litellm_call_id": None, - "completion_start_time": "2024-05-01 07:31:29.903685", - "temperature": 0.1, - "extra_body": {}, - "input": [{"role": "user", "content": "what's the weather in boston"}], - "api_key": "my-api-key", - "additional_args": { - "complete_input_dict": { - "model": "gpt-3.5-turbo", - "messages": [ - {"role": "user", "content": "what's the weather in boston"} - ], - "temperature": 0.1, - "extra_body": {}, - } - }, - "log_event_type": "successful_api_call", - "end_time": "2024-05-01 07:31:29.903685", - "cache_hit": None, - "response_cost": 6.25e-05, - }, - "start_time": datetime.datetime(2024, 5, 1, 7, 31, 27, 986164), - "end_time": datetime.datetime(2024, 5, 1, 7, 31, 29, 903685), - "user_id": None, - "print_verbose": litellm.print_verbose, - "level": "DEFAULT", - "status_message": None, - } - - langfuse_response_object = langfuse_Logger.log_event(**langfuse_args) - - import langfuse - - langfuse_client = langfuse.Langfuse( - public_key=os.getenv("LANGFUSE_PROJECT2_PUBLIC"), - secret_key=os.getenv("LANGFUSE_PROJECT2_SECRET"), - ) - - trace_id = langfuse_response_object["trace_id"] - - assert trace_id is not None - - langfuse_client.flush() - - time.sleep(2) - - print(langfuse_client.get_trace(id=trace_id)) - - initial_langfuse_trace = langfuse_client.get_trace(id=trace_id) - - ### EXISTING TRACE ### - - new_metadata = {"existing_trace_id": trace_id} - new_messages = [{"role": "user", "content": "What do you know?"}] - new_response_obj = litellm.ModelResponse( - id="chatcmpl-9K5HUAbVRqFrMZKXL0WoC295xhguY", - choices=[ - litellm.Choices( - finish_reason="stop", - index=0, - message=litellm.Message( - content="What do I know?", - role="assistant", - ), - ) - ], - created=1714573888, - model="gpt-3.5-turbo-0125", - object="chat.completion", - system_fingerprint="fp_3b956da36b", - usage=litellm.Usage(completion_tokens=37, prompt_tokens=14, total_tokens=51), - ) - langfuse_args = { - "response_obj": new_response_obj, - "kwargs": { - "model": "gpt-3.5-turbo", - "litellm_params": { - "acompletion": False, - "api_key": None, - "force_timeout": 600, - "logger_fn": None, - "verbose": False, - "custom_llm_provider": "openai", - "api_base": "https://api.openai.com/v1/", - "litellm_call_id": "508113a1-c6f1-48ce-a3e1-01c6cce9330e", - "model_alias_map": {}, - "completion_call_id": None, - "metadata": new_metadata, - "model_info": None, - "proxy_server_request": None, - "preset_cache_key": None, - "no-log": False, - "stream_response": {}, - }, - "messages": new_messages, - "optional_params": {"temperature": 0.1, "extra_body": {}}, - "start_time": "2024-05-01 07:31:27.986164", - "stream": False, - "user": None, - "call_type": "completion", - "litellm_call_id": "508113a1-c6f1-48ce-a3e1-01c6cce9330e", - "completion_start_time": "2024-05-01 07:31:29.903685", - "temperature": 0.1, - "extra_body": {}, - "input": [{"role": "user", "content": "what's the weather in boston"}], - "api_key": "my-api-key", - "additional_args": { - "complete_input_dict": { - "model": "gpt-3.5-turbo", - "messages": [ - {"role": "user", "content": "what's the weather in boston"} - ], - "temperature": 0.1, - "extra_body": {}, - } - }, - "log_event_type": "successful_api_call", - "end_time": "2024-05-01 07:31:29.903685", - "cache_hit": None, - "response_cost": 6.25e-05, - }, - "start_time": datetime.datetime(2024, 5, 1, 7, 31, 27, 986164), - "end_time": datetime.datetime(2024, 5, 1, 7, 31, 29, 903685), - "user_id": None, - "print_verbose": litellm.print_verbose, - "level": "DEFAULT", - "status_message": None, - } - - langfuse_response_object = langfuse_Logger.log_event(**langfuse_args) - - new_trace_id = langfuse_response_object["trace_id"] - - assert new_trace_id == trace_id - - langfuse_client.flush() - - time.sleep(2) - - print(langfuse_client.get_trace(id=trace_id)) - - new_langfuse_trace = langfuse_client.get_trace(id=trace_id) - - initial_langfuse_trace_dict = dict(initial_langfuse_trace) - initial_langfuse_trace_dict.pop("updatedAt") - initial_langfuse_trace_dict.pop("timestamp") - - new_langfuse_trace_dict = dict(new_langfuse_trace) - new_langfuse_trace_dict.pop("updatedAt") - new_langfuse_trace_dict.pop("timestamp") - - assert initial_langfuse_trace_dict == new_langfuse_trace_dict - - @pytest.mark.skipif( condition=not os.environ.get("OPENAI_API_KEY", False), reason="Authentication missing for openai", @@ -928,42 +701,6 @@ async def test_make_request(): ) -@pytest.mark.skip( - reason="local only test, use this to verify if dynamic langfuse logging works as expected" -) -def test_aaalangfuse_dynamic_logging(): - """ - pass in langfuse credentials via completion call - - assert call is logged. - - Covers the team-logging scenario. - """ - from litellm._uuid import uuid - - import langfuse - - trace_id = str(uuid.uuid4()) - _ = litellm.completion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hey"}], - mock_response="Hey! how's it going?", - langfuse_public_key=os.getenv("LANGFUSE_PROJECT2_PUBLIC"), - langfuse_secret_key=os.getenv("LANGFUSE_PROJECT2_SECRET"), - metadata={"trace_id": trace_id}, - success_callback=["langfuse"], - ) - - time.sleep(3) - - langfuse_client = langfuse.Langfuse( - public_key=os.getenv("LANGFUSE_PROJECT2_PUBLIC"), - secret_key=os.getenv("LANGFUSE_PROJECT2_SECRET"), - ) - - langfuse_client.get_trace(id=trace_id) - - import datetime generation_params = { diff --git a/tests/local_testing/test_embedding.py b/tests/local_testing/test_embedding.py index 2c265d444a..e6776e3329 100644 --- a/tests/local_testing/test_embedding.py +++ b/tests/local_testing/test_embedding.py @@ -549,7 +549,7 @@ def test_bedrock_embedding_cohere(): "good morning from litellm, attempting to embed data", "lets test a second string for good measure", ], - aws_region_name="os.environ/AWS_REGION_NAME_2", + aws_region_name="os.environ/AWS_REGION_NAME", ) assert isinstance( response["data"][0]["embedding"], list diff --git a/tests/local_testing/test_gcs_bucket.py b/tests/local_testing/test_gcs_bucket.py index e2f51eb76c..018a79b53f 100644 --- a/tests/local_testing/test_gcs_bucket.py +++ b/tests/local_testing/test_gcs_bucket.py @@ -22,6 +22,7 @@ from litellm.integrations.gcs_bucket.gcs_bucket import ( ) from litellm.types.utils import StandardCallbackDynamicParams from unittest.mock import patch + verbose_logger.setLevel(logging.DEBUG) @@ -52,8 +53,8 @@ def load_vertex_ai_credentials(): service_account_key_data = {} # Update the service_account_key_data with environment variables - private_key_id = os.environ.get("GCS_PRIVATE_KEY_ID", "") - private_key = os.environ.get("GCS_PRIVATE_KEY", "") + private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "") + private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "") private_key = private_key.replace("\\n", "\n") service_account_key_data["private_key_id"] = private_key_id service_account_key_data["private_key"] = private_key @@ -686,6 +687,7 @@ async def test_basic_gcs_logger_with_folder_in_bucket_name(): if old_bucket_name is not None: os.environ["GCS_BUCKET_NAME"] = old_bucket_name + @pytest.mark.skip(reason="This test is flaky on ci/cd") def test_create_file_e2e(): """ @@ -696,6 +698,7 @@ def test_create_file_e2e(): test_file = ("test.wav", test_file_content, "audio/wav") from litellm import create_file + response = create_file( file=test_file, purpose="user_data", @@ -704,6 +707,7 @@ def test_create_file_e2e(): print("response", response) assert response is not None + @pytest.mark.skip(reason="This test is flaky on ci/cd") def test_create_file_e2e_jsonl(): """ @@ -714,14 +718,41 @@ def test_create_file_e2e_jsonl(): client = HTTPHandler() - example_jsonl = [{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gemini-1.5-flash-001", "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": "gemini-1.5-flash-001", "messages": [{"role": "system", "content": "You are an unhelpful assistant."},{"role": "user", "content": "Hello world!"}],"max_tokens": 10}}] - + example_jsonl = [ + { + "custom_id": "request-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-1.5-flash-001", + "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": "gemini-1.5-flash-001", + "messages": [ + {"role": "system", "content": "You are an unhelpful assistant."}, + {"role": "user", "content": "Hello world!"}, + ], + "max_tokens": 10, + }, + }, + ] + # Create and write to the file file_path = "example.jsonl" with open(file_path, "w") as f: for item in example_jsonl: f.write(json.dumps(item) + "\n") - + # Verify file content with open(file_path, "r") as f: content = f.read() @@ -729,10 +760,11 @@ def test_create_file_e2e_jsonl(): assert len(content) > 0, "File is empty" from litellm import create_file + with patch.object(client, "post") as mock_create_file: - try: + try: response = create_file( - file=open(file_path, "rb"), + file=open(file_path, "rb"), purpose="user_data", custom_llm_provider="vertex_ai", client=client, @@ -744,4 +776,7 @@ def test_create_file_e2e_jsonl(): print(f"kwargs: {mock_create_file.call_args.kwargs}") - assert mock_create_file.call_args.kwargs["data"] is not None and len(mock_create_file.call_args.kwargs["data"]) > 0 \ No newline at end of file + assert ( + mock_create_file.call_args.kwargs["data"] is not None + and len(mock_create_file.call_args.kwargs["data"]) > 0 + ) From 1ba109079dd90ec682382261f810f3825fcb059c Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 21:22:04 -0700 Subject: [PATCH 15/16] test: update tests --- .../test_adding_passthrough_model.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/store_model_in_db_tests/test_adding_passthrough_model.py b/tests/store_model_in_db_tests/test_adding_passthrough_model.py index e901be5bd7..80cd7e1aab 100644 --- a/tests/store_model_in_db_tests/test_adding_passthrough_model.py +++ b/tests/store_model_in_db_tests/test_adding_passthrough_model.py @@ -1,12 +1,12 @@ """ Test adding a pass through assemblyai model + api key + api base to the db -wait 20 seconds -make request +wait 20 seconds +make request -Cases to cover -1. user points api base to /assemblyai +Cases to cover +1. user points api base to /assemblyai 2. user points api base to /asssemblyai/us -3. user points api base to /assemblyai/eu +3. user points api base to /assemblyai/eu 4. Bad API Key / credential - 401 """ @@ -21,7 +21,7 @@ TEST_MASTER_KEY = "sk-1234" PROXY_BASE_URL = "http://0.0.0.0:4000" US_BASE_URL = f"{PROXY_BASE_URL}/assemblyai" EU_BASE_URL = f"{PROXY_BASE_URL}/eu.assemblyai" -ASSEMBLYAI_API_KEY_ENV_VAR = "TEST_SPECIAL_ASSEMBLYAI_API_KEY" +ASSEMBLYAI_API_KEY_ENV_VAR = "ASSEMBLYAI_API_KEY" def _delete_all_assemblyai_models_from_db(): From 85cc7bc433b9dec89bbc4df876293bd7c5157fa7 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 28 Mar 2026 21:44:46 -0700 Subject: [PATCH 16/16] refactor: make unit test --- tests/agent_tests/test_a2a_agent.py | 82 ++++++++++++++++++++--------- 1 file changed, 57 insertions(+), 25 deletions(-) diff --git a/tests/agent_tests/test_a2a_agent.py b/tests/agent_tests/test_a2a_agent.py index ace1f8c54c..ed7a5ab982 100644 --- a/tests/agent_tests/test_a2a_agent.py +++ b/tests/agent_tests/test_a2a_agent.py @@ -1,38 +1,72 @@ """ Simple A2A agent tests - non-streaming and streaming. -These tests validate the localhost URL retry logic: if an A2A agent's card -contains a localhost/internal URL (e.g., http://0.0.0.0:8001/), the request -will fail with a connection error. LiteLLM detects this and automatically -retries using the original api_base URL instead. - -Requires A2A_AGENT_URL environment variable to be set. - -Run with: - A2A_AGENT_URL=https://your-agent.example.com pytest tests/agent_tests/test_a2a_agent.py -v -s +These tests use a mocked A2A client to avoid network/env dependencies. """ -import os - -import pytest +from types import SimpleNamespace from uuid import uuid4 +import pytest -def get_a2a_agent_url(): - """Get A2A agent URL from environment, skip test if not set.""" - url = os.environ.get("A2A_AGENT_URL") - return url + +class MockA2AResponse: + def __init__(self, text: str): + self._payload = { + "id": str(uuid4()), + "jsonrpc": "2.0", + "result": { + "message": { + "role": "agent", + "parts": [{"kind": "text", "text": text}], + "messageId": uuid4().hex, + } + }, + } + + def model_dump(self, mode="json", exclude_none=True): + return self._payload + + +class MockA2AStreamingChunk(MockA2AResponse): + def __init__(self, text: str, state: str): + super().__init__(text=text) + self._payload["result"]["status"] = {"state": state} + + +class MockA2AClient: + def __init__(self): + self._litellm_agent_card = SimpleNamespace( + name="mock-agent", url="http://mock-agent.local" + ) + + async def send_message(self, request): + return MockA2AResponse(text="hello") + + def send_message_streaming(self, request): + async def _stream(): + yield MockA2AStreamingChunk(text="hel", state="in_progress") + yield MockA2AStreamingChunk(text="hello", state="completed") + + return _stream() + + +@pytest.fixture +def mock_a2a_client(monkeypatch): + import litellm.a2a_protocol.main as a2a_main + + async def _fake_create_a2a_client(base_url, timeout=60.0, extra_headers=None): + return MockA2AClient() + + monkeypatch.setattr(a2a_main, "create_a2a_client", _fake_create_a2a_client) @pytest.mark.asyncio -@pytest.mark.flaky(retries=3, delay=5) -async def test_a2a_non_streaming(): +async def test_a2a_non_streaming(mock_a2a_client): """Test non-streaming A2A request.""" from a2a.types import MessageSendParams, SendMessageRequest from litellm.a2a_protocol import asend_message - api_base = get_a2a_agent_url() - request = SendMessageRequest( id=str(uuid4()), params=MessageSendParams( @@ -46,7 +80,7 @@ async def test_a2a_non_streaming(): response = await asend_message( request=request, - api_base=api_base, + api_base="http://mock", ) assert response is not None @@ -54,13 +88,11 @@ async def test_a2a_non_streaming(): @pytest.mark.asyncio -async def test_a2a_streaming(): +async def test_a2a_streaming(mock_a2a_client): """Test streaming A2A request.""" from a2a.types import MessageSendParams, SendStreamingMessageRequest from litellm.a2a_protocol import asend_message_streaming - api_base = get_a2a_agent_url() - request = SendStreamingMessageRequest( id=str(uuid4()), params=MessageSendParams( @@ -75,7 +107,7 @@ async def test_a2a_streaming(): chunks = [] async for chunk in asend_message_streaming( request=request, - api_base=api_base, + api_base="http://mock", ): chunks.append(chunk) print(f"\nStreaming chunk: {chunk}")