PR Feedback

This commit is contained in:
William Easton 2025-10-13 09:47:13 -05:00
commit 6ea6351d57
No known key found for this signature in database
2 changed files with 8 additions and 39 deletions

View file

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

View file

@ -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: