mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
fix(vertex_ai): forget a remembered context cache that generateContent rejects
A cache deleted out of band stayed remembered until its expireTime, so every request failed with 400 (Vertex) or 403 (Gemini) until then. Drop the name on a 400, 403, or 404 so the next request lists and recreates it. Move the registry into its own module so the Gemini handler can import it without the circular import through the context caching module's Cache().
This commit is contained in:
mateo-berri
parent
d6a2f203b3
commit
c8706c8b03
5 changed files with 229 additions and 160 deletions
139
litellm/llms/vertex_ai/context_caching/known_cached_contents.py
Normal file
139
litellm/llms/vertex_ai/context_caching/known_cached_contents.py
Normal file
|
|
@ -0,0 +1,139 @@
|
|||
import asyncio
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final, Literal
|
||||
|
||||
# A remembered cache stops being reused this long before Vertex expires it, so a request
|
||||
# resolved just before expiry does not reach generateContent with an expired cache.
|
||||
KNOWN_CACHE_EXPIRY_MARGIN_SECONDS: Final = 10.0
|
||||
MAX_KNOWN_CACHED_CONTENTS: Final = 10_000
|
||||
CACHE_REJECTED_STATUS_CODES: Final = frozenset({400, 403, 404})
|
||||
|
||||
_RFC3339_FRACTION: Final = re.compile(r"\.(\d{1,6})\d*")
|
||||
|
||||
|
||||
def parse_expire_time(expire_time: object) -> float | None:
|
||||
"""Parse a Vertex RFC 3339 ``expireTime`` (nanosecond precision, ``Z`` suffix) to epoch seconds."""
|
||||
if not isinstance(expire_time, str) or not expire_time:
|
||||
return None
|
||||
normalized: Final = _RFC3339_FRACTION.sub(lambda m: "." + m.group(1), expire_time.replace("Z", "+00:00"), count=1)
|
||||
try:
|
||||
parsed: Final = datetime.fromisoformat(normalized)
|
||||
except ValueError:
|
||||
return None
|
||||
if parsed.tzinfo is None:
|
||||
return parsed.replace(tzinfo=timezone.utc).timestamp()
|
||||
return parsed.timestamp()
|
||||
|
||||
|
||||
class KnownCachedContents:
|
||||
"""Process-local record of the cachedContents this process created or found.
|
||||
|
||||
Vertex meters every cachedContents list page against a per-minute, per-location quota,
|
||||
so listing on every cache-marked request exhausts it under load. A cache is only
|
||||
remembered with the ``expireTime`` Vertex reported, and is never reused within
|
||||
``KNOWN_CACHE_EXPIRY_MARGIN_SECONDS`` of it.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._entries: dict[tuple[tuple[str, str], str], tuple[str, float]] = {} # mutable-ok: guarded by _lock
|
||||
self._lock: Final = threading.Lock()
|
||||
self._inflight: weakref.WeakKeyDictionary[ # mutable-ok: private, guarded by _lock
|
||||
asyncio.AbstractEventLoop, dict[tuple[tuple[str, str], str], asyncio.Future[str | None]]
|
||||
] = weakref.WeakKeyDictionary()
|
||||
|
||||
@staticmethod
|
||||
def scope(
|
||||
base_url: str,
|
||||
custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
|
||||
headers: Mapping[str, str],
|
||||
) -> tuple[str, str]:
|
||||
"""Where a cache name is valid: the Vertex URL names project and location, a Gemini key names the project.
|
||||
|
||||
The Gemini key is read from the headers actually sent, so an ``extra_headers`` override scopes the
|
||||
cache to the caller's project. It is kept as-is in this in-memory scope and never logged or persisted.
|
||||
"""
|
||||
if custom_llm_provider != "gemini":
|
||||
return (base_url, "")
|
||||
sent_keys: Final = [value for name, value in headers.items() if name.lower() == "x-goog-api-key"]
|
||||
return (base_url, sent_keys[-1] if sent_keys else "")
|
||||
|
||||
@staticmethod
|
||||
def is_usable(expire_time: object) -> bool:
|
||||
expires_at: Final = parse_expire_time(expire_time)
|
||||
return expires_at is None or expires_at - KNOWN_CACHE_EXPIRY_MARGIN_SECONDS > time.time()
|
||||
|
||||
def get(self, scope: tuple[str, str], cache_key: str) -> str | None:
|
||||
with self._lock:
|
||||
entry: Final = self._entries.get((scope, cache_key))
|
||||
if entry is None:
|
||||
return None
|
||||
name, reuse_until = entry
|
||||
if reuse_until > time.time():
|
||||
return name
|
||||
del self._entries[(scope, cache_key)]
|
||||
return None
|
||||
|
||||
def remember(self, scope: tuple[str, str], cache_key: str, name: str | None, expire_time: object) -> None:
|
||||
expires_at: Final = parse_expire_time(expire_time)
|
||||
if not name or expires_at is None:
|
||||
return
|
||||
reuse_until: Final = expires_at - KNOWN_CACHE_EXPIRY_MARGIN_SECONDS
|
||||
now: Final = time.time()
|
||||
if reuse_until <= now:
|
||||
return
|
||||
with self._lock:
|
||||
if len(self._entries) >= MAX_KNOWN_CACHED_CONTENTS:
|
||||
self._entries = {k: v for k, v in self._entries.items() if v[1] > now}
|
||||
while len(self._entries) >= MAX_KNOWN_CACHED_CONTENTS:
|
||||
del self._entries[next(iter(self._entries))]
|
||||
self._entries[(scope, cache_key)] = (name, reuse_until)
|
||||
|
||||
def forget_rejected(self, cached_content: str | None, status_code: int) -> None:
|
||||
if cached_content is None or status_code not in CACHE_REJECTED_STATUS_CODES:
|
||||
return
|
||||
with self._lock:
|
||||
self._entries = {key: entry for key, entry in self._entries.items() if entry[0] != cached_content}
|
||||
|
||||
async def single_flight(
|
||||
self, scope: tuple[str, str], cache_key: str, resolve: Callable[[], Awaitable[str | None]]
|
||||
) -> str | None:
|
||||
"""Run ``resolve`` once per (scope, cache_key) at a time within this event loop.
|
||||
|
||||
Concurrent requests for the same prefix wait for the first one's lookup-or-create
|
||||
instead of each listing (and creating a duplicate cache). If that first attempt
|
||||
fails, each waiter resolves on its own.
|
||||
"""
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
own: Final[asyncio.Future[str | None]] = loop.create_future()
|
||||
with self._lock:
|
||||
inflight: Final = self._inflight.setdefault(loop, {})
|
||||
leader: Final = inflight.setdefault((scope, cache_key), own)
|
||||
if leader is not own:
|
||||
name: Final = await asyncio.shield(leader)
|
||||
if name:
|
||||
return name
|
||||
return await resolve()
|
||||
try:
|
||||
resolved: Final = await resolve()
|
||||
except BaseException:
|
||||
own.set_result(None)
|
||||
raise
|
||||
else:
|
||||
own.set_result(resolved)
|
||||
return resolved
|
||||
finally:
|
||||
with self._lock:
|
||||
inflight.pop((scope, cache_key), None)
|
||||
|
||||
def clear(self) -> None:
|
||||
with self._lock:
|
||||
self._entries.clear()
|
||||
self._inflight.clear()
|
||||
|
||||
|
||||
known_cached_contents: Final = KnownCachedContents()
|
||||
|
|
@ -1,10 +1,4 @@
|
|||
import asyncio
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from collections.abc import Sequence
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
|
|
@ -26,6 +20,7 @@ from litellm.utils import is_prompt_caching_valid_prompt
|
|||
|
||||
from ..common_utils import VertexAIError, get_vertex_base_url
|
||||
from ..vertex_llm_base import VertexBase
|
||||
from .known_cached_contents import KnownCachedContents, known_cached_contents
|
||||
from .transformation import (
|
||||
cached_messages_end_on_supported_turn,
|
||||
separate_cached_messages,
|
||||
|
|
@ -36,13 +31,6 @@ local_cache_obj: Final = Cache(type=LiteLLMCacheType.LOCAL) # only used for cal
|
|||
|
||||
MAX_PAGINATION_PAGES: Final = 100 # Reasonable upper bound for pagination
|
||||
|
||||
# A remembered cache stops being reused this long before Vertex expires it, so a request
|
||||
# resolved just before expiry does not reach generateContent with an expired cache.
|
||||
KNOWN_CACHE_EXPIRY_MARGIN_SECONDS: Final = 10.0
|
||||
MAX_KNOWN_CACHED_CONTENTS: Final = 10_000
|
||||
|
||||
_RFC3339_FRACTION: Final = re.compile(r"\.(\d{1,6})\d*")
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _CachedContentItem(TypedDict):
|
||||
|
|
@ -82,123 +70,6 @@ def _record_cache_creation_tokens(logging_obj: Logging, created: _CreatedCachedC
|
|||
logging_obj.vertex_ai_cache_creation_input_tokens = max(tokens, 0) # rebind-ok: set provider count
|
||||
|
||||
|
||||
def _parse_expire_time(expire_time: object) -> float | None:
|
||||
"""Parse a Vertex RFC 3339 ``expireTime`` (nanosecond precision, ``Z`` suffix) to epoch seconds."""
|
||||
if not isinstance(expire_time, str) or not expire_time:
|
||||
return None
|
||||
normalized: Final = _RFC3339_FRACTION.sub(lambda m: "." + m.group(1), expire_time.replace("Z", "+00:00"), count=1)
|
||||
try:
|
||||
parsed: Final = datetime.fromisoformat(normalized)
|
||||
except ValueError:
|
||||
return None
|
||||
if parsed.tzinfo is None:
|
||||
return parsed.replace(tzinfo=timezone.utc).timestamp()
|
||||
return parsed.timestamp()
|
||||
|
||||
|
||||
class _KnownCachedContents:
|
||||
"""Process-local record of the cachedContents this process created or found.
|
||||
|
||||
Vertex meters every cachedContents list page against a per-minute, per-location quota,
|
||||
so listing on every cache-marked request exhausts it under load. A cache is only
|
||||
remembered with the ``expireTime`` Vertex reported, and is never reused within
|
||||
``KNOWN_CACHE_EXPIRY_MARGIN_SECONDS`` of it.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._entries: dict[tuple[tuple[str, str], str], tuple[str, float]] = {} # mutable-ok: guarded by _lock
|
||||
self._lock: Final = threading.Lock()
|
||||
self._inflight: weakref.WeakKeyDictionary[ # mutable-ok: private, guarded by _lock
|
||||
asyncio.AbstractEventLoop, dict[tuple[tuple[str, str], str], asyncio.Future[str | None]]
|
||||
] = weakref.WeakKeyDictionary()
|
||||
|
||||
@staticmethod
|
||||
def scope(
|
||||
base_url: str,
|
||||
custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
|
||||
headers: Mapping[str, str],
|
||||
) -> tuple[str, str]:
|
||||
"""Where a cache name is valid: the Vertex URL names project and location, a Gemini key names the project.
|
||||
|
||||
The Gemini key is read from the headers actually sent, so an ``extra_headers`` override scopes the
|
||||
cache to the caller's project. It is kept as-is in this in-memory scope and never logged or persisted.
|
||||
"""
|
||||
if custom_llm_provider != "gemini":
|
||||
return (base_url, "")
|
||||
sent_keys: Final = [value for name, value in headers.items() if name.lower() == "x-goog-api-key"]
|
||||
return (base_url, sent_keys[-1] if sent_keys else "")
|
||||
|
||||
@staticmethod
|
||||
def is_usable(expire_time: object) -> bool:
|
||||
expires_at: Final = _parse_expire_time(expire_time)
|
||||
return expires_at is None or expires_at - KNOWN_CACHE_EXPIRY_MARGIN_SECONDS > time.time()
|
||||
|
||||
def get(self, scope: tuple[str, str], cache_key: str) -> str | None:
|
||||
with self._lock:
|
||||
entry: Final = self._entries.get((scope, cache_key))
|
||||
if entry is None:
|
||||
return None
|
||||
name, reuse_until = entry
|
||||
if reuse_until > time.time():
|
||||
return name
|
||||
del self._entries[(scope, cache_key)]
|
||||
return None
|
||||
|
||||
def remember(self, scope: tuple[str, str], cache_key: str, name: str | None, expire_time: object) -> None:
|
||||
expires_at: Final = _parse_expire_time(expire_time)
|
||||
if not name or expires_at is None:
|
||||
return
|
||||
reuse_until: Final = expires_at - KNOWN_CACHE_EXPIRY_MARGIN_SECONDS
|
||||
now: Final = time.time()
|
||||
if reuse_until <= now:
|
||||
return
|
||||
with self._lock:
|
||||
if len(self._entries) >= MAX_KNOWN_CACHED_CONTENTS:
|
||||
self._entries = {k: v for k, v in self._entries.items() if v[1] > now}
|
||||
while len(self._entries) >= MAX_KNOWN_CACHED_CONTENTS:
|
||||
del self._entries[next(iter(self._entries))]
|
||||
self._entries[(scope, cache_key)] = (name, reuse_until)
|
||||
|
||||
async def single_flight(
|
||||
self, scope: tuple[str, str], cache_key: str, resolve: Callable[[], Awaitable[str | None]]
|
||||
) -> str | None:
|
||||
"""Run ``resolve`` once per (scope, cache_key) at a time within this event loop.
|
||||
|
||||
Concurrent requests for the same prefix wait for the first one's lookup-or-create
|
||||
instead of each listing (and creating a duplicate cache). If that first attempt
|
||||
fails, each waiter resolves on its own.
|
||||
"""
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
own: Final[asyncio.Future[str | None]] = loop.create_future()
|
||||
with self._lock:
|
||||
inflight: Final = self._inflight.setdefault(loop, {})
|
||||
leader: Final = inflight.setdefault((scope, cache_key), own)
|
||||
if leader is not own:
|
||||
name: Final = await asyncio.shield(leader)
|
||||
if name:
|
||||
return name
|
||||
return await resolve()
|
||||
try:
|
||||
resolved: Final = await resolve()
|
||||
except BaseException:
|
||||
own.set_result(None)
|
||||
raise
|
||||
else:
|
||||
own.set_result(resolved)
|
||||
return resolved
|
||||
finally:
|
||||
with self._lock:
|
||||
inflight.pop((scope, cache_key), None)
|
||||
|
||||
def clear(self) -> None:
|
||||
with self._lock:
|
||||
self._entries.clear()
|
||||
self._inflight.clear()
|
||||
|
||||
|
||||
known_cached_contents: Final = _KnownCachedContents()
|
||||
|
||||
|
||||
class ContextCachingEndpoints(VertexBase):
|
||||
"""
|
||||
Covers context caching endpoints for Vertex AI + Google AI Studio
|
||||
|
|
@ -291,7 +162,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
vertex_auth_header=vertex_auth_header,
|
||||
model=model,
|
||||
)
|
||||
scope: Final = _KnownCachedContents.scope(base_url, custom_llm_provider, headers)
|
||||
scope: Final = KnownCachedContents.scope(base_url, custom_llm_provider, headers)
|
||||
known_name: Final = known_cached_contents.get(scope, cache_key)
|
||||
if known_name is not None:
|
||||
return known_name
|
||||
|
|
@ -340,7 +211,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
if display_name is None or display_name != cache_key:
|
||||
continue
|
||||
expire_time = cached_item.get("expireTime")
|
||||
if not _KnownCachedContents.is_usable(expire_time):
|
||||
if not KnownCachedContents.is_usable(expire_time):
|
||||
continue
|
||||
name = cached_item.get("name")
|
||||
known_cached_contents.remember(scope, cache_key, name, expire_time)
|
||||
|
|
@ -388,7 +259,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
vertex_auth_header=vertex_auth_header,
|
||||
model=model,
|
||||
)
|
||||
scope: Final = _KnownCachedContents.scope(base_url, custom_llm_provider, headers)
|
||||
scope: Final = KnownCachedContents.scope(base_url, custom_llm_provider, headers)
|
||||
known_name: Final = known_cached_contents.get(scope, cache_key)
|
||||
if known_name is not None:
|
||||
return known_name
|
||||
|
|
@ -437,7 +308,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
if display_name is None or display_name != cache_key:
|
||||
continue
|
||||
expire_time = cached_item.get("expireTime")
|
||||
if not _KnownCachedContents.is_usable(expire_time):
|
||||
if not KnownCachedContents.is_usable(expire_time):
|
||||
continue
|
||||
name = cached_item.get("name")
|
||||
known_cached_contents.remember(scope, cache_key, name, expire_time)
|
||||
|
|
@ -604,7 +475,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root
|
||||
_record_cache_creation_tokens(logging_obj, raw_response_cached)
|
||||
known_cached_contents.remember(
|
||||
_KnownCachedContents.scope(url, custom_llm_provider, headers),
|
||||
KnownCachedContents.scope(url, custom_llm_provider, headers),
|
||||
generated_cache_key,
|
||||
raw_response_cached.get("name"),
|
||||
raw_response_cached.get("expireTime"),
|
||||
|
|
@ -706,7 +577,7 @@ class ContextCachingEndpoints(VertexBase):
|
|||
generated_cache_key: Final = local_cache_obj.get_cache_key(
|
||||
messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model
|
||||
)
|
||||
scope: Final = _KnownCachedContents.scope(url, custom_llm_provider, headers)
|
||||
scope: Final = KnownCachedContents.scope(url, custom_llm_provider, headers)
|
||||
|
||||
async def lookup_or_create() -> str | None:
|
||||
google_cache_name: Final = await self.async_check_cache(
|
||||
|
|
|
|||
|
|
@ -96,6 +96,7 @@ from ..common_utils import (
|
|||
build_vertex_schema,
|
||||
supports_response_json_schema,
|
||||
)
|
||||
from ..context_caching.known_cached_contents import known_cached_contents
|
||||
from ..vertex_llm_base import VertexBase
|
||||
from .grounding_requests import calculate_grounding_requests
|
||||
from .long_decode import (
|
||||
|
|
@ -2831,6 +2832,7 @@ async def make_call(
|
|||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as e:
|
||||
known_cached_contents.forget_rejected(request_body.get("cachedContent"), e.response.status_code)
|
||||
exception_string: Final = str(await e.response.aread())
|
||||
raise VertexAIError(
|
||||
status_code=e.response.status_code,
|
||||
|
|
@ -2891,14 +2893,18 @@ def make_sync_call(
|
|||
if client is None:
|
||||
client = HTTPHandler() # Create a new client if none provided
|
||||
|
||||
response: Final = client.post(
|
||||
url_for(api_base, request_body),
|
||||
headers=headers,
|
||||
data=data,
|
||||
stream=True,
|
||||
logging_obj=logging_obj,
|
||||
timeout=timeout,
|
||||
)
|
||||
try:
|
||||
response: Final = client.post(
|
||||
url_for(api_base, request_body),
|
||||
headers=headers,
|
||||
data=data,
|
||||
stream=True,
|
||||
logging_obj=logging_obj,
|
||||
timeout=timeout,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
known_cached_contents.forget_rejected(request_body.get("cachedContent"), e.response.status_code)
|
||||
raise
|
||||
|
||||
if response.status_code != 200 and response.status_code != 201:
|
||||
raise VertexAIError(
|
||||
|
|
@ -3175,6 +3181,7 @@ class VertexLLM(VertexBase):
|
|||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
known_cached_contents.forget_rejected(body.get("cachedContent"), err.response.status_code)
|
||||
raise VertexAIError(
|
||||
status_code=err.response.status_code,
|
||||
message=err.response.text,
|
||||
|
|
@ -3385,6 +3392,7 @@ class VertexLLM(VertexBase):
|
|||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
known_cached_contents.forget_rejected(body.get("cachedContent"), err.response.status_code)
|
||||
raise VertexAIError(
|
||||
status_code=err.response.status_code,
|
||||
message=err.response.text,
|
||||
|
|
|
|||
|
|
@ -1,8 +1,6 @@
|
|||
import pytest
|
||||
|
||||
from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import (
|
||||
known_cached_contents,
|
||||
)
|
||||
from litellm.llms.vertex_ai.context_caching.known_cached_contents import known_cached_contents
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
|
|
|
|||
|
|
@ -1,21 +1,24 @@
|
|||
import asyncio
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError
|
||||
from litellm.llms.vertex_ai.context_caching import vertex_ai_context_caching
|
||||
from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import (
|
||||
ContextCachingEndpoints,
|
||||
_parse_expire_time,
|
||||
known_cached_contents,
|
||||
)
|
||||
from litellm.llms.vertex_ai.context_caching.known_cached_contents import known_cached_contents, parse_expire_time
|
||||
from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ContextCachingEndpoints
|
||||
|
||||
_MODULE: Final = "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching"
|
||||
_REGISTRY_MODULE: Final = "litellm.llms.vertex_ai.context_caching.known_cached_contents"
|
||||
_SCOPE: Final = ("https://vertex.test/v1/projects/p/locations/us/cachedContents", "")
|
||||
|
||||
|
||||
|
|
@ -163,7 +166,7 @@ def test_remembered_cache_is_listed_again_once_its_reuse_window_closes():
|
|||
client.post.return_value = _response(_created("cachedContents/1", _expire_time(60)))
|
||||
assert _sync_call(client)[2] == "cachedContents/1"
|
||||
|
||||
with patch(f"{_MODULE}.time.time", return_value=datetime.now(timezone.utc).timestamp() + 55):
|
||||
with patch(f"{_REGISTRY_MODULE}.time.time", return_value=datetime.now(timezone.utc).timestamp() + 55):
|
||||
assert _sync_call(client)[2] == "cachedContents/listed"
|
||||
|
||||
|
||||
|
|
@ -290,18 +293,18 @@ async def test_waiters_resolve_on_their_own_when_the_first_request_fails():
|
|||
|
||||
|
||||
def test_parse_expire_time_accepts_nanosecond_precision():
|
||||
assert _parse_expire_time("2026-10-14T18:06:21.123456789Z") == pytest.approx(
|
||||
assert parse_expire_time("2026-10-14T18:06:21.123456789Z") == pytest.approx(
|
||||
datetime(2026, 10, 14, 18, 6, 21, 123456, tzinfo=timezone.utc).timestamp()
|
||||
)
|
||||
assert (
|
||||
_parse_expire_time("2026-10-14T18:06:21Z") == datetime(2026, 10, 14, 18, 6, 21, tzinfo=timezone.utc).timestamp()
|
||||
parse_expire_time("2026-10-14T18:06:21Z") == datetime(2026, 10, 14, 18, 6, 21, tzinfo=timezone.utc).timestamp()
|
||||
)
|
||||
assert _parse_expire_time("not a time") is None
|
||||
assert _parse_expire_time(None) is None
|
||||
assert parse_expire_time("not a time") is None
|
||||
assert parse_expire_time(None) is None
|
||||
|
||||
|
||||
def test_parse_expire_time_reads_a_missing_offset_as_utc():
|
||||
assert _parse_expire_time("2026-10-14T18:06:21") == datetime(2026, 10, 14, 18, 6, 21, tzinfo=timezone.utc).timestamp()
|
||||
assert parse_expire_time("2026-10-14T18:06:21") == datetime(2026, 10, 14, 18, 6, 21, tzinfo=timezone.utc).timestamp()
|
||||
|
||||
|
||||
def test_cache_expiring_within_the_margin_is_not_remembered():
|
||||
|
|
@ -311,9 +314,59 @@ def test_cache_expiring_within_the_margin_is_not_remembered():
|
|||
|
||||
|
||||
def test_registry_is_bounded():
|
||||
with patch(f"{_MODULE}.MAX_KNOWN_CACHED_CONTENTS", 3):
|
||||
with patch(f"{_REGISTRY_MODULE}.MAX_KNOWN_CACHED_CONTENTS", 3):
|
||||
for index in range(5):
|
||||
known_cached_contents.remember(_SCOPE, f"key-{index}", f"cachedContents/{index}", _expire_time(60))
|
||||
|
||||
assert known_cached_contents.get(_SCOPE, "key-0") is None
|
||||
assert known_cached_contents.get(_SCOPE, "key-4") == "cachedContents/4"
|
||||
|
||||
|
||||
async def _complete(sync_mode: bool, stream: bool) -> None:
|
||||
if sync_mode:
|
||||
response = litellm.completion(model="gemini/gemini-3.1-pro", messages=_messages(), stream=stream)
|
||||
for _chunk in response if stream else ():
|
||||
pass
|
||||
return
|
||||
response = await litellm.acompletion(model="gemini/gemini-3.1-pro", messages=_messages(), stream=stream)
|
||||
if stream:
|
||||
async for _chunk in response:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.parametrize("stream", [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
("status_code", "relisted"),
|
||||
[(400, True), (403, True), (404, True), (429, False), (500, False)],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_rejected_by_generate_content_is_looked_up_again(
|
||||
respx_mock: respx.MockRouter,
|
||||
fake_provider_credentials: None,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sync_mode: bool,
|
||||
stream: bool,
|
||||
status_code: int,
|
||||
relisted: bool,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
list_route: Final = respx_mock.get(url__regex=r".*/cachedContents.*").mock(
|
||||
return_value=httpx.Response(200, json={})
|
||||
)
|
||||
respx_mock.post(url__regex=r".*/cachedContents$").mock(
|
||||
return_value=httpx.Response(200, json=_created("cachedContents/deleted", _expire_time(300)))
|
||||
)
|
||||
generate_route: Final = respx_mock.post(url__regex=r".*:(stream)?[gG]enerateContent.*").mock(
|
||||
return_value=httpx.Response(status_code, json={"error": {"code": status_code, "message": "rejected"}})
|
||||
)
|
||||
|
||||
for _ in range(2):
|
||||
with pytest.raises(openai.APIStatusError):
|
||||
await _complete(sync_mode, stream)
|
||||
|
||||
assert [json.loads(call.request.content)["cachedContent"] for call in generate_route.calls] == [
|
||||
"cachedContents/deleted",
|
||||
"cachedContents/deleted",
|
||||
]
|
||||
assert list_route.call_count == (2 if relisted else 1)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue