fix(mcp): support YAML OpenAPI specs

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:

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)