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))