diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 5b23695d06d..de02bbfd723 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -13,6 +13,7 @@ from typing import Any, Final, TypedDict from urllib.parse import quote import httpx +import yaml from typing_extensions import ReadOnly, Required from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError @@ -165,6 +166,18 @@ def load_openapi_spec(filepath: str) -> dict[str, Any]: return asyncio.run(load_openapi_spec_async(filepath)) +def _parse_openapi_spec(text: str) -> dict[str, Any]: + """Parse JSON or YAML OpenAPI documents into a mapping.""" + try: + parsed: Any = json.loads(text) + except json.JSONDecodeError: + parsed = yaml.safe_load(text) + + if not isinstance(parsed, dict): + raise TypeError("OpenAPI spec must be a JSON or YAML object") + return parsed + + async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None) -> dict[str, Any]: if filepath.startswith("http://") or filepath.startswith("https://"): client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) @@ -174,14 +187,14 @@ 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 _parse_openapi_spec(r.text) # fallback: local file # Local filesystem path 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: Mapping[str, Any], spec_path: str | None = None) -> str: diff --git a/tests/unit/proxy/_experimental/mcp_server/test_openapi_spec_path_url.py b/tests/unit/proxy/_experimental/mcp_server/test_openapi_spec_path_url.py index 17a0022046e..26854dfa44f 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_openapi_spec_path_url.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_openapi_spec_path_url.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Dict +from typing import Any import httpx import pytest @@ -33,7 +33,7 @@ class _FakeAsyncHTTPHandler: def test_load_openapi_spec_supports_http_url(monkeypatch: pytest.MonkeyPatch) -> None: url = "http://example.local/openapi.json" - expected: Dict[str, Any] = { + expected: dict[str, Any] = { "openapi": "3.0.0", "info": {"title": "Test API", "version": "1.0.0"}, "paths": {}, @@ -44,7 +44,7 @@ def test_load_openapi_spec_supports_http_url(monkeypatch: pytest.MonkeyPatch) -> resp = httpx.Response(status_code=200, json=expected, request=req) calls = {"get_async_httpx_client": 0} - handler_holder: Dict[str, Any] = {} + handler_holder: dict[str, Any] = {} def fake_get_async_httpx_client(*args, **kwargs): calls["get_async_httpx_client"] += 1 @@ -56,9 +56,7 @@ def test_load_openapi_spec_supports_http_url(monkeypatch: pytest.MonkeyPatch) -> monkeypatch.setattr(gen, "get_async_httpx_client", fake_get_async_httpx_client) # Bypass SSRF validation in test (example.local doesn't resolve) - monkeypatch.setattr( - gen, "async_safe_get", lambda client, url, **kw: client.get(url) - ) + monkeypatch.setattr(gen, "async_safe_get", lambda client, url, **kw: client.get(url)) # Fail loudly if someone reintroduces direct httpx.get() def boom(*args, **kwargs): @@ -73,10 +71,35 @@ def test_load_openapi_spec_supports_http_url(monkeypatch: pytest.MonkeyPatch) -> assert handler_holder["handler"].calls == 1 -def test_load_openapi_spec_supports_local_file_path( - tmp_path, monkeypatch: pytest.MonkeyPatch -) -> None: - expected: Dict[str, Any] = { +def test_load_openapi_spec_supports_yaml_http_url(monkeypatch: pytest.MonkeyPatch) -> None: + url = "http://example.local/openapi.yaml" + yaml_spec = """ +openapi: 3.0.0 +info: + title: YAML API + version: 1.0.0 +paths: {} +""" + expected: dict[str, Any] = { + "openapi": "3.0.0", + "info": {"title": "YAML API", "version": "1.0.0"}, + "paths": {}, + } + + req = httpx.Request("GET", url) + resp = httpx.Response(status_code=200, text=yaml_spec, request=req) + + def fake_get_async_httpx_client(*args, **kwargs): + return _FakeAsyncHTTPHandler(resp, expected_url=url) + + monkeypatch.setattr(gen, "get_async_httpx_client", fake_get_async_httpx_client) + monkeypatch.setattr(gen, "async_safe_get", lambda client, url, **kw: client.get(url)) + + assert gen.load_openapi_spec(url) == expected + + +def test_load_openapi_spec_supports_local_file_path(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None: + expected: dict[str, Any] = { "openapi": "3.0.0", "info": {"title": "Local API", "version": "1.0.0"}, "paths": {}, @@ -90,9 +113,7 @@ def test_load_openapi_spec_supports_local_file_path( # For local files, shared client must NOT be used. def boom_client(*args, **kwargs): - raise AssertionError( - "get_async_httpx_client() must not be called for local file paths" - ) + raise AssertionError("get_async_httpx_client() must not be called for local file paths") monkeypatch.setattr(gen, "get_async_httpx_client", boom_client)