phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 03:04:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(caching.py): add s3 path as a top-level param

Browse source
This commit is contained in:
Krrish Dholakia 2024-03-06 18:07:28 -08:00
parent 32adb4e80d
commit 726dad5756
2 changed files with 6 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

3
litellm/caching.py
View file

@ -572,6 +572,7 @@ class S3Cache(BaseCache):
self.bucket_name = s3_bucket_name self.bucket_name = s3_bucket_name
self.key_prefix = s3_path.rstrip("/") + "/" if s3_path else "" self.key_prefix = s3_path.rstrip("/") + "/" if s3_path else ""
# Create an S3 client with custom endpoint URL # Create an S3 client with custom endpoint URL
self.s3_client = boto3.client( self.s3_client = boto3.client(
"s3", "s3",
region_name=s3_region_name, region_name=s3_region_name,
@ -776,6 +777,7 @@ class Cache:
s3_aws_secret_access_key: Optional[str] = None, s3_aws_secret_access_key: Optional[str] = None,
s3_aws_session_token: Optional[str] = None, s3_aws_session_token: Optional[str] = None,
s3_config: Optional[Any] = None, s3_config: Optional[Any] = None,
s3_path: Optional[str] = None,
redis_semantic_cache_use_async=False, redis_semantic_cache_use_async=False,
redis_semantic_cache_embedding_model="text-embedding-ada-002", redis_semantic_cache_embedding_model="text-embedding-ada-002",
**kwargs, **kwargs,
@ -825,6 +827,7 @@ class Cache:
s3_aws_secret_access_key=s3_aws_secret_access_key, s3_aws_secret_access_key=s3_aws_secret_access_key,
s3_aws_session_token=s3_aws_session_token, s3_aws_session_token=s3_aws_session_token,
s3_config=s3_config, s3_config=s3_config,
s3_path=s3_path,
**kwargs, **kwargs,
) )
if "cache" not in litellm.input_callback: if "cache" not in litellm.input_callback:

5
litellm/tests/test_caching.py
View file

@ -698,7 +698,6 @@ def test_s3_cache_acompletion_stream_azure():
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.skip(reason="AWS Suspended Account")
async def test_s3_cache_acompletion_azure(): async def test_s3_cache_acompletion_azure():
import asyncio import asyncio
import logging import logging
@ -717,7 +716,9 @@ async def test_s3_cache_acompletion_azure():
} }
] ]
litellm.cache = Cache( litellm.cache = Cache(
type="s3", s3_bucket_name="cache-bucket-litellm", s3_region_name="us-west-2" type="s3",
s3_bucket_name="litellm-my-test-bucket-2",
s3_region_name="us-east-1",
) )
print("s3 Cache: test for caching, streaming + completion") print("s3 Cache: test for caching, streaming + completion")