core[minor]: Add maxsize for InMemoryCache (#23405)

This PR introduces a maxsize parameter for the InMemoryCache class,
allowing users to specify the maximum number of items to store in the
cache. If the cache exceeds the specified maximum size, the oldest items
are removed. Additionally, comprehensive unit tests have been added to
ensure all functionalities are thoroughly tested. The tests are written
using pytest and cover both synchronous and asynchronous methods.

Twitter: @spyrosavl

---------

Co-authored-by: Eugene Yurtsev <eyurtsev@gmail.com>
This commit is contained in:
Spyros Avlonitis 2024-07-01 21:21:21 +03:00 committed by GitHub
parent 96af8f31ae
commit 8cfb2fa1b7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 134 additions and 2 deletions

View File

@ -145,9 +145,18 @@ class BaseCache(ABC):
class InMemoryCache(BaseCache):
"""Cache that stores things in memory."""
def __init__(self) -> None:
"""Initialize with empty cache."""
def __init__(self, *, maxsize: Optional[int] = None) -> None:
"""Initialize with empty cache.
Args:
maxsize: The maximum number of items to store in the cache.
If None, the cache has no maximum size.
If the cache exceeds the maximum size, the oldest items are removed.
"""
self._cache: Dict[Tuple[str, str], RETURN_VAL_TYPE] = {}
if maxsize is not None and maxsize <= 0:
raise ValueError("maxsize must be greater than 0")
self._maxsize = maxsize
def lookup(self, prompt: str, llm_string: str) -> Optional[RETURN_VAL_TYPE]:
"""Look up based on prompt and llm_string.
@ -174,6 +183,8 @@ class InMemoryCache(BaseCache):
return_val: The value to be cached. The value is a list of Generations
(or subclasses).
"""
if self._maxsize is not None and len(self._cache) == self._maxsize:
del self._cache[next(iter(self._cache))]
self._cache[(prompt, llm_string)] = return_val
def clear(self, **kwargs: Any) -> None:

View File

@ -0,0 +1,121 @@
from typing import Tuple
import pytest
from langchain_core.caches import RETURN_VAL_TYPE, InMemoryCache
from langchain_core.outputs import Generation
@pytest.fixture
def cache() -> InMemoryCache:
"""Fixture to provide an instance of InMemoryCache."""
return InMemoryCache()
def cache_item(item_id: int) -> Tuple[str, str, RETURN_VAL_TYPE]:
"""Generate a valid cache item."""
prompt = f"prompt{item_id}"
llm_string = f"llm_string{item_id}"
generations = [Generation(text=f"text{item_id}")]
return prompt, llm_string, generations
def test_initialization() -> None:
"""Test the initialization of InMemoryCache."""
cache = InMemoryCache()
assert cache._cache == {}
assert cache._maxsize is None
cache_with_maxsize = InMemoryCache(maxsize=2)
assert cache_with_maxsize._cache == {}
assert cache_with_maxsize._maxsize == 2
with pytest.raises(ValueError):
InMemoryCache(maxsize=0)
def test_lookup(
cache: InMemoryCache,
) -> None:
"""Test the lookup method of InMemoryCache."""
prompt, llm_string, generations = cache_item(1)
cache.update(prompt, llm_string, generations)
assert cache.lookup(prompt, llm_string) == generations
assert cache.lookup("prompt2", "llm_string2") is None
def test_update_with_no_maxsize(cache: InMemoryCache) -> None:
"""Test the update method of InMemoryCache with no maximum size."""
prompt, llm_string, generations = cache_item(1)
cache.update(prompt, llm_string, generations)
assert cache.lookup(prompt, llm_string) == generations
def test_update_with_maxsize() -> None:
"""Test the update method of InMemoryCache with a maximum size."""
cache = InMemoryCache(maxsize=2)
prompt1, llm_string1, generations1 = cache_item(1)
cache.update(prompt1, llm_string1, generations1)
assert cache.lookup(prompt1, llm_string1) == generations1
prompt2, llm_string2, generations2 = cache_item(2)
cache.update(prompt2, llm_string2, generations2)
assert cache.lookup(prompt2, llm_string2) == generations2
prompt3, llm_string3, generations3 = cache_item(3)
cache.update(prompt3, llm_string3, generations3)
assert cache.lookup(prompt1, llm_string1) is None # 'prompt1' should be evicted
assert cache.lookup(prompt2, llm_string2) == generations2
assert cache.lookup(prompt3, llm_string3) == generations3
def test_clear(cache: InMemoryCache) -> None:
"""Test the clear method of InMemoryCache."""
prompt, llm_string, generations = cache_item(1)
cache.update(prompt, llm_string, generations)
cache.clear()
assert cache.lookup(prompt, llm_string) is None
async def test_alookup(cache: InMemoryCache) -> None:
"""Test the asynchronous lookup method of InMemoryCache."""
prompt, llm_string, generations = cache_item(1)
await cache.aupdate(prompt, llm_string, generations)
assert await cache.alookup(prompt, llm_string) == generations
assert await cache.alookup("prompt2", "llm_string2") is None
async def test_aupdate_with_no_maxsize(cache: InMemoryCache) -> None:
"""Test the asynchronous update method of InMemoryCache with no maximum size."""
prompt, llm_string, generations = cache_item(1)
await cache.aupdate(prompt, llm_string, generations)
assert await cache.alookup(prompt, llm_string) == generations
async def test_aupdate_with_maxsize() -> None:
"""Test the asynchronous update method of InMemoryCache with a maximum size."""
cache = InMemoryCache(maxsize=2)
prompt, llm_string, generations = cache_item(1)
await cache.aupdate(prompt, llm_string, generations)
assert await cache.alookup(prompt, llm_string) == generations
prompt2, llm_string2, generations2 = cache_item(2)
await cache.aupdate(prompt2, llm_string2, generations2)
assert await cache.alookup(prompt2, llm_string2) == generations2
prompt3, llm_string3, generations3 = cache_item(3)
await cache.aupdate(prompt3, llm_string3, generations3)
assert await cache.alookup(prompt, llm_string) is None
assert await cache.alookup(prompt2, llm_string2) == generations2
assert await cache.alookup(prompt3, llm_string3) == generations3
async def test_aclear(cache: InMemoryCache) -> None:
"""Test the asynchronous clear method of InMemoryCache."""
prompt, llm_string, generations = cache_item(1)
await cache.aupdate(prompt, llm_string, generations)
await cache.aclear()
assert await cache.alookup(prompt, llm_string) is None