mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-20 00:11:50 +00:00
215bf03373
Narrows reportAny / reportExplicitAny hot spots in provider transformations, proxy endpoints, integrations and secret managers by introducing TypedDicts, Protocols and object-typed boundaries instead of Any, then ratchets the budget ceilings down to match. reportAny 14765 -> 14076, reportExplicitAny 4493 -> 4128, ANN401 387 -> 307
504 lines
19 KiB
Python
504 lines
19 KiB
Python
"""
|
|
Qdrant Semantic Cache implementation
|
|
|
|
Has 4 methods:
|
|
- set_cache
|
|
- get_cache
|
|
- async_set_cache
|
|
- async_get_cache
|
|
"""
|
|
|
|
import ast
|
|
import asyncio
|
|
import json
|
|
import os
|
|
from typing import TYPE_CHECKING, Any, Final, Protocol, cast
|
|
|
|
import litellm
|
|
from litellm._logging import print_verbose
|
|
from litellm.constants import (
|
|
QDRANT_SCALAR_QUANTILE,
|
|
QDRANT_VECTOR_SIZE,
|
|
SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS,
|
|
)
|
|
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|
get_str_from_messages,
|
|
)
|
|
from litellm.types.utils import EmbeddingResponse
|
|
|
|
from ._embedding_router import (
|
|
build_router_embedding_metadata,
|
|
resolve_embedding_max_input_tokens,
|
|
resolve_embedding_router,
|
|
resolve_embedding_timeout,
|
|
truncate_embedding_input,
|
|
)
|
|
from .base_cache import BaseCache
|
|
|
|
if TYPE_CHECKING:
|
|
from litellm.router import Router
|
|
|
|
|
|
class _QdrantCollectionDetailsResponse(Protocol):
|
|
"""The qdrant `/collections/{name}` response, whose body is kept as an opaque JSON object."""
|
|
|
|
def json(self) -> dict[str, object]: ...
|
|
|
|
|
|
class QdrantSemanticCache(BaseCache):
|
|
CACHE_KEY_FIELD_NAME = "litellm_cache_key"
|
|
embedding_max_input_tokens: int | None = None
|
|
embedding_timeout: float = SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
|
|
|
|
def __init__(
|
|
self,
|
|
qdrant_api_base=None,
|
|
qdrant_api_key=None,
|
|
collection_name=None,
|
|
similarity_threshold=None,
|
|
quantization_config=None,
|
|
embedding_model="text-embedding-ada-002",
|
|
host_type=None,
|
|
vector_size=None,
|
|
embedding_max_input_tokens: int | None = None,
|
|
embedding_timeout: float | None = None,
|
|
):
|
|
from litellm.llms.custom_httpx.http_handler import (
|
|
_get_httpx_client,
|
|
get_async_httpx_client,
|
|
httpxSpecialProvider,
|
|
)
|
|
from litellm.secret_managers.main import get_secret_str
|
|
|
|
if collection_name is None:
|
|
raise Exception("collection_name must be provided, passed None")
|
|
|
|
self.collection_name = collection_name
|
|
print_verbose(f"qdrant semantic-cache initializing COLLECTION - {self.collection_name}")
|
|
|
|
if similarity_threshold is None:
|
|
raise Exception("similarity_threshold must be provided, passed None")
|
|
self.similarity_threshold = similarity_threshold
|
|
self.embedding_model = embedding_model
|
|
self.embedding_max_input_tokens = embedding_max_input_tokens
|
|
self.embedding_timeout = resolve_embedding_timeout(embedding_timeout)
|
|
self.vector_size = vector_size if vector_size is not None else QDRANT_VECTOR_SIZE
|
|
headers = {}
|
|
|
|
# check if defined as os.environ/ variable
|
|
if qdrant_api_base:
|
|
if isinstance(qdrant_api_base, str) and qdrant_api_base.startswith("os.environ/"):
|
|
qdrant_api_base = get_secret_str(qdrant_api_base)
|
|
if qdrant_api_key:
|
|
if isinstance(qdrant_api_key, str) and qdrant_api_key.startswith("os.environ/"):
|
|
qdrant_api_key = get_secret_str(qdrant_api_key)
|
|
|
|
qdrant_api_base = qdrant_api_base or os.getenv("QDRANT_URL") or os.getenv("QDRANT_API_BASE")
|
|
qdrant_api_key = qdrant_api_key or os.getenv("QDRANT_API_KEY")
|
|
headers = {"Content-Type": "application/json"}
|
|
if qdrant_api_key:
|
|
headers["api-key"] = qdrant_api_key
|
|
|
|
if qdrant_api_base is None:
|
|
raise ValueError("Qdrant url must be provided")
|
|
|
|
self.qdrant_api_base = qdrant_api_base
|
|
self.qdrant_api_key = qdrant_api_key
|
|
print_verbose(f"qdrant semantic-cache qdrant_api_base: {self.qdrant_api_base}")
|
|
|
|
self.headers = headers
|
|
|
|
self.sync_client = _get_httpx_client()
|
|
self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Caching)
|
|
|
|
if quantization_config is None:
|
|
print_verbose("Quantization config is not provided. Default binary quantization will be used.")
|
|
collection_exists: Final = self.sync_client.get(
|
|
url=f"{self.qdrant_api_base}/collections/{self.collection_name}/exists",
|
|
headers=self.headers,
|
|
)
|
|
if collection_exists.status_code != 200:
|
|
raise ValueError(f"Error from qdrant checking if /collections exist {collection_exists.text}")
|
|
|
|
if collection_exists.json()["result"]["exists"]:
|
|
collection_details: _QdrantCollectionDetailsResponse = self.sync_client.get(
|
|
url=f"{self.qdrant_api_base}/collections/{self.collection_name}",
|
|
headers=self.headers,
|
|
)
|
|
self.collection_info: dict[str, object] = collection_details.json()
|
|
print_verbose(f"Collection already exists.\nCollection details:{self.collection_info}")
|
|
self._ensure_cache_key_payload_index()
|
|
else:
|
|
quantization_params: dict[str, dict[str, object]]
|
|
if quantization_config is None or quantization_config == "binary":
|
|
quantization_params = {
|
|
"binary": {
|
|
"always_ram": False,
|
|
}
|
|
}
|
|
elif quantization_config == "scalar":
|
|
quantization_params = {
|
|
"scalar": {
|
|
"type": "int8",
|
|
"quantile": QDRANT_SCALAR_QUANTILE,
|
|
"always_ram": False,
|
|
}
|
|
}
|
|
elif quantization_config == "product":
|
|
quantization_params = {"product": {"compression": "x16", "always_ram": False}}
|
|
else:
|
|
raise Exception("Quantization config must be one of 'scalar', 'binary' or 'product'")
|
|
|
|
new_collection_status: Final = self.sync_client.put(
|
|
url=f"{self.qdrant_api_base}/collections/{self.collection_name}",
|
|
json={
|
|
"vectors": {"size": self.vector_size, "distance": "Cosine"},
|
|
"quantization_config": quantization_params,
|
|
},
|
|
headers=self.headers,
|
|
)
|
|
if new_collection_status.json()["result"]:
|
|
collection_details = self.sync_client.get(
|
|
url=f"{self.qdrant_api_base}/collections/{self.collection_name}",
|
|
headers=self.headers,
|
|
)
|
|
self.collection_info = collection_details.json()
|
|
print_verbose(f"New collection created.\nCollection details:{self.collection_info}")
|
|
self._ensure_cache_key_payload_index()
|
|
else:
|
|
raise Exception("Error while creating new collection")
|
|
|
|
def _get_cache_logic(self, cached_response: Any):
|
|
if cached_response is None:
|
|
return cached_response
|
|
try:
|
|
cached_response = json.loads(cached_response) # Convert string to dictionary
|
|
except Exception:
|
|
cached_response = ast.literal_eval(cached_response)
|
|
return cached_response
|
|
|
|
def _get_qdrant_cache_key_filter(self, key: str) -> dict:
|
|
return {
|
|
"must": [
|
|
{
|
|
"key": self.CACHE_KEY_FIELD_NAME,
|
|
"match": {"value": str(key)},
|
|
}
|
|
]
|
|
}
|
|
|
|
def _add_cache_key_filter_to_search_data(self, data: dict, key: str) -> None:
|
|
data["filter"] = self._get_qdrant_cache_key_filter(key)
|
|
|
|
def _ensure_cache_key_payload_index(self) -> None:
|
|
try:
|
|
response: Final = self.sync_client.put(
|
|
url=f"{self.qdrant_api_base}/collections/{self.collection_name}/index",
|
|
headers=self.headers,
|
|
json={
|
|
"field_name": self.CACHE_KEY_FIELD_NAME,
|
|
"field_schema": "keyword",
|
|
},
|
|
)
|
|
if response.status_code not in (200, 201):
|
|
print_verbose(f"Qdrant semantic-cache could not create cache-key payload index: {response.text}")
|
|
except Exception as exc:
|
|
print_verbose(f"Qdrant semantic-cache could not create cache-key payload index: {exc}")
|
|
|
|
def _payload_matches_cache_key(self, payload: dict, key: str) -> bool:
|
|
# Pre-isolation points stored only prompt + response with no cache-key
|
|
# payload field. Reassigning them to a caller's key would risk
|
|
# cross-scope hits, so they're treated as misses and re-populated on
|
|
# the next set_cache.
|
|
cached_key: Final = payload.get(self.CACHE_KEY_FIELD_NAME)
|
|
return cached_key is not None and str(cached_key) == str(key)
|
|
|
|
def _embedding_input(self, prompt: str, router: "Router | None") -> str:
|
|
return truncate_embedding_input(
|
|
prompt,
|
|
self.embedding_model,
|
|
resolve_embedding_max_input_tokens(self.embedding_max_input_tokens, self.embedding_model, router),
|
|
)
|
|
|
|
def _get_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> EmbeddingResponse:
|
|
"""Embed via the proxy Router when it serves the model, else direct."""
|
|
try:
|
|
from litellm.proxy.proxy_server import llm_model_list, llm_router
|
|
except ImportError:
|
|
llm_model_list = None
|
|
llm_router = None
|
|
|
|
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
|
|
embedding_input: Final = self._embedding_input(prompt, router)
|
|
if router is not None:
|
|
return router.embedding(
|
|
model=self.embedding_model,
|
|
input=embedding_input,
|
|
cache={"no-store": True, "no-cache": True},
|
|
metadata=build_router_embedding_metadata(metadata),
|
|
timeout=self.embedding_timeout,
|
|
num_retries=0,
|
|
)
|
|
return litellm.embedding(
|
|
model=self.embedding_model,
|
|
input=embedding_input,
|
|
cache={"no-store": True, "no-cache": True},
|
|
timeout=self.embedding_timeout,
|
|
num_retries=0,
|
|
)
|
|
|
|
async def _get_async_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> EmbeddingResponse:
|
|
try:
|
|
from litellm.proxy.proxy_server import llm_model_list, llm_router
|
|
except ImportError:
|
|
llm_model_list = None
|
|
llm_router = None
|
|
|
|
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
|
|
embedding_input: Final = self._embedding_input(prompt, router)
|
|
embedding_call: Final = (
|
|
router.aembedding(
|
|
model=self.embedding_model,
|
|
input=embedding_input,
|
|
cache={"no-store": True, "no-cache": True},
|
|
metadata=build_router_embedding_metadata(metadata),
|
|
timeout=self.embedding_timeout,
|
|
num_retries=0,
|
|
)
|
|
if router is not None
|
|
else litellm.aembedding(
|
|
model=self.embedding_model,
|
|
input=embedding_input,
|
|
cache={"no-store": True, "no-cache": True},
|
|
timeout=self.embedding_timeout,
|
|
num_retries=0,
|
|
)
|
|
)
|
|
return await asyncio.wait_for(embedding_call, self.embedding_timeout)
|
|
|
|
def set_cache(self, key, value, **kwargs):
|
|
print_verbose(f"qdrant semantic-cache set_cache, kwargs: {kwargs}")
|
|
from litellm._uuid import uuid
|
|
|
|
# get the prompt
|
|
messages: Final = kwargs["messages"]
|
|
prompt: Final = get_str_from_messages(messages)
|
|
|
|
# create an embedding for prompt
|
|
embedding_response: Final = cast(
|
|
EmbeddingResponse,
|
|
self._get_embedding(prompt, metadata=kwargs.get("metadata")),
|
|
)
|
|
|
|
# get the embedding
|
|
embedding: Final = embedding_response["data"][0]["embedding"]
|
|
|
|
value = str(value)
|
|
assert isinstance(value, str)
|
|
|
|
data: Final = {
|
|
"points": [
|
|
{
|
|
"id": str(uuid.uuid4()),
|
|
"vector": embedding,
|
|
"payload": {
|
|
self.CACHE_KEY_FIELD_NAME: str(key),
|
|
"text": prompt,
|
|
"response": value,
|
|
},
|
|
},
|
|
]
|
|
}
|
|
self.sync_client.put(
|
|
url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points",
|
|
headers=self.headers,
|
|
json=data,
|
|
)
|
|
|
|
def get_cache(self, key, **kwargs):
|
|
print_verbose(f"sync qdrant semantic-cache get_cache, kwargs: {kwargs}")
|
|
|
|
# get the messages
|
|
messages: Final = kwargs["messages"]
|
|
prompt: Final = get_str_from_messages(messages)
|
|
|
|
# convert to embedding
|
|
embedding_response: Final = cast(
|
|
EmbeddingResponse,
|
|
self._get_embedding(prompt, metadata=kwargs.get("metadata")),
|
|
)
|
|
|
|
# get the embedding
|
|
embedding: Final = embedding_response["data"][0]["embedding"]
|
|
|
|
data: Final = {
|
|
"vector": embedding,
|
|
"params": {
|
|
"quantization": {
|
|
"ignore": False,
|
|
"rescore": True,
|
|
"oversampling": 3.0,
|
|
}
|
|
},
|
|
"limit": 1,
|
|
"with_payload": True,
|
|
}
|
|
self._add_cache_key_filter_to_search_data(data=data, key=key)
|
|
|
|
search_response: Final = self.sync_client.post(
|
|
url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points/search",
|
|
headers=self.headers,
|
|
json=data,
|
|
)
|
|
results: Final = search_response.json()["result"]
|
|
|
|
if results is None:
|
|
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
|
return None
|
|
if isinstance(results, list):
|
|
if len(results) == 0:
|
|
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
|
return None
|
|
|
|
similarity: Final = results[0]["score"]
|
|
payload: Final = results[0]["payload"]
|
|
if not self._payload_matches_cache_key(payload=payload, key=key):
|
|
print_verbose("Qdrant semantic-cache hit did not match cache key scope")
|
|
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
|
return None
|
|
|
|
cached_prompt: Final = payload["text"]
|
|
|
|
# check similarity, if more than self.similarity_threshold, return results
|
|
print_verbose(
|
|
f"semantic cache: similarity threshold: {self.similarity_threshold}, similarity: {similarity}, prompt: {prompt}, closest_cached_prompt: {cached_prompt}"
|
|
)
|
|
|
|
# update kwargs["metadata"] with similarity, don't rewrite the original metadata
|
|
kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity
|
|
|
|
if similarity >= self.similarity_threshold:
|
|
# cache hit !
|
|
cached_value: Final = payload["response"]
|
|
print_verbose(
|
|
f"got a cache hit, similarity: {similarity}, Current prompt: {prompt}, cached_prompt: {cached_prompt}"
|
|
)
|
|
return self._get_cache_logic(cached_response=cached_value)
|
|
else:
|
|
# cache miss !
|
|
return None
|
|
|
|
async def async_set_cache(self, key, value, **kwargs):
|
|
from litellm._uuid import uuid
|
|
|
|
print_verbose(f"async qdrant semantic-cache set_cache, kwargs: {kwargs}")
|
|
|
|
# get the prompt
|
|
messages: Final = kwargs["messages"]
|
|
prompt: Final = get_str_from_messages(messages)
|
|
embedding_response: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))
|
|
|
|
# get the embedding
|
|
embedding: Final = embedding_response["data"][0]["embedding"]
|
|
|
|
value = str(value)
|
|
assert isinstance(value, str)
|
|
|
|
data: Final = {
|
|
"points": [
|
|
{
|
|
"id": str(uuid.uuid4()),
|
|
"vector": embedding,
|
|
"payload": {
|
|
self.CACHE_KEY_FIELD_NAME: str(key),
|
|
"text": prompt,
|
|
"response": value,
|
|
},
|
|
},
|
|
]
|
|
}
|
|
|
|
await self.async_client.put(
|
|
url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points",
|
|
headers=self.headers,
|
|
json=data,
|
|
)
|
|
|
|
async def async_get_cache(self, key, **kwargs):
|
|
print_verbose(f"async qdrant semantic-cache get_cache, kwargs: {kwargs}")
|
|
|
|
# get the messages
|
|
messages: Final = kwargs["messages"]
|
|
prompt: Final = get_str_from_messages(messages)
|
|
|
|
embedding_response: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))
|
|
|
|
# get the embedding
|
|
embedding: Final = embedding_response["data"][0]["embedding"]
|
|
|
|
data: Final = {
|
|
"vector": embedding,
|
|
"params": {
|
|
"quantization": {
|
|
"ignore": False,
|
|
"rescore": True,
|
|
"oversampling": 3.0,
|
|
}
|
|
},
|
|
"limit": 1,
|
|
"with_payload": True,
|
|
}
|
|
self._add_cache_key_filter_to_search_data(data=data, key=key)
|
|
|
|
search_response: Final = await self.async_client.post(
|
|
url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points/search",
|
|
headers=self.headers,
|
|
json=data,
|
|
)
|
|
|
|
results: Final = search_response.json()["result"]
|
|
|
|
if results is None:
|
|
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
|
return None
|
|
if isinstance(results, list):
|
|
if len(results) == 0:
|
|
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
|
return None
|
|
|
|
similarity: Final = results[0]["score"]
|
|
payload: Final = results[0]["payload"]
|
|
if not self._payload_matches_cache_key(payload=payload, key=key):
|
|
print_verbose("Qdrant semantic-cache hit did not match cache key scope")
|
|
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
|
return None
|
|
|
|
cached_prompt: Final = payload["text"]
|
|
|
|
# check similarity, if more than self.similarity_threshold, return results
|
|
print_verbose(
|
|
f"semantic cache: similarity threshold: {self.similarity_threshold}, similarity: {similarity}, prompt: {prompt}, closest_cached_prompt: {cached_prompt}"
|
|
)
|
|
|
|
# update kwargs["metadata"] with similarity, don't rewrite the original metadata
|
|
kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity
|
|
|
|
if similarity >= self.similarity_threshold:
|
|
# cache hit !
|
|
cached_value: Final = payload["response"]
|
|
print_verbose(
|
|
f"got a cache hit, similarity: {similarity}, Current prompt: {prompt}, cached_prompt: {cached_prompt}"
|
|
)
|
|
return self._get_cache_logic(cached_response=cached_value)
|
|
else:
|
|
# cache miss !
|
|
return None
|
|
|
|
async def _collection_info(self):
|
|
return self.collection_info
|
|
|
|
async def async_set_cache_pipeline(self, cache_list, **kwargs):
|
|
tasks: Final = []
|
|
for val in cache_list:
|
|
tasks.append(self.async_set_cache(val[0], val[1], **kwargs))
|
|
await asyncio.gather(*tasks)
|