feat: boost ranking for tested functions and enable reference graph

- Add existing_unit_test_count() with parametrized test deduplication
- Stable-sort ranked functions so tested ones come first
- Enable reference graph resolver (was disabled) for non-CI runs
- Add per-function logging with ref count and test count
- Auto-upgrade top N functions to high effort when user hasn't set --effort
- Add CallGraph model with traversal (BFS, topological sort, subgraph)
- Add get_call_graph() to DependencyResolver protocol and ReferenceGraph
- Refactor get_callees() to delegate through get_call_graph()

CF-1660

Author: KRRT7

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

@@ -35,6 +35,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 = func.qualified_name_with_modules_from_root(project_root)
+    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.

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/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())

codeflash/languages/python/reference_graph.py

@@ -3,20 +3,21 @@
 import hashlib
 import os
 import sqlite3
-from collections import defaultdict
 from pathlib import Path
 from typing import TYPE_CHECKING
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_utils import get_qualified_name, path_belongs_to_site_packages
 from codeflash.languages.base import IndexResult
-from codeflash.models.models import FunctionSource
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterable
 
     from jedi.api.classes import Name
 
+    from codeflash.models.call_graph import CallGraph
+    from codeflash.models.models import FunctionSource
+
 
 # ---------------------------------------------------------------------------
 # Module-level helpers (must be top-level for ProcessPoolExecutor pickling)
@@ -262,49 +263,10 @@ def _init_schema(self) -> None:
     def get_callees(
         self, file_path_to_qualified_names: dict[Path, set[str]]
     ) -> tuple[dict[Path, set[FunctionSource]], list[FunctionSource]]:
-        file_path_to_function_source: dict[Path, set[FunctionSource]] = defaultdict(set)
-        function_source_list: list[FunctionSource] = []
-
-        all_caller_keys: list[tuple[str, str]] = []
-        for file_path, qualified_names in file_path_to_qualified_names.items():
-            resolved = str(file_path.resolve())
-            self.ensure_file_indexed(file_path, resolved)
-            all_caller_keys.extend((resolved, qn) for qn in qualified_names)
+        from codeflash.models.call_graph import callees_from_graph
 
