universal-mcp 0.1.12__py3-none-any.whl → 0.1.13__py3-none-any.whl

universal_mcp/utils/openapi.py

@@ -1,9 +1,65 @@
 import json
 import re
+from functools import cache
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal
 
 import yaml
+from pydantic import BaseModel
+
+
+class Parameters(BaseModel):
+    name: str
+    identifier: str
+    description: str = ""
+    type: str = "string"
+    where: Literal["path", "query", "header", "body"]
+    required: bool
+    example: str | None = None
+
+    def __str__(self):
+        return f"{self.name}: ({self.type})"
+
+
+class Method(BaseModel):
+    name: str
+    summary: str
+    tags: list[str]
+    path: str
+    method: str
+    path_params: list[Parameters]
+    query_params: list[Parameters]
+    body_params: list[Parameters]
+    return_type: str
+
+    def deduplicate_params(self):
+        """
+        Deduplicate parameters by name.
+        Sometimes the same parameter is defined in multiple places, we only want to include it once.
+        """
+        # TODO: Implement this
+        pass
+
+    def render(self, template_dir: str, template_name: str = "method.jinja2") -> str:
+        """
+        Render this Method instance into source code using a Jinja2 template.
+
+        Args:
+            template_dir (str): Directory where the Jinja2 templates are located.
+            template_name (str): Filename of the method template.
+
+        Returns:
+            str: The rendered method source code.
+        """
+        from jinja2 import Environment, FileSystemLoader
+
+        env = Environment(
+            loader=FileSystemLoader(template_dir),
+            trim_blocks=True,
+            lstrip_blocks=True,
+        )
+        template = env.get_template(template_name)
+        return template.render(method=self)
 
 
 def convert_to_snake_case(identifier: str) -> str:
@@ -24,16 +80,68 @@ def convert_to_snake_case(identifier: str) -> str:
     return result.lower()
 
 
-def load_schema(path: Path):
+@cache
+def _resolve_schema_reference(reference, schema):
+    """
+    Resolve a JSON schema reference to its target schema.
+
+    Args:
+        reference (str): The reference string (e.g., '#/components/schemas/User')
+        schema (dict): The complete OpenAPI schema that contains the reference
+
+    Returns:
+        dict: The resolved schema, or None if not found
+    """
+    if not reference.startswith("#/"):
+        return None
+
+    # Split the reference path and navigate through the schema
+    parts = reference[2:].split("/")
+    current = schema
+
+    for part in parts:
+        if part in current:
+            current = current[part]
+        else:
+            return None
+
+    return current
+
+
+def _resolve_references(schema: dict[str, Any]):
+    """
+    Recursively walk the OpenAPI schema and inline all JSON Schema $ref references.
+    """
+
+    def _resolve(node):
+        if isinstance(node, dict):
+            # If this dict is a reference, replace it with the resolved schema
+            if "$ref" in node:
+                ref = node["$ref"]
+                resolved = _resolve_schema_reference(ref, schema)
+                # If resolution fails, leave the ref dict as-is
+                return _resolve(resolved) if resolved is not None else node
+            # Otherwise, recurse into each key/value
+            return {key: _resolve(value) for key, value in node.items()}
+        elif isinstance(node, list):
+            # Recurse into list elements
+            return [_resolve(item) for item in node]
+        # Primitive value, return as-is
+        return node
+
+    return _resolve(schema)
+
+
+def _load_and_resolve_references(path: Path):
+    # Load the schema
     type = "yaml" if path.suffix == ".yaml" else "json"
     with open(path) as f:
