Merge pull request #1660 from codeflash-ai/unstructured-inference

feat: improve function ranking with reference graph and test-based boosting

Author: KRRT7

codeflash/api/aiservice.py

@@ -14,7 +14,7 @@
 from codeflash.code_utils.env_utils import get_codeflash_api_key
 from codeflash.code_utils.git_utils import get_last_commit_author_if_pr_exists, get_repo_owner_and_name
 from codeflash.code_utils.time_utils import humanize_runtime
-from codeflash.languages import Language, current_language
+from codeflash.languages import Language, current_language, current_language_support
 from codeflash.models.ExperimentMetadata import ExperimentMetadata
 from codeflash.models.models import (
     AIServiceRefinerRequest,
@@ -58,6 +58,8 @@ def add_language_metadata(
         payload: dict[str, Any], language_version: str | None = None, module_system: str | None = None
     ) -> None:
         """Add language version and module system metadata to an API payload."""
+        if language_version is None:
+            language_version = current_language_support().language_version
         payload["language_version"] = language_version
         payload["python_version"] = language_version if current_language() == Language.PYTHON else None
 

codeflash/code_utils/config_consts.py

@@ -50,6 +50,8 @@
 
 MAX_CONTEXT_LEN_REVIEW = 1000
 
+HIGH_EFFORT_TOP_N = 15
+
 
 class EffortLevel(str, Enum):
     LOW = "low"

codeflash/discovery/discover_unit_tests.py

@@ -11,6 +11,7 @@
 import subprocess
 import unittest
 from collections import defaultdict
+from functools import lru_cache
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable, Optional, final
 
@@ -35,6 +36,21 @@
     from codeflash.verification.verification_utils import TestConfig
 
 
+def existing_unit_test_count(
+    func: FunctionToOptimize, project_root: Path, function_to_tests: dict[str, set[FunctionCalledInTest]]
+) -> int:
+    key = f"{module_name_from_file_path_cached(func.file_path, project_root)}.{func.qualified_name}"
+    tests = function_to_tests.get(key, set())
+    seen: set[tuple[Path, str | None, str]] = set()
+    for t in tests:
+        if t.tests_in_file.test_type != TestType.EXISTING_UNIT_TEST:
+            continue
+        tif = t.tests_in_file
+        base_name = tif.test_function.split("[", 1)[0]
+        seen.add((tif.test_file, tif.test_class, base_name))
+    return len(seen)
+
+
 @final
 class PytestExitCode(enum.IntEnum):  # don't need to import entire pytest just for this
     #: Tests passed.
@@ -1079,3 +1095,9 @@ def process_test_files(
     tests_cache.close()
 
     return dict(function_to_test_map), num_discovered_tests, num_discovered_replay_tests
+
+
+# Cache module name resolution to avoid repeated Path.resolve()/relative_to() calls
+@lru_cache(maxsize=128)
+def module_name_from_file_path_cached(file_path: Path, project_root: Path) -> str:
+    return module_name_from_file_path(file_path, project_root)

codeflash/languages/base.py

@@ -18,6 +18,7 @@
     from pathlib import Path
 
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+    from codeflash.models.call_graph import CallGraph
     from codeflash.models.models import FunctionSource, GeneratedTestsList, InvocationId, ValidCode
     from codeflash.verification.verification_utils import TestConfig
 
@@ -250,6 +251,12 @@ def count_callees_per_function(
         """Return the number of callees for each (file_path, qualified_name) pair."""
         ...
 
+    def get_call_graph(
+        self, file_path_to_qualified_names: dict[Path, set[str]], *, include_metadata: bool = False
+    ) -> CallGraph:
+        """Return a CallGraph with full caller→callee edges for the given functions."""
+        ...
+
     def close(self) -> None:
         """Release resources (e.g. database connections)."""
         ...

codeflash/languages/code_replacer.py

@@ -6,13 +6,15 @@
 
 from __future__ import annotations
 
-from pathlib import Path
+import os
 from typing import TYPE_CHECKING
 
 from codeflash.cli_cmds.console import logger
 from codeflash.languages.base import FunctionFilterCriteria, Language
 
 if TYPE_CHECKING:
+    from pathlib import Path
+
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
     from codeflash.languages.base import LanguageSupport
     from codeflash.models.models import CodeStringsMarkdown
@@ -38,7 +40,7 @@ def get_optimized_code_for_module(relative_path: Path, optimized_code: CodeStrin
     # directory prefix but the correct filename
     target_name = relative_path.name
     basename_matches = [
-        code for path, code in file_to_code_context.items() if path != "None" and Path(path).name == target_name
+        code for path, code in file_to_code_context.items() if path != "None" and os.path.basename(path) == target_name
     ]
     if len(basename_matches) == 1:
         logger.debug(f"Using basename-matched code block for {relative_path}")

codeflash/languages/function_optimizer.py

@@ -425,6 +425,7 @@ def __init__(
         args: Namespace | None = None,
         replay_tests_dir: Path | None = None,
         call_graph: DependencyResolver | None = None,
+        effort_override: str | None = None,
     ) -> None:
         self.project_root = test_cfg.project_root_path.resolve()
         self.test_cfg = test_cfg
@@ -451,7 +452,8 @@ def __init__(
         self.local_aiservice_client = LocalAiServiceClient() if self.experiment_id else None
         self.test_files = TestFiles(test_files=[])
 
-        self.effort = getattr(args, "effort", EffortLevel.MEDIUM.value) if args else EffortLevel.MEDIUM.value
+        default_effort = getattr(args, "effort", EffortLevel.MEDIUM.value) if args else EffortLevel.MEDIUM.value
+        self.effort = effort_override or default_effort
 
         self.args = args  # Check defaults for these
         self.function_trace_id: str = str(uuid.uuid4())