universal-mcp 0.1.24rc2__py3-none-any.whl → 0.1.24rc3__py3-none-any.whl

universal_mcp/utils/openapi/openapi.py

@@ -1,3 +1,4 @@
+import hashlib
 import json
 import re
 import textwrap
@@ -11,6 +12,177 @@ from pydantic import BaseModel
 
 from .filters import load_filter_config, should_process_operation
 
+# Schema registry for tracking unique response schemas to avoid duplicates
+_schema_registry: dict[str, str] = {}  # schema_hash -> model_class_name
+_generated_models: dict[str, str] = {}  # model_class_name -> model_source_code
+
+
+def _get_schema_hash(schema: dict[str, Any]) -> str:
+    """Generate a hash for a schema to identify unique schemas."""
+    try:
+        schema_str = json.dumps(schema, sort_keys=True, default=str)
+        return hashlib.md5(schema_str.encode()).hexdigest()[:8]
+    except (TypeError, ValueError):
+        # Fallback to string representation if JSON serialization fails
+        schema_str = str(sorted(schema.items())) if isinstance(schema, dict) else str(schema)
+        return hashlib.md5(schema_str.encode()).hexdigest()[:8]
+
+
+def _generate_model_name(operation: dict[str, Any], path: str, method: str, schema: dict[str, Any]) -> str:
+    """Generate a meaningful model name for a response schema."""
+    if "title" in schema:
+        name = schema["title"]
+    else:
+        # Generate name from operation info
+        if "operationId" in operation:
+            name = operation["operationId"] + "Response"
+        else:
+            # Generate from path and method
+            path_parts = [
+                part for part in path.strip("/").split("/") if not (part.startswith("{") and part.endswith("}"))
+            ]
+            if path_parts:
+                name = f"{method.capitalize()}{path_parts[-1].capitalize()}Response"
+            else:
+                name = f"{method.capitalize()}Response"
+
+    name = "".join(word.capitalize() for word in re.split(r"[^a-zA-Z0-9]", name) if word)
+
+    if name and name[0].isdigit():
+        name = "Response" + name
+
+    return name or "Response"
+
+
+def _generate_response_model_class(schema: dict[str, Any], model_name: str) -> str:
+    """Generate Pydantic model source code from OpenAPI response schema."""
+    if not schema:
+        return ""
+
+    # Handle array responses
+    if schema.get("type") == "array":
+        items_schema = schema.get("items", {})
+        if items_schema and (items_schema.get("properties") or items_schema.get("type") == "object"):
+            # Generate model for array items if it's an object
+            item_model_name = f"{model_name}Item"
+            item_model_code = _generate_response_model_class(items_schema, item_model_name)
+
+            # Create collection model
+            collection_model = f"""
+class {model_name}(BaseModel):
+    value: List[{item_model_name}]
+"""
+            return item_model_code + collection_model
+        else:
+            # Fallback for arrays with simple items or no schema
+            item_type = "Any"
+            if items_schema:
+                if items_schema.get("type") == "string":
+                    item_type = "str"
+                elif items_schema.get("type") == "integer":
+                    item_type = "int"
+                elif items_schema.get("type") == "number":
+                    item_type = "float"
+                elif items_schema.get("type") == "boolean":
+                    item_type = "bool"
+
+            return f"""
+class {model_name}(BaseModel):
+    value: List[{item_type}]
+"""
+
+    # Handle object responses
+    if schema.get("type") == "object" or "properties" in schema:
+        properties, required_fields = _extract_properties_from_schema(schema)
+
+        if not properties:
+            return f"""
+class {model_name}(BaseModel):
+    pass
+"""
+
+        field_definitions = []
+        for prop_name, prop_schema in properties.items():
+            field_name = _sanitize_identifier(prop_name)
+            is_required = prop_name in required_fields
+
+            # Handle arrays with object items specially
+            if prop_schema.get("type") == "array" and prop_schema.get("items", {}).get("properties"):
+                # Generate a model for the array items
+                item_model_name = f"{model_name}{field_name.capitalize()}Item"
+                items_schema = prop_schema.get("items", {})
+
+                # Generate the item model and store it globally
+                item_model_code = _generate_response_model_class(items_schema, item_model_name)
+                if item_model_code and item_model_name not in _generated_models:
+                    _generated_models[item_model_name] = item_model_code
+
+                python_type = f"List[{item_model_name}]" if is_required else f"Optional[List[{item_model_name}]]"
+            else:
+                python_type = _openapi_type_to_python_type(prop_schema, required=is_required)
+
+            # Handle field aliases for special characters like @odata.context
+            if prop_name != field_name or prop_name.startswith("@"):
+                if is_required:
+                    field_definitions.append(f"    {field_name}: {python_type} = Field(alias='{prop_name}')")
+                else:
+                    field_definitions.append(f"    {field_name}: {python_type} = Field(None, alias='{prop_name}')")
+            else:
+                if is_required:
+                    field_definitions.append(f"    {field_name}: {python_type}")
+                else:
+                    field_definitions.append(f"    {field_name}: {python_type} = None")
+
+        model_code = f"""
+class {model_name}(BaseModel):
+{chr(10).join(field_definitions)}
+"""
+        return model_code
+
+    # Fallback for other schema types
+    return ""
+
+
+def _get_or_create_response_model(
+    operation: dict[str, Any], path: str, method: str, schema: dict[str, Any]
+) -> str | None:
+    """Get or create a response model for a given schema, avoiding duplicates."""
+    if not schema:
+        return None
+
+    try:
+        # Generate hash for this schema
+        schema_hash = _get_schema_hash(schema)
+
+        # Check if we already have a model for this schema
+        if schema_hash in _schema_registry:
+            return _schema_registry[schema_hash]
+
+        # Generate new model
+        model_name = _generate_model_name(operation, path, method, schema)
+
+        # Ensure unique model name
+        base_name = model_name
+        counter = 1
+        while model_name in _generated_models:
+            model_name = f"{base_name}{counter}"
+            counter += 1
+
+        # Generate model source code
+        model_code = _generate_response_model_class(schema, model_name)
+
+        if model_code:
+            # Register the model
+            _schema_registry[schema_hash] = model_name
+            _generated_models[model_name] = model_code
+            return model_name
+
+    except Exception as e:
+        # If model generation fails, log and continue with fallback
+        print(f"Warning: Could not generate model for {method.upper()} {path}: {e}")
+
+    return None
+
 
 class Parameters(BaseModel):
     name: str
@@ -220,15 +392,20 @@ def _load_and_resolve_references(path: Path):
     return replace_refs(schema)
 
 
-def _determine_return_type(operation: dict[str, Any]) -> str:
+def _determine_return_type(operation: dict[str, Any], path: str, method: str) -> str:
     """
     Determine the return type from the response schema.
 
+    Now generates specific Pydantic model classes for response schemas where possible,
+    falling back to generic types for complex or missing schemas.
+
     Args:
         operation (dict): The operation details from the schema.
+        path (str): The API path (e.g., '/users/{user_id}').
+        method (str): The HTTP method (e.g., 'get').
 
     Returns:
-        str: The appropriate return type annotation (list[Any], dict[str, Any], or Any)
+        str: The appropriate return type annotation (specific model class name or generic type)
     """
     responses = operation.get("responses", {})
     # Find successful response (2XX)
@@ -247,7 +424,12 @@ def _determine_return_type(operation: dict[str, Any]) -> str:
             if content_type.startswith("application/json") and "schema" in content_info:
                 schema = content_info["schema"]
 
-                # Only determine if it's a list, dict, or unknown (Any)
+                # generate a specific model class for this schema
+                model_name = _get_or_create_response_model(operation, path, method, schema)
+
+                if model_name:
+                    return model_name
+
                 if schema.get("type") == "array":
                     return "list[Any]"
                 elif schema.get("type") == "object" or "$ref" in schema:
@@ -536,7 +718,7 @@ def _generate_method_code(path, method, operation):
     # --- End Alias duplicate parameter names ---
 
     # --- Determine Return Type and Body Characteristics ---
-    return_type = _determine_return_type(operation)
+    return_type = _determine_return_type(operation, path, method)
 
     body_required = has_body and operation["requestBody"].get("required", False)  # Remains useful
 
@@ -751,7 +933,7 @@ def _generate_method_code(path, method, operation):
     # openapi_path_comment_for_docstring = f"# openapi_path: {path}"
     # docstring_parts.append(openapi_path_comment_for_docstring)
 
-    return_type = _determine_return_type(operation)
+    return_type = _determine_return_type(operation, path, method)
 
     # Summary
     summary = operation.get("summary", "").strip()