-        if type == "yaml":
-            return yaml.safe_load(f)
-        else:
-            return json.load(f)
+        schema = yaml.safe_load(f) if type == "yaml" else json.load(f)
+    # Resolve references
+    return _resolve_references(schema)
 
 
-def determine_return_type(operation: dict[str, Any]) -> str:
+def _determine_return_type(operation: dict[str, Any]) -> str:
     """
     Determine the return type from the response schema.
 
@@ -70,144 +178,10 @@ def determine_return_type(operation: dict[str, Any]) -> str:
     return "Any"
 
 
-def resolve_schema_reference(reference, schema):
-    """
-    Resolve a JSON schema reference to its target schema.
-
-    Args:
-        reference (str): The reference string (e.g., '#/components/schemas/User')
-        schema (dict): The complete OpenAPI schema that contains the reference
-
-    Returns:
-        dict: The resolved schema, or None if not found
-    """
-    if not reference.startswith("#/"):
-        return None
-
-    # Split the reference path and navigate through the schema
-    parts = reference[2:].split("/")
-    current = schema
-
-    for part in parts:
-        if part in current:
-            current = current[part]
-        else:
-            return None
-
-    return current
-
-
-def generate_api_client(schema):
-    """
-    Generate a Python API client class from an OpenAPI schema.
-
-    Args:
-        schema (dict): The OpenAPI schema as a dictionary.
-
-    Returns:
-        str: A string containing the Python code for the API client class.
-    """
-    methods = []
-    method_names = []
-
-    # Extract API info for naming and base URL
-    info = schema.get("info", {})
-    api_title = info.get("title", "API")
-
-    # Get base URL from servers array if available
-    base_url = ""
-    servers = schema.get("servers", [])
-    if servers and isinstance(servers, list) and "url" in servers[0]:
-        base_url = servers[0]["url"].rstrip("/")
-
-    # Create a clean class name from API title
-    if api_title:
-        # Convert API title to a clean class name
-        base_name = "".join(word.capitalize() for word in api_title.split())
-        clean_name = "".join(c for c in base_name if c.isalnum())
-        class_name = f"{clean_name}App"
-
-        # Extract tool name - remove spaces and convert to lowercase
-        tool_name = api_title.lower()
-
-        # Remove version numbers (like 3.0, v1, etc.)
-        tool_name = re.sub(r"\s*v?\d+(\.\d+)*", "", tool_name)
-
-        # Remove common words that aren't needed
-        common_words = ["api", "openapi", "open", "swagger", "spec", "specification"]
-        for word in common_words:
-            tool_name = tool_name.replace(word, "")
-
-        # Remove spaces, hyphens, underscores
-        tool_name = tool_name.replace(" ", "").replace("-", "").replace("_", "")
-
-        # Remove any non-alphanumeric characters
-        tool_name = "".join(c for c in tool_name if c.isalnum())
-
-        # If empty (after cleaning), use generic name
-        if not tool_name:
-            tool_name = "api"
-    else:
-        class_name = "APIClient"
-        tool_name = "api"
-
-    # Iterate over paths and their operations
-    for path, path_info in schema.get("paths", {}).items():
-        for method in path_info:
-            if method in ["get", "post", "put", "delete", "patch", "options", "head"]:
-                operation = path_info[method]
-                method_code, func_name = generate_method_code(
-                    path, method, operation, schema, tool_name
-                )
-                methods.append(method_code)
-                method_names.append(func_name)
-
-    # Generate list_tools method with all the function names
-    tools_list = ",\n            ".join([f"self.{name}" for name in method_names])
-    list_tools_method = f"""    def list_tools(self):
-        return [
-            {tools_list}
-        ]"""
-
-    # Generate class imports
-    imports = [
-        "from typing import Any",
-        "from universal_mcp.applications import APIApplication",
-        "from universal_mcp.integrations import Integration",
-    ]
-
-    # Construct the class code
-    class_code = (
-        "\n".join(imports) + "\n\n"
-        f"class {class_name}(APIApplication):\n"
-        f"    def __init__(self, integration: Integration = None, **kwargs) -> None:\n"
-        f"        super().__init__(name='{class_name.lower()}', integration=integration, **kwargs)\n"
-        f'        self.base_url = "{base_url}"\n\n'
-        + "\n\n".join(methods)
-        + "\n\n"
-        + list_tools_method
-        + "\n"
-    )
-    return class_code
-
-
-def generate_method_code(path, method, operation, full_schema, tool_name=None):
+def _determine_function_name(operation: dict[str, Any], path: str, method: str) -> str:
     """
