universal-mcp 0.1.7rc1__py3-none-any.whl → 0.1.8__py3-none-any.whl

universal_mcp/utils/docgen.py

@@ -5,7 +5,12 @@ using LLMs with structured output
 """
 
 import ast
+import json
 import os
+import re
+import sys
+import textwrap
+import traceback
 
 import litellm
 from pydantic import BaseModel, Field
@@ -21,6 +26,14 @@ class DocstringOutput(BaseModel):
         description="Dictionary mapping parameter names to their descriptions"
     )
     returns: str = Field(description="Description of what the function returns")
+    raises: dict[str, str] = Field(
+        default_factory=dict,
+        description="Dictionary mapping potential exception types/reasons to their descriptions",
+    )
+    tags: list[str] = Field(
+        default_factory=list,
+        description="List of relevant tags for the function (e.g., action, job type, async status, importance)",
+    )
 
 
 class FunctionExtractor(ast.NodeVisitor):
@@ -47,17 +60,23 @@ class FunctionExtractor(ast.NodeVisitor):
             return None
 
     def visit_FunctionDef(self, node: ast.FunctionDef):
-        """Visits a regular function definition."""
-        source_code = self._get_source_segment(node)
-        if source_code:
-            self.functions.append((node.name, source_code))
+        """Visits a regular function definition and collects it if not excluded."""
+        # Add the exclusion logic here
+        if not node.name.startswith("_") and node.name != "list_tools":
+            source_code = self._get_source_segment(node)
+            if source_code:
+                self.functions.append((node.name, source_code))
+        # Continue traversing the AST for nested functions/classes
         self.generic_visit(node)
 
     def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
-        """Visits an asynchronous function definition."""
-        source_code = self._get_source_segment(node)
-        if source_code:
-            self.functions.append((node.name, source_code))
+        """Visits an asynchronous function definition and collects it if not excluded."""
+        # Add the exclusion logic here
+        if not node.name.startswith("_") and node.name != "list_tools":
+            source_code = self._get_source_segment(node)
+            if source_code:
+                self.functions.append((node.name, source_code))
+        # Continue traversing the AST for nested functions/classes
         self.generic_visit(node)
 
 
@@ -115,8 +134,49 @@ def extract_functions_from_script(file_path: str) -> list[tuple[str, str]]:
         return []
 
 
+def extract_json_from_text(text):
+    """Extract valid JSON from text that might contain additional content.
+
+    Args:
+        text: Raw text response from the model
+
+    Returns:
+        Dict containing the extracted JSON data
+
+    Raises:
+        ValueError: If no valid JSON could be extracted
+    """
+    # Try to find JSON between triple backticks (common markdown pattern)
+    json_match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", text)
+    if json_match:
+        try:
+            return json.loads(json_match.group(1))
+        except json.JSONDecodeError:
+            pass
+
+    # Try to find the first { and last } for a complete JSON object
+    try:
+        start = text.find("{")
+        if start >= 0:
+            brace_count = 0
+            for i in range(start, len(text)):
+                if text[i] == "{":
+                    brace_count += 1
+                elif text[i] == "}":
+                    brace_count -= 1
+                    if brace_count == 0:
+                        return json.loads(text[start : i + 1])
+    except json.JSONDecodeError:
+        pass
+
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError as e:
+        raise ValueError("Could not extract valid JSON from the response") from e
+
+
 def generate_docstring(
-    function_code: str, model: str = "openai/gpt-4o"
+    function_code: str, model: str = "perplexity/sonar-pro"
 ) -> DocstringOutput:
     """
     Generate a docstring for a Python function using litellm with structured output.
