-
Notifications
You must be signed in to change notification settings - Fork 1.5k
feat: add schema dereference middleware #1641
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
c7188ef to
b5ec2be
Compare
ea7ea72 to
5858dbe
Compare
|
Hi @jlowin could you review this PR |
5858dbe to
75978f6
Compare
|
Since I have a similar use-case and have my own hack locally, I tested this PR. I have comments on the PR below. However, since I wanted slightly different behaviour, I also tried re-implementing it with First, here it is with jsonref -- it removes tools which fail dereferencing (to avoid LLM misuse) but logging a warning (i.e., fix the server): from jsonref import JsonRefError, replace_refs
class DereferenceSchemaMiddleware(Middleware):
def _dereference_tool(self, tool, remove_defs=True):
"""Dereference a tool's schema, returning the modified tool or None on failure.
Args:
tool: Tool object with parameters schema to dereference
remove_defs: If True, remove $defs from schema after dereferencing (default: True)
"""
try:
# proxies=False returns plain dicts to ensure compatibility with FastMCP
# lazy_load=False resolves immediately
dereferenced = replace_refs(tool.parameters, proxies=False, lazy_load=False)
# If requested, remove $defs (since all references have been resolved)
if remove_defs and isinstance(dereferenced, dict) and "$defs" in dereferenced:
dereferenced = {k: v for k, v in dereferenced.items() if k != "$defs"}
tool.parameters = dereferenced
return tool
except JsonRefError as e:
logger.warning(
f"Failed to dereference schema for tool '{tool.name}': {e}. "
"Tool will be hidden from client."
)
return None
except Exception as e:
logger.error(
f"Unexpected error dereferencing schema for tool '{tool.name}': {e}. "
"Tool will be hidden from client."
)
return None
async def on_list_tools(self, context: MiddlewareContext, call_next):
tools = await call_next(context)
return [
dereferenced_tool
for tool in tools
if (dereferenced_tool := self._dereference_tool(tool)) is not None
]Even more minimal if tolerating a lingering from jsonref import JsonRefError, replace_refs
class DereferenceSchemaMiddleware(Middleware):
def _dereference_tool(self, tool):
"""Dereference a tool's schema, returning the unmodified tool if dereferencing fails.
Args:
tool: Tool object with parameters schema to dereference
"""
try:
# proxies=False returns plain dicts to ensure compatibility with FastMCP
# lazy_load=False resolves immediately
tool.parameters = replace_refs(tool.parameters, proxies=False, lazy_load=False)
except JsonRefError:
pass
return tool
async def on_list_tools(self, context: MiddlewareContext, call_next):
tools = await call_next(context)
return [ self._dereference_tool(tool) for tool in tools ]Comments on the PR
The middleware is not conditional, so I tried both subclass and separate class importing the deref function; neither seem great options. Since the deref function is module-level, a subclass can't If not subclassing, the deref function can be imported, but this essentially side-stepping this middleware entirely, and just using its "internals" as a utility function. I think these could be helped by defining the deref function in the class, improving encapsulation and permitting a subclass to call it whenever needed. Both mean re-implementing the hooks, adding the condition. Conditionality could be provided by adding a callable to the constructor for simple cases, and also a template function (e.g.,
I'm curious about the thinking here. The PR leaves alone refs in the corner-cases. I think this means tools failing to deref all arguments will still be potentially misused by the LLM (if it uses those args) -- making the problem less frequent, but still present. Also, there is no logging when this occurs.
This is debatable, i.e., a general vs. pragmatic and sufficient solution. How were the specific hooks needing deref determined? I don't think the deref in So a simple solution, for now, only needs to deref in OTOH, a general approach would alter FastMCPs schema generation, and all future responses would be deref'd. (Tempting, to avoid MCP Schema interpretation.) So my thought is that these knobs and design-decisions seem unnecessary, given the simplicity of the jsonref implementation. 🤷 I have other comments about the code itself, which I can add later if going forward. |
Description
Contributors Checklist
Review Checklist
Follow up with #1192