-    Generate the code for a single API method.
-
-    Args:
-        path (str): The API path (e.g., '/users/{user_id}').
-        method (str): The HTTP method (e.g., 'get').
-        operation (dict): The operation details from the schema.
-        full_schema (dict): The complete OpenAPI schema, used for reference resolution.
-        tool_name (str, optional): The name of the tool/app to prefix the function name with.
-
-    Returns:
-        tuple: (method_code, func_name) - The Python code for the method and its name.
+    Determine the function name from the operation.
     """
-    # Extract path parameters from the URL path
-    path_params_in_url = re.findall(r"{([^}]+)}", path)
-
     # Determine function name
     if "operationId" in operation:
         raw_name = operation["operationId"]
@@ -237,25 +211,108 @@ def generate_method_code(path, method, operation, full_schema, tool_name=None):
     func_name = re.sub(
         r"_an$", r"_an", func_name
     )  # Don't change if 'an' is at the end of the name
+    return func_name
 
-    # Get parameters and request body
-    # Resolve parameter references before processing
-    resolved_parameters = []
-    for param in operation.get("parameters", []):
-        if "$ref" in param:
-            # Resolve reference to actual parameter object
-            ref_param = resolve_schema_reference(param["$ref"], full_schema)
-            if ref_param:
-                resolved_parameters.append(ref_param)
-            else:
-                print(
-                    f"Warning: Could not resolve parameter reference: {param['$ref']}"
+
+def _generate_path_params(path: str) -> list[Parameters]:
+    path_params_in_url = re.findall(r"{([^}]+)}", path)
+    parameters = []
+    for param in path_params_in_url:
+        try:
+            parameters.append(
+                Parameters(
+                    name=param.replace("-", "_"),
+                    identifier=param,
+                    description=param,
+                    type="string",
+                    where="path",
+                    required=True,
                 )
-        else:
-            resolved_parameters.append(param)
+            )
+        except Exception as e:
+            print(f"Error generating path parameters {param}: {e}")
+            raise e
+    return parameters
+
+
+def _generate_url(path: str, path_params: list[Parameters]):
+    formatted_path = path
+    for param in path_params:
+        formatted_path = formatted_path.replace(
+            f"{{{param.identifier}}}", f"{{{param.name}}}"
+        )
+    return formatted_path
+
+
+def _generate_query_params(operation: dict[str, Any]) -> list[Parameters]:
+    query_params = []
+    for param in operation.get("parameters", []):
+        name = param.get("name")
+        description = param.get("description", "")
+        type = param.get("type")
+        where = param.get("in")
+        required = param.get("required")
+        if where == "query":
+            parameter = Parameters(
+                name=name.replace("-", "_"),
+                identifier=name,
+                description=description,
+                type=type,
+                where=where,
+                required=required,
+            )
+            query_params.append(parameter)
+    return query_params
+
+
+def _generate_body_params(operation: dict[str, Any]) -> list[Parameters]:
+    body_params = []
+    request_body = operation.get("requestBody", {})
+    required = request_body.get("required", False)
+    content = request_body.get("content", {})
+    json_content = content.get("application/json", {})
+    schema = json_content.get("schema", {})
+    properties = schema.get("properties", {})
+    for param in properties:
+        body_params.append(
+            Parameters(
+                name=param,
+                identifier=param,
+                description=param,
+                type="string",
+                where="body",
+                required=required,
+            )
+        )
+    return body_params
+
+
+def _generate_method_code(path, method, operation):
+    """
+    Generate the code for a single API method.
+
+    Args:
+        path (str): The API path (e.g., '/users/{user_id}').
+        method (str): The HTTP method (e.g., 'get').
+        operation (dict): The operation details from the schema.
+        full_schema (dict): The complete OpenAPI schema, used for reference resolution.
+        tool_name (str, optional): The name of the tool/app to prefix the function name with.
+
+    Returns:
+        tuple: (method_code, func_name) - The Python code for the method and its name.
+    """
 
-    # Filter out header parameters from the resolved parameters
-    parameters = [param for param in resolved_parameters if param.get("in") != "header"]
+    func_name = _determine_function_name(operation, path, method)
+    operation.get("summary", "")
+    operation.get("tags", [])
+    # Extract path parameters from the URL path
+    path_params = _generate_path_params(path)
+    query_params = _generate_query_params(operation)
+    _generate_body_params(operation)
+    return_type = _determine_return_type(operation)
+    # gen_method   = Method(name=func_name, summary=summary, tags=tags, path=path, method=method, path_params=path_params, query_params=query_params, body_params=body_params, return_type=return_type)
+    # logger.info(f"Generated method: {gen_method.model_dump()}")
+    # return method.render(template_dir="templates", template_name="method.jinja2")
 
     has_body = "requestBody" in operation
     body_required = has_body and operation["requestBody"].get("required", False)
@@ -274,14 +331,6 @@ def generate_method_code(path, method, operation, full_schema, tool_name=None):
                 if content_type.startswith("application/json") and "schema" in content:
                     schema = content["schema"]
 
-                    # Resolve schema reference if present
-                    if "$ref" in schema:
-                        ref_schema = resolve_schema_reference(
-                            schema["$ref"], full_schema
-                        )
-                        if ref_schema:
-                            schema = ref_schema
-
                     # Check if properties is empty and additionalProperties is true
                     if (
                         schema.get("type") == "object"
@@ -295,7 +344,6 @@ def generate_method_code(path, method, operation, full_schema, tool_name=None):
     required_fields = []
     request_body_properties = {}
     is_array_body = False
-    array_items_schema = None
 
     if has_body:
         for content_type, content in (
@@ -304,21 +352,11 @@ def generate_method_code(path, method, operation, full_schema, tool_name=None):
             if content_type.startswith("application/json") and "schema" in content:
                 schema = content["schema"]
 
-                # Resolve schema reference if present
-                if "$ref" in schema:
-                    ref_schema = resolve_schema_reference(schema["$ref"], full_schema)
-                    if ref_schema:
-                        schema = ref_schema
-
                 # Check if the schema is an array type
                 if schema.get("type") == "array":
                     is_array_body = True
-                    array_items_schema = schema.get("items", {})
-                    # Try to resolve any reference in items
-                    if "$ref" in array_items_schema:
-                        array_items_schema = resolve_schema_reference(
-                            array_items_schema["$ref"], full_schema
-                        )
+                    schema.get("items", {})
+
                 else:
                     # Extract required fields from schema
                     if "required" in schema:
@@ -327,15 +365,6 @@ def generate_method_code(path, method, operation, full_schema, tool_name=None):
                     if "properties" in schema:
                         request_body_properties = schema["properties"]
 
-                        # Check for nested references in properties
-                        for prop_name, prop_schema in request_body_properties.items():
-                            if "$ref" in prop_schema:
-                                ref_prop_schema = resolve_schema_reference(
-                                    prop_schema["$ref"], full_schema
-                                )
-                                if ref_prop_schema:
-                                    request_body_properties[prop_name] = ref_prop_schema
-
                 # Handle schemas with empty properties but additionalProperties: true
                 # by treating them similar to empty bodies
                 if (
@@ -348,18 +377,23 @@ def generate_method_code(path, method, operation, full_schema, tool_name=None):
     optional_args = []
 
     # Add path parameters
-    for param_name in path_params_in_url:
-        if param_name not in required_args:
-            required_args.append(param_name)
+    for param in path_params:
+        if param.name not in required_args:
+            required_args.append(param.name)
 
-    # Add query parameters
-    for param in parameters:
+    for param in query_params:
         param_name = param["name"]
-        if param_name not in required_args:
+        # Handle parameters with square brackets and hyphens by converting to valid Python identifiers
+        param_identifier = (
+            param_name.replace("[", "_").replace("]", "").replace("-", "_")
+        )
+        if param_identifier not in required_args and param_identifier not in [
+            p.split("=")[0] for p in optional_args
+        ]:
             if param.get("required", False):
-                required_args.append(param_name)
+                required_args.append(param_identifier)
             else:
-                optional_args.append(f"{param_name}=None")
+                optional_args.append(f"{param_identifier}=None")
 
     # Handle array type request body differently
     request_body_params = []
@@ -401,18 +435,19 @@ def generate_method_code(path, method, operation, full_schema, tool_name=None):
     args = required_args + optional_args
 
     # Determine return type
-    return_type = determine_return_type(operation)
-
-    signature = f"    def {func_name}(self, {', '.join(args)}) -> {return_type}:"
+    return_type = _determine_return_type(operation)
+    if args:
+        signature = f"    def {func_name}(self, {', '.join(args)}) -> {return_type}:"
+    else:
+        signature = f"    def {func_name}(self) -> {return_type}:"
 
     # Build method body
     body_lines = []
 
-    # Validate required parameters including path parameters
-    for param_name in required_args:
-        body_lines.append(f"        if {param_name} is None:")
+    for param in path_params:
+        body_lines.append(f"        if {param.name} is None:")
         body_lines.append(
-            f"            raise ValueError(\"Missing required parameter '{param_name}'\")"
+            f"            raise ValueError(\"Missing required parameter '{param.identifier}'\")"  # Use original name in error
         )
 
     # Build request body (handle array and object types differently)
@@ -437,17 +472,21 @@ def generate_method_code(path, method, operation, full_schema, tool_name=None):
             )
 
     # Format URL directly with path parameters
-    url_line = f'        url = f"{{self.base_url}}{path}"'
+    url = _generate_url(path, path_params)
+    url_line = f'        url = f"{{self.base_url}}{url}"'
     body_lines.append(url_line)
 
-    # Query parameters
-    query_params = [p for p in parameters if p["in"] == "query"]
+    # Build query parameters, handling square brackets in parameter names
     if query_params:
-        query_params_items = ", ".join(
-            [f"('{p['name']}', {p['name']})" for p in query_params]
-        )
+        query_params_items = []
+        for param in query_params:
+            param_name = param.name
+            param_identifier = (
+                param_name.replace("[", "_").replace("]", "").replace("-", "_")
+            )
+            query_params_items.append(f"('{param_name}', {param_identifier})")
         body_lines.append(
-            f"        query_params = {{k: v for k, v in [{query_params_items}] if v is not None}}"
+            f"        query_params = {{k: v for k, v in [{', '.join(query_params_items)}] if v is not None}}"
         )
     else:
         body_lines.append("        query_params = {}")
@@ -492,6 +531,109 @@ def generate_method_code(path, method, operation, full_schema, tool_name=None):
     return method_code, func_name
 
 
+def load_schema(path: Path):
+    return _load_and_resolve_references(path)
+
+
+def generate_api_client(schema, class_name: str | None = None):
+    """
+    Generate a Python API client class from an OpenAPI schema.
+
+    Args:
+        schema (dict): The OpenAPI schema as a dictionary.
+
+    Returns:
+        str: A string containing the Python code for the API client class.
+    """
+    methods = []
+    method_names = []
+
+    # Extract API info for naming and base URL
+    info = schema.get("info", {})
+    api_title = info.get("title", "API")
+
+    # Get base URL from servers array if available
+    base_url = ""
+    servers = schema.get("servers", [])
+    if servers and isinstance(servers, list) and "url" in servers[0]:
+        base_url = servers[0]["url"].rstrip("/")
+
+    # Create a clean class name from API title
+    if api_title:
+        # Convert API title to a clean class name
+        if class_name:
+            clean_name = (
+                class_name.capitalize()[:-3]
+                if class_name.endswith("App")
+                else class_name.capitalize()
+            )
+        else:
+            base_name = "".join(word.capitalize() for word in api_title.split())
+            clean_name = "".join(c for c in base_name if c.isalnum())
+        class_name = f"{clean_name}App"
+
+        # Extract tool name - remove spaces and convert to lowercase
+        tool_name = api_title.lower()
+
+        # Remove version numbers (like 3.0, v1, etc.)
+        tool_name = re.sub(r"\s*v?\d+(\.\d+)*", "", tool_name)
+
+        # Remove common words that aren't needed
+        common_words = ["api", "openapi", "open", "swagger", "spec", "specification"]
+        for word in common_words:
+            tool_name = tool_name.replace(word, "")
+
+        # Remove spaces, hyphens, underscores
+        tool_name = tool_name.replace(" ", "").replace("-", "").replace("_", "")
+
+        # Remove any non-alphanumeric characters
+        tool_name = "".join(c for c in tool_name if c.isalnum())
+
+        # If empty (after cleaning), use generic name
+        if not tool_name:
+            tool_name = "api"
+    else:
+        class_name = "APIClient"
+        tool_name = "api"
+
+    # Iterate over paths and their operations
+    for path, path_info in schema.get("paths", {}).items():
+        for method in path_info:
+            if method in ["get", "post", "put", "delete", "patch", "options", "head"]:
+                operation = path_info[method]
+                method_code, func_name = _generate_method_code(path, method, operation)
+                methods.append(method_code)
+                method_names.append(func_name)
+
+    # Generate list_tools method with all the function names
+    tools_list = ",\n            ".join([f"self.{name}" for name in method_names])
+    list_tools_method = f"""    def list_tools(self):
+        return [
+            {tools_list}
+        ]"""
+
+    # Generate class imports
+    imports = [
+        "from typing import Any",
+        "from universal_mcp.applications import APIApplication",
+        "from universal_mcp.integrations import Integration",
+    ]
+
+    # Construct the class code
+    class_code = (
+        "\n".join(imports) + "\n\n"
+        f"class {class_name}(APIApplication):\n"
+        f"    def __init__(self, integration: Integration = None, **kwargs) -> None:\n"
+        f"        super().__init__(name='{class_name.lower()}', integration=integration, **kwargs)\n"
+        f'        self.base_url = "{base_url}"\n\n'
+        + "\n\n".join(methods)
+        + "\n\n"
+        + list_tools_method
+        + "\n"
+    )
+    return class_code
+
+
 # Example usage
 if __name__ == "__main__":
     # Sample OpenAPI schema