diff --git a/src/fastmcp/contrib/middleware/caching/elasticsearch/elasticsearch_cache.py b/src/fastmcp/contrib/middleware/caching/elasticsearch/elasticsearch_cache.py index 4267acaf2..973b4fc00 100644 --- a/src/fastmcp/contrib/middleware/caching/elasticsearch/elasticsearch_cache.py +++ b/src/fastmcp/contrib/middleware/caching/elasticsearch/elasticsearch_cache.py @@ -31,18 +31,13 @@ DEFAULT_MAPPING = { "value": { "type": "keyword", "index": False, + "doc_values": False, + "ignore_above": 256, }, }, } -class SetEncoder(json.JSONEncoder): - def default(self, obj): - if isinstance(obj, set): - return list(obj) - return json.JSONEncoder.default(self, obj) - - class ElasticsearchCache(CacheProtocol): """A cache client that uses Elasticsearch.""" @@ -111,7 +106,7 @@ class ElasticsearchCache(CacheProtocol): collection=cache_entry.collection, key=cache_entry.key ) - document = json.loads(cache_entry.model_dump_json()) + document = json.loads(cache_entry.model_dump_json(serialize_as_any=True)) document["value"] = json.dumps(document["value"]) diff --git a/src/fastmcp/server/middleware/caching.py b/src/fastmcp/server/middleware/caching.py index d4a223a58..d2cd0b4e6 100644 --- a/src/fastmcp/server/middleware/caching.py +++ b/src/fastmcp/server/middleware/caching.py @@ -65,9 +65,7 @@ class BaseCacheEntry(BaseModel, ABC): ) created_at: datetime = Field(default_factory=lambda: datetime.now(tz=timezone.utc)) - value: CachableValueTypes - collection: str key: str ttl: int @@ -265,12 +263,12 @@ class CacheProtocol(Protocol): self, collection: str, key: str, - ) -> BaseCacheEntry | None: + ) -> CacheEntryTypes | None: """Get a cache entry from the cache.""" async def set_entry( self, - cache_entry: BaseCacheEntry, + cache_entry: CacheEntryTypes, ) -> None: """Set a value in the cache using the collection and key.""" @@ -307,7 +305,7 @@ class DiskCache(CacheProtocol): directory=path, size_limit=size_limit ) - async def get_entry(self, collection: str, key: str) -> BaseCacheEntry | None: + async def get_entry(self, collection: str, key: str) -> CacheEntryTypes | None: collection_key: str = self.make_collection_key(collection=collection, key=key) cache_entry = self._cache.get(key=collection_key) @@ -319,7 +317,7 @@ class DiskCache(CacheProtocol): async def set_entry( self, - cache_entry: BaseCacheEntry, + cache_entry: CacheEntryTypes, ) -> None: collection_key: str = self.make_collection_key( collection=cache_entry.collection, key=cache_entry.key @@ -356,20 +354,22 @@ class InMemoryCache(CacheProtocol): max_entries: The maximum number of entries to store in the cache. Defaults to 1000. """ self.max_entries = max_entries - self._cache = MemoryCacheClient( + self._cache: MemoryCacheClient[Any, CacheEntryTypes] = MemoryCacheClient[ + Any, CacheEntryTypes + ]( maxsize=max_entries, ttu=_memory_cache_ttu, getsizeof=_memory_cache_getsizeof, ) - async def get_entry(self, collection: str, key: str) -> BaseCacheEntry | None: - collection_key = self.make_collection_key(collection=collection, key=key) + async def get_entry(self, collection: str, key: str) -> CacheEntryTypes | None: + collection_key: str = self.make_collection_key(collection=collection, key=key) return self._cache.get(collection_key) async def set_entry( self, - cache_entry: BaseCacheEntry, + cache_entry: CacheEntryTypes, ) -> None: collection_key: str = self.make_collection_key( collection=cache_entry.collection, key=cache_entry.key @@ -483,7 +483,7 @@ class MethodSettings(TypedDict): MethodSettingsType = TypeVar("MethodSettingsType", bound=SharedMethodSettings) -MCP_METHOD_TO_CACHE_ENTRY_TYPE: dict[str, type[BaseCacheEntry]] = { +MCP_METHOD_TO_CACHE_ENTRY_TYPE: dict[str, type[CacheEntryTypes]] = { "tools/list": ListToolsCacheEntry, "tools/call": ToolResultCacheEntry, "resources/list": ListResourcesCacheEntry, @@ -746,10 +746,10 @@ class ResponseCachingMiddleware(Middleware): ttl: int = self._get_cache_ttl(context=context) - cache_entry: BaseCacheEntry = MCP_METHOD_TO_CACHE_ENTRY_TYPE[collection]( - collection=collection, + cache_entry: CacheEntryTypes = MCP_METHOD_TO_CACHE_ENTRY_TYPE[collection]( + collection=collection, # pyright: ignore[reportArgumentType] key=key, - value=value, + value=value, # pyright: ignore[reportArgumentType] ttl=ttl, )