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")




class DiskCacheStoreConfig(BaseModel):
    """DiskCacheStore 的配置"""

    directory: str  # Path where cache is stored

    # Could add other diskcache.Cache parameters like size_limit, etc.




class DiskCacheStore(CacheStore[T], Component[DiskCacheStoreConfig]):
    """
    一个使用 diskcache 作为底层存储的类型化 CacheStore 实现。
    使用示例请参考 :class:`~autogen_ext.models.cache.ChatCompletionCache`。

    Args:
        cache_instance: diskcache.Cache 的实例。
                        用户需自行管理 DiskCache 实例的生命周期。
    """

    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



    def get(self, key: str, default: Optional[T] = None) -> Optional[T]:
        return cast(Optional[T], self.cache.get(key, default))  # type: ignore[reportUnknownMemberType]




    def set(self, key: str, value: T) -> None:
        self.cache.set(key, cast(Any, value))  # type: ignore[reportUnknownMemberType]




    def _to_config(self) -> DiskCacheStoreConfig:
        # Get directory from cache instance
        return DiskCacheStoreConfig(directory=self.cache.directory)




    @classmethod
    def _from_config(cls, config: DiskCacheStoreConfig) -> Self:
        return cls(cache_instance=diskcache.Cache(config.directory))  # type: ignore[no-any-return]