Source code for autogen_ext.cache_store.diskcache

from typing import Any, Optional, TypeVar, cast

import diskcache
from autogen_core import CacheStore, Component
from pydantic import BaseModel
from typing_extensions import Self

T = TypeVar("T")


[docs] class DiskCacheStoreConfig(BaseModel): """Configuration for DiskCacheStore""" directory: str # Path where cache is stored
# Could add other diskcache.Cache parameters like size_limit, etc.
[docs] class DiskCacheStore(CacheStore[T], Component[DiskCacheStoreConfig]): """ A typed CacheStore implementation that uses diskcache as the underlying storage. See :class:`~autogen_ext.models.cache.ChatCompletionCache` for an example of usage. Args: cache_instance: An instance of diskcache.Cache. The user is responsible for managing the DiskCache instance's lifetime. """ component_config_schema = DiskCacheStoreConfig component_provider_override = "autogen_ext.cache_store.diskcache.DiskCacheStore" def __init__(self, cache_instance: diskcache.Cache): # type: ignore[no-any-unimported] self.cache = cache_instance
[docs] def get(self, key: str, default: Optional[T] = None) -> Optional[T]: return cast(Optional[T], self.cache.get(key, default)) # type: ignore[reportUnknownMemberType]
[docs] def set(self, key: str, value: T) -> None: self.cache.set(key, cast(Any, value)) # type: ignore[reportUnknownMemberType]
[docs] def _to_config(self) -> DiskCacheStoreConfig: # Get directory from cache instance return DiskCacheStoreConfig(directory=self.cache.directory)
[docs] @classmethod def _from_config(cls, config: DiskCacheStoreConfig) -> Self: return cls(cache_instance=diskcache.Cache(config.directory)) # type: ignore[no-any-return]