litellm/litellm/caching/qdrant_semantic_cache.py

"""
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 types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol, cast

from pydantic import TypeAdapter

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.asyncify import asyncify
from litellm.litellm_core_utils.prompt_templates.common_utils import (
    get_semantic_cache_prompt_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

_WAIT_FOR_INDEXING: Final = MappingProxyType({"wait": "true"})
_EMBEDDING_VECTOR: Final = TypeAdapter(object)

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_async_httpx_client,
            get_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 = await asyncify(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_semantic_cache_prompt_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_VECTOR.validate_python(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,
            params=_WAIT_FOR_INDEXING,
            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_semantic_cache_prompt_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_VECTOR.validate_python(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_semantic_cache_prompt_from_messages(messages)
        embedding_response: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))

        # get the embedding
        embedding: Final = _EMBEDDING_VECTOR.validate_python(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,
            params=_WAIT_FOR_INDEXING,
            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_semantic_cache_prompt_from_messages(messages)

        embedding_response: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))

        # get the embedding
        embedding: Final = _EMBEDDING_VECTOR.validate_python(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)