mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-23 14:04:18 +02:00
PR Feedback
This commit is contained in:
parent
82dc3bdff9
commit
6ea6351d57
2 changed files with 8 additions and 39 deletions
|
|
@ -1,11 +1,11 @@
|
|||
"""A middleware for response caching."""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, TypedDict, TypeVar, cast
|
||||
|
||||
import mcp.types
|
||||
import pydantic_core
|
||||
from key_value.aio.adapters.pydantic import PydanticAdapter
|
||||
from key_value.aio.protocols.key_value import AsyncKeyValue
|
||||
from key_value.aio.stores.memory import MemoryStore
|
||||
|
|
@ -598,22 +598,19 @@ class ResponseCachingMiddleware(Middleware):
|
|||
def _make_call_tool_cache_key(msg: mcp.types.CallToolRequestParams) -> str:
|
||||
"""Make a cache key for a tool call by hashing the tool name and its arguments."""
|
||||
|
||||
raw = f"{msg.name}:{_get_arguments_str(msg.arguments)}"
|
||||
return hashlib.sha256(raw.encode("utf-8")).hexdigest()
|
||||
return f"{msg.name}:{_get_arguments_str(msg.arguments)}"
|
||||
|
||||
|
||||
def _make_read_resource_cache_key(msg: mcp.types.ReadResourceRequestParams) -> str:
|
||||
"""Make a cache key for a resource read by hashing the resource URI."""
|
||||
|
||||
raw = f"{msg.uri}"
|
||||
return hashlib.sha256(raw.encode("utf-8")).hexdigest()
|
||||
return f"{msg.uri}"
|
||||
|
||||
|
||||
def _make_get_prompt_cache_key(msg: mcp.types.GetPromptRequestParams) -> str:
|
||||
"""Make a cache key for a prompt get by hashing the prompt name and its arguments."""
|
||||
|
||||
raw = f"{msg.name}:{_get_arguments_str(msg.arguments)}"
|
||||
return hashlib.sha256(raw.encode("utf-8")).hexdigest()
|
||||
return f"{msg.name}:{_get_arguments_str(msg.arguments)}"
|
||||
|
||||
|
||||
def _get_arguments_str(arguments: dict[str, Any] | None) -> str:
|
||||
|
|
@ -623,7 +620,7 @@ def _get_arguments_str(arguments: dict[str, Any] | None) -> str:
|
|||
return "null"
|
||||
|
||||
try:
|
||||
return json.dumps(arguments, sort_keys=True, separators=(",", ":"))
|
||||
return pydantic_core.to_json(value=arguments, fallback=str).decode()
|
||||
|
||||
except TypeError:
|
||||
return repr(arguments)
|
||||
|
|
@ -651,21 +648,3 @@ def _get_size_of_tool_result(value: ToolResult) -> int:
|
|||
)
|
||||
|
||||
return content_size + structured_content_size
|
||||
|
||||
|
||||
# def get_size_of_one_value(value: BaseModel | ToolResult | ReadResourceContents) -> int:
|
||||
# """Get the size of an mcp type."""
|
||||
|
||||
# if isinstance(value, ToolResult):
|
||||
# return get_size_of_tool_result(value)
|
||||
# if isinstance(value, ReadResourceContents):
|
||||
# return len(value.content)
|
||||
# return len(value.model_dump_json())
|
||||
|
||||
|
||||
# def get_size_of_value(value: CachableListTypes) -> int:
|
||||
# """Get the size of a cache entry."""
|
||||
# if isinstance(value, (BaseModel | ToolResult | ReadResourceContents)):
|
||||
# return get_size_of_one_value(value)
|
||||
|
||||
# return sum(get_size_of_one_value(item) for item in value)
|
||||
|
|
|
|||
|
|
@ -367,11 +367,7 @@ class TestResponseCachingMiddleware:
|
|||
|
||||
key = _make_call_tool_cache_key(msg)
|
||||
|
||||
# Should be a SHA256 hash
|
||||
assert len(key) == 64
|
||||
assert key == snapshot(
|
||||
"7fa3d5c7967a202457eeca0731709fd87ec98546ddaee829ee86ca54b1858c59"
|
||||
)
|
||||
assert key == snapshot('test_tool:{"param1":"value1","param2":42}')
|
||||
|
||||
msg = mcp.types.ReadResourceRequestParams(
|
||||
uri=AnyUrl("https://test_uri"),
|
||||
|
|
@ -379,10 +375,7 @@ class TestResponseCachingMiddleware:
|
|||
|
||||
key = _make_read_resource_cache_key(msg)
|
||||
|
||||
assert len(key) == 64
|
||||
assert key == snapshot(
|
||||
"e34cc47c03ed1ad54f02501d95ecc463b65646568961c97ca4b730cb274e9d42"
|
||||
)
|
||||
assert key == snapshot("https://test_uri/")
|
||||
|
||||
msg = mcp.types.GetPromptRequestParams(
|
||||
name="test_prompt", arguments={"param1": "value1"}
|
||||
|
|
@ -390,10 +383,7 @@ class TestResponseCachingMiddleware:
|
|||
|
||||
key = _make_get_prompt_cache_key(msg)
|
||||
|
||||
assert len(key) == 64
|
||||
assert key == snapshot(
|
||||
"6306ff84fd3ff247a4bd91271e9d727d7f051bba53fb2e3bf80958988c4baf57"
|
||||
)
|
||||
assert key == snapshot('test_prompt:{"param1":"value1"}')
|
||||
|
||||
|
||||
class TestResponseCachingMiddlewareIntegration:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue