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 c0375e9437e..44a954d2a70 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -57,6 +57,19 @@ from litellm.proxy._experimental.mcp_server.tool_registry import ( from litellm.types.mcp import credential_redirect_hook, custom_credential_slot +def _import_yaml(): + """Import and return the yaml module, raising a clear error if missing.""" + try: + import yaml as _yaml + + return _yaml + except ImportError: + raise ImportError( + "PyYAML is required to parse YAML OpenAPI specs. " + "Install it with: pip install pyyaml" + ) from None + + class _OpenAPIJSONSchema(TypedDict, total=False): properties: Mapping[str, object] type: ReadOnly[str] @@ -187,17 +200,13 @@ async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None content_type = r.headers.get("content-type", "") if _is_yaml_content(filepath, content_type): - import yaml - - return yaml.safe_load(r.text) + return _import_yaml().safe_load(r.text) # Try JSON first; fall back to YAML for specs served without # proper Content-Type headers (common with raw GitHub URLs). try: return r.json() except Exception: - import yaml - - return yaml.safe_load(r.text) + return _import_yaml().safe_load(r.text) # fallback: local file # Local filesystem path @@ -205,19 +214,15 @@ async def load_openapi_spec_async(filepath: str, *, max_bytes: int | None = None raise FileNotFoundError(f"OpenAPI spec not found at {filepath}") if _is_yaml_content(filepath): - import yaml - with open(filepath, "r", encoding="utf-8") as f: - return yaml.safe_load(f) + return _import_yaml().safe_load(f) with open(filepath, "r", encoding="utf-8") as f: try: return json.load(f) except Exception: - import yaml - f.seek(0) - return yaml.safe_load(f) + return _import_yaml().safe_load(f) def get_base_url(spec: Mapping[str, Any], spec_path: str | None = None) -> str: