codekingpro/portable-devtools

1"""Optional caching layer for language models.2 3Distinct from provider-based [prompt caching](https://docs.langchain.com/oss/python/langchain/models#prompt-caching).4 5!!! warning "Beta feature"6 7    This is a beta feature. Please be wary of deploying experimental code to production8    unless you've taken appropriate precautions.9 10A cache is useful for two reasons:11 121. It can save you money by reducing the number of API calls you make to the LLM13    provider if you're often requesting the same completion multiple times.142. It can speed up your application by reducing the number of API calls you make to the15    LLM provider.16"""17 18from __future__ import annotations19 20from abc import ABC, abstractmethod21from collections.abc import Sequence22from typing import Any23 24from typing_extensions import override25 26from langchain_core.outputs import Generation27from langchain_core.runnables import run_in_executor28 29RETURN_VAL_TYPE = Sequence[Generation]30 31 32class BaseCache(ABC):33    """Interface for a caching layer for LLMs and Chat models.34 35    The cache interface consists of the following methods:36 37    - lookup: Look up a value based on a prompt and `llm_string`.38    - update: Update the cache based on a prompt and `llm_string`.39    - clear: Clear the cache.40 41    In addition, the cache interface provides an async version of each method.42 43    The default implementation of the async methods is to run the synchronous44    method in an executor. It's recommended to override the async methods45    and provide async implementations to avoid unnecessary overhead.46    """47 48    @abstractmethod49    def lookup(self, prompt: str, llm_string: str) -> RETURN_VAL_TYPE | None:50        """Look up based on `prompt` and `llm_string`.51 52        A cache implementation is expected to generate a key from the 2-tuple53        of `prompt` and `llm_string` (e.g., by concatenating them with a delimiter).54 55        Args:56            prompt: A string representation of the prompt.57 58                In the case of a chat model, the prompt is a non-trivial59                serialization of the prompt into the language model.60            llm_string: A string representation of the LLM configuration.61 62                This is used to capture the invocation parameters of the LLM63                (e.g., model name, temperature, stop tokens, max tokens, etc.).64 65                These invocation parameters are serialized into a string representation.66 67        Returns:68            On a cache miss, return `None`. On a cache hit, return the cached value.69                The cached value is a list of `Generation` (or subclasses).70        """71 72    @abstractmethod73    def update(self, prompt: str, llm_string: str, return_val: RETURN_VAL_TYPE) -> None:74        """Update cache based on `prompt` and `llm_string`.75 76        The `prompt` and `llm_string` are used to generate a key for the cache. The key77        should match that of the lookup method.78 79        Args:80            prompt: A string representation of the prompt.81 82                In the case of a chat model, the prompt is a non-trivial83                serialization of the prompt into the language model.84            llm_string: A string representation of the LLM configuration.85 86                This is used to capture the invocation parameters of the LLM87                (e.g., model name, temperature, stop tokens, max tokens, etc.).88 89                These invocation parameters are serialized into a string90                representation.91            return_val: The value to be cached.92 93                The value is a list of `Generation` (or subclasses).94        """95 96    @abstractmethod97    def clear(self, **kwargs: Any) -> None:98        """Clear cache that can take additional keyword arguments."""99 100    async def alookup(self, prompt: str, llm_string: str) -> RETURN_VAL_TYPE | None:101        """Async look up based on `prompt` and `llm_string`.102 103        A cache implementation is expected to generate a key from the 2-tuple104        of `prompt` and `llm_string` (e.g., by concatenating them with a delimiter).105 106        Args:107            prompt: A string representation of the prompt.108 109                In the case of a chat model, the prompt is a non-trivial110                serialization of the prompt into the language model.111            llm_string: A string representation of the LLM configuration.112 113                This is used to capture the invocation parameters of the LLM114                (e.g., model name, temperature, stop tokens, max tokens, etc.).115 116                These invocation parameters are serialized into a string117                representation.118 119        Returns:120            On a cache miss, return `None`. On a cache hit, return the cached value.121                The cached value is a list of `Generation` (or subclasses).122        """123        return await run_in_executor(None, self.lookup, prompt, llm_string)124 125    async def aupdate(126        self, prompt: str, llm_string: str, return_val: RETURN_VAL_TYPE127    ) -> None:128        """Async update cache based on `prompt` and `llm_string`.129 130        The prompt and llm_string are used to generate a key for the cache.131        The key should match that of the look up method.132 133        Args:134            prompt: A string representation of the prompt.135 136                In the case of a chat model, the prompt is a non-trivial137                serialization of the prompt into the language model.138            llm_string: A string representation of the LLM configuration.139 140                This is used to capture the invocation parameters of the LLM141                (e.g., model name, temperature, stop tokens, max tokens, etc.).142 143                These invocation parameters are serialized into a string144                representation.145            return_val: The value to be cached. The value is a list of `Generation`146                (or subclasses).147        """148        return await run_in_executor(None, self.update, prompt, llm_string, return_val)149 150    async def aclear(self, **kwargs: Any) -> None:151        """Async clear cache that can take additional keyword arguments."""152        return await run_in_executor(None, self.clear, **kwargs)153 154 155class InMemoryCache(BaseCache):156    """Cache that stores things in memory.157 158    Example:159        ```python160        from langchain_core.caches import InMemoryCache161        from langchain_core.outputs import Generation162 163        # Initialize cache164        cache = InMemoryCache()165 166        # Update cache167        cache.update(168            prompt="What is the capital of France?",169            llm_string="model='gpt-5.4-mini',170            return_val=[Generation(text="Paris")],171        )172 173        # Lookup cache174        result = cache.lookup(175            prompt="What is the capital of France?",176            llm_string="model='gpt-5.4-mini',177        )178        # result is [Generation(text="Paris")]179        ```180    """181 182    def __init__(self, *, maxsize: int | None = None) -> None:183        """Initialize with empty cache.184 185        Args:186            maxsize: The maximum number of items to store in the cache.187 188                If `None`, the cache has no maximum size.189 190                If the cache exceeds the maximum size, the oldest items are removed.191 192        Raises:193            ValueError: If `maxsize` is less than or equal to `0`.194        """195        self._cache: dict[tuple[str, str], RETURN_VAL_TYPE] = {}196        if maxsize is not None and maxsize <= 0:197            msg = "maxsize must be greater than 0"198            raise ValueError(msg)199        self._maxsize = maxsize200 201    def lookup(self, prompt: str, llm_string: str) -> RETURN_VAL_TYPE | None:202        """Look up based on `prompt` and `llm_string`.203 204        Args:205            prompt: A string representation of the prompt.206 207                In the case of a chat model, the prompt is a non-trivial208                serialization of the prompt into the language model.209            llm_string: A string representation of the LLM configuration.210 211        Returns:212            On a cache miss, return `None`. On a cache hit, return the cached value.213        """214        return self._cache.get((prompt, llm_string), None)215 216    def update(self, prompt: str, llm_string: str, return_val: RETURN_VAL_TYPE) -> None:217        """Update cache based on `prompt` and `llm_string`.218 219        Args:220            prompt: A string representation of the prompt.221 222                In the case of a chat model, the prompt is a non-trivial223                serialization of the prompt into the language model.224            llm_string: A string representation of the LLM configuration.225            return_val: The value to be cached.226 227                The value is a list of `Generation` (or subclasses).228        """229        if self._maxsize is not None and len(self._cache) == self._maxsize:230            del self._cache[next(iter(self._cache))]231        self._cache[prompt, llm_string] = return_val232 233    @override234    def clear(self, **kwargs: Any) -> None:235        """Clear cache."""236        self._cache = {}237 238    async def alookup(self, prompt: str, llm_string: str) -> RETURN_VAL_TYPE | None:239        """Async look up based on `prompt` and `llm_string`.240 241        Args:242            prompt: A string representation of the prompt.243 244                In the case of a chat model, the prompt is a non-trivial245                serialization of the prompt into the language model.246            llm_string: A string representation of the LLM configuration.247 248        Returns:249            On a cache miss, return `None`. On a cache hit, return the cached value.250        """251        return self.lookup(prompt, llm_string)252 253    async def aupdate(254        self, prompt: str, llm_string: str, return_val: RETURN_VAL_TYPE255    ) -> None:256        """Async update cache based on `prompt` and `llm_string`.257 258        Args:259            prompt: A string representation of the prompt.260 261                In the case of a chat model, the prompt is a non-trivial262                serialization of the prompt into the language model.263            llm_string: A string representation of the LLM configuration.264            return_val: The value to be cached. The value is a list of `Generation`265                (or subclasses).266        """267        self.update(prompt, llm_string, return_val)268 269    @override270    async def aclear(self, **kwargs: Any) -> None:271        """Async clear cache."""272        self.clear()273