Decompose CodeMode into composable discovery tools (#3354)

* Decompose CodeMode into composable discovery tools

CodeMode now owns only the execute sandbox. Discovery tools (search,
get_schema, etc.) are composable via the discovery_tools parameter.
Each is a Callable[[GetToolCatalog], Tool] factory.

Ships SearchTool (lightweight name+description results) and SchemaTool
(brief markdown or full JSON schemas by tool name) as built-in defaults.

* Rename to Search/GetSchemas/Tags, add tag filtering, fix bugs

- Rename SearchTool→Search, SchemaTool→GetSchemas, Categories→Tags
- Add tags parameter to Search for LLM-side tag filtering
- Add Tags discovery tool for browsing tools by tag
- Fix shared singleton default factories (now per-instance)
- Fix get_schema full mode returning invalid JSON on partial matches
- Fix "untagged" filter inconsistency between Tags and Search
- Split serialization tests to comply with loq line limit
- Rewrite docs for conceptual clarity

* Add three-tier detail levels, remove default_arguments, rename Tags→GetTags, rewrite docs

* Clean up __all__ exports, return valid JSON for empty full-detail results
This commit is contained in:
Jeremiah Lowin 2026-03-02 16:35:55 -05:00 committed by GitHub
commit 59da3e4ed4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 1316 additions and 675 deletions

View file

@ -163,10 +163,10 @@
"servers/transforms/namespace",
"servers/transforms/tool-transformation",
"servers/visibility",
"servers/transforms/code-mode",
"servers/transforms/tool-search",
"servers/transforms/resources-as-tools",
"servers/transforms/prompts-as-tools",
"servers/transforms/code-mode"
"servers/transforms/prompts-as-tools"
]
},
{

View file

@ -1,26 +1,32 @@
---
title: Code Mode (Experimental)
title: Code Mode
sidebarTitle: Code Mode
description: Let LLMs write Python to orchestrate tools in a sandbox
icon: flask
tag: EXPERIMENTAL
tag: NEW
---
import { VersionBadge } from '/snippets/version-badge.mdx'
<VersionBadge version="3.1.0" />
<Warning>
CodeMode is experimental. The core interface is stable, but the specific discovery tools and their parameters may evolve as we learn more about what works best in practice.
</Warning>
Standard MCP tool usage has two scaling problems. First, every tool in the catalog is loaded into the LLM's context upfront — with hundreds of tools, that's tens of thousands of tokens spent before the LLM even reads the user's request. Second, every tool call is a round-trip: the LLM calls a tool, the result passes back through the context window, the LLM reasons about it, calls another tool, and so on. Intermediate results that only exist to feed the next step still burn tokens flowing through the model.
`CodeMode` solves both problems by replacing the tool catalog with two meta-tools — `search` for discovering tools by keyword, and `execute` for running Python scripts that chain tool calls in a sandbox. The LLM discovers what it needs, writes a script, and gets back only the final answer. One round-trip instead of ten; tool definitions loaded only when needed instead of all at once.
CodeMode solves both problems. Instead of seeing your entire tool catalog, the LLM gets meta-tools for discovering what's available and for writing and executing code that calls the tools it needs. It discovers on demand, writes a script that chains tool calls in a sandbox, and gets back only the final answer.
The approach was introduced by Cloudflare in [Code Mode](https://blog.cloudflare.com/code-mode/) and explored further by Anthropic in [Code Execution with MCP](https://www.anthropic.com/engineering/code-execution-with-mcp).
<Note>
CodeMode requires a sandbox to execute LLM-generated code safely. The default sandbox uses [pydantic-monty](https://github.com/pydantic/pydantic-monty), installed via `pip install "fastmcp[code-mode]"`. You can also provide your own sandbox — see [Custom Sandbox Providers](#custom-sandbox-providers).
</Note>
## Getting Started
## Basic Usage
<Tip>
CodeMode requires the `code-mode` extra for sandbox support. Install it with `pip install "fastmcp[code-mode]"`.
</Tip>
You take a normal server with normally registered tools and add a `CodeMode` transform. The transform wraps your existing tools in the code mode machinery — your tool functions don't change at all:
```python
from fastmcp import FastMCP
@ -39,157 +45,225 @@ def multiply(x: int, y: int) -> int:
return x * y
```
Clients now see only two tools. The LLM discovers the real tools through `search`, then orchestrates them through `execute`:
Clients connecting to this server no longer see `add` and `multiply` directly. Instead, they see the meta-tools that CodeMode provides — tools for discovering what's available and executing code against it. The original tools are still there, but they're accessed through the CodeMode layer.
```python
# Discover tools with a keyword search
result = await client.call_tool("search", {"query": "add multiply numbers"})
# [{"name": "add", "inputSchema": {...}, ...}, {"name": "multiply", ...}]
## Discovery
# Chain tool calls in a single round-trip
result = await client.call_tool("execute", {
"code": """
a = await call_tool("add", {"x": 3, "y": 4})
b = await call_tool("multiply", {"x": a["result"], "y": 2})
return b
"""
})
# {"result": 14}
```
Before the LLM can write code that calls your tools, it needs to know what tools exist and how to call them. This is the **discovery** process — the LLM uses meta-tools to learn about your tool catalog, then writes code against what it finds.
## Search
The fundamental tradeoff is **tokens vs. round-trips**. Each discovery step is an LLM round-trip: the model calls a tool, waits for the response, reasons about it, then decides what to do next. More steps mean less wasted context (each step is targeted) but more latency and API calls. Fewer steps mean the LLM gets information upfront but pays for detail it might not need.
The `search` meta-tool takes a `query` string and returns matching tools ranked by relevance, including their full `inputSchema` and `outputSchema`. The LLM uses these schemas to understand how to call each tool and what to expect back.
By default, CodeMode gives the LLM three tools — `search`, `get_schema`, and `execute` — creating a three-stage discovery flow:
Search uses BM25 ranking by default, matching against tool names and descriptions. You can swap in any [search transform](/servers/transforms/tool-search):
<Steps>
<Step title="Search for tools">
First, the LLM uses the `search` meta-tool to find tools by keyword.
```python
from fastmcp.server.transforms.search import RegexSearchTransform
mcp = FastMCP(
"Server",
transforms=[CodeMode(search_transform=RegexSearchTransform())],
)
```
### Search Result Format
By default, `search` returns tool definitions in the same JSON format as `list_tools` — full `inputSchema`, `outputSchema`, and metadata. This is complete, but verbose: each tool definition carries significant structural overhead that the LLM doesn't need just to decide which tool to call next.
`serialize_tools_for_output_markdown` is a built-in alternative that renders the same information as compact markdown. In benchmarks across typical tool catalogs it reduces search result tokens by ~65-70%.
```python
from fastmcp import FastMCP
from fastmcp.experimental.transforms import CodeMode
from fastmcp.server.transforms.search import serialize_tools_for_output_markdown
mcp = FastMCP(
"Server",
transforms=[CodeMode(search_result_serializer=serialize_tools_for_output_markdown)],
)
```
For a tool like `create_document`, the output looks like:
For example, it might do `search(query="math numbers")` and receive the following response:
```
### create_document
- add: Add two numbers.
- multiply: Multiply two numbers.
```
Create a new document in the workspace with the given metadata.
This lets the LLM know which tools are available and what they do, significantly reducing the surface area it needs to consider.
</Step>
<Step title="Get parameter details for the tools">
Next, the LLM calls `get_schema` to get parameter details for the tools it found in the previous step.
For example, it might do `get_schema(tools=["add", "multiply"])` and receive the following response:
```
### add
Add two numbers.
**Parameters**
- `title` (string, required)
- `content` (string, required)
- `tags` (string[], required)
- `author` (string, required)
- `published` (boolean)
- `parent_id` (string?)
- `x` (integer, required)
- `y` (integer, required)
### multiply
Multiply two numbers.
**Parameters**
- `x` (integer, required)
- `y` (integer, required)
```
You can also supply any callable that takes a sequence of tools and returns a string or a list of dicts. Async callables are supported too:
Now the LLM knows the parameters for the tools it found, and can write code that chains the tool calls. If it needed more detail, it could have called `get_schema` with `detail="full"` to get the complete JSON schema.
</Step>
<Step title="Write and execute code that chains the tool calls">
Finally, the LLM writes and executes code that chains the tool calls in a Python sandbox. Inside the sandbox, `call_tool(name, params)` is the only function available. The LLM uses this to compose tools into a workflow and return a final result.
For example, it might write the following code and call the `execute` tool with it:
```python
from fastmcp.server.transforms.search import SearchResultSerializer
def my_serializer(tools) -> str:
return "\n".join(f"{t.name}: {t.description}" for t in tools)
mcp = FastMCP(
"Server",
transforms=[CodeMode(search_result_serializer=my_serializer)],
)
a = await call_tool("add", {"x": 3, "y": 4})
b = await call_tool("multiply", {"x": a, "y": 2})
return b
```
The default JSON format is unchanged when `search_result_serializer` is not set — this is purely opt-in.
The result is returned to the LLM.
</Step>
</Steps>
## Execute
This three-stage flow works well for most servers — each step pulls in only the information needed for the next one, keeping context usage minimal. But CodeMode's discovery surface is fully configurable. The sections below explain each built-in discovery tool and how to combine them into different patterns.
The `execute` meta-tool takes a `code` string containing async Python. Inside the sandbox, one function is available:
## Discovery Tools
```python
await call_tool(tool_name, params) # -> dict | str
CodeMode ships with three built-in discovery tools: `Search`, `GetSchemas`, and `GetTags`. By default, only `Search` and `GetSchemas` are enabled. Each tool supports a `default_detail` parameter that sets the default verbosity level, and the LLM can override the detail level on any individual call.
### Detail Levels
`Search` and `GetSchemas` share the same three detail levels, so the same `detail` value produces the same output format regardless of which tool the LLM calls:
| Level | Output | Token cost |
|---|---|---|
| `"brief"` | Tool names and one-line descriptions | Cheapest — good for scanning |
| `"detailed"` | Compact markdown with parameter names, types, and required markers | Medium — often enough to write code |
| `"full"` | Complete JSON schema | Most expensive — everything |
`Search` defaults to `"brief"` and `GetSchemas` defaults to `"detailed"`.
### Search
`Search` finds tools by natural-language query using BM25 ranking. At its default `"brief"` detail, results include just tool names and descriptions — enough to decide which tools are worth inspecting further. The LLM can request `"detailed"` to get parameter schemas inline, or `"full"` for the complete JSON.
If your tools use [tags](/servers/tools#tags), Search also accepts a `tags` parameter so the LLM can narrow results to specific categories before searching.
### GetSchemas
`GetSchemas` returns parameter details for specific tools by name. At its default `"detailed"` level, it renders compact markdown with parameter names, types, and required markers. At `"full"`, it returns the complete JSON schema — useful when tools have deeply nested parameters that the compact format doesn't capture.
### GetTags
`GetTags` lets the LLM browse tools by category using [tag](/servers/tools#tags) metadata. At brief detail, the LLM sees tag names with counts. At full detail, it sees tools listed under each tag:
```
- math (3 tools)
- text (2 tools)
- untagged (1 tool)
```
The return type depends on whether the tool declares an output schema. When it does, `call_tool` returns the structured content dict exactly as the schema describes — for example, `{"result": 42}` for a tool returning `int`. When there's no output schema, `call_tool` returns the text content as a string. The LLM can tell which to expect from the `outputSchema` field in search results.
`GetTags` isn't included in the defaults — add it when browsing by category would help the LLM orient itself in a large catalog. The LLM can browse tags first, then pass specific tags into Search to narrow results.
Use `return` to produce the final output from the script.
## Discovery Patterns
### Default Arguments
The right discovery configuration depends on your server — how many tools you have and how complex their parameters are. It may be tempting to minimize round-trips by collapsing everything into fewer steps, but for the complex servers that benefit most from CodeMode, our experience is that staged discovery leads to better results. Flooding the LLM with detailed schemas for tools it doesn't end up using can hurt more than the extra round-trip costs. Each pattern below is a complete, copyable configuration.
When tools share common parameters (like workspace IDs or API keys), `default_arguments` injects them automatically:
### Three-Stage
The default. The LLM searches for candidates, inspects schemas for the ones it wants, then writes code. Best for **large or complex tool sets** where you want to minimize context usage — the LLM only pays for schemas it actually needs.
```python
mcp = FastMCP(
"Server",
transforms=[CodeMode(default_arguments={"workspace_id": "ws-123"})],
)
```
Defaults are only injected when the tool actually accepts the parameter and the LLM hasn't provided it explicitly. Parameters that a tool doesn't accept are silently skipped.
## OpenAPI Integration
`CodeMode` pairs naturally with OpenAPI-backed providers, where a single API spec can expose hundreds of endpoints as tools:
```python
import httpx
from fastmcp import FastMCP
from fastmcp.experimental.transforms import CodeMode
from fastmcp.server.providers.openapi import OpenAPIProvider
openapi_spec = httpx.get("https://api.example.com/openapi.json").json()
api_client = httpx.AsyncClient(base_url="https://api.example.com")
provider = OpenAPIProvider(
openapi_spec=openapi_spec,
client=api_client,
)
mcp = FastMCP("API Code Mode", providers=[provider], transforms=[CodeMode()])
mcp = FastMCP("Server", transforms=[CodeMode()])
```
## Configuration
### Custom Tool Names
The default `search` and `execute` names can be changed:
If your tools use [tags](/servers/tools#tags), add `GetTags` so the LLM can browse by category before searching — giving it four stages of progressive disclosure:
```python
mcp = FastMCP(
"Server",
transforms=[
CodeMode(
search_tool_name="find_tools",
execute_tool_name="run_workflow",
execute_description="Run multi-step API workflows",
)
],
from fastmcp import FastMCP
from fastmcp.experimental.transforms import CodeMode
from fastmcp.experimental.transforms.code_mode import GetTags, Search, GetSchemas
code_mode = CodeMode(
discovery_tools=[GetTags(), Search(), GetSchemas()],
)
mcp = FastMCP("Server", transforms=[code_mode])
```
### Two-Stage
Search returns parameter schemas inline, so the LLM can go straight from search to execute. Best for **smaller catalogs** where the extra tokens per search result are a reasonable price for one fewer round-trip.
```python
from fastmcp import FastMCP
from fastmcp.experimental.transforms import CodeMode
from fastmcp.experimental.transforms.code_mode import Search, GetSchemas
code_mode = CodeMode(
discovery_tools=[Search(default_detail="detailed"), GetSchemas()],
)
mcp = FastMCP("Server", transforms=[code_mode])
```
`GetSchemas` is still available as a fallback — the LLM can call it with `detail="full"` if it encounters a tool with complex nested parameters where the compact markdown isn't enough.
### Single-Stage
Skip discovery entirely and bake tool instructions into the execute tool's description. Best for **very simple servers** where the LLM already knows what tools are available — maybe there are only a few, or they're described in the system prompt.
```python
from fastmcp import FastMCP
from fastmcp.experimental.transforms import CodeMode
code_mode = CodeMode(
discovery_tools=[],
execute_description=(
"Available tools:\n"
"- add(x: int, y: int) -> int: Add two numbers\n"
"- multiply(x: int, y: int) -> int: Multiply two numbers\n\n"
"Write Python using `await call_tool(name, params)` and `return` the result."
),
)
mcp = FastMCP("Server", transforms=[code_mode])
```
## Custom Discovery Tools
Discovery tools are composable — you can mix the built-ins with your own. Each discovery tool is a callable that receives catalog access and returns a `Tool`. The catalog accessor is a function (not the catalog itself) because the catalog is request-scoped — different users may see different tools based on auth.
Here's a minimal example:
```python
from fastmcp.experimental.transforms import CodeMode
from fastmcp.experimental.transforms.code_mode import GetToolCatalog, GetSchemas
from fastmcp.server.context import Context
from fastmcp.tools.tool import Tool
def list_all_tools(get_catalog: GetToolCatalog) -> Tool:
async def list_tools(ctx: Context) -> str:
"""List all available tool names."""
tools = await get_catalog(ctx)
return ", ".join(t.name for t in tools)
return Tool.from_function(fn=list_tools, name="list_tools")
code_mode = CodeMode(discovery_tools=[list_all_tools, GetSchemas()])
```
The LLM sees the docstring of each discovery tool's inner function as its description — that's how it learns what each tool does and when to use it. Write docstrings that explain what the tool returns and when the LLM should call it.
Discovery tools and the execute tool can also have custom names:
```python
from fastmcp.experimental.transforms.code_mode import Search, GetSchemas
code_mode = CodeMode(
discovery_tools=[
Search(name="find_tools"),
GetSchemas(name="describe"),
],
execute_tool_name="run_workflow",
)
mcp = FastMCP("Server", transforms=[code_mode])
```
## Sandbox Configuration
### Resource Limits
The default `MontySandboxProvider` can enforce execution limits on sandboxed code — timeouts, memory caps, recursion depth, and more. Without limits, LLM-generated scripts can run indefinitely.
The default `MontySandboxProvider` can enforce execution limits — timeouts, memory caps, recursion depth, and more. Without limits, LLM-generated scripts can run indefinitely.
```python
from fastmcp.experimental.transforms import CodeMode, MontySandboxProvider
@ -213,7 +287,7 @@ All keys are optional — omit any to leave that dimension uncapped:
### Custom Sandbox Providers
The default `MontySandboxProvider` uses [pydantic-monty](https://github.com/pydantic/pydantic-monty) for sandboxed execution. You can replace it with any object implementing the `SandboxProvider` protocol:
You can replace the default sandbox with any object implementing the `SandboxProvider` protocol:
```python
from collections.abc import Callable
@ -232,7 +306,10 @@ class RemoteSandboxProvider:
# Send code to your remote sandbox runtime
...
mcp = FastMCP("Server", transforms=[CodeMode(sandbox_provider=RemoteSandboxProvider())])
mcp = FastMCP(
"Server",
transforms=[CodeMode(sandbox_provider=RemoteSandboxProvider())],
)
```
The `external_functions` dict contains async callables injected into the sandbox scope — `execute` uses this to provide `call_tool`.

View file

@ -1,8 +1,8 @@
import asyncio
import importlib
import logging
from collections.abc import Callable, Sequence
from typing import Annotated, Any, Protocol
import json
from collections.abc import Awaitable, Callable, Sequence
from typing import Annotated, Any, Literal, Protocol
from mcp.types import TextContent
from pydantic import Field
@ -12,16 +12,29 @@ from fastmcp.server.context import Context
from fastmcp.server.transforms import GetToolNext
from fastmcp.server.transforms.catalog import CatalogTransform
from fastmcp.server.transforms.search.base import (
BaseSearchTransform,
SearchResultSerializer,
_invoke_serializer,
serialize_tools_for_output_json,
serialize_tools_for_output_markdown,
)
from fastmcp.tools.tool import Tool, ToolResult
from fastmcp.utilities.versions import VersionSpec
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Type aliases
# ---------------------------------------------------------------------------
GetToolCatalog = Callable[[Context], Awaitable[Sequence[Tool]]]
"""Async callable that returns the auth-filtered tool catalog."""
SearchFn = Callable[[Sequence[Tool], str], Awaitable[Sequence[Tool]]]
"""Async callable that searches a tool sequence by query string."""
DiscoveryToolFactory = Callable[[GetToolCatalog], Tool]
"""Factory that receives catalog access and returns a synthetic Tool."""
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _ensure_async(fn: Callable[..., Any]) -> Callable[..., Any]:
@ -52,6 +65,11 @@ def _unwrap_tool_result(result: ToolResult) -> dict[str, Any] | str:
return "\n".join(parts)
# ---------------------------------------------------------------------------
# Sandbox providers
# ---------------------------------------------------------------------------
class SandboxProvider(Protocol):
"""Interface for executing LLM-generated Python code in a sandbox.
@ -122,45 +140,297 @@ class MontySandboxProvider:
return await pydantic_monty.run_monty_async(monty, **run_kwargs)
class CodeMode(CatalogTransform):
"""Transform that collapses all tools into `search` + `execute` meta-tools."""
# ---------------------------------------------------------------------------
# Built-in discovery tools
# ---------------------------------------------------------------------------
ToolDetailLevel = Literal["brief", "detailed", "full"]
"""Detail level for discovery tool output.
- ``"brief"``: tool names and one-line descriptions
- ``"detailed"``: compact markdown with parameter names, types, and required markers
- ``"full"``: complete JSON schema
"""
def _render_tools(tools: Sequence[Tool], detail: ToolDetailLevel) -> str:
"""Render tools at the requested detail level.
The same detail value produces the same output format regardless of
which discovery tool calls this, so ``detail="detailed"`` on Search
gives identical formatting to ``detail="detailed"`` on GetSchemas.
"""
if not tools:
if detail == "full":
return json.dumps([], indent=2)
return "No tools matched the query."
if detail == "full":
return json.dumps(serialize_tools_for_output_json(tools), indent=2)
if detail == "detailed":
return serialize_tools_for_output_markdown(tools)
# brief
lines: list[str] = []
for tool in tools:
desc = f": {tool.description}" if tool.description else ""
lines.append(f"- {tool.name}{desc}")
return "\n".join(lines)
class Search:
"""Discovery tool factory that searches the catalog by query.
Args:
search_fn: Async callable ``(tools, query) -> matching_tools``.
Defaults to BM25 ranking.
name: Name of the synthetic tool exposed to the LLM.
default_detail: Default detail level for search results.
``"brief"`` returns tool names and descriptions only.
``"detailed"`` returns compact markdown with parameter schemas.
``"full"`` returns complete JSON tool definitions.
"""
def __init__(
self,
*,
search_fn: SearchFn | None = None,
name: str = "search",
default_detail: ToolDetailLevel | None = None,
) -> None:
if search_fn is None:
from fastmcp.server.transforms.search.bm25 import BM25SearchTransform
_bm25 = BM25SearchTransform()
search_fn = _bm25._search
self._search_fn = search_fn
self._name = name
self._default_detail: ToolDetailLevel = default_detail or "brief"
def __call__(self, get_catalog: GetToolCatalog) -> Tool:
search_fn = self._search_fn
default_detail = self._default_detail
async def search(
query: Annotated[str, "Search query to find available tools"],
tags: Annotated[
list[str] | None,
"Filter to tools with any of these tags before searching",
] = None,
detail: Annotated[
ToolDetailLevel,
"'brief' for names and descriptions, 'detailed' for parameter schemas as markdown, 'full' for complete JSON schemas",
] = default_detail,
ctx: Context = None, # type: ignore[assignment]
) -> str:
"""Search for available tools by query.
Returns matching tools ranked by relevance.
"""
tools = await get_catalog(ctx)
if tags:
tag_set = set(tags)
has_untagged = "untagged" in tag_set
real_tags = tag_set - {"untagged"}
tools = [
t
for t in tools
if (t.tags & real_tags) or (has_untagged and not t.tags)
]
results = await search_fn(tools, query)
return _render_tools(results, detail)
return Tool.from_function(fn=search, name=self._name)
class GetSchemas:
"""Discovery tool factory that returns schemas for tools by name.
Args:
name: Name of the synthetic tool exposed to the LLM.
default_detail: Default detail level for schema results.
``"brief"`` returns tool names and descriptions only.
``"detailed"`` renders compact markdown with parameter names,
types, and required markers.
``"full"`` returns the complete JSON schema.
"""
def __init__(
self,
*,
name: str = "get_schema",
default_detail: ToolDetailLevel | None = None,
) -> None:
self._name = name
self._default_detail: ToolDetailLevel = default_detail or "detailed"
def __call__(self, get_catalog: GetToolCatalog) -> Tool:
default_detail = self._default_detail
async def get_schema(
tools: Annotated[
list[str],
"List of tool names to get schemas for",
],
detail: Annotated[
ToolDetailLevel,
"'brief' for names and descriptions, 'detailed' for parameter schemas as markdown, 'full' for complete JSON schemas",
] = default_detail,
ctx: Context = None, # type: ignore[assignment]
) -> str:
"""Get parameter schemas for specific tools.
Use after searching to get the detail needed to call a tool.
"""
catalog = await get_catalog(ctx)
catalog_by_name = {t.name: t for t in catalog}
matched = [catalog_by_name[n] for n in tools if n in catalog_by_name]
not_found = [n for n in tools if n not in catalog_by_name]
if not matched and not_found:
return f"Tools not found: {', '.join(not_found)}"
if detail == "full":
data = serialize_tools_for_output_json(matched)
if not_found:
data.append({"not_found": not_found})
return json.dumps(data, indent=2)
result = _render_tools(matched, detail)
if not_found:
result += f"\n\nTools not found: {', '.join(not_found)}"
return result
return Tool.from_function(fn=get_schema, name=self._name)
class GetTags:
"""Discovery tool factory that lists tool tags from the catalog.
Reads ``tool.tags`` from the catalog and groups tools by tag. Tools
without tags appear under ``"untagged"``.
Args:
name: Name of the synthetic tool exposed to the LLM.
default_detail: Default detail level.
``"brief"`` returns tag names with tool counts.
``"full"`` lists all tools under each tag.
"""
def __init__(
self,
*,
name: str = "tags",
default_detail: Literal["brief", "full"] | None = None,
) -> None:
self._name = name
self._default_detail: Literal["brief", "full"] = default_detail or "brief"
def __call__(self, get_catalog: GetToolCatalog) -> Tool:
default_detail = self._default_detail
async def tags(
detail: Annotated[
Literal["brief", "full"],
"Level of detail: 'brief' for tag names and counts, 'full' for tools listed under each tag",
] = default_detail,
ctx: Context = None, # type: ignore[assignment]
) -> str:
"""List available tool tags.
Use to browse available tools by tag before searching.
"""
catalog = await get_catalog(ctx)
by_tag: dict[str, list[Tool]] = {}
for tool in catalog:
if tool.tags:
for tag in tool.tags:
by_tag.setdefault(tag, []).append(tool)
else:
by_tag.setdefault("untagged", []).append(tool)
if not by_tag:
return "No tools available."
if detail == "brief":
lines = [
f"- {tag} ({len(tools)} tool{'s' if len(tools) != 1 else ''})"
for tag, tools in sorted(by_tag.items())
]
return "\n".join(lines)
blocks: list[str] = []
for tag, tools in sorted(by_tag.items()):
lines = [f"### {tag}"]
for tool in tools:
desc = f": {tool.description}" if tool.description else ""
lines.append(f"- {tool.name}{desc}")
blocks.append("\n".join(lines))
return "\n\n".join(blocks)
return Tool.from_function(fn=tags, name=self._name)
# ---------------------------------------------------------------------------
# CodeMode
# ---------------------------------------------------------------------------
def _default_discovery_tools() -> list[DiscoveryToolFactory]:
return [Search(), GetSchemas()]
class CodeMode(CatalogTransform):
"""Transform that collapses all tools into discovery + execute meta-tools.
Discovery tools are composable via the ``discovery_tools`` parameter.
Each is a callable that receives catalog access and returns a ``Tool``.
By default, ``Search`` and ``GetSchemas`` are included for
progressive disclosure: search finds candidates, get_schema retrieves
parameter details, and execute runs code.
The ``execute`` tool is always present and provides a sandboxed Python
environment with ``call_tool(name, params)`` in scope.
"""
def __init__(
self,
*,
default_arguments: dict[str, Any] | None = None,
sandbox_provider: SandboxProvider | None = None,
search_transform: BaseSearchTransform | None = None,
search_tool_name: str = "search",
discovery_tools: list[DiscoveryToolFactory] | None = None,
execute_tool_name: str = "execute",
execute_description: str | None = None,
search_result_serializer: SearchResultSerializer | None = None,
) -> None:
if search_tool_name == execute_tool_name:
raise ValueError(
"search_tool_name and execute_tool_name must be different."
)
super().__init__()
self._default_arguments = default_arguments or {}
self.search_tool_name = search_tool_name
self.execute_tool_name = execute_tool_name
self.execute_description = execute_description
self.sandbox_provider = sandbox_provider or MontySandboxProvider()
self._cached_search_tool: Tool | None = None
self._cached_execute_tool: Tool | None = None
self._search_result_serializer: SearchResultSerializer | None = (
search_result_serializer
self._discovery_factories = (
discovery_tools
if discovery_tools is not None
else _default_discovery_tools()
)
self._built_discovery_tools: list[Tool] | None = None
self._cached_execute_tool: Tool | None = None
if search_transform is None:
from fastmcp.server.transforms.search.bm25 import BM25SearchTransform
search_transform = BM25SearchTransform()
self._search_transform = search_transform
def _build_discovery_tools(self) -> list[Tool]:
if self._built_discovery_tools is None:
tools = [
factory(self.get_tool_catalog) for factory in self._discovery_factories
]
names = {t.name for t in tools}
if self.execute_tool_name in names:
raise ValueError(
f"Discovery tool name '{self.execute_tool_name}' "
f"collides with execute_tool_name."
)
if len(names) != len(tools):
raise ValueError("Discovery tools must have unique names.")
self._built_discovery_tools = tools
return self._built_discovery_tools
async def transform_tools(self, tools: Sequence[Tool]) -> Sequence[Tool]:
return [self._get_search_tool(), self._get_execute_tool()]
return [*self._build_discovery_tools(), self._get_execute_tool()]
async def get_tool(
self,
@ -169,8 +439,9 @@ class CodeMode(CatalogTransform):
*,
version: VersionSpec | None = None,
) -> Tool | None:
if name == self.search_tool_name:
return self._get_search_tool()
for tool in self._build_discovery_tools():
if tool.name == name:
return tool
if name == self.execute_tool_name:
return self._get_execute_tool()
return await call_next(name, version=version)
@ -193,41 +464,11 @@ class CodeMode(CatalogTransform):
return tool
return None
def _get_search_tool(self) -> Tool:
if self._cached_search_tool is None:
self._cached_search_tool = self._make_search_tool()
return self._cached_search_tool
def _get_execute_tool(self) -> Tool:
if self._cached_execute_tool is None:
self._cached_execute_tool = self._make_execute_tool()
return self._cached_execute_tool
def _make_search_tool(self) -> Tool:
transform = self
async def search(
query: Annotated[
str,
"Search query to find available tools",
],
ctx: Context = None, # type: ignore[assignment]
) -> str | list[dict[str, Any]]:
"""Search for available tools by query.
Returns matching tool definitions ranked by relevance,
in the same format as list_tools.
"""
tools = await transform.get_tool_catalog(ctx)
results = await transform._search_transform._search(tools, query)
if transform._search_result_serializer is not None:
return await _invoke_serializer(
transform._search_result_serializer, results
)
return await transform._search_transform._render_results(results)
return Tool.from_function(fn=search, name=self.search_tool_name)
def _make_execute_tool(self) -> Tool:
transform = self
@ -243,9 +484,6 @@ class CodeMode(CatalogTransform):
ctx: Context = None, # type: ignore[assignment]
) -> Any:
"""Execute tool calls using Python code."""
defaults = transform._default_arguments
# Cache the tool catalog for the duration of this execute block
# so multiple call_tool() invocations don't each trigger list_tools().
cached_tools: Sequence[Tool] | None = None
async def _get_cached_tools() -> Sequence[Tool]:
@ -260,26 +498,7 @@ class CodeMode(CatalogTransform):
if tool is None:
raise NotFoundError(f"Unknown tool: {tool_name}")
accepted_args = set(tool.parameters.get("properties", {}).keys())
skipped = {
key
for key in defaults
if key not in params and key not in accepted_args
}
if skipped:
logger.debug(
"default_arguments keys %s not accepted by tool %r, skipping",
skipped,
tool.name,
)
merged = {
key: value
for key, value in defaults.items()
if key not in params and key in accepted_args
}
merged.update(params)
result = await ctx.fastmcp.call_tool(tool.name, merged)
result = await ctx.fastmcp.call_tool(tool.name, params)
return _unwrap_tool_result(result)
return await transform.sandbox_provider.run(
@ -296,9 +515,10 @@ class CodeMode(CatalogTransform):
__all__ = [
"CodeMode",
"GetSchemas",
"GetTags",
"GetToolCatalog",
"MontySandboxProvider",
"SandboxProvider",
"SearchResultSerializer",
"serialize_tools_for_output_json",
"serialize_tools_for_output_markdown",
"Search",
]

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,197 @@
from typing import Any
import pytest
from fastmcp import FastMCP
from fastmcp.server.transforms.search.base import (
_schema_section,
_schema_type,
serialize_tools_for_output_markdown,
)
# ---------------------------------------------------------------------------
# _schema_type unit tests
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"schema,expected",
[
({"type": "string"}, "string"),
({"type": "integer"}, "integer"),
({"type": "boolean"}, "boolean"),
({"type": "null"}, "null"),
({"type": "array", "items": {"type": "string"}}, "string[]"),
({"type": "array", "items": {"type": "integer"}}, "integer[]"),
({"type": "array"}, "any[]"),
({"$ref": "#/$defs/Foo"}, "object"),
({"properties": {"x": {"type": "int"}}}, "object"),
({}, "any"),
(None, "any"),
("not a dict", "any"),
],
)
def test_schema_type_basic(schema: Any, expected: str) -> None:
assert _schema_type(schema) == expected
@pytest.mark.parametrize(
"schema,expected",
[
({"anyOf": [{"type": "string"}, {"type": "null"}]}, "string?"),
({"anyOf": [{"type": "string"}, {"type": "integer"}]}, "string | integer"),
(
{"anyOf": [{"type": "string"}, {"type": "integer"}, {"type": "null"}]},
"string | integer?",
),
({"anyOf": [{"type": "null"}]}, "null"),
({"anyOf": []}, "any"),
({"oneOf": [{"type": "string"}, {"type": "null"}]}, "string?"),
({"oneOf": [{"type": "string"}, {"type": "integer"}]}, "string | integer"),
({"allOf": [{"type": "object"}]}, "object"),
({"allOf": [{"$ref": "#/$defs/Foo"}, {"$ref": "#/$defs/Bar"}]}, "object"),
],
)
def test_schema_type_unions(schema: Any, expected: str) -> None:
assert _schema_type(schema) == expected
# ---------------------------------------------------------------------------
# _schema_section unit tests
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"schema,expected_lines",
[
(None, ["**Parameters**", "- `value` (any)"]),
("string", ["**Parameters**", "- `value` (any)"]),
({"type": "string"}, ["**Parameters**", "- `value` (string)"]),
(
{"type": "object", "properties": {}},
["**Parameters**", "*(no parameters)*"],
),
],
)
def test_schema_section_fallbacks(schema: Any, expected_lines: list[str]) -> None:
assert _schema_section(schema, "Parameters") == expected_lines
def test_schema_section_lists_fields_with_required_marker() -> None:
schema = {
"type": "object",
"properties": {
"name": {"type": "string"},
"age": {"type": "integer"},
},
"required": ["name"],
}
lines = _schema_section(schema, "Parameters")
assert lines[0] == "**Parameters**"
assert "- `name` (string, required)" in lines
assert "- `age` (integer)" in lines
# ---------------------------------------------------------------------------
# serialize_tools_for_output_markdown unit tests
# ---------------------------------------------------------------------------
def test_serialize_tools_for_output_markdown_empty_list() -> None:
assert serialize_tools_for_output_markdown([]) == "No tools matched the query."
async def test_serialize_tools_for_output_markdown_basic_tool() -> None:
mcp = FastMCP("MD Basic")
@mcp.tool
def square(x: int) -> int:
"""Compute the square of a number."""
return x * x
tools = await mcp.list_tools()
result = serialize_tools_for_output_markdown(tools)
assert "### square" in result
assert "Compute the square of a number." in result
assert "**Parameters**" in result
assert "`x` (integer, required)" in result
async def test_serialize_tools_for_output_markdown_omits_output_section_when_no_schema() -> (
None
):
mcp = FastMCP("MD No Output")
@mcp.tool
def ping() -> None:
pass
tools = await mcp.list_tools()
result = serialize_tools_for_output_markdown(tools)
assert "**Returns**" not in result
async def test_serialize_tools_for_output_markdown_includes_output_section_when_schema_present() -> (
None
):
mcp = FastMCP("MD With Output")
@mcp.tool
def double(x: int) -> int:
return x * 2
tools = await mcp.list_tools()
result = serialize_tools_for_output_markdown(tools)
assert "**Returns**" in result
async def test_serialize_tools_for_output_markdown_omits_description_when_absent() -> (
None
):
mcp = FastMCP("MD No Desc")
@mcp.tool
def ping() -> None:
pass
tools = await mcp.list_tools()
result = serialize_tools_for_output_markdown(tools)
assert "### ping" in result
async def test_serialize_tools_for_output_markdown_optional_field_uses_question_mark() -> (
None
):
mcp = FastMCP("MD Optional")
@mcp.tool
def greet(name: str, greeting: str | None = None) -> str:
return f"{greeting or 'Hello'}, {name}!"
tools = await mcp.list_tools()
result = serialize_tools_for_output_markdown(tools)
assert "`greeting` (string?)" in result
async def test_serialize_tools_for_output_markdown_multiple_tools_separated() -> None:
mcp = FastMCP("MD Multi")
@mcp.tool
def add(a: int, b: int) -> int:
return a + b
@mcp.tool
def subtract(a: int, b: int) -> int:
return a - b
tools = await mcp.list_tools()
result = serialize_tools_for_output_markdown(tools)
assert "### add" in result
assert "### subtract" in result
assert "\n\n" in result