"""The cache interface. The cache interface is used to cache LLM results and embedding results. Maybe we can cache more server results in the future. """ from abc import ABC, abstractmethod from dataclasses import dataclass from enum import Enum from typing import Any, Generic, Optional, TypeVar from dbgpt.core.interface.serialization import Serializable K = TypeVar("K") V = TypeVar("V") class RetrievalPolicy(str, Enum): """The retrieval policy of the cache.""" EXACT_MATCH = "exact_match" SIMILARITY_MATCH = "similarity_match" class CachePolicy(str, Enum): """The cache policy of the cache.""" LRU = "lru" FIFO = "fifo" @dataclass class CacheConfig: """The cache config.""" retrieval_policy: Optional[RetrievalPolicy] = RetrievalPolicy.EXACT_MATCH cache_policy: Optional[CachePolicy] = CachePolicy.LRU class CacheKey(Serializable, ABC, Generic[K]): """The key of the cache. Must be hashable and comparable. Supported cache keys: - The LLM cache key: Include user prompt and the parameters to LLM. - The embedding model cache key: Include the texts to embedding and the parameters to embedding model. """ @abstractmethod def __hash__(self) -> int: """Return the hash value of the key.""" @abstractmethod def __eq__(self, other: Any) -> bool: """Check equality with another key.""" @abstractmethod def get_hash_bytes(self) -> bytes: """Return the byte array of hash value.""" @abstractmethod def get_value(self) -> K: """Get the underlying value of the cache key. Returns: K: The real object of current cache key """ class CacheValue(Serializable, ABC, Generic[V]): """Cache value abstract class.""" @abstractmethod def get_value(self) -> V: """Get the underlying real value.""" class CacheClient(ABC, Generic[K, V]): """The cache client interface.""" @abstractmethod async def get( self, key: CacheKey[K], cache_config: Optional[CacheConfig] = None ) -> Optional[CacheValue[V]]: """Retrieve a value from the cache using the provided key. Args: key (CacheKey[K]): The key to get cache cache_config (Optional[CacheConfig]): Cache config Returns: Optional[CacheValue[V]]: The value retrieved according to key. If cache key not exist, return None. """ @abstractmethod async def set( self, key: CacheKey[K], value: CacheValue[V], cache_config: Optional[CacheConfig] = None, ) -> None: """Set a value in the cache for the provided key. Args: key (CacheKey[K]): The key to set to cache value (CacheValue[V]): The value to set to cache cache_config (Optional[CacheConfig]): Cache config """ @abstractmethod async def exists( self, key: CacheKey[K], cache_config: Optional[CacheConfig] = None ) -> bool: """Check if a key exists in the cache. Args: key (CacheKey[K]): The key to set to cache cache_config (Optional[CacheConfig]): Cache config Return: bool: True if the key in the cache, otherwise is False """ @abstractmethod def new_key(self, **kwargs) -> CacheKey[K]: """Create a cache key with params.""" @abstractmethod def new_value(self, **kwargs) -> CacheValue[K]: """Create a cache key with params."""