ceos/fastmcp
1
0
Fork 0
mirror of https://github.com/PrefectHQ/fastmcp.git synced 2026-08-22 21:44:18 +02:00
Code Issues Projects Releases Packages Wiki Activity

small fix for caching

Browse source
This commit is contained in:
William Easton 2025-09-21 12:35:04 -05:00
commit e84f9446ae
No known key found for this signature in database
2 changed files with 17 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

11
src/fastmcp/contrib/middleware/caching/elasticsearch/elasticsearch_cache.py
View file

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

28
src/fastmcp/server/middleware/caching.py
View file

@ -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,
)