mirror of https://github.com/microsoft/autogen.git
Add in memory cache class (#2266)
* Add in memory cache class * formatting
This commit is contained in:
parent
2053dd9f3d
commit
0644a2bc0c
|
@ -1,3 +1,4 @@
|
|||
from .cache import Cache, AbstractCache
|
||||
from .abstract_cache_base import AbstractCache
|
||||
from .cache import Cache
|
||||
|
||||
__all__ = ["Cache", "AbstractCache"]
|
||||
|
|
|
@ -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()
|
|
@ -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"
|
Loading…
Reference in New Issue