mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(mcp): accept YAML OpenAPI documents (#48622)
This commit is contained in:
tin-berri
GitHub
parent
c23cff7e45
commit
ad4b35d791
3 changed files with 284 additions and 6 deletions
|
|
@ -5,15 +5,29 @@ This module is used to generate MCP tools from OpenAPI specs.
|
|||
import asyncio
|
||||
import contextvars
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
from pathlib import PurePosixPath
|
||||
from typing import Final, Literal, Protocol, TypeAlias, TypedDict, overload
|
||||
from typing import (
|
||||
ClassVar,
|
||||
Final,
|
||||
Literal,
|
||||
Protocol,
|
||||
TypeAlias,
|
||||
TypedDict,
|
||||
cast, # noqa: TID251 # existing OpenAPI Protocol fields have no runtime schema
|
||||
overload,
|
||||
)
|
||||
from urllib.parse import quote
|
||||
|
||||
import httpx
|
||||
import yaml
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, Required
|
||||
from yaml.events import AliasEvent, ScalarEvent
|
||||
from yaml.nodes import MappingNode, Node, ScalarNode
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError
|
||||
from litellm.proxy.mcp_server.exceptions import (
|
||||
|
|
@ -30,6 +44,18 @@ _OPENAPI_TOOL_NAME_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]")
|
|||
OPENAPI_TOOL_NAME_MAX_LEN: Final = 128
|
||||
|
||||
_OPENAPI_TOOL_NAME_MAX_LEN: Final = OPENAPI_TOOL_NAME_MAX_LEN
|
||||
_OPENAPI_DOCUMENT: Final = TypeAdapter(dict[str, object])
|
||||
_MAX_YAML_INT_LENGTH: Final = 1024
|
||||
# YAML 1.2 core scalars; isolate these from PyYAML's global YAML 1.1 resolver.
|
||||
_YAML_SCALARS: Final = {
|
||||
"tag:yaml.org,2002:null": re.compile(r"^(?:null|Null|NULL|~|)$"),
|
||||
"tag:yaml.org,2002:bool": re.compile(r"^(?:true|True|TRUE|false|False|FALSE)$"),
|
||||
"tag:yaml.org,2002:int": re.compile(r"^(?:[-+]?[0-9]+|0o[0-7]+|0x[0-9a-fA-F]+)$"),
|
||||
"tag:yaml.org,2002:float": re.compile(
|
||||
r"^(?:[-+]?(?:\.[0-9]+|[0-9]+(?:\.[0-9]*)?)(?:[eE][-+]?[0-9]+)?"
|
||||
r"|[-+]?\.(?:inf|Inf|INF)|\.(?:nan|NaN|NAN))$"
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def sanitize_openapi_tool_name(raw_name: str) -> str:
|
||||
|
|
@ -199,6 +225,70 @@ def load_openapi_spec(filepath: str) -> _OpenAPISpec:
|
|||
return asyncio.run(load_openapi_spec_async(filepath))
|
||||
|
||||
|
||||
class _OpenAPISafeLoader(yaml.SafeLoader):
|
||||
yaml_implicit_resolvers: ClassVar[Mapping[str | None, Sequence[tuple[str, re.Pattern[str]]]]] = { # pyright: ignore[reportIncompatibleVariableOverride] # PyYAML resolver tables are class state
|
||||
None: list(_YAML_SCALARS.items())
|
||||
}
|
||||
|
||||
def compose_node(self, parent: Node | None, index: int) -> Node | None:
|
||||
event: Final[object] = self.peek_event() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # incomplete PyYAML event stubs
|
||||
if isinstance(event, AliasEvent):
|
||||
raise ValueError("YAML aliases are not supported in OpenAPI specs") # noqa: TRY004 # callers handle malformed documents as ValueError
|
||||
node: Final = super().compose_node(parent, index)
|
||||
# PyYAML otherwise implicitly resolves even an explicit nonspecific ! tag.
|
||||
if isinstance(event, ScalarEvent) and event.tag == "!" and isinstance(node, ScalarNode):
|
||||
node.tag = "tag:yaml.org,2002:str"
|
||||
return node
|
||||
|
||||
def construct_object(self, node: Node, deep: bool = False) -> object:
|
||||
pattern: Final = _YAML_SCALARS.get(node.tag)
|
||||
if pattern is not None:
|
||||
if not isinstance(node, ScalarNode):
|
||||
raise ValueError("Invalid YAML OpenAPI spec")
|
||||
scalar: Final[str] = node.value
|
||||
if node.tag == "tag:yaml.org,2002:int" and len(scalar) > _MAX_YAML_INT_LENGTH:
|
||||
raise ValueError("YAML integer is too long in OpenAPI spec")
|
||||
if pattern.fullmatch(scalar) is None:
|
||||
raise ValueError("Invalid YAML OpenAPI spec")
|
||||
if node.tag == "tag:yaml.org,2002:int":
|
||||
# YAML 1.2 decimal values with leading zeroes are not octal.
|
||||
return int(scalar, 0 if scalar.startswith(("0o", "0x")) else 10)
|
||||
if node.tag == "tag:yaml.org,2002:float":
|
||||
value: Final = self.construct_yaml_float(node)
|
||||
if not math.isfinite(value):
|
||||
raise ValueError("Invalid YAML OpenAPI spec")
|
||||
return value
|
||||
elif node.tag not in {"tag:yaml.org,2002:str", "tag:yaml.org,2002:seq", "tag:yaml.org,2002:map"} or (
|
||||
node.tag == "tag:yaml.org,2002:str" and not isinstance(node, ScalarNode)
|
||||
):
|
||||
raise ValueError("Invalid YAML OpenAPI spec")
|
||||
if isinstance(node, MappingNode):
|
||||
pairs: Final[Sequence[tuple[Node, Node]]] = node.value
|
||||
if any(not isinstance(key, ScalarNode) or key.tag != "tag:yaml.org,2002:str" for key, _ in pairs):
|
||||
raise ValueError("Invalid YAML OpenAPI spec")
|
||||
return super().construct_object( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # incomplete PyYAML stubs
|
||||
node, deep
|
||||
)
|
||||
|
||||
|
||||
def _validate_openapi_spec(value: object) -> _OpenAPISpec:
|
||||
try:
|
||||
document: Final = _OPENAPI_DOCUMENT.validate_python(value)
|
||||
except ValidationError as exc:
|
||||
raise ValueError("OpenAPI spec must be a JSON or YAML object") from exc
|
||||
return cast(_OpenAPISpec, document) # cast-ok: preserve existing lazy OpenAPI Protocol validation
|
||||
|
||||
|
||||
def _parse_openapi_spec(content: str | bytes) -> _OpenAPISpec:
|
||||
try:
|
||||
return _validate_openapi_spec(json.loads(content))
|
||||
except json.JSONDecodeError:
|
||||
try:
|
||||
return _validate_openapi_spec(yaml.load(content, Loader=_OpenAPISafeLoader))
|
||||
except (yaml.YAMLError, RecursionError) as exc:
|
||||
raise ValueError("Invalid YAML OpenAPI spec") from exc
|
||||
|
||||
|
||||
async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None) -> _OpenAPISpec:
|
||||
if filepath.startswith("http://") or filepath.startswith("https://"):
|
||||
client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
|
||||
|
|
@ -208,14 +298,17 @@ async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None
|
|||
else await async_safe_get(client, filepath, max_response_bytes=max_bytes)
|
||||
)
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
return await asyncio.to_thread(_parse_openapi_spec, r.content)
|
||||
|
||||
# fallback: local file
|
||||
# Local filesystem path
|
||||
return await asyncio.to_thread(_load_openapi_file, filepath)
|
||||
|
||||
|
||||
def _load_openapi_file(filepath: str) -> _OpenAPISpec:
|
||||
if not os.path.exists(filepath):
|
||||
raise FileNotFoundError(f"OpenAPI spec not found at {filepath}")
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
return _parse_openapi_spec(f.read())
|
||||
|
||||
|
||||
def get_base_url(spec: _OpenAPISpec, spec_path: str | None = None) -> str:
|
||||
|
|
|
|||
|
|
@ -5405,7 +5405,42 @@ class TestMCPServerManager:
|
|||
(
|
||||
httpx.Response(200, text="secret invalid JSON body"),
|
||||
"unhealthy",
|
||||
"OpenAPI specification could not be loaded (JSONDecodeError)",
|
||||
"OpenAPI specification could not be loaded (ValueError)",
|
||||
),
|
||||
(
|
||||
httpx.Response(200, text="paths: [secret invalid YAML body"),
|
||||
"unhealthy",
|
||||
"OpenAPI specification could not be loaded (ValueError)",
|
||||
),
|
||||
(
|
||||
httpx.Response(200, text="paths: {}\nx: \x01"),
|
||||
"unhealthy",
|
||||
"OpenAPI specification could not be loaded (ValueError)",
|
||||
),
|
||||
(
|
||||
httpx.Response(200, text="x: &a []\ny: *a"),
|
||||
"unhealthy",
|
||||
"OpenAPI specification could not be loaded (ValueError)",
|
||||
),
|
||||
(
|
||||
httpx.Response(200, text="x: !!bool invalid"),
|
||||
"unhealthy",
|
||||
"OpenAPI specification could not be loaded (ValueError)",
|
||||
),
|
||||
(
|
||||
httpx.Response(200, text="x: !!int ''"),
|
||||
"unhealthy",
|
||||
"OpenAPI specification could not be loaded (ValueError)",
|
||||
),
|
||||
(
|
||||
httpx.Response(200, text="x: !!float ''"),
|
||||
"unhealthy",
|
||||
"OpenAPI specification could not be loaded (ValueError)",
|
||||
),
|
||||
(
|
||||
httpx.Response(200, text="x: !!timestamp invalid"),
|
||||
"unhealthy",
|
||||
"OpenAPI specification could not be loaded (ValueError)",
|
||||
),
|
||||
],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,9 +1,15 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict
|
||||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from respx import MockRouter
|
||||
|
||||
from litellm.proxy.mcp_server import openapi_to_mcp_generator as gen
|
||||
|
||||
|
|
@ -98,3 +104,147 @@ def test_load_openapi_spec_supports_local_file_path(
|
|||
|
||||
spec = gen.load_openapi_spec(str(p))
|
||||
assert spec == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("source", ["url", "file"])
|
||||
@pytest.mark.parametrize("suffix", [".yaml", ".json", ""])
|
||||
@pytest.mark.parametrize("document_format", ["json", "yaml"])
|
||||
@pytest.mark.parametrize("use_async", [False, True])
|
||||
def test_openapi_formats_preserve_document_across_sources(
|
||||
tmp_path: Path, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch,
|
||||
source: str, suffix: str, document_format: str, use_async: bool,
|
||||
) -> None:
|
||||
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
||||
expected: Final = {
|
||||
"openapi": "3.0.0",
|
||||
"paths": {"/ping": {"get": {
|
||||
"operationId": "on",
|
||||
"parameters": [
|
||||
{"name": "on", "in": "query", "schema": {"type": "string"}},
|
||||
{"name": "off", "in": "query", "schema": {"type": "string"}},
|
||||
],
|
||||
}}},
|
||||
"x-values": [100.0, "on", "off", "yes", "no", "2026-01-01", "1:20", True, False, None, 12, 10, 10],
|
||||
}
|
||||
content: Final = (
|
||||
json.dumps(expected).replace("100.0", "1e2")
|
||||
if document_format == "json"
|
||||
else "openapi: 3.0.0\npaths:\n /ping:\n get:\n operationId: on\n"
|
||||
" parameters:\n"
|
||||
" - {name: on, in: query, schema: {type: string}}\n"
|
||||
" - {name: off, in: query, schema: {type: string}}\n"
|
||||
"x-values: [1e2, on, off, yes, no, 2026-01-01, 1:20, TRUE, False, null, 012, 0o12, 0xA]\n"
|
||||
)
|
||||
path: Final = tmp_path / ("spec" + suffix)
|
||||
path.write_text(content, encoding="utf-8")
|
||||
target: Final = str(path) if source == "file" else "https://93.184.216.34/spec" + suffix
|
||||
if source == "url":
|
||||
respx_mock.get(target).respond(200, text=content, headers={"content-type": "text/plain"})
|
||||
|
||||
spec: Final = asyncio.run(gen.load_openapi_spec_async(target)) if use_async else gen.load_openapi_spec(target)
|
||||
|
||||
assert spec == expected
|
||||
operation: Final = spec["paths"]["/ping"]["get"]
|
||||
assert gen.sanitize_openapi_tool_name(operation["operationId"]) == "on"
|
||||
assert gen.build_input_schema(operation) == {
|
||||
"type": "object",
|
||||
"properties": {"on": {"type": "string", "description": ""}, "off": {"type": "string", "description": ""}},
|
||||
"required": [],
|
||||
}
|
||||
|
||||
|
||||
_LONG_YAML_INTEGER: Final = "1" * 1025
|
||||
|
||||
|
||||
@pytest.mark.parametrize("content,message", [
|
||||
("", "OpenAPI spec must be a JSON or YAML object"),
|
||||
("null", "OpenAPI spec must be a JSON or YAML object"),
|
||||
("[]", "OpenAPI spec must be a JSON or YAML object"),
|
||||
("- item", "OpenAPI spec must be a JSON or YAML object"),
|
||||
("plain text", "OpenAPI spec must be a JSON or YAML object"),
|
||||
("paths: [", "Invalid YAML OpenAPI spec"),
|
||||
("paths: {}\nx: \x01", "Invalid YAML OpenAPI spec"),
|
||||
("paths: " + "[" * 600 + "0" + "]" * 600, "Invalid YAML OpenAPI spec"),
|
||||
("x: !!python/object/apply:builtins.str [unsafe]", "Invalid YAML OpenAPI spec"),
|
||||
("x: &a []\ny: *a", "YAML aliases are not supported in OpenAPI specs"),
|
||||
("x: " + _LONG_YAML_INTEGER, "YAML integer is too long in OpenAPI spec"),
|
||||
("x: !!int " + _LONG_YAML_INTEGER, "YAML integer is too long in OpenAPI spec"),
|
||||
("x: !!int {=: '" + _LONG_YAML_INTEGER + "'}", "Invalid YAML OpenAPI spec"),
|
||||
("x: !!int [1]", "Invalid YAML OpenAPI spec"),
|
||||
("x: !!bool invalid", "Invalid YAML OpenAPI spec"),
|
||||
("x: !!bool []", "Invalid YAML OpenAPI spec"),
|
||||
("x: !!str []", "Invalid YAML OpenAPI spec"),
|
||||
("x: !!int ''", "Invalid YAML OpenAPI spec"),
|
||||
("x: !!float ''", "Invalid YAML OpenAPI spec"),
|
||||
("x: !!float {=: 1.2}", "Invalid YAML OpenAPI spec"),
|
||||
("x: !!timestamp invalid", "Invalid YAML OpenAPI spec"),
|
||||
("x: !!timestamp 2026-01-01", "Invalid YAML OpenAPI spec"),
|
||||
("x: !!set {a: null}", "Invalid YAML OpenAPI spec"),
|
||||
("x: !!binary YQ==", "Invalid YAML OpenAPI spec"),
|
||||
("x: {1: value}", "Invalid YAML OpenAPI spec"),
|
||||
("x: {[key]: value}", "Invalid YAML OpenAPI spec"),
|
||||
("x: .inf", "Invalid YAML OpenAPI spec"),
|
||||
("x: .nan", "Invalid YAML OpenAPI spec"),
|
||||
("x: 1e999", "Invalid YAML OpenAPI spec"),
|
||||
], ids=[
|
||||
"empty", "null", "json-array", "yaml-sequence", "scalar", "syntax", "reader", "nesting",
|
||||
"python-tag", "alias", "implicit-int", "explicit-int", "mapping-int", "sequence-int",
|
||||
"invalid-bool", "sequence-bool", "sequence-str", "empty-int", "empty-float", "mapping-float", "invalid-timestamp",
|
||||
"timestamp", "set", "binary", "non-string-key", "complex-key", "infinity", "nan", "float-overflow",
|
||||
])
|
||||
def test_openapi_loader_rejects_invalid_or_unsafe_documents(
|
||||
tmp_path: Path, content: str, message: str,
|
||||
) -> None:
|
||||
path: Final = tmp_path / "spec.yaml"
|
||||
path.write_text(content, encoding="utf-8")
|
||||
|
||||
with pytest.raises(ValueError, match="^" + message + "$"):
|
||||
gen.load_openapi_spec(str(path))
|
||||
|
||||
|
||||
def test_openapi_yaml_integer_boundary_preserves_quoted_strings(tmp_path: Path) -> None:
|
||||
integer: Final = "1" * 1024
|
||||
path: Final = tmp_path / "spec.yaml"
|
||||
path.write_text("x-integer: " + integer + "\nx-string: '" + integer + "1'", encoding="utf-8")
|
||||
|
||||
assert gen.load_openapi_spec(str(path)) == {"x-integer": int(integer), "x-string": integer + "1"}
|
||||
|
||||
|
||||
def test_openapi_json_integer_preserved_even_with_yaml_extension(tmp_path: Path) -> None:
|
||||
path: Final = tmp_path / "spec.yaml"
|
||||
path.write_text('{"x-integer":' + _LONG_YAML_INTEGER + "}", encoding="utf-8")
|
||||
|
||||
assert gen.load_openapi_spec(str(path)) == {"x-integer": int(_LONG_YAML_INTEGER)}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["! " + _LONG_YAML_INTEGER, "! '" + _LONG_YAML_INTEGER + "'"], ids=["plain", "quoted"])
|
||||
def test_openapi_nonspecific_tag_disables_scalar_resolution(tmp_path: Path, value: str) -> None:
|
||||
path: Final = tmp_path / "spec.yaml"
|
||||
path.write_text("x: " + value, encoding="utf-8")
|
||||
|
||||
assert gen.load_openapi_spec(str(path)) == {"x": _LONG_YAML_INTEGER}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("source", ["url", "file"])
|
||||
async def test_openapi_yaml_parses_off_the_event_loop(
|
||||
tmp_path: Path, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, source: str,
|
||||
) -> None:
|
||||
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
||||
event_loop_thread: Final = threading.get_ident()
|
||||
parser_threads: Final[list[int]] = []
|
||||
original_load: Final = yaml.load
|
||||
|
||||
def observed_load(stream: str | bytes, Loader: type[yaml.SafeLoader]) -> object:
|
||||
parser_threads.append(threading.get_ident())
|
||||
return original_load(stream, Loader=Loader)
|
||||
|
||||
monkeypatch.setattr(yaml, "load", observed_load)
|
||||
path: Final = tmp_path / "spec.yaml"
|
||||
path.write_text("paths: {}", encoding="utf-8")
|
||||
target: Final = str(path) if source == "file" else "https://93.184.216.34/spec.yaml"
|
||||
if source == "url":
|
||||
respx_mock.get(target).respond(200, text="paths: {}")
|
||||
|
||||
assert await gen.load_openapi_spec_async(target) == {"paths": {}}
|
||||
assert len(parser_threads) == 1 and parser_threads[0] != event_loop_thread
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue