Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions litellm/caching/gcs_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import json
import asyncio
from typing import Optional
from urllib.parse import quote

from litellm._logging import print_verbose, verbose_logger
from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase
Expand Down Expand Up @@ -48,7 +49,7 @@ def set_cache(self, key, value, **kwargs):
headers = self._construct_headers()
object_name = self.key_prefix + key
bucket_name = self.bucket_name
url = f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={object_name}"
url = f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={quote(object_name, safe='')}"
data = json.dumps(value)
self.sync_client.post(url=url, data=data, headers=headers)
except Exception as e:
Expand All @@ -59,7 +60,7 @@ async def async_set_cache(self, key, value, **kwargs):
headers = self._construct_headers()
object_name = self.key_prefix + key
bucket_name = self.bucket_name
url = f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={object_name}"
url = f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={quote(object_name, safe='')}"
data = json.dumps(value)
await self.async_client.post(url=url, data=data, headers=headers)
except Exception as e:
Expand All @@ -72,7 +73,7 @@ def get_cache(self, key, **kwargs):
headers = self._construct_headers()
object_name = self.key_prefix + key
bucket_name = self.bucket_name
url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{object_name}?alt=media"
url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{quote(object_name, safe='')}?alt=media"
response = self.sync_client.get(url=url, headers=headers)
if response.status_code == 200:
cached_response = json.loads(response.text)
Expand All @@ -91,7 +92,7 @@ async def async_get_cache(self, key, **kwargs):
headers = self._construct_headers()
object_name = self.key_prefix + key
bucket_name = self.bucket_name
url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{object_name}?alt=media"
url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{quote(object_name, safe='')}?alt=media"
response = await self.async_client.get(url=url, headers=headers)
if response.status_code == 200:
return json.loads(response.text)
Expand Down
61 changes: 61 additions & 0 deletions tests/test_litellm/caching/test_gcs_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,3 +44,64 @@ async def test_gcs_cache_async_set_and_get(mock_gcs_dependencies):
mock_gcs_dependencies["async_client"].get.return_value.text = '{"foo": "bar"}'
result = await cache.async_get_cache("key")
assert result == {"foo": "bar"}


@pytest.mark.asyncio
async def test_gcs_cache_async_get_encodes_object_name_in_path(mock_gcs_dependencies):
"""
Regression test for https://github.com/BerriAI/litellm/issues/30377

When gcs_path is set, the object name contains a '/' (e.g. "my_cache/<hash>").
The GCS JSON API requires the object name in the GET path to be URL-encoded,
so the '/' must be sent as '%2F'. Otherwise GCS returns 404 and every read
silently misses.
"""
cache = GCSCache(bucket_name="test-bucket", gcs_path="my_cache/")

mock_gcs_dependencies["async_client"].get.return_value.status_code = 200
mock_gcs_dependencies["async_client"].get.return_value.text = '{"foo": "bar"}'

result = await cache.async_get_cache("abc123")
assert result == {"foo": "bar"}

called_url = mock_gcs_dependencies["async_client"].get.call_args.kwargs["url"]
# The slash from gcs_path must be percent-encoded in the path segment.
assert "/o/my_cache%2Fabc123?alt=media" in called_url
assert "/o/my_cache/abc123" not in called_url


def test_gcs_cache_get_encodes_object_name_in_path(mock_gcs_dependencies):
"""Sync counterpart of the regression test for issue #30377."""
cache = GCSCache(bucket_name="test-bucket", gcs_path="my_cache/")

mock_gcs_dependencies["sync_client"].get.return_value.status_code = 200
mock_gcs_dependencies["sync_client"].get.return_value.text = '{"foo": "bar"}'

result = cache.get_cache("abc123")
assert result == {"foo": "bar"}

called_url = mock_gcs_dependencies["sync_client"].get.call_args.kwargs["url"]
assert "/o/my_cache%2Fabc123?alt=media" in called_url
assert "/o/my_cache/abc123" not in called_url


def test_gcs_cache_set_encodes_object_name_in_query(mock_gcs_dependencies):
"""
The set path uses the object name as a query parameter. Encoding it keeps
both sides symmetric so the key written matches the key read back.
"""
cache = GCSCache(bucket_name="test-bucket", gcs_path="my_cache/")
cache.set_cache("abc123", {"foo": "bar"})

called_url = mock_gcs_dependencies["sync_client"].post.call_args.kwargs["url"]
assert "name=my_cache%2Fabc123" in called_url
Comment thread
darktheorys marked this conversation as resolved.
Comment thread
greptile-apps[bot] marked this conversation as resolved.


@pytest.mark.asyncio
async def test_gcs_cache_async_set_encodes_object_name_in_query(mock_gcs_dependencies):
"""Async counterpart of test_gcs_cache_set_encodes_object_name_in_query."""
cache = GCSCache(bucket_name="test-bucket", gcs_path="my_cache/")
await cache.async_set_cache("abc123", {"foo": "bar"})

called_url = mock_gcs_dependencies["async_client"].post.call_args.kwargs["url"]
assert "name=my_cache%2Fabc123" in called_url
Loading