From 744dca1dc312bc66c3f05521f8f42c2c69ff1364 Mon Sep 17 00:00:00 2001 From: Michal Otmianowski Date: Thu, 21 Aug 2025 10:22:54 +0200 Subject: [PATCH 1/6] init asyc implementation --- litellm/caching/caching.py | 4 +- litellm/caching/caching_handler.py | 8 +- litellm/caching/s3_cache.py | 38 +++- .../caching/test_s3_cache_async.py | 207 ++++++++++++++++++ 4 files changed, 240 insertions(+), 17 deletions(-) create mode 100644 tests/test_litellm/caching/test_s3_cache_async.py diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 0a8b6ef18d..82fc37e0cb 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -774,11 +774,9 @@ class Cache: """ Internal method to check if the cache type supports async get/set operations - Only S3 Cache Does NOT support async operations + All cache types now support async operations """ - if self.type and self.type == LiteLLMCacheType.S3: - return False return True diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 7580752c30..1dcc0f1fdb 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -599,7 +599,7 @@ class LLMCachingHandler: cached_result = await litellm.cache.async_get_cache( dynamic_cache_object=self.dual_cache, **new_kwargs ) - else: # for s3 caching. [NOT RECOMMENDED IN PROD - this will slow down responses since boto3 is sync] + else: # fallback for caches that don't support async cached_result = litellm.cache.get_cache( dynamic_cache_object=self.dual_cache, **new_kwargs ) @@ -806,12 +806,6 @@ class LLMCachingHandler: result, dynamic_cache_object=self.dual_cache, **new_kwargs ) ) - elif isinstance(litellm.cache.cache, S3Cache): - threading.Thread( - target=litellm.cache.add_cache, - args=(result,), - kwargs=new_kwargs, - ).start() else: asyncio.create_task( litellm.cache.async_add_cache( diff --git a/litellm/caching/s3_cache.py b/litellm/caching/s3_cache.py index 6095d547aa..929de6493c 100644 --- a/litellm/caching/s3_cache.py +++ b/litellm/caching/s3_cache.py @@ -1,17 +1,17 @@ """ S3 Cache implementation -WARNING: DO NOT USE THIS IN PRODUCTION - This is not ASYNC Has 4 methods: - set_cache - get_cache - - async_set_cache - - async_get_cache + - async_set_cache (uses run_in_executor) + - async_get_cache (uses run_in_executor) """ import ast import asyncio import json +from functools import partial from typing import Optional from litellm._logging import print_verbose, verbose_logger @@ -72,7 +72,7 @@ class S3Cache(BaseCache): import datetime # Calculate expiration time - expiration_time = datetime.datetime.now() + ttl + expiration_time = datetime.datetime.now() + datetime.timedelta(seconds=ttl) # Upload the data to S3 with the calculated expiration time self.s3_client.put_object( @@ -102,7 +102,18 @@ class S3Cache(BaseCache): print_verbose(f"S3 Caching: set_cache() - Got exception from S3: {e}") async def async_set_cache(self, key, value, **kwargs): - self.set_cache(key=key, value=value, **kwargs) + """ + Asynchronously set cache using run_in_executor to avoid blocking the event loop. + Compatible with Python 3.8+. + """ + try: + verbose_logger.debug(f"LiteLLM ASYNC SET Cache - S3. Key={key}. Value={value}") + loop = asyncio.get_event_loop() + func = partial(self.set_cache, key, value, **kwargs) + await loop.run_in_executor(None, func) + except Exception as e: + # NON blocking - notify users S3 is throwing an exception + verbose_logger.error(f"S3 Caching: async_set_cache() - Got exception from S3: {e}") def get_cache(self, key, **kwargs): import botocore @@ -142,13 +153,26 @@ class S3Cache(BaseCache): return None except Exception as e: - # NON blocking - notify users S3 is throwing an exception verbose_logger.error( f"S3 Caching: get_cache() - Got exception from S3: {e}" ) async def async_get_cache(self, key, **kwargs): - return self.get_cache(key=key, **kwargs) + """ + Asynchronously get cache using run_in_executor to avoid blocking the event loop. + Compatible with Python 3.8+. + """ + try: + verbose_logger.debug(f"Get ASYNC S3 Cache: key: {key}") + loop = asyncio.get_event_loop() + func = partial(self.get_cache, key, **kwargs) + result = await loop.run_in_executor(None, func) + return result + except Exception as e: + verbose_logger.error( + f"S3 Caching: async_get_cache() - Got exception from S3: {e}" + ) + return None def flush_cache(self): pass diff --git a/tests/test_litellm/caching/test_s3_cache_async.py b/tests/test_litellm/caching/test_s3_cache_async.py new file mode 100644 index 0000000000..4c18a5f295 --- /dev/null +++ b/tests/test_litellm/caching/test_s3_cache_async.py @@ -0,0 +1,207 @@ +import os +import sys +from unittest.mock import MagicMock, patch +import json +import datetime +import asyncio + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +from litellm.caching.s3_cache import S3Cache + + +@pytest.fixture +def mock_s3_dependencies(): + mock_s3_client = MagicMock() + + with patch("boto3.client", return_value=mock_s3_client): + yield {"s3_client": mock_s3_client} + + +@pytest.mark.asyncio +async def test_s3_cache_async_set_cache(mock_s3_dependencies): + """Test async_set_cache functionality using run_in_executor""" + cache = S3Cache("test-bucket") + test_value = {"key": "value", "number": 42} + + await cache.async_set_cache("test_key", test_value) + + cache.s3_client.put_object.assert_called_once() + call_args = cache.s3_client.put_object.call_args + + assert call_args[1]["Bucket"] == "test-bucket" + assert call_args[1]["Key"] == "test_key" + assert call_args[1]["Body"] == json.dumps(test_value) + assert call_args[1]["ContentType"] == "application/json" + assert call_args[1]["ContentLanguage"] == "en" + assert call_args[1]["ContentDisposition"] == 'inline; filename="test_key.json"' + + +@pytest.mark.asyncio +async def test_s3_cache_async_set_cache_with_ttl(mock_s3_dependencies): + """Test async_set_cache with TTL functionality""" + cache = S3Cache("test-bucket") + test_value = {"key": "value"} + ttl = datetime.timedelta(seconds=3600) # 1 hour + + await cache.async_set_cache("test_key", test_value, ttl=ttl) + + cache.s3_client.put_object.assert_called_once() + call_args = cache.s3_client.put_object.call_args + + assert "Expires" in call_args[1] + assert "CacheControl" in call_args[1] + assert "max-age=1:00:00" in call_args[1]["CacheControl"] + + +@pytest.mark.asyncio +async def test_s3_cache_async_get_cache(mock_s3_dependencies): + """Test async_get_cache functionality using run_in_executor""" + cache = S3Cache("test-bucket") + + mock_response = { + "Body": MagicMock() + } + mock_response["Body"].read.return_value = b'{"key": "value", "number": 42}' + cache.s3_client.get_object.return_value = mock_response + + result = await cache.async_get_cache("test_key") + + cache.s3_client.get_object.assert_called_once_with( + Bucket="test-bucket", + Key="test_key" + ) + + assert result == {"key": "value", "number": 42} + + +@pytest.mark.asyncio +async def test_s3_cache_async_get_cache_not_found(mock_s3_dependencies): + """Test async_get_cache when key is not found""" + import botocore.exceptions + + cache = S3Cache("test-bucket") + + error_response = {"Error": {"Code": "NoSuchKey"}} + cache.s3_client.get_object.side_effect = botocore.exceptions.ClientError( + error_response, "GetObject" + ) + + result = await cache.async_get_cache("nonexistent_key") + + cache.s3_client.get_object.assert_called_once_with( + Bucket="test-bucket", + Key="nonexistent_key" + ) + assert result is None + + +@pytest.mark.asyncio +async def test_s3_cache_async_set_cache_pipeline(mock_s3_dependencies): + """Test async_set_cache_pipeline functionality""" + cache = S3Cache("test-bucket") + + cache_list = [ + ("key1", {"data": "value1"}), + ("key2", {"data": "value2"}), + ("key3", {"data": "value3"}), + ] + + await cache.async_set_cache_pipeline(cache_list) + + # Should have called put_object 3 times + assert cache.s3_client.put_object.call_count == 3 + + # Verify each call + calls = cache.s3_client.put_object.call_args_list + for i, (key, value) in enumerate(cache_list): + call_args = calls[i][1] + assert call_args["Bucket"] == "test-bucket" + assert call_args["Key"] == key + assert call_args["Body"] == json.dumps(value) + + +@pytest.mark.asyncio +async def test_s3_cache_concurrent_async_operations(mock_s3_dependencies): + """Test concurrent async operations to ensure they don't block each other""" + cache = S3Cache("test-bucket") + + # Create multiple concurrent set operations + tasks = [] + for i in range(5): + key = f"concurrent_key_{i}" + value = {"id": i, "data": f"test_data_{i}"} + tasks.append(cache.async_set_cache(key, value)) + + # Execute all tasks concurrently + await asyncio.gather(*tasks) + + # Verify all operations were called + assert cache.s3_client.put_object.call_count == 5 + + # Verify each call had correct parameters + calls = cache.s3_client.put_object.call_args_list + for i, call in enumerate(calls): + call_args = call[1] + assert call_args["Bucket"] == "test-bucket" + assert f"concurrent_key_{i}" == call_args["Key"] + + +@pytest.mark.asyncio +async def test_s3_cache_async_error_handling(mock_s3_dependencies): + """Test that async methods handle errors gracefully""" + cache = S3Cache("test-bucket") + + # Test async_set_cache error handling + cache.s3_client.put_object.side_effect = Exception("S3 Error") + + # Should not raise exception, just log it + await cache.async_set_cache("error_key", {"data": "value"}) + + # Test async_get_cache error handling + cache.s3_client.get_object.side_effect = Exception("S3 Error") + + result = await cache.async_get_cache("error_key") + assert result is None + + +@pytest.mark.asyncio +async def test_s3_cache_async_with_key_prefix(mock_s3_dependencies): + """Test async operations with s3_path prefix""" + cache = S3Cache("test-bucket", s3_path="cache/data") + test_value = {"key": "value"} + + await cache.async_set_cache("namespace:key", test_value) + + cache.s3_client.put_object.assert_called_once() + call_args = cache.s3_client.put_object.call_args + + # Should transform key with prefix and colon replacement + assert call_args[1]["Key"] == "cache/data/namespace/key" + + +def test_s3_cache_supports_async(): + """Test that S3Cache now supports async operations""" + from litellm.caching.caching import Cache, LiteLLMCacheType + + cache = Cache(type=LiteLLMCacheType.S3, s3_bucket_name="test-bucket") + + # Should now return True for async support + assert cache._supports_async() is True + + +@pytest.mark.asyncio +async def test_s3_cache_async_disconnect(mock_s3_dependencies): + """Test async disconnect method""" + cache = S3Cache("test-bucket") + + # Should not raise any exceptions + await cache.disconnect() + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) From 08c942c30699dfdbf79369ae7d5b071bcc142461 Mon Sep 17 00:00:00 2001 From: Michal Otmianowski Date: Thu, 21 Aug 2025 10:38:43 +0200 Subject: [PATCH 2/6] refactor tests --- tests/test_litellm/caching/test_s3_cache.py | 191 +++++++++++++++- .../caching/test_s3_cache_async.py | 207 ------------------ 2 files changed, 187 insertions(+), 211 deletions(-) delete mode 100644 tests/test_litellm/caching/test_s3_cache_async.py diff --git a/tests/test_litellm/caching/test_s3_cache.py b/tests/test_litellm/caching/test_s3_cache.py index a37f7f5706..9ab49987be 100644 --- a/tests/test_litellm/caching/test_s3_cache.py +++ b/tests/test_litellm/caching/test_s3_cache.py @@ -3,6 +3,7 @@ import sys from unittest.mock import MagicMock, patch import json import datetime +import asyncio import pytest @@ -43,7 +44,7 @@ def test_s3_cache_set_cache_with_ttl(mock_s3_dependencies): """Test set_cache with TTL functionality""" cache = S3Cache("test-bucket") test_value = {"key": "value"} - ttl = datetime.timedelta(seconds=3600) # 1 hour + ttl = 3600 # 1 hour in seconds cache.set_cache("test_key", test_value, ttl=ttl) @@ -52,7 +53,7 @@ def test_s3_cache_set_cache_with_ttl(mock_s3_dependencies): assert "Expires" in call_args[1] assert "CacheControl" in call_args[1] - assert "max-age=1:00:00" in call_args[1]["CacheControl"] + assert "max-age=3600" in call_args[1]["CacheControl"] def test_s3_cache_get_cache(mock_s3_dependencies): @@ -120,7 +121,189 @@ def test_s3_cache_initialization(): cache = S3Cache("test-bucket") assert cache.bucket_name == "test-bucket" assert cache.key_prefix == "" - + # Test with s3_path cache_with_path = S3Cache("test-bucket", s3_path="my/cache/path") - assert cache_with_path.key_prefix == "my/cache/path/" \ No newline at end of file + assert cache_with_path.key_prefix == "my/cache/path/" + + +# ============================================================================ +# ASYNC TESTS +# ============================================================================ + +@pytest.mark.asyncio +async def test_s3_cache_async_set_cache(mock_s3_dependencies): + cache = S3Cache("test-bucket") + test_value = {"key": "value", "number": 42} + + await cache.async_set_cache("test_key", test_value) + + cache.s3_client.put_object.assert_called_once() + call_args = cache.s3_client.put_object.call_args + + assert call_args[1]["Bucket"] == "test-bucket" + assert call_args[1]["Key"] == "test_key" + assert call_args[1]["Body"] == json.dumps(test_value) + assert call_args[1]["ContentType"] == "application/json" + assert call_args[1]["ContentLanguage"] == "en" + assert call_args[1]["ContentDisposition"] == 'inline; filename="test_key.json"' + + +@pytest.mark.asyncio +async def test_s3_cache_async_set_cache_with_ttl(mock_s3_dependencies): + cache = S3Cache("test-bucket") + test_value = {"key": "value"} + ttl = 3600 # 1 hour in seconds + + await cache.async_set_cache("test_key", test_value, ttl=ttl) + + cache.s3_client.put_object.assert_called_once() + call_args = cache.s3_client.put_object.call_args + + assert "Expires" in call_args[1] + assert "CacheControl" in call_args[1] + assert "max-age=3600" in call_args[1]["CacheControl"] + + +@pytest.mark.asyncio +async def test_s3_cache_async_get_cache(mock_s3_dependencies): + cache = S3Cache("test-bucket") + + mock_response = { + "Body": MagicMock() + } + mock_response["Body"].read.return_value = b'{"key": "value", "number": 42}' + cache.s3_client.get_object.return_value = mock_response + + result = await cache.async_get_cache("test_key") + + cache.s3_client.get_object.assert_called_once_with( + Bucket="test-bucket", + Key="test_key" + ) + + assert result == {"key": "value", "number": 42} + + +@pytest.mark.asyncio +async def test_s3_cache_async_get_cache_not_found(mock_s3_dependencies): + """Test async_get_cache when key is not found""" + import botocore.exceptions + + cache = S3Cache("test-bucket") + + error_response = {"Error": {"Code": "NoSuchKey"}} + cache.s3_client.get_object.side_effect = botocore.exceptions.ClientError( + error_response, "GetObject" + ) + + result = await cache.async_get_cache("nonexistent_key") + + cache.s3_client.get_object.assert_called_once_with( + Bucket="test-bucket", + Key="nonexistent_key" + ) + assert result is None + + +@pytest.mark.asyncio +async def test_s3_cache_async_set_cache_pipeline(mock_s3_dependencies): + """Test async_set_cache_pipeline functionality""" + cache = S3Cache("test-bucket") + + cache_list = [ + ("key1", {"data": "value1"}), + ("key2", {"data": "value2"}), + ("key3", {"data": "value3"}), + ] + + await cache.async_set_cache_pipeline(cache_list) + + # Should have called put_object 3 times + assert cache.s3_client.put_object.call_count == 3 + + # Verify each call + calls = cache.s3_client.put_object.call_args_list + for i, (key, value) in enumerate(cache_list): + call_args = calls[i][1] + assert call_args["Bucket"] == "test-bucket" + assert call_args["Key"] == key + assert call_args["Body"] == json.dumps(value) + + +@pytest.mark.asyncio +async def test_s3_cache_concurrent_async_operations(mock_s3_dependencies): + """Test concurrent async operations to ensure they don't block each other""" + cache = S3Cache("test-bucket") + + # Create multiple concurrent set operations + tasks = [] + for i in range(5): + key = f"concurrent_key_{i}" + value = {"id": i, "data": f"test_data_{i}"} + tasks.append(cache.async_set_cache(key, value)) + + # Execute all tasks concurrently + await asyncio.gather(*tasks) + + # Verify all operations were called + assert cache.s3_client.put_object.call_count == 5 + + # Verify each call had correct parameters + calls = cache.s3_client.put_object.call_args_list + for i, call in enumerate(calls): + call_args = call[1] + assert call_args["Bucket"] == "test-bucket" + assert f"concurrent_key_{i}" == call_args["Key"] + + +@pytest.mark.asyncio +async def test_s3_cache_async_error_handling(mock_s3_dependencies): + """Test that async methods handle errors gracefully""" + cache = S3Cache("test-bucket") + + # Test async_set_cache error handling + cache.s3_client.put_object.side_effect = Exception("S3 Error") + + # Should not raise exception, just log it + await cache.async_set_cache("error_key", {"data": "value"}) + + # Test async_get_cache error handling + cache.s3_client.get_object.side_effect = Exception("S3 Error") + + result = await cache.async_get_cache("error_key") + assert result is None + + +@pytest.mark.asyncio +async def test_s3_cache_async_with_key_prefix(mock_s3_dependencies): + """Test async operations with s3_path prefix""" + cache = S3Cache("test-bucket", s3_path="cache/data") + test_value = {"key": "value"} + + await cache.async_set_cache("namespace:key", test_value) + + cache.s3_client.put_object.assert_called_once() + call_args = cache.s3_client.put_object.call_args + + # Should transform key with prefix and colon replacement + assert call_args[1]["Key"] == "cache/data/namespace/key" + + +def test_s3_cache_supports_async(): + """Test that S3Cache now supports async operations""" + from litellm.caching.caching import Cache, LiteLLMCacheType + + cache = Cache(type=LiteLLMCacheType.S3, s3_bucket_name="test-bucket") + + # Should now return True for async support + assert cache._supports_async() is True + + +@pytest.mark.asyncio +async def test_s3_cache_async_disconnect(mock_s3_dependencies): + """Test async disconnect method""" + cache = S3Cache("test-bucket") + + # Should not raise any exceptions + await cache.disconnect() \ No newline at end of file diff --git a/tests/test_litellm/caching/test_s3_cache_async.py b/tests/test_litellm/caching/test_s3_cache_async.py deleted file mode 100644 index 4c18a5f295..0000000000 --- a/tests/test_litellm/caching/test_s3_cache_async.py +++ /dev/null @@ -1,207 +0,0 @@ -import os -import sys -from unittest.mock import MagicMock, patch -import json -import datetime -import asyncio - -import pytest - -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path - -from litellm.caching.s3_cache import S3Cache - - -@pytest.fixture -def mock_s3_dependencies(): - mock_s3_client = MagicMock() - - with patch("boto3.client", return_value=mock_s3_client): - yield {"s3_client": mock_s3_client} - - -@pytest.mark.asyncio -async def test_s3_cache_async_set_cache(mock_s3_dependencies): - """Test async_set_cache functionality using run_in_executor""" - cache = S3Cache("test-bucket") - test_value = {"key": "value", "number": 42} - - await cache.async_set_cache("test_key", test_value) - - cache.s3_client.put_object.assert_called_once() - call_args = cache.s3_client.put_object.call_args - - assert call_args[1]["Bucket"] == "test-bucket" - assert call_args[1]["Key"] == "test_key" - assert call_args[1]["Body"] == json.dumps(test_value) - assert call_args[1]["ContentType"] == "application/json" - assert call_args[1]["ContentLanguage"] == "en" - assert call_args[1]["ContentDisposition"] == 'inline; filename="test_key.json"' - - -@pytest.mark.asyncio -async def test_s3_cache_async_set_cache_with_ttl(mock_s3_dependencies): - """Test async_set_cache with TTL functionality""" - cache = S3Cache("test-bucket") - test_value = {"key": "value"} - ttl = datetime.timedelta(seconds=3600) # 1 hour - - await cache.async_set_cache("test_key", test_value, ttl=ttl) - - cache.s3_client.put_object.assert_called_once() - call_args = cache.s3_client.put_object.call_args - - assert "Expires" in call_args[1] - assert "CacheControl" in call_args[1] - assert "max-age=1:00:00" in call_args[1]["CacheControl"] - - -@pytest.mark.asyncio -async def test_s3_cache_async_get_cache(mock_s3_dependencies): - """Test async_get_cache functionality using run_in_executor""" - cache = S3Cache("test-bucket") - - mock_response = { - "Body": MagicMock() - } - mock_response["Body"].read.return_value = b'{"key": "value", "number": 42}' - cache.s3_client.get_object.return_value = mock_response - - result = await cache.async_get_cache("test_key") - - cache.s3_client.get_object.assert_called_once_with( - Bucket="test-bucket", - Key="test_key" - ) - - assert result == {"key": "value", "number": 42} - - -@pytest.mark.asyncio -async def test_s3_cache_async_get_cache_not_found(mock_s3_dependencies): - """Test async_get_cache when key is not found""" - import botocore.exceptions - - cache = S3Cache("test-bucket") - - error_response = {"Error": {"Code": "NoSuchKey"}} - cache.s3_client.get_object.side_effect = botocore.exceptions.ClientError( - error_response, "GetObject" - ) - - result = await cache.async_get_cache("nonexistent_key") - - cache.s3_client.get_object.assert_called_once_with( - Bucket="test-bucket", - Key="nonexistent_key" - ) - assert result is None - - -@pytest.mark.asyncio -async def test_s3_cache_async_set_cache_pipeline(mock_s3_dependencies): - """Test async_set_cache_pipeline functionality""" - cache = S3Cache("test-bucket") - - cache_list = [ - ("key1", {"data": "value1"}), - ("key2", {"data": "value2"}), - ("key3", {"data": "value3"}), - ] - - await cache.async_set_cache_pipeline(cache_list) - - # Should have called put_object 3 times - assert cache.s3_client.put_object.call_count == 3 - - # Verify each call - calls = cache.s3_client.put_object.call_args_list - for i, (key, value) in enumerate(cache_list): - call_args = calls[i][1] - assert call_args["Bucket"] == "test-bucket" - assert call_args["Key"] == key - assert call_args["Body"] == json.dumps(value) - - -@pytest.mark.asyncio -async def test_s3_cache_concurrent_async_operations(mock_s3_dependencies): - """Test concurrent async operations to ensure they don't block each other""" - cache = S3Cache("test-bucket") - - # Create multiple concurrent set operations - tasks = [] - for i in range(5): - key = f"concurrent_key_{i}" - value = {"id": i, "data": f"test_data_{i}"} - tasks.append(cache.async_set_cache(key, value)) - - # Execute all tasks concurrently - await asyncio.gather(*tasks) - - # Verify all operations were called - assert cache.s3_client.put_object.call_count == 5 - - # Verify each call had correct parameters - calls = cache.s3_client.put_object.call_args_list - for i, call in enumerate(calls): - call_args = call[1] - assert call_args["Bucket"] == "test-bucket" - assert f"concurrent_key_{i}" == call_args["Key"] - - -@pytest.mark.asyncio -async def test_s3_cache_async_error_handling(mock_s3_dependencies): - """Test that async methods handle errors gracefully""" - cache = S3Cache("test-bucket") - - # Test async_set_cache error handling - cache.s3_client.put_object.side_effect = Exception("S3 Error") - - # Should not raise exception, just log it - await cache.async_set_cache("error_key", {"data": "value"}) - - # Test async_get_cache error handling - cache.s3_client.get_object.side_effect = Exception("S3 Error") - - result = await cache.async_get_cache("error_key") - assert result is None - - -@pytest.mark.asyncio -async def test_s3_cache_async_with_key_prefix(mock_s3_dependencies): - """Test async operations with s3_path prefix""" - cache = S3Cache("test-bucket", s3_path="cache/data") - test_value = {"key": "value"} - - await cache.async_set_cache("namespace:key", test_value) - - cache.s3_client.put_object.assert_called_once() - call_args = cache.s3_client.put_object.call_args - - # Should transform key with prefix and colon replacement - assert call_args[1]["Key"] == "cache/data/namespace/key" - - -def test_s3_cache_supports_async(): - """Test that S3Cache now supports async operations""" - from litellm.caching.caching import Cache, LiteLLMCacheType - - cache = Cache(type=LiteLLMCacheType.S3, s3_bucket_name="test-bucket") - - # Should now return True for async support - assert cache._supports_async() is True - - -@pytest.mark.asyncio -async def test_s3_cache_async_disconnect(mock_s3_dependencies): - """Test async disconnect method""" - cache = S3Cache("test-bucket") - - # Should not raise any exceptions - await cache.disconnect() - - -if __name__ == "__main__": - pytest.main([__file__, "-v", "-s"]) From 5dd7e001274c7157d9d772d89e7657dad4497235 Mon Sep 17 00:00:00 2001 From: Michal Otmianowski Date: Thu, 21 Aug 2025 11:27:12 +0200 Subject: [PATCH 3/6] formatting adjustments --- tests/test_litellm/caching/test_s3_cache.py | 54 ++++++++++----------- 1 file changed, 25 insertions(+), 29 deletions(-) diff --git a/tests/test_litellm/caching/test_s3_cache.py b/tests/test_litellm/caching/test_s3_cache.py index 9ab49987be..85c62cc462 100644 --- a/tests/test_litellm/caching/test_s3_cache.py +++ b/tests/test_litellm/caching/test_s3_cache.py @@ -1,3 +1,4 @@ +from litellm.caching.s3_cache import S3Cache import os import sys from unittest.mock import MagicMock, patch @@ -11,13 +12,11 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -from litellm.caching.s3_cache import S3Cache - @pytest.fixture def mock_s3_dependencies(): mock_s3_client = MagicMock() - + with patch("boto3.client", return_value=mock_s3_client): yield {"s3_client": mock_s3_client} @@ -26,12 +25,12 @@ def test_s3_cache_set_cache(mock_s3_dependencies): """Test basic set_cache functionality""" cache = S3Cache("test-bucket") test_value = {"key": "value", "number": 42} - + cache.set_cache("test_key", test_value) - + cache.s3_client.put_object.assert_called_once() call_args = cache.s3_client.put_object.call_args - + assert call_args[1]["Bucket"] == "test-bucket" assert call_args[1]["Key"] == "test_key" assert call_args[1]["Body"] == json.dumps(test_value) @@ -44,7 +43,7 @@ def test_s3_cache_set_cache_with_ttl(mock_s3_dependencies): """Test set_cache with TTL functionality""" cache = S3Cache("test-bucket") test_value = {"key": "value"} - ttl = 3600 # 1 hour in seconds + ttl = datetime.timedelta(seconds=3600) # 1 hour cache.set_cache("test_key", test_value, ttl=ttl) @@ -53,44 +52,44 @@ def test_s3_cache_set_cache_with_ttl(mock_s3_dependencies): assert "Expires" in call_args[1] assert "CacheControl" in call_args[1] - assert "max-age=3600" in call_args[1]["CacheControl"] + assert "max-age=1:00:00" in call_args[1]["CacheControl"] def test_s3_cache_get_cache(mock_s3_dependencies): """Test basic get_cache functionality""" cache = S3Cache("test-bucket") - + mock_response = { "Body": MagicMock() } mock_response["Body"].read.return_value = b'{"key": "value", "number": 42}' cache.s3_client.get_object.return_value = mock_response - + result = cache.get_cache("test_key") - + cache.s3_client.get_object.assert_called_once_with( - Bucket="test-bucket", + Bucket="test-bucket", Key="test_key" ) - + assert result == {"key": "value", "number": 42} def test_s3_cache_get_cache_not_found(mock_s3_dependencies): """Test get_cache when key is not found""" import botocore.exceptions - + cache = S3Cache("test-bucket") - + error_response = {"Error": {"Code": "NoSuchKey"}} cache.s3_client.get_object.side_effect = botocore.exceptions.ClientError( error_response, "GetObject" ) - + result = cache.get_cache("nonexistent_key") - + cache.s3_client.get_object.assert_called_once_with( - Bucket="test-bucket", + Bucket="test-bucket", Key="nonexistent_key" ) assert result is None @@ -99,16 +98,16 @@ def test_s3_cache_get_cache_not_found(mock_s3_dependencies): def test_s3_key_transformation(): """Test the _to_s3_key method for key transformation""" cache = S3Cache("test-bucket") - + # Test basic key transformation (colon to slash) result = cache._to_s3_key("user:123:session:456") assert result == "user/123/session/456" - + # Test with s3_path prefix cache_with_prefix = S3Cache("test-bucket", s3_path="cache/data") result = cache_with_prefix._to_s3_key("namespace:key") assert result == "cache/data/namespace/key" - + # Test with s3_path that has trailing slash cache_with_slash = S3Cache("test-bucket", s3_path="cache/data/") result = cache_with_slash._to_s3_key("namespace:key") @@ -131,6 +130,7 @@ def test_s3_cache_initialization(): # ASYNC TESTS # ============================================================================ + @pytest.mark.asyncio async def test_s3_cache_async_set_cache(mock_s3_dependencies): cache = S3Cache("test-bucket") @@ -169,17 +169,14 @@ async def test_s3_cache_async_set_cache_with_ttl(mock_s3_dependencies): async def test_s3_cache_async_get_cache(mock_s3_dependencies): cache = S3Cache("test-bucket") - mock_response = { - "Body": MagicMock() - } + mock_response = {"Body": MagicMock()} mock_response["Body"].read.return_value = b'{"key": "value", "number": 42}' cache.s3_client.get_object.return_value = mock_response result = await cache.async_get_cache("test_key") cache.s3_client.get_object.assert_called_once_with( - Bucket="test-bucket", - Key="test_key" + Bucket="test-bucket", Key="test_key" ) assert result == {"key": "value", "number": 42} @@ -200,8 +197,7 @@ async def test_s3_cache_async_get_cache_not_found(mock_s3_dependencies): result = await cache.async_get_cache("nonexistent_key") cache.s3_client.get_object.assert_called_once_with( - Bucket="test-bucket", - Key="nonexistent_key" + Bucket="test-bucket", Key="nonexistent_key" ) assert result is None @@ -306,4 +302,4 @@ async def test_s3_cache_async_disconnect(mock_s3_dependencies): cache = S3Cache("test-bucket") # Should not raise any exceptions - await cache.disconnect() \ No newline at end of file + await cache.disconnect() From feae09da5d94fbc16d8acf65ca40601cdaaa9560 Mon Sep 17 00:00:00 2001 From: Michal Otmianowski Date: Thu, 21 Aug 2025 11:33:11 +0200 Subject: [PATCH 4/6] minor adjustments --- tests/test_litellm/caching/test_s3_cache.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/caching/test_s3_cache.py b/tests/test_litellm/caching/test_s3_cache.py index 85c62cc462..7fbb7f316c 100644 --- a/tests/test_litellm/caching/test_s3_cache.py +++ b/tests/test_litellm/caching/test_s3_cache.py @@ -1,4 +1,3 @@ -from litellm.caching.s3_cache import S3Cache import os import sys from unittest.mock import MagicMock, patch @@ -12,6 +11,8 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path +from litellm.caching.s3_cache import S3Cache + @pytest.fixture def mock_s3_dependencies(): From d53dd58c4d4942589aa6fc6d53e51f4d33754d56 Mon Sep 17 00:00:00 2001 From: Michal Otmianowski Date: Thu, 21 Aug 2025 11:39:08 +0200 Subject: [PATCH 5/6] remove unnecessary comments --- litellm/caching/s3_cache.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/litellm/caching/s3_cache.py b/litellm/caching/s3_cache.py index 929de6493c..dbf5164b08 100644 --- a/litellm/caching/s3_cache.py +++ b/litellm/caching/s3_cache.py @@ -98,7 +98,6 @@ class S3Cache(BaseCache): ContentDisposition=f'inline; filename="{key}.json"', ) except Exception as e: - # NON blocking - notify users S3 is throwing an exception print_verbose(f"S3 Caching: set_cache() - Got exception from S3: {e}") async def async_set_cache(self, key, value, **kwargs): @@ -112,7 +111,6 @@ class S3Cache(BaseCache): func = partial(self.set_cache, key, value, **kwargs) await loop.run_in_executor(None, func) except Exception as e: - # NON blocking - notify users S3 is throwing an exception verbose_logger.error(f"S3 Caching: async_set_cache() - Got exception from S3: {e}") def get_cache(self, key, **kwargs): From a6f052882b825880316158dbce4000b3de1da5ee Mon Sep 17 00:00:00 2001 From: Michal Otmianowski Date: Thu, 21 Aug 2025 12:10:29 +0200 Subject: [PATCH 6/6] fix test --- litellm/caching/s3_cache.py | 2 +- tests/test_litellm/caching/test_s3_cache.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/caching/s3_cache.py b/litellm/caching/s3_cache.py index dbf5164b08..15f7a5c1e1 100644 --- a/litellm/caching/s3_cache.py +++ b/litellm/caching/s3_cache.py @@ -106,7 +106,7 @@ class S3Cache(BaseCache): Compatible with Python 3.8+. """ try: - verbose_logger.debug(f"LiteLLM ASYNC SET Cache - S3. Key={key}. Value={value}") + verbose_logger.debug(f"Set ASYNC S3 Cache: Key={key}. Value={value}") loop = asyncio.get_event_loop() func = partial(self.set_cache, key, value, **kwargs) await loop.run_in_executor(None, func) diff --git a/tests/test_litellm/caching/test_s3_cache.py b/tests/test_litellm/caching/test_s3_cache.py index 7fbb7f316c..dce3f7d585 100644 --- a/tests/test_litellm/caching/test_s3_cache.py +++ b/tests/test_litellm/caching/test_s3_cache.py @@ -44,7 +44,7 @@ def test_s3_cache_set_cache_with_ttl(mock_s3_dependencies): """Test set_cache with TTL functionality""" cache = S3Cache("test-bucket") test_value = {"key": "value"} - ttl = datetime.timedelta(seconds=3600) # 1 hour + ttl = 3600 # 1 hour in seconds cache.set_cache("test_key", test_value, ttl=ttl) @@ -53,7 +53,7 @@ def test_s3_cache_set_cache_with_ttl(mock_s3_dependencies): assert "Expires" in call_args[1] assert "CacheControl" in call_args[1] - assert "max-age=1:00:00" in call_args[1]["CacheControl"] + assert "max-age=3600" in call_args[1]["CacheControl"] def test_s3_cache_get_cache(mock_s3_dependencies):