-        if not all_caller_keys:
-            return file_path_to_function_source, function_source_list
-
-        cur = self.conn.cursor()
-        cur.execute("CREATE TEMP TABLE IF NOT EXISTS _caller_keys (caller_file TEXT, caller_qualified_name TEXT)")
-        cur.execute("DELETE FROM _caller_keys")
-        cur.executemany("INSERT INTO _caller_keys VALUES (?, ?)", all_caller_keys)
-
-        rows = cur.execute(
-            """
-            SELECT ce.callee_file, ce.callee_qualified_name, ce.callee_fully_qualified_name,
-                   ce.callee_only_function_name, ce.callee_definition_type, ce.callee_source_line
-            FROM call_edges ce
-            INNER JOIN _caller_keys ck
-                ON ce.caller_file = ck.caller_file AND ce.caller_qualified_name = ck.caller_qualified_name
-            WHERE ce.project_root = ? AND ce.language = ?
-            """,
-            (self.project_root_str, self.language),
-        ).fetchall()
-
-        for callee_file, callee_qn, callee_fqn, callee_name, callee_type, callee_src in rows:
-            callee_path = Path(callee_file)
-            fs = FunctionSource(
-                file_path=callee_path,
-                qualified_name=callee_qn,
-                fully_qualified_name=callee_fqn,
-                only_function_name=callee_name,
-                source_code=callee_src,
-                definition_type=callee_type,
-            )
-            file_path_to_function_source[callee_path].add(fs)
-            function_source_list.append(fs)
-
-        return file_path_to_function_source, function_source_list
+        graph = self.get_call_graph(file_path_to_qualified_names, include_metadata=True)
+        return callees_from_graph(graph)
 
     def count_callees_per_function(
         self, file_path_to_qualified_names: dict[Path, set[str]]
@@ -540,5 +502,90 @@ def _fallback_sequential_index(
             result = self.index_file(file_path, file_hash, resolved)
             self._report_progress(on_progress, result)
 
+    def get_call_graph(
+        self, file_path_to_qualified_names: dict[Path, set[str]], *, include_metadata: bool = False
+    ) -> CallGraph:
+        from codeflash.models.call_graph import CallEdge, CalleeMetadata, CallGraph, FunctionNode
+
+        all_caller_keys: list[tuple[Path, str, str]] = []
+        for file_path, qualified_names in file_path_to_qualified_names.items():
+            resolved = str(file_path.resolve())
+            self.ensure_file_indexed(file_path, resolved)
+            all_caller_keys.extend((file_path, resolved, qn) for qn in qualified_names)
+
+        if not all_caller_keys:
+            return CallGraph(edges=[])
+
+        cur = self.conn.cursor()
+        cur.execute("CREATE TEMP TABLE IF NOT EXISTS _graph_keys (caller_file TEXT, caller_qualified_name TEXT)")
+        cur.execute("DELETE FROM _graph_keys")
+        cur.executemany(
+            "INSERT INTO _graph_keys VALUES (?, ?)", [(resolved, qn) for _, resolved, qn in all_caller_keys]
+        )
+
+        if include_metadata:
+            rows = cur.execute(
+                """
+                SELECT ce.caller_file, ce.caller_qualified_name,
+                       ce.callee_file, ce.callee_qualified_name,
+                       ce.callee_fully_qualified_name, ce.callee_only_function_name,
+                       ce.callee_definition_type, ce.callee_source_line
+                FROM call_edges ce
+                INNER JOIN _graph_keys gk
+                    ON ce.caller_file = gk.caller_file AND ce.caller_qualified_name = gk.caller_qualified_name
+                WHERE ce.project_root = ? AND ce.language = ?
+                """,
+                (self.project_root_str, self.language),
+            ).fetchall()
+
+            edges: list[CallEdge] = []
+            for (
+                caller_file,
+                caller_qn,
+                callee_file,
+                callee_qn,
+                callee_fqn,
+                callee_name,
+                callee_type,
+                callee_src,
+            ) in rows:
+                edges.append(
+                    CallEdge(
+                        caller=FunctionNode(file_path=Path(caller_file), qualified_name=caller_qn),
+                        callee=FunctionNode(file_path=Path(callee_file), qualified_name=callee_qn),
+                        is_cross_file=caller_file != callee_file,
+                        callee_metadata=CalleeMetadata(
+                            fully_qualified_name=callee_fqn,
+                            only_function_name=callee_name,
+                            definition_type=callee_type,
+                            source_line=callee_src,
+                        ),
+                    )
+                )
+        else:
+            rows = cur.execute(
+                """
+                SELECT ce.caller_file, ce.caller_qualified_name,
+                       ce.callee_file, ce.callee_qualified_name
+                FROM call_edges ce
+                INNER JOIN _graph_keys gk
+                    ON ce.caller_file = gk.caller_file AND ce.caller_qualified_name = gk.caller_qualified_name
+                WHERE ce.project_root = ? AND ce.language = ?
+                """,
+                (self.project_root_str, self.language),
+            ).fetchall()
+
+            edges = []
+            for caller_file, caller_qn, callee_file, callee_qn in rows:
+                edges.append(
+                    CallEdge(
+                        caller=FunctionNode(file_path=Path(caller_file), qualified_name=caller_qn),
+                        callee=FunctionNode(file_path=Path(callee_file), qualified_name=callee_qn),
+                        is_cross_file=caller_file != callee_file,
+                    )
+                )
+
+        return CallGraph(edges=edges)
+
     def close(self) -> None:
         self.conn.close()

codeflash/languages/python/support.py

@@ -957,7 +957,7 @@ def create_dependency_resolver(self, project_root: Path) -> DependencyResolver |
         try:
             return ReferenceGraph(project_root, language=self.language.value)
         except Exception:
-            logger.debug("Failed to initialize ReferenceGraph, falling back to per-function Jedi analysis")
+            logger.info("Failed to initialize ReferenceGraph, falling back to per-function Jedi analysis")
             return None
 
     def instrument_existing_test(