@@ -129,31 +189,42 @@ def generate_docstring(
         A DocstringOutput object containing the structured docstring components
     """
     system_prompt = """You are a helpful AI assistant specialized in writing high-quality Google-style Python docstrings.
-    You MUST ALWAYS include an Args section, even if there are no arguments (in which case mention 'None')."""
+    You MUST ALWAYS include an Args section, even if there are no arguments (in which case mention 'None').
+    You should also generate a list of tags describing the function's purpose and characteristics."""
+
+    user_prompt = f"""Generate a high-quality Google-style docstring for the following Python function.
+    Analyze the function's name, parameters, return values, potential exceptions, and functionality to create a comprehensive docstring.
 
-    user_prompt = f"""Generate a high-quality Google-style docstring for the following Python function. 
-    Analyze the function's name, parameters, return values, and functionality to create a comprehensive docstring.
-    
     The docstring MUST:
     1. Start with a clear, concise summary of what the function does
     2. ALWAYS include Args section with description of each parameter (or 'None' if no parameters)
-    3. Include Returns section describing the return value
-    4. Be formatted according to Google Python Style Guide
-    
+    3. Include Returns section describing the return value (or 'None' if nothing is explicitly returned)
+    4. **Optionally include a Raises section if the function might raise exceptions, describing the exception type/reason and when it's raised.**
+    5. **Include a Tags section with a list of strings describing the function's purpose, characteristics, or keywords.** Tags should be lowercase and single words or hyphenated phrases. Include tags like:
+        - The main action (e.g., 'scrape', 'search', 'start', 'check', 'cancel', 'list')
+        - The type of job ('async_job', 'batch')
+        - The stage of an asynchronous job ('start', 'status', 'cancel')
+        - Related domain/feature ('ai', 'management')
+        - **Significance: Add the tag 'important' to functions that represent core capabilities or primary interaction points of the class (e.g., initiating actions like scrape, search, or starting async jobs).**
+    6. Be formatted according to Google Python Style Guide
+
     Here is the function:
-    
+
     {function_code}
-    
-    Respond in JSON format with the following structure:
+
+    Respond ONLY in JSON format with the following structure. **Include the 'raises' field only if the function is likely to raise exceptions.** **Include the 'tags' field as a list of strings.**
     {{
-      "summary": "A clear, concise summary of what the function does",
-      "args": {{"param_name": "param description", "param_name2": "param description"}},
-      "returns": "Description of what the function returns"
+    "summary": "A clear, concise summary of what the function does",
+    "args": {{"param_name": "param description", "param_name2": "param description"}},
+    "returns": "Description of what the function returns",
+    "raises": {{
+        "ExceptionType": "Description of when/why this exception is raised"
+    }},
+    "tags": ["tag1", "tag2", "important"]
     }}
     """
 
     try:
-        # Use regular completion and parse the JSON ourselves instead of using response_model
         response = litellm.completion(
             model=model,
             messages=[
@@ -162,149 +233,283 @@ def generate_docstring(
             ],
         )
 
-        # Get the response content
         response_text = response.choices[0].message.content
 
-        # Simple JSON extraction in case the model includes extra text
-        import json
-        import re
-
-        # Find JSON object in the response using regex
-        json_match = re.search(r"({.*})", response_text.replace("\n", " "), re.DOTALL)
-        if json_match:
-            json_str = json_match.group(1)
-            parsed_data = json.loads(json_str)
-        else:
-            # Try to parse the whole response as JSON
-            parsed_data = json.loads(response_text)
-
-        # Ensure args is never empty
-        if not parsed_data.get("args"):
+        try:
+            parsed_data = extract_json_from_text(response_text)
+        except ValueError as e:
+            print(f"JSON extraction failed: {e}")
+            print(
+                f"Raw response: {response_text[:100]}..."
+            )  # Log first 100 chars for debugging
+            # Return a default structure if extraction fails
+            return DocstringOutput(
+                summary="Failed to extract docstring information",
+                args={"None": "This function takes no arguments"},
+                returns="Unknown return value",
+            )
+        model_args = parsed_data.get("args")
+        if not model_args:
             parsed_data["args"] = {"None": "This function takes no arguments"}
 
-        # Create DocstringOutput from parsed data
         return DocstringOutput(
-            summary=parsed_data.get("summary", ""),
+            summary=parsed_data.get("summary", "No documentation available"),
             args=parsed_data.get("args", {"None": "This function takes no arguments"}),
-            returns=parsed_data.get("returns", ""),
+            returns=parsed_data.get("returns", "None"),
+            raises=parsed_data.get("raises", {}),
+            tags=parsed_data.get("tags", []),  # Get tags, default to empty list
         )
 
     except Exception as e:
-        print(f"Error generating docstring: {e}")
-        # Return a docstring object with default values
+        print(f"Error generating docstring: {e}", file=sys.stderr)
+        traceback.print_exc(file=sys.stderr)
         return DocstringOutput(
-            summary="No documentation available",
+            summary=f"Error generating docstring: {e}",
             args={"None": "This function takes no arguments"},
             returns="None",
+            raises={},
+            tags=["generation-error"],
         )
 
 
 def format_docstring(docstring: DocstringOutput) -> str:
     """
-    Format a DocstringOutput object into a properly formatted docstring string.
+    Format a DocstringOutput object into the content string for a docstring.
+    This function produces the content *between* the triple quotes, without
+    the leading/trailing triple quotes or the main indentation.
 
     Args:
         docstring: The DocstringOutput object to format
 
     Returns:
-        A formatted docstring string ready to be inserted into code
+        A formatted docstring content string ready to be indented and wrapped
+        in triple quotes for insertion into code.
     """
-    formatted_docstring = f"{docstring.summary}\n\n"
-
-    if docstring.args:
-        formatted_docstring += "Args:\n"
-        for arg_name, arg_desc in docstring.args.items():
-            formatted_docstring += f"    {arg_name}: {arg_desc}\n"
-        formatted_docstring += "\n"
-
-    if docstring.returns:
-        formatted_docstring += f"Returns:\n    {docstring.returns}\n"
-
-    return formatted_docstring.strip()
+    parts = []
+
+    summary = docstring.summary.strip()
+    if summary:
+        parts.append(summary)
+
+    filtered_args = {
+        name: desc
+        for name, desc in docstring.args.items()
+        if name not in ("self", "cls")
+    }
+    args_lines = []
+    if filtered_args:
+        args_lines.append("Args:")
+        for arg_name, arg_desc in filtered_args.items():
+            arg_desc_cleaned = arg_desc.strip()
+            args_lines.append(f"    {arg_name}: {arg_desc_cleaned}")
+    elif docstring.args.get(
+        "None"
+    ):  # Include the 'None' placeholder if it was generated
+        args_lines.append("Args:")
+        none_desc_cleaned = docstring.args["None"].strip()
+        args_lines.append(f"    None: {none_desc_cleaned}")
+
+    if args_lines:
+        parts.append("\n".join(args_lines))
+
+    returns_desc_cleaned = docstring.returns.strip()
+    if returns_desc_cleaned and returns_desc_cleaned.lower() not in ("none", ""):
+        parts.append(f"Returns:\n    {returns_desc_cleaned}")
+
+    raises_lines = []
+    if docstring.raises:
+        raises_lines.append("Raises:")
+        for exception_type, exception_desc in docstring.raises.items():
+            exception_desc_cleaned = exception_desc.strip()
+            if (
+                exception_type.strip() and exception_desc_cleaned
+            ):  # Ensure type and desc are not empty
+                raises_lines.append(
+                    f"    {exception_type.strip()}: {exception_desc_cleaned}"
+                )
+    if raises_lines:
+        parts.append("\n".join(raises_lines))
+
+    cleaned_tags = [tag.strip() for tag in docstring.tags if tag and tag.strip()]
+    if cleaned_tags:
+        tags_string = ", ".join(cleaned_tags)
+        parts.append(f"Tags:\n    {tags_string}")
+
+    return "\n\n".join(parts)
 
 
 def insert_docstring_into_function(function_code: str, docstring: str) -> str:
     """
-    Insert a docstring into a function's code.
+    Insert a docstring into a function's code, replacing an existing one if present
+    at the correct location, and attempting to remove misplaced string literals
+    from the body.
+
+    This version handles multiline function definitions and existing docstrings
+    by carefully splicing lines based on AST node positions. It also tries to
+    clean up old, misplaced string literals that might have been interpreted
+    as docstrings previously.
 
     Args:
-        function_code: The source code of the function
-        docstring: The formatted docstring string to insert
+        function_code: The source code of the function snippet. This snippet is
+                       expected to contain exactly one function definition.
+        docstring: The formatted docstring string content (without triple quotes or
+                   leading/trailing newlines within the content itself).
 
     Returns:
-        The updated function code with the docstring inserted
+        The updated function code with the docstring inserted, or the original
+        code if an error occurs during processing or parsing.
     """
     try:
-        function_ast = ast.parse(function_code)
-        if not function_ast.body or not hasattr(function_ast.body[0], "body"):
-            return function_code
+        lines = function_code.splitlines(keepends=True)
 
-        function_lines = function_code.splitlines()
+        tree = ast.parse(function_code)
+        if not tree.body or not isinstance(
+            tree.body[0], ast.FunctionDef | ast.AsyncFunctionDef
+        ):
+            print(
+                "Warning: Could not parse function definition from code snippet. Returning original code.",
+                file=sys.stderr,
+            )
+            return function_code  # Return original code if parsing fails or isn't a function
 
-        # Find the function definition line (ends with ':')
-        func_def_line = None
-        for i, line in enumerate(function_lines):
-            if "def " in line and line.strip().endswith(":"):
-                func_def_line = i
-                break
+        func_node = tree.body[0]
+        func_name = getattr(func_node, "name", "unknown_function")
+
+        insert_idx = func_node.end_lineno
 
-        if func_def_line is None:
-            return function_code
+        if func_node.body:
+            insert_idx = func_node.body[0].lineno - 1
 
-        # Determine indentation from the first non-empty line after the function definition
-        body_indent = ""
-        for line in function_lines[func_def_line + 1 :]:
-            if line.strip():
-                body_indent = " " * (len(line) - len(line.lstrip()))
+        body_indent = "    "  # Default indentation (PEP 8)
+
+        indent_source_idx = insert_idx
+        actual_first_body_line_idx = -1
+        for i in range(indent_source_idx, len(lines)):
+            line = lines[i]
+            stripped = line.lstrip()
+            if stripped and not stripped.startswith("#"):
+                actual_first_body_line_idx = i
                 break
 
-        # Check if the function already has a docstring
-        first_element = (
-            function_ast.body[0].body[0] if function_ast.body[0].body else None
-        )
-        has_docstring = (
-            isinstance(first_element, ast.Expr)
-            and isinstance(first_element.value, ast.Constant)
-            and isinstance(first_element.value.value, str)
+        # If a meaningful line was found at or after insertion point, use its indentation
+        if actual_first_body_line_idx != -1:
+            body_line = lines[actual_first_body_line_idx]
+            body_indent = body_line[: len(body_line) - len(body_line.lstrip())]
+        else:
+            if func_node.lineno - 1 < len(lines):  # Ensure def line exists
+                def_line = lines[func_node.lineno - 1]
+                def_line_indent = def_line[: len(def_line) - len(def_line.lstrip())]
+                body_indent = (
+                    def_line_indent + "    "
+                )  # Standard 4 spaces relative indent
+
+        # Format the new docstring lines with the calculated indentation
+        new_docstring_lines_formatted = [f'{body_indent}"""\n']
+        new_docstring_lines_formatted.extend(
+            [f"{body_indent}{line}\n" for line in docstring.splitlines()]
         )
-
-        docstring_lines = [
-            f'{body_indent}"""',
-            *[f"{body_indent}{line}" for line in docstring.split("\n")],
-            f'{body_indent}"""',
-        ]
-
-        if has_docstring:
-            # Find the existing docstring in the source and replace it
-            for i in range(func_def_line + 1, len(function_lines)):
-                if '"""' in function_lines[i] or "'''" in function_lines[i]:
-                    docstring_start = i
-                    # Find end of docstring
-                    for j in range(docstring_start + 1, len(function_lines)):
-                        if '"""' in function_lines[j] or "'''" in function_lines[j]:
-                            docstring_end = j
-                            # Replace the existing docstring
-                            return "\n".join(
-                                function_lines[:docstring_start]
-                                + docstring_lines
-                                + function_lines[docstring_end + 1 :]
-                            )
+        new_docstring_lines_formatted.append(f'{body_indent}"""\n')
+
+        output_lines = []
+        output_lines.extend(lines[:insert_idx])
+
+        # 2. Insert the new docstring
+        output_lines.extend(new_docstring_lines_formatted)
+        remaining_body_lines = lines[insert_idx:]
+
+        remaining_body_code = "".join(remaining_body_lines)
+
+        if remaining_body_code.strip():  # Only parse if there's non-whitespace content
+            try:
+                dummy_code = f"def _dummy_func():\n{textwrap.indent(remaining_body_code, body_indent)}"
+                dummy_tree = ast.parse(dummy_code)
+                dummy_body_statements = (
+                    dummy_tree.body[0].body
+                    if dummy_tree.body
+                    and isinstance(
+                        dummy_tree.body[0], ast.FunctionDef | ast.AsyncFunctionDef
+                    )
+                    else []
+                )
+                cleaned_body_parts = []
+                for _node in dummy_body_statements:
+                    break  # Exit this loop, we'll process func_node.body instead
+                cleaned_body_parts = []
+                start_stmt_index = (
+                    1
+                    if func_node.body
+                    and isinstance(func_node.body[0], ast.Expr)
+                    and isinstance(func_node.body[0].value, ast.Constant)
+                    and isinstance(func_node.body[0].value.value, str)
+                    else 0
+                )
+
+                for i in range(start_stmt_index, len(func_node.body)):
+                    stmt_node = func_node.body[i]
+
+                    is_just_string_stmt = (
+                        isinstance(stmt_node, ast.Expr)
+                        and isinstance(stmt_node.value, ast.Constant)
+                        and isinstance(stmt_node.value.value, str)
+                    )
+
+                    if not is_just_string_stmt:
+                        stmt_start_idx = stmt_node.lineno - 1
+                        stmt_end_idx = (
+                            stmt_node.end_lineno - 1
+                        )  # Inclusive end line index
+
+                        cleaned_body_parts.extend(
+                            lines[stmt_start_idx : stmt_end_idx + 1]
+                        )
+
+                if func_node.body:
+                    last_stmt_end_idx = func_node.body[-1].end_lineno - 1
+                    for line in lines[last_stmt_end_idx + 1 :]:
+                        if line.strip():
+                            cleaned_body_parts.append(line)
+                cleaned_body_lines = cleaned_body_parts
+
+            except SyntaxError as parse_e:
+                print(
+                    f"WARNING: Could not parse function body for cleaning, keeping all body lines: {parse_e}",
+                    file=sys.stderr,
+                )
+                traceback.print_exc(file=sys.stderr)
+                cleaned_body_lines = remaining_body_lines
+            except Exception as other_e:
+                print(
+                    f"WARNING: Unexpected error processing function body for cleaning, keeping all body lines: {other_e}",
+                    file=sys.stderr,
+                )
+                traceback.print_exc(file=sys.stderr)
+                cleaned_body_lines = remaining_body_lines
         else:
-            # Insert new docstring after function definition
-            return "\n".join(
-                function_lines[: func_def_line + 1]
-                + docstring_lines
-                + function_lines[func_def_line + 1 :]
-            )
+            cleaned_body_lines = []
+            output_lines.extend(lines[func_node.end_lineno :])
+
+        if func_node.body or not remaining_body_code.strip():
+            output_lines.extend(cleaned_body_lines)
 
-        # Default return if insertion logic fails
+        final_code = "".join(output_lines)
+        ast.parse(final_code)
+        return final_code
+
+    except SyntaxError as e:
+        print(
+            f"WARNING: Generated code snippet for '{func_name}' has syntax error: {e}",
+            file=sys.stderr,
+        )
+        traceback.print_exc(file=sys.stderr)
         return function_code
     except Exception as e:
-        print(f"Error inserting docstring: {e}")
+        print(f"Error processing function snippet for insertion: {e}", file=sys.stderr)
+        traceback.print_exc(file=sys.stderr)
+
         return function_code
 
 
-def process_file(file_path: str, model: str = "openai/gpt-4o") -> int:
+def process_file(file_path: str, model: str = "perplexity/sonar-pro") -> int:
     """
     Process a Python file and add docstrings to all functions in it.
 
@@ -355,6 +560,7 @@ def process_file(file_path: str, model: str = "openai/gpt-4o") -> int:
             f.write(updated_content)
         print(f"Updated {count} functions in {file_path}")
     else:
+        print(updated_function, "formatted docstring", formatted_docstring)
         print(f"No changes made to {file_path}")
 
     return count