diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index c246094c81..9367170b7a 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -166,39 +166,6 @@ def test_completion_bedrock_guardrails(streaming): pytest.fail(f"Error occurred: {e}") -def test_completion_bedrock_claude_2_1_completion_auth(): - print("calling bedrock claude 2.1 completion params auth") - import os - - aws_access_key_id = os.environ["AWS_ACCESS_KEY_ID"] - aws_secret_access_key = os.environ["AWS_SECRET_ACCESS_KEY"] - aws_region_name = os.environ["AWS_REGION_NAME"] - - os.environ.pop("AWS_ACCESS_KEY_ID", None) - os.environ.pop("AWS_SECRET_ACCESS_KEY", None) - os.environ.pop("AWS_REGION_NAME", None) - try: - response = completion( - model="bedrock/anthropic.claude-v2:1", - messages=messages, - max_tokens=10, - temperature=0.1, - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - aws_region_name=aws_region_name, - ) - # Add any assertions here to check the response - print(response) - - os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id - os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key - os.environ["AWS_REGION_NAME"] = aws_region_name - except RateLimitError: - pass - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - # test_completion_bedrock_claude_2_1_completion_auth() diff --git a/tests/local_testing/test_caching.py b/tests/local_testing/test_caching.py index 8a0956958b..d1abec5ff7 100644 --- a/tests/local_testing/test_caching.py +++ b/tests/local_testing/test_caching.py @@ -314,6 +314,7 @@ async def test_caching_with_cache_controls(sync_flag): # test_caching_with_cache_controls() + @pytest.mark.flaky(retries=3, delay=1) def test_caching_with_models_v2(): messages = [ @@ -449,6 +450,7 @@ def test_embedding_caching(): # test_embedding_caching() + @pytest.mark.asyncio async def test_embedding_caching_individual_items_and_then_list(): litellm._turn_on_debug() @@ -473,7 +475,7 @@ async def test_embedding_caching_individual_items_and_then_list(): assert embedding3["data"][0]["embedding"] == embedding1["data"][0]["embedding"] assert embedding3["data"][1]["embedding"] == embedding2["data"][0]["embedding"] assert embedding3._hidden_params["cache_hit"] == True - assert embedding3.usage.prompt_tokens != 0 + assert embedding3.usage.prompt_tokens != 0 ## with new input, check that prompt tokens increase additional_text = "this is a new text" @@ -483,6 +485,7 @@ async def test_embedding_caching_individual_items_and_then_list(): ) assert embedding4.usage.prompt_tokens > embedding3.usage.prompt_tokens + @pytest.mark.asyncio async def test_embedding_caching_individual_items(): litellm.cache = Cache() @@ -500,7 +503,7 @@ async def test_embedding_caching_individual_items(): assert embedding3["data"][0]["embedding"] == embedding1["data"][0]["embedding"] assert len(embedding3.data) == 1 assert embedding3._hidden_params["cache_hit"] == True - assert embedding3.usage.prompt_tokens != 0 + assert embedding3.usage.prompt_tokens != 0 def test_embedding_caching_azure(): @@ -1156,7 +1159,7 @@ async def test_redis_cache_acompletion_stream_bedrock(): response_2_content = "" response1 = await litellm.acompletion( - model="bedrock/anthropic.claude-v2", + model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", messages=messages, max_tokens=40, temperature=1, @@ -1171,7 +1174,7 @@ async def test_redis_cache_acompletion_stream_bedrock(): print("\n\n Response 1 content: ", response_1_content, "\n\n") response2 = await litellm.acompletion( - model="bedrock/anthropic.claude-v2", + model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", messages=messages, max_tokens=40, temperature=1, @@ -1883,9 +1886,6 @@ def test_caching_redis_simple(caplog, capsys): assert "async success_callback: reaches cache for logging" not in captured.out - - - @pytest.mark.asyncio() async def test_cache_default_off_acompletion(): litellm.set_verbose = True @@ -2417,7 +2417,7 @@ async def test_redis_increment_pipeline(): results = await redis_cache.async_increment_pipeline(increment_list) # Verify results - assert len(results) == 4 + assert len(results) == 4 # Verify the values were actually set in Redis value1 = await redis_cache.async_get_cache("test_key1") @@ -2502,116 +2502,136 @@ def test_redis_caching_multiple_namespaces(): # Use a fixed uuid to ensure consistent cache keys test_uuid = "12345678-1234-1234-1234-123456789abc" messages = [{"role": "user", "content": f"what is litellm? {test_uuid}"}] - + # Mock the Redis client creation from the _redis module - with patch('litellm._redis.get_redis_client') as mock_get_redis_client, \ - patch('litellm._redis.get_redis_connection_pool') as mock_get_redis_connection_pool: + with patch("litellm._redis.get_redis_client") as mock_get_redis_client, patch( + "litellm._redis.get_redis_connection_pool" + ) as mock_get_redis_connection_pool: # Create a mock Redis client that simulates real Redis behavior mock_redis_client = MagicMock() mock_get_redis_client.return_value = mock_redis_client - + # Mock the connection pool mock_connection_pool = MagicMock() mock_get_redis_connection_pool.return_value = mock_connection_pool - + # Dictionary to simulate Redis storage with namespace support redis_storage = {} - + def mock_redis_get(key): print(f"Redis GET: {key}") value = redis_storage.get(key, None) # Convert to bytes to match real Redis behavior if value is not None: import json - return json.dumps(value).encode('utf-8') + + return json.dumps(value).encode("utf-8") return None - + def mock_redis_set(name, value, ex=None, **kwargs): print(f"Redis SET: {name} = {value}") redis_storage[name] = value return True - + def mock_redis_ping(): return True - + def mock_redis_info(): return {"redis_version": "7.0.0"} - + mock_redis_client.get = mock_redis_get mock_redis_client.set = mock_redis_set mock_redis_client.ping = mock_redis_ping mock_redis_client.info = mock_redis_info - + # Initialize the cache litellm.cache = Cache(type="redis") - + namespace_1 = "org-id1" namespace_2 = "org-id2" # Use mock_response to ensure deterministic responses without external API calls response_1 = completion( - model="gpt-3.5-turbo", - messages=messages, + model="gpt-3.5-turbo", + messages=messages, cache={"namespace": namespace_1}, - mock_response="Response for namespace 1" + mock_response="Response for namespace 1", ) response_2 = completion( - model="gpt-3.5-turbo", - messages=messages, + model="gpt-3.5-turbo", + messages=messages, cache={"namespace": namespace_2}, - mock_response="Response for namespace 2" + mock_response="Response for namespace 2", ) response_3 = completion( - model="gpt-3.5-turbo", - messages=messages, + model="gpt-3.5-turbo", + messages=messages, cache={"namespace": namespace_1}, - mock_response="This should be cached" + mock_response="This should be cached", ) response_4 = completion( - model="gpt-3.5-turbo", + model="gpt-3.5-turbo", messages=messages, - mock_response="Response without namespace" + mock_response="Response without namespace", + ) + + print( + f"Response 1 type: {type(response_1)} - ID: {getattr(response_1, 'id', 'N/A')}" + ) + print( + f"Response 2 type: {type(response_2)} - ID: {getattr(response_2, 'id', 'N/A')}" + ) + print( + f"Response 3 type: {type(response_3)} - Cache hit: {isinstance(response_3, str)}" + ) + print( + f"Response 4 type: {type(response_4)} - ID: {getattr(response_4, 'id', 'N/A')}" ) - print(f"Response 1 type: {type(response_1)} - ID: {getattr(response_1, 'id', 'N/A')}") - print(f"Response 2 type: {type(response_2)} - ID: {getattr(response_2, 'id', 'N/A')}") - print(f"Response 3 type: {type(response_3)} - Cache hit: {isinstance(response_3, str)}") - print(f"Response 4 type: {type(response_4)} - ID: {getattr(response_4, 'id', 'N/A')}") - print(f"Redis storage keys: {list(redis_storage.keys())}") # Verify that different namespaces created different cache keys cache_keys = list(redis_storage.keys()) namespace_1_keys = [k for k in cache_keys if k.startswith(f"{namespace_1}:")] namespace_2_keys = [k for k in cache_keys if k.startswith(f"{namespace_2}:")] - no_namespace_keys = [k for k in cache_keys if not k.startswith(f"{namespace_1}:") and not k.startswith(f"{namespace_2}:")] - + no_namespace_keys = [ + k + for k in cache_keys + if not k.startswith(f"{namespace_1}:") + and not k.startswith(f"{namespace_2}:") + ] + print(f"Namespace 1 keys: {namespace_1_keys}") print(f"Namespace 2 keys: {namespace_2_keys}") print(f"No namespace keys: {no_namespace_keys}") - + # Should have at least one key for each namespace assert len(namespace_1_keys) > 0, "Should have cache keys for namespace 1" assert len(namespace_2_keys) > 0, "Should have cache keys for namespace 2" assert len(no_namespace_keys) > 0, "Should have cache keys for no namespace" - + # The main test: response 3 should be a cache hit (string) because it uses same namespace as response 1 - assert isinstance(response_3, str), "Response 3 should be a cache hit (string) for same namespace" - + assert isinstance( + response_3, str + ), "Response 3 should be a cache hit (string) for same namespace" + # response 1 & 2 should be ModelResponse objects (cache misses) - assert hasattr(response_1, 'id'), "Response 1 should be a ModelResponse object" - assert hasattr(response_2, 'id'), "Response 2 should be a ModelResponse object" - assert hasattr(response_4, 'id'), "Response 4 should be a ModelResponse object" - + assert hasattr(response_1, "id"), "Response 1 should be a ModelResponse object" + assert hasattr(response_2, "id"), "Response 2 should be a ModelResponse object" + assert hasattr(response_4, "id"), "Response 4 should be a ModelResponse object" + # response 1 & 2 should have different IDs (different namespaces) - assert response_1.id != response_2.id, f"Expected different response ID for different namespace. Got {response_1.id} and {response_2.id}" - + assert ( + response_1.id != response_2.id + ), f"Expected different response ID for different namespace. Got {response_1.id} and {response_2.id}" + # response 1 & 4 should have different IDs (different namespaces) - assert response_1.id != response_4.id, f"Expected different response ID for no namespace vs namespaced. Got {response_1.id} and {response_4.id}" - + assert ( + response_1.id != response_4.id + ), f"Expected different response ID for no namespace vs namespaced. Got {response_1.id} and {response_4.id}" def test_caching_with_reasoning_content(): @@ -2643,12 +2663,22 @@ def test_caching_with_reasoning_content(): def test_caching_reasoning_args_miss(): # test in memory cache try: - #litellm._turn_on_debug() + # litellm._turn_on_debug() litellm.set_verbose = True - litellm.cache = Cache( + litellm.cache = Cache() + response1 = completion( + model="claude-3-7-sonnet-latest", + messages=messages, + caching=True, + reasoning_effort="low", + mock_response="My response", + ) + response2 = completion( + model="claude-3-7-sonnet-latest", + messages=messages, + caching=True, + mock_response="My response", ) - response1 = completion(model="claude-3-7-sonnet-latest", messages=messages, caching=True, reasoning_effort="low", mock_response="My response") - response2 = completion(model="claude-3-7-sonnet-latest", messages=messages, caching=True, mock_response="My response") print(f"response1: {response1}") print(f"response2: {response2}") assert response1.id != response2.id @@ -2656,29 +2686,52 @@ def test_caching_reasoning_args_miss(): # test in memory cache print(f"error occurred: {traceback.format_exc()}") pytest.fail(f"Error occurred: {e}") + def test_caching_reasoning_args_hit(): # test in memory cache try: - #litellm._turn_on_debug() + # litellm._turn_on_debug() litellm.set_verbose = True - litellm.cache = Cache( + litellm.cache = Cache() + response1 = completion( + model="claude-3-7-sonnet-latest", + messages=messages, + caching=True, + reasoning_effort="low", + mock_response="My response", + ) + response2 = completion( + model="claude-3-7-sonnet-latest", + messages=messages, + caching=True, + reasoning_effort="low", + mock_response="My response", ) - response1 = completion(model="claude-3-7-sonnet-latest", messages=messages, caching=True, reasoning_effort="low", mock_response="My response") - response2 = completion(model="claude-3-7-sonnet-latest", messages=messages, caching=True, reasoning_effort="low", mock_response="My response") print(f"response1: {response1}") print(f"response2: {response2}") assert response1.id == response2.id except Exception as e: print(f"error occurred: {traceback.format_exc()}") pytest.fail(f"Error occurred: {e}") - + + def test_caching_thinking_args_miss(): # test in memory cache try: - #litellm._turn_on_debug() + # litellm._turn_on_debug() litellm.set_verbose = True - litellm.cache = Cache( + litellm.cache = Cache() + response1 = completion( + model="claude-3-7-sonnet-latest", + messages=messages, + caching=True, + thinking={"type": "enabled", "budget_tokens": 1024}, + mock_response="My response", + ) + response2 = completion( + model="claude-3-7-sonnet-latest", + messages=messages, + caching=True, + mock_response="My response", ) - response1 = completion(model="claude-3-7-sonnet-latest", messages=messages, caching=True, thinking={"type": "enabled", "budget_tokens": 1024}, mock_response="My response") - response2 = completion(model="claude-3-7-sonnet-latest", messages=messages, caching=True, mock_response="My response") print(f"response1: {response1}") print(f"response2: {response2}") assert response1.id != response2.id @@ -2686,18 +2739,29 @@ def test_caching_thinking_args_miss(): # test in memory cache print(f"error occurred: {traceback.format_exc()}") pytest.fail(f"Error occurred: {e}") + def test_caching_thinking_args_hit(): # test in memory cache try: - #litellm._turn_on_debug() + # litellm._turn_on_debug() litellm.set_verbose = True - litellm.cache = Cache( + litellm.cache = Cache() + response1 = completion( + model="claude-3-7-sonnet-latest", + messages=messages, + caching=True, + thinking={"type": "enabled", "budget_tokens": 1024}, + mock_response="My response", + ) + response2 = completion( + model="claude-3-7-sonnet-latest", + messages=messages, + caching=True, + thinking={"type": "enabled", "budget_tokens": 1024}, + mock_response="My response", ) - response1 = completion(model="claude-3-7-sonnet-latest", messages=messages, caching=True, thinking={"type": "enabled", "budget_tokens": 1024}, mock_response="My response" ) - response2 = completion(model="claude-3-7-sonnet-latest", messages=messages, caching=True, thinking={"type": "enabled", "budget_tokens": 1024}, mock_response="My response") print(f"response1: {response1}") print(f"response2: {response2}") assert response1.id == response2.id except Exception as e: print(f"error occurred: {traceback.format_exc()}") pytest.fail(f"Error occurred: {e}") - diff --git a/tests/local_testing/test_custom_callback_input.py b/tests/local_testing/test_custom_callback_input.py index 5da87d2be6..90d1d15c5a 100644 --- a/tests/local_testing/test_custom_callback_input.py +++ b/tests/local_testing/test_custom_callback_input.py @@ -554,91 +554,6 @@ async def test_async_chat_openai_stream_options(): pytest.fail(f"An exception occurred: {str(e)}") -## Test Bedrock + sync -def test_chat_bedrock_stream(): - try: - customHandler = CompletionCustomHandler() - litellm.callbacks = [customHandler] - response = litellm.completion( - model="bedrock/anthropic.claude-v2", - messages=[{"role": "user", "content": "Hi 👋 - i'm sync bedrock"}], - ) - # test streaming - response = litellm.completion( - model="bedrock/anthropic.claude-v2", - messages=[{"role": "user", "content": "Hi 👋 - i'm sync bedrock"}], - stream=True, - ) - for chunk in response: - continue - # test failure callback - try: - response = litellm.completion( - model="bedrock/anthropic.claude-v2", - messages=[{"role": "user", "content": "Hi 👋 - i'm sync bedrock"}], - aws_region_name="my-bad-region", - stream=True, - ) - for chunk in response: - continue - except Exception: - pass - time.sleep(1) - print(f"customHandler.errors: {customHandler.errors}") - assert len(customHandler.errors) == 0 - litellm.callbacks = [] - except Exception as e: - pytest.fail(f"An exception occurred: {str(e)}") - - -# test_chat_bedrock_stream() - - -## Test Bedrock + Async -@pytest.mark.asyncio -async def test_async_chat_bedrock_stream(): - try: - litellm.set_verbose = True - customHandler = CompletionCustomHandler() - litellm.callbacks = [customHandler] - response = await litellm.acompletion( - model="bedrock/anthropic.claude-v2", - messages=[{"role": "user", "content": "Hi 👋 - i'm async bedrock"}], - ) - # test streaming - response = await litellm.acompletion( - model="bedrock/anthropic.claude-v2", - messages=[{"role": "user", "content": "Hi 👋 - i'm async bedrock"}], - stream=True, - ) - print(f"response: {response}") - async for chunk in response: - print(f"chunk: {chunk}") - continue - - await asyncio.sleep(1) - ## test failure callback - try: - response = await litellm.acompletion( - model="bedrock/anthropic.claude-v2", - messages=[{"role": "user", "content": "Hi 👋 - i'm async bedrock"}], - aws_region_name="my-bad-key", - stream=True, - ) - async for chunk in response: - continue - - await asyncio.sleep(1) - except Exception: - pass - await asyncio.sleep(1) - print(f"customHandler.errors: {customHandler.errors}") - assert len(customHandler.errors) == 0 - litellm.callbacks = [] - except Exception as e: - pytest.fail(f"An exception occurred: {str(e)}") - - # asyncio.run(test_async_chat_bedrock_stream())