Skip to content

Commit 4e02e7e

Browse files
mateo-berrimateo-berri
authored
Merge pull request #37742 from BerriAI/litellm_lit5879_semantic_cache_embedding_timeout
fix(caching): bound the semantic cache embedding lookup so a dead embedding endpoint can't block requests
2 parents 66a6a09 + d57715b commit 4e02e7e

9 files changed

Lines changed: 296 additions & 24 deletions

File tree

‎litellm/caching/_embedding_router.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
from typing import TYPE_CHECKING, Any, Final
1717

1818
import litellm
19+
from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
1920

2021
if TYPE_CHECKING:
2122
from litellm.router import Router
@@ -60,6 +61,13 @@ def resolve_embedding_max_input_tokens(
6061
return deployment_max_input_tokens
6162

6263

64+
def resolve_embedding_timeout(configured_timeout: float | None) -> float:
65+
"""Explicit cache setting first, else the short semantic-cache default."""
66+
if configured_timeout is not None:
67+
return configured_timeout
68+
return SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
69+
70+
6371
def truncate_embedding_input(prompt: str, embedding_model: str, max_input_tokens: int | None) -> str:
6472
"""Keep only the first ``max_input_tokens`` tokens of ``prompt`` for the embedding call."""
6573
if max_input_tokens is None:

‎litellm/caching/caching.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@ def __init__(
9898
qdrant_semantic_cache_embedding_model: str = "text-embedding-ada-002",
9999
qdrant_semantic_cache_vector_size: int | None = None,
100100
semantic_cache_embedding_max_input_tokens: int | None = None,
101+
semantic_cache_embedding_timeout: float | None = None,
101102
# GCP IAM authentication parameters
102103
gcp_service_account: str | None = None,
103104
gcp_ssl_ca_certs: str | None = None,
@@ -124,6 +125,7 @@ def __init__(
124125
qdrant_collection_name (str, optional): The name for your qdrant collection. Required if type is "qdrant-semantic".
125126
similarity_threshold (float, optional): The similarity threshold for semantic-caching, Required if type is "redis-semantic" or "qdrant-semantic".
126127
semantic_cache_embedding_max_input_tokens (int, optional): Truncate prompts to this many tokens before embedding them for semantic caching. Defaults to the embedding deployment's configured max_input_tokens.
128+
semantic_cache_embedding_timeout (float, optional): Seconds a semantic-cache lookup may spend embedding the prompt before it gives up and lets the request continue to the LLM. Defaults to SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS.
127129
128130
# Disk Cache Args
129131
disk_cache_dir (str, optional): The directory for the disk cache. Defaults to None.
@@ -195,6 +197,7 @@ def __init__(
195197
embedding_model=redis_semantic_cache_embedding_model,
196198
index_name=redis_semantic_cache_index_name,
197199
embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens,
200+
embedding_timeout=semantic_cache_embedding_timeout,
198201
**kwargs,
199202
)
200203
elif type == LiteLLMCacheType.VALKEY_SEMANTIC:
@@ -211,6 +214,7 @@ def __init__(
211214
index_name=valkey_semantic_cache_index_name,
212215
startup_nodes=redis_startup_nodes,
213216
embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens,
217+
embedding_timeout=semantic_cache_embedding_timeout,
214218
**kwargs,
215219
)
216220
elif type == LiteLLMCacheType.QDRANT_SEMANTIC:
@@ -223,6 +227,7 @@ def __init__(
223227
embedding_model=qdrant_semantic_cache_embedding_model,
224228
vector_size=qdrant_semantic_cache_vector_size,
225229
embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens,
230+
embedding_timeout=semantic_cache_embedding_timeout,
226231
)
227232
elif type == LiteLLMCacheType.LOCAL:
228233
self.cache = InMemoryCache()

‎litellm/caching/qdrant_semantic_cache.py‎

Lines changed: 26 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,11 @@
1616

1717
import litellm
1818
from litellm._logging import print_verbose
19-
from litellm.constants import QDRANT_SCALAR_QUANTILE, QDRANT_VECTOR_SIZE
19+
from litellm.constants import (
20+
QDRANT_SCALAR_QUANTILE,
21+
QDRANT_VECTOR_SIZE,
22+
SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS,
23+
)
2024
from litellm.litellm_core_utils.prompt_templates.common_utils import (
2125
get_str_from_messages,
2226
)
@@ -26,6 +30,7 @@
2630
build_router_embedding_metadata,
2731
resolve_embedding_max_input_tokens,
2832
resolve_embedding_router,
33+
resolve_embedding_timeout,
2934
truncate_embedding_input,
3035
)
3136
from .base_cache import BaseCache
@@ -37,6 +42,7 @@
3742
class QdrantSemanticCache(BaseCache):
3843
CACHE_KEY_FIELD_NAME = "litellm_cache_key"
3944
embedding_max_input_tokens: int | None = None
45+
embedding_timeout: float = SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
4046

4147
def __init__(
4248
self,
@@ -49,6 +55,7 @@ def __init__(
4955
host_type=None,
5056
vector_size=None,
5157
embedding_max_input_tokens: int | None = None,
58+
embedding_timeout: float | None = None,
5259
):
5360
from litellm.llms.custom_httpx.http_handler import (
5461
_get_httpx_client,
@@ -68,6 +75,7 @@ def __init__(
6875
self.similarity_threshold = similarity_threshold
6976
self.embedding_model = embedding_model
7077
self.embedding_max_input_tokens = embedding_max_input_tokens
78+
self.embedding_timeout = resolve_embedding_timeout(embedding_timeout)
7179
self.vector_size = vector_size if vector_size is not None else QDRANT_VECTOR_SIZE
7280
headers = {}
7381

@@ -222,11 +230,15 @@ def _get_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) ->
222230
input=embedding_input,
223231
cache={"no-store": True, "no-cache": True},
224232
metadata=build_router_embedding_metadata(metadata),
233+
timeout=self.embedding_timeout,
234+
num_retries=0,
225235
)
226236
return litellm.embedding(
227237
model=self.embedding_model,
228238
input=embedding_input,
229239
cache={"no-store": True, "no-cache": True},
240+
timeout=self.embedding_timeout,
241+
num_retries=0,
230242
)
231243

232244
async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse:
@@ -238,19 +250,25 @@ async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | Non
238250

239251
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
240252
embedding_input: Final = self._embedding_input(prompt, router)
241-
if router is not None:
242-
return await router.aembedding(
253+
embedding_call: Final = (
254+
router.aembedding(
243255
model=self.embedding_model,
244256
input=embedding_input,
245257
cache={"no-store": True, "no-cache": True},
246258
metadata=build_router_embedding_metadata(metadata),
259+
timeout=self.embedding_timeout,
260+
num_retries=0,
261+
)
262+
if router is not None
263+
else litellm.aembedding(
264+
model=self.embedding_model,
265+
input=embedding_input,
266+
cache={"no-store": True, "no-cache": True},
267+
timeout=self.embedding_timeout,
268+
num_retries=0,
247269
)
248-
249-
return await litellm.aembedding(
250-
model=self.embedding_model,
251-
input=embedding_input,
252-
cache={"no-store": True, "no-cache": True},
253270
)
271+
return await asyncio.wait_for(embedding_call, self.embedding_timeout)
254272

255273
def set_cache(self, key, value, **kwargs):
256274
print_verbose(f"qdrant semantic-cache set_cache, kwargs: {kwargs}")

‎litellm/caching/redis_semantic_cache.py‎

Lines changed: 30 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import litellm
2020
from litellm._logging import print_verbose, verbose_logger
21+
from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
2122
from litellm.litellm_core_utils.prompt_templates.common_utils import (
2223
get_str_from_messages,
2324
)
@@ -27,6 +28,7 @@
2728
build_router_embedding_metadata,
2829
resolve_embedding_max_input_tokens,
2930
resolve_embedding_router,
31+
resolve_embedding_timeout,
3032
truncate_embedding_input,
3133
)
3234
from .base_cache import BaseCache
@@ -47,6 +49,7 @@ class RedisSemanticCache(BaseCache):
4749
DEFAULT_REDIS_INDEX_NAME: str = "litellm_semantic_cache_index"
4850
CACHE_KEY_FIELD_NAME: str = "litellm_cache_key"
4951
embedding_max_input_tokens: int | None = None
52+
embedding_timeout: float = SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
5053

5154
def __init__(
5255
self,
@@ -58,6 +61,7 @@ def __init__(
5861
embedding_model: str = "text-embedding-ada-002",
5962
index_name: str | None = None,
6063
embedding_max_input_tokens: int | None = None,
64+
embedding_timeout: float | None = None,
6165
**kwargs: object,
6266
):
6367
"""
@@ -74,6 +78,8 @@ def __init__(
7478
index_name: Name for the Redis index
7579
embedding_max_input_tokens: Truncate prompts to this many tokens before
7680
embedding; defaults to the Router deployment's configured max_input_tokens
81+
embedding_timeout: Seconds a cache lookup may spend embedding the prompt before it
82+
gives up and lets the request continue to the LLM
7783
ttl: Default time-to-live for cache entries in seconds
7884
**kwargs: Additional arguments passed to the Redis client
7985
@@ -99,6 +105,7 @@ def __init__(
99105
self.distance_threshold = 1 - similarity_threshold
100106
self.embedding_model = embedding_model
101107
self.embedding_max_input_tokens = embedding_max_input_tokens
108+
self.embedding_timeout = resolve_embedding_timeout(embedding_timeout)
102109

103110
# Set up Redis connection
104111
if redis_url is None:
@@ -349,6 +356,8 @@ def _get_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) ->
349356
input=embedding_input,
350357
cache={"no-store": True, "no-cache": True},
351358
metadata=build_router_embedding_metadata(metadata),
359+
timeout=self.embedding_timeout,
360+
num_retries=0,
352361
),
353362
)
354363
else:
@@ -358,6 +367,8 @@ def _get_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) ->
358367
model=self.embedding_model,
359368
input=embedding_input,
360369
cache={"no-store": True, "no-cache": True},
370+
timeout=self.embedding_timeout,
371+
num_retries=0,
361372
),
362373
)
363374
return embedding_response["data"][0]["embedding"]
@@ -512,20 +523,26 @@ async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | Non
512523

513524
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
514525
embedding_input: Final = self._embedding_input(prompt, router)
526+
embedding_call: Final = (
527+
router.aembedding(
528+
model=self.embedding_model,
529+
input=embedding_input,
530+
cache={"no-store": True, "no-cache": True},
531+
metadata=build_router_embedding_metadata(metadata),
532+
timeout=self.embedding_timeout,
533+
num_retries=0,
534+
)
535+
if router is not None
536+
else litellm.aembedding(
537+
model=self.embedding_model,
538+
input=embedding_input,
539+
cache={"no-store": True, "no-cache": True},
540+
timeout=self.embedding_timeout,
541+
num_retries=0,
542+
)
543+
)
515544
try:
516-
if router is not None:
517-
embedding_response = await router.aembedding(
518-
model=self.embedding_model,
519-
input=embedding_input,
520-
cache={"no-store": True, "no-cache": True},
521-
metadata=build_router_embedding_metadata(metadata),
522-
)
523-
else:
524-
embedding_response = await litellm.aembedding(
525-
model=self.embedding_model,
526-
input=embedding_input,
527-
cache={"no-store": True, "no-cache": True},
528-
)
545+
embedding_response: Final = await asyncio.wait_for(embedding_call, self.embedding_timeout)
529546
return embedding_response["data"][0]["embedding"]
530547
except Exception as e:
531548
print_verbose(f"Error generating async embedding: {e}")

‎litellm/caching/valkey_semantic_cache.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
from litellm._uuid import uuid
3131
from litellm.llms.valkey.common_utils import build_valkey_url, pack_vector
3232

33+
from ._embedding_router import resolve_embedding_timeout
3334
from .redis_semantic_cache import RedisSemanticCache
3435

3536

@@ -62,6 +63,7 @@ def __init__(
6263
sync_client: Redis | None = None,
6364
async_client: AsyncRedis | None = None,
6465
embedding_max_input_tokens: int | None = None,
66+
embedding_timeout: float | None = None,
6567
**kwargs: Any,
6668
):
6769
if similarity_threshold is None:
@@ -80,6 +82,7 @@ def __init__(
8082
self.similarity_threshold = similarity_threshold
8183
self.embedding_model = embedding_model
8284
self.embedding_max_input_tokens = embedding_max_input_tokens
85+
self.embedding_timeout = resolve_embedding_timeout(embedding_timeout)
8386
self.index_name = index_name or self.DEFAULT_VALKEY_INDEX_NAME
8487
self.key_prefix = f"{self.index_name}:"
8588
self._index_dim: int | None = None

‎litellm/constants.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -436,6 +436,9 @@
436436
# deadline and connect handshake (see ``http_handler`` cached handler paths).
437437
COMPLETION_HTTP_FALLBACK_SECONDS: Final[float] = 600.0
438438
HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS: Final[float] = 5.0
439+
SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS: Final[float] = float(
440+
os.getenv("SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS", "5.0")
441+
)
439442
request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", str(int(DEFAULT_REQUEST_TIMEOUT_SECONDS))))
440443
request_timeout_explicitly_set: bool = "REQUEST_TIMEOUT" in os.environ
441444
DEFAULT_A2A_AGENT_TIMEOUT: Final[float] = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes

‎litellm/main.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5974,7 +5974,7 @@ def embedding(
59745974
# Optional params
59755975
dimensions: int | None = None,
59765976
encoding_format: str | None = None,
5977-
timeout=600, # default to 10 minutes
5977+
timeout: float = 600, # default to 10 minutes
59785978
# set api_base, api_version, api_key
59795979
api_base: str | None = None,
59805980
api_version: str | None = None,
@@ -6000,7 +6000,7 @@ def embedding(
60006000
# Optional params
60016001
dimensions: int | None = None,
60026002
encoding_format: str | None = None,
6003-
timeout=600, # default to 10 minutes
6003+
timeout: float = 600, # default to 10 minutes
60046004
# set api_base, api_version, api_key
60056005
api_base: str | None = None,
60066006
api_version: str | None = None,
@@ -6027,7 +6027,7 @@ def embedding(
60276027
# Optional params
60286028
dimensions: int | None = None,
60296029
encoding_format: str | None = None,
6030-
timeout=600, # default to 10 minutes
6030+
timeout: float = 600, # default to 10 minutes
60316031
# set api_base, api_version, api_key
60326032
api_base: str | None = None,
60336033
api_version: str | None = None,

‎tests/test_litellm/caching/test_qdrant_semantic_cache.py‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -966,3 +966,67 @@ async def test_qdrant_async_embedding_explicit_limit_beats_deployment_limit(monk
966966

967967
sent_input = router.aembedding.call_args.kwargs["input"]
968968
assert _token_count("sem-embed", sent_input) == 3
969+
970+
971+
@pytest.mark.asyncio
972+
async def test_qdrant_async_embedding_call_is_bounded(monkeypatch):
973+
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
974+
975+
cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
976+
cache.embedding_model = "sem-embed"
977+
cache.embedding_max_input_tokens = None
978+
cache.embedding_timeout = 1.5
979+
980+
router = MagicMock()
981+
router.get_configured_token_limits.return_value = (None, None)
982+
router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]})
983+
monkeypatch.setitem(
984+
sys.modules,
985+
"litellm.proxy.proxy_server",
986+
_router_proxy_module(router, "sem-embed"),
987+
)
988+
989+
await cache._get_async_embedding("What is the capital of France?")
990+
991+
assert router.aembedding.call_args.kwargs["timeout"] == 1.5
992+
assert router.aembedding.call_args.kwargs["num_retries"] == 0
993+
994+
995+
@pytest.mark.asyncio
996+
async def test_qdrant_async_embedding_gives_up_on_unresponsive_endpoint(monkeypatch):
997+
import asyncio
998+
import time
999+
1000+
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
1001+
1002+
cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
1003+
cache.embedding_model = "sem-embed"
1004+
cache.embedding_max_input_tokens = None
1005+
cache.embedding_timeout = 0.05
1006+
1007+
async def never_responds(**kwargs):
1008+
await asyncio.sleep(3)
1009+
return {"data": [{"embedding": [0.1, 0.2]}]}
1010+
1011+
router = MagicMock()
1012+
router.get_configured_token_limits.return_value = (None, None)
1013+
router.aembedding = never_responds
1014+
monkeypatch.setitem(
1015+
sys.modules,
1016+
"litellm.proxy.proxy_server",
1017+
_router_proxy_module(router, "sem-embed"),
1018+
)
1019+
1020+
started = time.monotonic()
1021+
with pytest.raises(asyncio.TimeoutError):
1022+
await cache._get_async_embedding("What is the capital of France?")
1023+
assert time.monotonic() - started < 1.0
1024+
1025+
1026+
def test_qdrant_semantic_cache_defaults_embedding_timeout():
1027+
from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache
1028+
from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
1029+
1030+
cache = QdrantSemanticCache.__new__(QdrantSemanticCache)
1031+
assert cache.embedding_timeout == SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
1032+
assert SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS < 60

0 commit comments

Comments
 (0)