Add in memory cache class (#2266)

* Add in memory cache class

* formatting
This commit is contained in:
Jack Gerrits 2024-04-03 18:18:40 -04:00 committed by GitHub
parent 2053dd9f3d
commit 0644a2bc0c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 102 additions and 1 deletions

View File

@ -1,3 +1,4 @@
from .cache import Cache, AbstractCache
from .abstract_cache_base import AbstractCache
from .cache import Cache
__all__ = ["Cache", "AbstractCache"]

54
autogen/cache/in_memory_cache.py vendored Normal file
View File

@ -0,0 +1,54 @@
from types import TracebackType
from typing import Any, Dict, Optional, Type, Union
import sys
from .abstract_cache_base import AbstractCache
if sys.version_info >= (3, 11):
from typing import Self
else:
from typing_extensions import Self
class InMemoryCache(AbstractCache):
def __init__(self, seed: Union[str, int] = ""):
self._seed = str(seed)
self._cache: Dict[str, Any] = {}
def _prefixed_key(self, key: str) -> str:
separator = "_" if self._seed else ""
return f"{self._seed}{separator}{key}"
def get(self, key: str, default: Optional[Any] = None) -> Optional[Any]:
result = self._cache.get(self._prefixed_key(key))
if result is None:
return default
return result
def set(self, key: str, value: Any) -> None:
self._cache[self._prefixed_key(key)] = value
def close(self) -> None:
pass
def __enter__(self) -> Self:
"""
Enter the runtime context related to the object.
Returns:
self: The instance itself.
"""
return self
def __exit__(
self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
) -> None:
"""
Exit the runtime context related to the object.
Args:
exc_type: The exception type if an exception was raised in the context.
exc_value: The exception value if an exception was raised in the context.
traceback: The traceback if an exception was raised in the context.
"""
self.close()

46
test/cache/test_in_memory_cache.py vendored Normal file
View File

@ -0,0 +1,46 @@
from autogen.cache.in_memory_cache import InMemoryCache
def test_prefixed_key():
cache = InMemoryCache(seed="test")
assert cache._prefixed_key("key") == "test_key"
def test_get_with_default_value():
cache = InMemoryCache()
assert cache.get("key", "default_value") == "default_value"
def test_get_without_default_value():
cache = InMemoryCache()
assert cache.get("key") is None
def test_get_with_set_value():
cache = InMemoryCache()
cache.set("key", "value")
assert cache.get("key") == "value"
def test_get_with_set_value_and_seed():
cache = InMemoryCache(seed="test")
cache.set("key", "value")
assert cache.get("key") == "value"
def test_set():
cache = InMemoryCache()
cache.set("key", "value")
assert cache._cache["key"] == "value"
def test_set_with_seed():
cache = InMemoryCache(seed="test")
cache.set("key", "value")
assert cache._cache["test_key"] == "value"
def test_context_manager():
with InMemoryCache() as cache:
cache.set("key", "value")
assert cache.get("key") == "value"