Merge branch 'main' of github.com:codeflash-ai/codeflash into lsp/verbose-quiet-logs

Author: mohammedahmed18

codeflash/api/aiservice.py

@@ -349,6 +349,51 @@ def get_new_explanation(  # noqa: D417
         console.rule()
         return ""
 
+    def generate_ranking(  # noqa: D417
+        self, trace_id: str, diffs: list[str], optimization_ids: list[str], speedups: list[float]
+    ) -> list[int] | None:
+        """Optimize the given python code for performance by making a request to the Django endpoint.
+
+        Parameters
+        ----------
+        - trace_id : unique uuid of function
+        - diffs : list of unified diff strings of opt candidates
+        - speedups : list of speedups of opt candidates
+
+        Returns
+        -------
+        - List[int]: Ranking of opt candidates in decreasing order
+
+        """
+        payload = {
+            "trace_id": trace_id,
+            "diffs": diffs,
+            "speedups": speedups,
+            "optimization_ids": optimization_ids,
+            "python_version": platform.python_version(),
+        }
+        logger.info("Generating ranking")
+        console.rule()
+        try:
+            response = self.make_ai_service_request("/rank", payload=payload, timeout=60)
+        except requests.exceptions.RequestException as e:
+            logger.exception(f"Error generating ranking: {e}")
+            ph("cli-optimize-error-caught", {"error": str(e)})
+            return None
+
+        if response.status_code == 200:
+            ranking: list[int] = response.json()["ranking"]
+            console.rule()
+            return ranking
+        try:
+            error = response.json()["error"]
+        except Exception:
+            error = response.text
+        logger.error(f"Error generating ranking: {response.status_code} - {error}")
+        ph("cli-optimize-error-response", {"response_status_code": response.status_code, "error": error})
+        console.rule()
+        return None
+
     def log_results(  # noqa: D417
         self,
         function_trace_id: str,

codeflash/code_utils/code_utils.py

@@ -20,6 +20,23 @@
 ImportErrorPattern = re.compile(r"ModuleNotFoundError.*$", re.MULTILINE)
 
 
+def unified_diff_strings(code1: str, code2: str, fromfile: str = "original", tofile: str = "modified") -> str:
+    """Return the unified diff between two code strings as a single string.
+
+    :param code1: First code string (original).
+    :param code2: Second code string (modified).
+    :param fromfile: Label for the first code string.
+    :param tofile: Label for the second code string.
+    :return: Unified diff as a string.
+    """
+    code1_lines = code1.splitlines(keepends=True)
+    code2_lines = code2.splitlines(keepends=True)
+
+    diff = difflib.unified_diff(code1_lines, code2_lines, fromfile=fromfile, tofile=tofile, lineterm="")
+
+    return "".join(diff)
+
+
 def diff_length(a: str, b: str) -> int:
     """Compute the length (in characters) of the unified diff between two strings.
 

codeflash/code_utils/deduplicate_code.py

@@ -0,0 +1,247 @@
+import ast
+import hashlib
+from typing import Dict, Set
+
+
+class VariableNormalizer(ast.NodeTransformer):
+    """Normalizes only local variable names in AST to canonical forms like var_0, var_1, etc.
+    Preserves function names, class names, parameters, built-ins, and imported names.
+    """
+
+    def __init__(self):
+        self.var_counter = 0
+        self.var_mapping: Dict[str, str] = {}
+        self.scope_stack = []
+        self.builtins = set(dir(__builtins__))
+        self.imports: Set[str] = set()
+        self.global_vars: Set[str] = set()
+        self.nonlocal_vars: Set[str] = set()
+        self.parameters: Set[str] = set()  # Track function parameters
+
+    def enter_scope(self):
+        """Enter a new scope (function/class)"""
+        self.scope_stack.append(
+            {"var_mapping": dict(self.var_mapping), "var_counter": self.var_counter, "parameters": set(self.parameters)}
+        )
+
+    def exit_scope(self):
+        """Exit current scope and restore parent scope"""
+        if self.scope_stack:
+            scope = self.scope_stack.pop()
+            self.var_mapping = scope["var_mapping"]
+            self.var_counter = scope["var_counter"]
+            self.parameters = scope["parameters"]
+
+    def get_normalized_name(self, name: str) -> str:
+        """Get or create normalized name for a variable"""
+        # Don't normalize if it's a builtin, import, global, nonlocal, or parameter
+        if (
+            name in self.builtins
+            or name in self.imports
+            or name in self.global_vars
+            or name in self.nonlocal_vars
+            or name in self.parameters
+        ):
+            return name
+
+        # Only normalize local variables
+        if name not in self.var_mapping:
+            self.var_mapping[name] = f"var_{self.var_counter}"
+            self.var_counter += 1
+        return self.var_mapping[name]
+
+    def visit_Import(self, node):
+        """Track imported names"""
+        for alias in node.names:
+            name = alias.asname if alias.asname else alias.name
+            self.imports.add(name.split(".")[0])
+        return node
+
+    def visit_ImportFrom(self, node):
+        """Track imported names from modules"""
+        for alias in node.names:
+            name = alias.asname if alias.asname else alias.name
+            self.imports.add(name)
+        return node
+
+    def visit_Global(self, node):
+        """Track global variable declarations"""
+        # Avoid repeated .add calls by using set.update with list
+        self.global_vars.update(node.names)
+        return node
+
+    def visit_Nonlocal(self, node):
+        """Track nonlocal variable declarations"""
+        for name in node.names:
+            self.nonlocal_vars.add(name)
+        return node
+
+    def visit_FunctionDef(self, node):
+        """Process function but keep function name and parameters unchanged"""
+        self.enter_scope()
+
+        # Track all parameters (don't modify them)
+        for arg in node.args.args:
+            self.parameters.add(arg.arg)
+        if node.args.vararg:
+            self.parameters.add(node.args.vararg.arg)
+        if node.args.kwarg:
+            self.parameters.add(node.args.kwarg.arg)
+        for arg in node.args.kwonlyargs:
+            self.parameters.add(arg.arg)
+
+        # Visit function body
+        node = self.generic_visit(node)
+        self.exit_scope()
+        return node
+
+    def visit_AsyncFunctionDef(self, node):
+        """Handle async functions same as regular functions"""
+        return self.visit_FunctionDef(node)
+
+    def visit_ClassDef(self, node):
+        """Process class but keep class name unchanged"""
+        self.enter_scope()
+        node = self.generic_visit(node)
+        self.exit_scope()
+        return node
+
+    def visit_Name(self, node):
+        """Normalize variable names in Name nodes"""
+        if isinstance(node.ctx, (ast.Store, ast.Del)):
+            # For assignments and deletions, check if we should normalize
+            if (
+                node.id not in self.builtins
+                and node.id not in self.imports
+                and node.id not in self.parameters
+                and node.id not in self.global_vars
+                and node.id not in self.nonlocal_vars
+            ):
+                node.id = self.get_normalized_name(node.id)
+        elif isinstance(node.ctx, ast.Load):
+            # For loading, use existing mapping if available
+            if node.id in self.var_mapping:
+                node.id = self.var_mapping[node.id]
+        return node
+
+    def visit_ExceptHandler(self, node):
+        """Normalize exception variable names"""
+        if node.name:
+            node.name = self.get_normalized_name(node.name)
+        return self.generic_visit(node)
+
+    def visit_comprehension(self, node):
+        """Normalize comprehension target variables"""
+        # Create new scope for comprehension
+        old_mapping = dict(self.var_mapping)
+        old_counter = self.var_counter
+
+        # Process the comprehension
+        node = self.generic_visit(node)
+
+        # Restore scope
+        self.var_mapping = old_mapping
+        self.var_counter = old_counter
+        return node
+
+    def visit_For(self, node):
+        """Handle for loop target variables"""
+        # The target in a for loop is a local variable that should be normalized
+        return self.generic_visit(node)
+
+    def visit_With(self, node):
+        """Handle with statement as variables"""
+        return self.generic_visit(node)
+
+
+def normalize_code(code: str, remove_docstrings: bool = True, return_ast_dump: bool = False) -> str:
+    """Normalize Python code by parsing, cleaning, and normalizing only variable names.
+    Function names, class names, and parameters are preserved.
+
+    Args:
+        code: Python source code as string
+        remove_docstrings: Whether to remove docstrings
+
+    Returns:
+        Normalized code as string
+
+    """
+    try:
+        # Parse the code
+        tree = ast.parse(code)
+
+        # Remove docstrings if requested
+        if remove_docstrings:
+            remove_docstrings_from_ast(tree)
+
+        # Normalize variable names
+        normalizer = VariableNormalizer()
+        normalized_tree = normalizer.visit(tree)
+        if return_ast_dump:
+            # This is faster than unparsing etc
+            return ast.dump(normalized_tree, annotate_fields=False, include_attributes=False)
+
+        # Fix missing locations in the AST
+        ast.fix_missing_locations(normalized_tree)
+
+        # Unparse back to code
+        return ast.unparse(normalized_tree)
+    except SyntaxError as e:
+        msg = f"Invalid Python syntax: {e}"
+        raise ValueError(msg) from e
+
+
+def remove_docstrings_from_ast(node):
+    """Remove docstrings from AST nodes."""
+    # Only FunctionDef, AsyncFunctionDef, ClassDef, and Module can contain docstrings in their body[0]
+    node_types = (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Module)
+    # Use our own stack-based DFS instead of ast.walk for efficiency
+    stack = [node]
+    while stack:
+        current_node = stack.pop()
+        if isinstance(current_node, node_types):
+            # Remove docstring if it's the first stmt in body
+            body = current_node.body
+            if (
+                body
+                and isinstance(body[0], ast.Expr)
+                and isinstance(body[0].value, ast.Constant)
+                and isinstance(body[0].value.value, str)
+            ):
+                current_node.body = body[1:]
+            # Only these nodes can nest more docstring-containing nodes
+            # Add their body elements to stack, avoiding unnecessary traversal
+            stack.extend([child for child in body if isinstance(child, node_types)])
+
+
+def get_code_fingerprint(code: str) -> str:
+    """Generate a fingerprint for normalized code.
+
+    Args:
+        code: Python source code
+
+    Returns:
+        SHA-256 hash of normalized code
+
+    """
+    normalized = normalize_code(code)
+    return hashlib.sha256(normalized.encode()).hexdigest()
+
+
+def are_codes_duplicate(code1: str, code2: str) -> bool:
+    """Check if two code segments are duplicates after normalization.
+
+    Args:
+        code1: First code segment
+        code2: Second code segment
+
+    Returns:
+        True if codes are structurally identical (ignoring local variable names)
+
+    """
+    try:
+        normalized1 = normalize_code(code1, return_ast_dump=True)
+        normalized2 = normalize_code(code2, return_ast_dump=True)
+        return normalized1 == normalized2
+    except Exception:
+        return False