@@ -978,9 +1160,15 @@ def _generate_method_code(path, method, operation):
     # using the prepared URL, query parameters, request body data, files, and content type.
     # Use convenience methods that automatically handle responses and errors
 
+    # Determine the appropriate return statement based on return type
+    if return_type in ["Any", "dict[str, Any]", "list[Any]"]:
+        return_statement = "        return self._handle_response(response)"
+    else:
+        return_statement = f"        return {return_type}.model_validate(self._handle_response(response))"
+
     if method_lower == "get":
         body_lines.append("        response = self._get(url, params=query_params)")
-        body_lines.append("        return self._handle_response(response)")
+        body_lines.append(return_statement)
     elif method_lower == "post":
         if selected_content_type == "multipart/form-data":
             body_lines.append(
@@ -990,7 +1178,7 @@ def _generate_method_code(path, method, operation):
             body_lines.append(
                 f"        response = self._post(url, data=request_body_data, params=query_params, content_type='{final_content_type_for_api_call}')"
             )
-        body_lines.append("        return self._handle_response(response)")
+        body_lines.append(return_statement)
     elif method_lower == "put":
         if selected_content_type == "multipart/form-data":
             body_lines.append(
@@ -1000,16 +1188,16 @@ def _generate_method_code(path, method, operation):
             body_lines.append(
                 f"        response = self._put(url, data=request_body_data, params=query_params, content_type='{final_content_type_for_api_call}')"
             )
-        body_lines.append("        return self._handle_response(response)")
+        body_lines.append(return_statement)
     elif method_lower == "patch":
         body_lines.append("        response = self._patch(url, data=request_body_data, params=query_params)")
-        body_lines.append("        return self._handle_response(response)")
+        body_lines.append(return_statement)
     elif method_lower == "delete":
         body_lines.append("        response = self._delete(url, params=query_params)")
-        body_lines.append("        return self._handle_response(response)")
+        body_lines.append(return_statement)
     else:
         body_lines.append(f"        response = self._{method_lower}(url, data=request_body_data, params=query_params)")
-        body_lines.append("        return self._handle_response(response)")
+        body_lines.append(return_statement)
 
     # --- Combine Signature, Docstring, and Body for Final Method Code ---
     method_code = signature + formatted_docstring + "\n" + "\n".join(body_lines)
@@ -1020,9 +1208,69 @@ def load_schema(path: Path):
     return _load_and_resolve_references(path)
 
 
+def generate_schemas_file(schema, class_name: str | None = None, filter_config_path: str | None = None):
+    """
+    Generate a Python file containing only the response schema classes from an OpenAPI schema.
+
+    Args:
+        schema (dict): The OpenAPI schema as a dictionary.
+        class_name (str | None): Optional class name for context.
+        filter_config_path (str | None): Optional path to JSON filter configuration file.
+
+    Returns:
+        str: A string containing the Python code for the response schema classes.
+    """
+    global _schema_registry, _generated_models
+    _schema_registry.clear()
+    _generated_models.clear()
+
+    # Load filter configuration if provided
+    filter_config = None
+    if filter_config_path:
+        filter_config = load_filter_config(filter_config_path)
+
+    # Generate response models by processing all 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"]:
+                # Apply filter configuration
+                if not should_process_operation(path, method, filter_config):
+                    continue
+
+                operation = path_info[method]
+                # Generate response model for this operation
+                _determine_return_type(operation, path, method)
+
+    # Generate the schemas file content
+    imports = [
+        "from typing import Any, Optional, List",
+        "from pydantic import BaseModel, Field",
+    ]
+
+    imports_section = "\n".join(imports)
+    models_section = "\n".join(_generated_models.values()) if _generated_models else ""
+
+    if not models_section:
+        # If no models were generated, create a minimal file
+        schemas_code = f"""{imports_section}
+
+# No response models were generated for this OpenAPI schema
+"""
+    else:
+        schemas_code = f"""{imports_section}
+
+# Generated Response Models
+
+{models_section}
+"""
+
+    return schemas_code
+
+
 def generate_api_client(schema, class_name: str | None = None, filter_config_path: str | None = None):
     """
     Generate a Python API client class from an OpenAPI schema.
+    Models are not included - they should be generated separately using generate_schemas_file.
 
     Args:
         schema (dict): The OpenAPI schema as a dictionary.
@@ -1032,6 +1280,10 @@ def generate_api_client(schema, class_name: str | None = None, filter_config_pat
     Returns:
         str: A string containing the Python code for the API client class.
     """
+    global _schema_registry, _generated_models
+    _schema_registry.clear()
+    _generated_models.clear()
+
     # Load filter configuration if provided
     filter_config = None
     if filter_config_path:
@@ -1057,7 +1309,7 @@ def generate_api_client(schema, class_name: str | None = None, filter_config_pat
     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()
+            clean_name = class_name[:-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())
@@ -1114,21 +1366,32 @@ def generate_api_client(schema, class_name: str | None = None, filter_config_pat
             {tools_list}
         ]"""
 
-    # Generate class imports
+    # Generate class imports - import from separate schemas file
     imports = [
         "from typing import Any, Optional, List",
         "from universal_mcp.applications import APIApplication",
         "from universal_mcp.integrations import Integration",
+        "from .schemas import *",
     ]
 
-    # 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"
-    )
+    # Construct the class code (no model classes since they're in separate file)
+    imports_section = "\n".join(imports)
+
+    class_code_parts = [
+        imports_section,
+        "",
+        f"class {class_name}(APIApplication):",
+        "    def __init__(self, integration: Integration = None, **kwargs) -> None:",
+        f"        super().__init__(name='{class_name.lower()}', integration=integration, **kwargs)",
+        f'        self.base_url = "{base_url}"',
+        "",
+        "\n\n".join(methods),
+        "",
+        list_tools_method,
+        "",
+    ]
+
+    class_code = "\n".join(class_code_parts)
     return class_code
 
 

universal_mcp/utils/openapi/postprocessor.py

@@ -0,0 +1,275 @@
+import ast
+import re
+
+import litellm
+
+
+def add_hint_tags_to_docstrings(input_path: str, output_path: str):
+    """
+    Reads a Python API client file, inspects each function, and adds appropriate tags to the docstring:
+    - 'readOnlyHint': Tool does not modify its environment (fetching, reading, etc.)
+    - 'destructiveHint': Tool may perform destructive updates
+    - 'openWorldHint': Tool interacts with external entities (3rd party APIs)
+
+    Functions can have multiple tags (e.g., 'readOnlyHint, openWorldHint').
+    Does not alter other tags in the docstring.
+    Writes the modified code to output_path.
+    """
+    with open(input_path, encoding="utf-8") as f:
+        source = f.read()
+    tree = ast.parse(source)
+
+    # Initialize counters
+    total_functions = 0
+    functions_with_http_methods = 0
+    functions_processed_by_llm = 0
+    functions_tagged = 0
+    llm_failures = 0
+
+    class DocstringTagAdder(ast.NodeTransformer):
+        def _find_http_method(self, node):
+            """Find the HTTP method used in the function body."""
+            http_methods = []
+
+            def visit_node(n):
+                if (
+                    isinstance(n, ast.Call)
+                    and isinstance(n.func, ast.Attribute)
+                    and isinstance(n.func.value, ast.Name)
+                    and n.func.value.id == "self"
+                    and n.func.attr in ["_get", "_post", "_put", "_patch", "_delete"]
+                ):
+                    http_methods.append(n.func.attr[1:])
+                for child in ast.iter_child_nodes(n):
+                    visit_node(child)
+
+            visit_node(node)
+            return http_methods[0] if http_methods else None
+
+        def visit_FunctionDef(self, node):
+            nonlocal \
+                total_functions, \
+                functions_with_http_methods, \
+                functions_processed_by_llm, \
+                functions_tagged, \
+                llm_failures
+
+            total_functions += 1
+            print(f"\n[{total_functions}] Processing function: {node.name}")
+
+            http_method = self._find_http_method(node)
+            tag_to_add = None
+
+            if http_method:
+                functions_with_http_methods += 1
+                print(f"  └─ Found HTTP method: {http_method.upper()}")
+
+                # Use simple agent to decide tag
+                print("  └─ Calling LLM to determine tag...")
+                tag_to_add = self._get_tag_suggestion_from_agent(node, http_method)
+
+                if tag_to_add:
+                    functions_processed_by_llm += 1
+                    print(f"  └─ LLM suggested tags: {tag_to_add}")
+                else:
+                    print("  └─ LLM failed or returned invalid response")
+            else:
+                print("  └─ No HTTP method found - skipping")
+
+            if tag_to_add:
+                docstring = ast.get_docstring(node, clean=False)
+                if docstring is not None:
+                    # Look for Tags: section in the docstring
+                    tags_match = re.search(r"Tags:\s*(.+)", docstring, re.DOTALL)
+                    if tags_match:
+                        tags_line = tags_match.group(1).strip()
+                        # Parse existing tags
+                        existing_tags = [tag.strip() for tag in tags_line.split(",")]
+
+                        # Parse new tags to add
+                        new_tags_to_add = [tag.strip() for tag in tag_to_add.split(",")]
+                        tags_to_add = [tag for tag in new_tags_to_add if tag not in existing_tags]
+
+                        if tags_to_add:
+                            # Add the new tags to the existing list
+                            new_tags_line = tags_line.rstrip() + f", {', '.join(tags_to_add)}"
+                            new_docstring = re.sub(r"(Tags:\s*)(.+)", r"\1" + new_tags_line, docstring, flags=re.DOTALL)
+                            # Replace docstring
+                            if isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Constant):
+                                node.body[0].value.value = new_docstring
+                                functions_tagged += 1
+                                print(f"  └─ ✅ Tags '{', '.join(tags_to_add)}' added successfully")
+                        else:
+                            print(f"  └─ ⚠️  All tags '{tag_to_add}' already exist - skipping")
+                    else:
+                        print("  └─ ⚠️  No 'Tags:' section found in docstring - skipping")
+                else:
+                    print("  └─ ⚠️  No docstring found - skipping")
+            return node
+
+        def _get_tag_suggestion_from_agent(self, node, http_method):
+            """Use a simple agent to decide which tag to add based on function context."""
+
+            function_name = node.name
+            docstring = ast.get_docstring(node, clean=False) or ""
+            parameters = [arg.arg for arg in node.args.args if arg.arg != "self"]
+
+            system_prompt = """You are an expert at analyzing API functions and determining their characteristics.
+
+            Your task is to analyze each function and decide which tags to add:
+            - 'readOnlyHint': Tool does not modify its environment (fetching, reading, etc.)
+            - 'destructiveHint': Tool may perform destructive updates
+            - 'openWorldHint': Tool interacts with external entities (3rd party APIs)
+
+            IMPORTANT:
+            - HTTP method alone is NOT enough to determine the tags. You must analyze the function's actual purpose.
+            - Since these are all API client functions, MOST functions should have 'openWorldHint' (they interact with external APIs).
+            - Only functions that are purely local operations (like reading local files) should NOT have 'openWorldHint'.
+
+            Functions can have multiple tags. For example:
+            - A function that reads from Gmail API: 'readOnlyHint, openWorldHint'
+            - A function that deletes from GitHub API: 'destructiveHint, openWorldHint'
+            - A function that only reads local files: 'readOnlyHint' (no openWorldHint)
+
+            Respond with comma-separated tags (e.g., 'readOnlyHint, openWorldHint') or 'none' if no tags apply."""
+
+            user_prompt = f"""Analyze this API function and decide which tags to add:
+
+Function Name: {function_name}
+HTTP Method: {http_method}
+Parameters: {", ".join(parameters)}
+Docstring: {docstring[:1000]}...
+
+Based on this information, which tags should this function get?
+
+Think through:
+1. What does this function actually do? (from name and docstring)
+2. Does it modify its environment or just read/fetch?
+3. Does it interact with external entities (3rd party APIs)?
+4. Could it be potentially destructive?
+
+GUIDELINES for readOnlyHint (does not modify environment):
+- Functions that only READ or FETCH data
+- Functions that VALIDATE or CHECK things without saving
+- Functions that EXPORT or DOWNLOAD data
+- Functions that perform HEALTH CHECKS or PING operations
+- Functions that REFRESH tokens or sessions
+- Functions that SEARCH or FILTER data
+- Functions that GET information without changing anything
+- Functions that LIST or RETRIEVE data
+
+GUIDELINES for destructiveHint (DESTROYS or DELETES things):
+- Functions that DELETE resources or data
+- Functions that REMOVE or ERASE things
+- Functions that DESTROY or TERMINATE resources
+- Functions that CANCEL or ABORT operations
+- Functions that REVOKE or INVALIDATE things
+
+IMPORTANT:
+- A function should NOT have both readOnlyHint and destructiveHint - they are mutually exclusive.
+- Creating, sending, or updating things is NOT destructive - only deleting/destroying is destructive.
+- Functions that CREATE, SEND, UPDATE, or MODIFY should NOT get destructiveHint.
+
+GUIDELINES for openWorldHint (interacts with external entities):
+- Functions that interact with 3rd party APIs (Gmail, Outlook, Reddit, GitHub, etc.)
+- Functions that make external HTTP requests
+- Functions that connect to external services
+- Functions that interact with cloud services
+- Functions that communicate with external databases
+- Functions that call external webhooks
+- MOST API client functions will have this tag since they interact with external APIs
+
+NOT openWorldHint (local operations):
+- Functions that only read local files
+- Functions that process local data
+- Functions that work with local databases
+- Functions that manipulate local variables
+- Functions that only work with local system resources
+
+Examples:
+- Gmail API read function: 'readOnlyHint, openWorldHint'
+- Gmail API send email: 'openWorldHint' (not destructive, just sending)
+- Gmail API create draft: 'openWorldHint' (not destructive, just creating)
+- GitHub API delete repository: 'destructiveHint, openWorldHint'
+- Local file reader: 'readOnlyHint' (no openWorldHint)
+- Local data processor: 'none' (no tags)
+
+Focus on the FUNCTION'S PURPOSE, not just the HTTP method.
+
+Your answer (comma-separated tags or 'none'):"""
+
+            try:
+                response = litellm.completion(
+                    model="perplexity/sonar-pro",
+                    messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}],
+                    temperature=0.1,
+                    max_tokens=50,
+                )
+
+                suggested_tags = response.choices[0].message.content.strip().lower()
+
+                if suggested_tags == "none":
+                    return None
+
+                # Parse comma-separated tags
+                tag_list = [tag.strip() for tag in suggested_tags.split(",")]
+                valid_tags = []
+
+                for tag in tag_list:
+                    if tag == "readonlyhint":
+                        valid_tags.append("readOnlyHint")
+                    elif tag == "destructivehint":
+                        valid_tags.append("destructiveHint")
+                    elif tag == "openworldhint":
+                        valid_tags.append("openWorldHint")
+
+                if valid_tags:
+                    return ", ".join(valid_tags)
+                else:
+                    # If LLM gives unexpected response, return None (no tag added)
+                    return None
+
+            except Exception as e:
+                nonlocal llm_failures
+                llm_failures += 1
+                print(f"  └─ ❌ LLM failed for function {function_name}: {e}")
+                # If LLM fails, return None (no tag added)
+                return None
+
+    new_tree = DocstringTagAdder().visit(tree)
+    ast.fix_missing_locations(new_tree)
+    new_source = ast.unparse(new_tree)
+
+    # Print summary statistics
+    print(f"\n{'=' * 60}")
+    print("📊 PROCESSING SUMMARY")
+    print(f"{'=' * 60}")
+    print(f"Total functions processed: {total_functions}")
+    print(f"Functions with HTTP methods: {functions_with_http_methods}")
+    print(f"Functions processed by LLM: {functions_processed_by_llm}")
+    print(f"Functions successfully tagged: {functions_tagged}")
+    print(f"LLM failures: {llm_failures}")
+    if functions_with_http_methods > 0:
+        print(
+            f"LLM success rate: {(functions_processed_by_llm / functions_with_http_methods * 100):.1f}% of HTTP functions"
+        )
+    print(f"{'=' * 60}")
+
+    # Format with Black in memory
+    try:
+        import black
+
+        formatted_content = black.format_file_contents(new_source, fast=False, mode=black.FileMode())
+        with open(output_path, "w", encoding="utf-8") as f:
+            f.write(formatted_content)
+        print(f"Black formatting applied successfully to: {output_path}")
+    except ImportError:
+        print(f"Black not installed. Skipping formatting for: {output_path}")
+        # Write unformatted version if Black is not available
+        with open(output_path, "w", encoding="utf-8") as f:
+            f.write(new_source)
+    except Exception as e:
+        print(f"Black formatting failed for {output_path}: {e}")
+        # Write unformatted version if Black formatting fails
+        with open(output_path, "w", encoding="utf-8") as f:
+            f.write(new_source)