47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
|
from typing import Any, Optional, TypeVar, cast
|
||
|
|
|
||
|
|
import diskcache
|
||
|
|
from agentdhal_core import CacheStore, Component
|
||
|
|
from pydantic import BaseModel
|
||
|
|
from typing_extensions import Self
|
||
|
|
|
||
|
|
T = TypeVar("T")
|
||
|
|
|
||
|
|
|
||
|
|
class DiskCacheStoreConfig(BaseModel):
|
||
|
|
"""Configuration for DiskCacheStore"""
|
||
|
|
|
||
|
|
directory: str # Path where cache is stored
|
||
|
|
# Could add other diskcache.Cache parameters like size_limit, etc.
|
||
|
|
|
||
|
|
|
||
|
|
class DiskCacheStore(CacheStore[T], Component[DiskCacheStoreConfig]):
|
||
|
|
"""
|
||
|
|
A typed CacheStore implementation that uses diskcache as the underlying storage.
|
||
|
|
See :class:`~agentdhal_extensions.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 = "agentdhal_extensions.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]
|