Merge pull request #1054 from codeflash-ai/let's-see-what-we're-sending

display RO / RW context (--verbose)

Author: KRRT7

codeflash/code_utils/config_consts.py

@@ -4,6 +4,8 @@
 from typing import Any, Union
 
 MAX_TEST_RUN_ITERATIONS = 5
+OPTIMIZATION_CONTEXT_TOKEN_LIMIT = 16000
+TESTGEN_CONTEXT_TOKEN_LIMIT = 16000
 INDIVIDUAL_TESTCASE_TIMEOUT = 15
 MAX_FUNCTION_TEST_SECONDS = 60
 MIN_IMPROVEMENT_THRESHOLD = 0.05

codeflash/context/code_context_extractor.py

@@ -12,6 +12,7 @@
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_extractor import add_needed_imports_from_module, find_preexisting_objects
 from codeflash.code_utils.code_utils import encoded_tokens_len, get_qualified_name, path_belongs_to_site_packages
+from codeflash.code_utils.config_consts import OPTIMIZATION_CONTEXT_TOKEN_LIMIT, TESTGEN_CONTEXT_TOKEN_LIMIT
 from codeflash.context.unused_definition_remover import (
     collect_top_level_defs_with_usages,
     extract_names_from_targets,
@@ -39,8 +40,8 @@
 def get_code_optimization_context(
     function_to_optimize: FunctionToOptimize,
     project_root_path: Path,
-    optim_token_limit: int = 16000,
-    testgen_token_limit: int = 16000,
+    optim_token_limit: int = OPTIMIZATION_CONTEXT_TOKEN_LIMIT,
+    testgen_token_limit: int = TESTGEN_CONTEXT_TOKEN_LIMIT,
 ) -> CodeOptimizationContext:
     # Get FunctionSource representation of helpers of FTO
     helpers_of_fto_dict, helpers_of_fto_list = get_function_sources_from_jedi(

codeflash/discovery/discover_unit_tests.py

@@ -508,7 +508,7 @@ def analyze_imports_in_test_file(test_file_path: Path | str, target_functions: s
         return True
 
     if analyzer.found_any_target_function:
-        logger.debug(f"Test file {test_file_path} imports target function: {analyzer.found_qualified_name}")
+        # logger.debug(f"Test file {test_file_path} imports target function: {analyzer.found_qualified_name}")
         return True
 
     # Be conservative with dynamic imports - if __import__ is used and a target function
@@ -517,10 +517,10 @@ def analyze_imports_in_test_file(test_file_path: Path | str, target_functions: s
         # Check if any target function name appears as a string literal or direct usage
         for target_func in target_functions:
             if target_func in source_code:
-                logger.debug(f"Test file {test_file_path} has dynamic imports and references {target_func}")
+                # logger.debug(f"Test file {test_file_path} has dynamic imports and references {target_func}")
                 return True
 
-    logger.debug(f"Test file {test_file_path} does not import any target functions.")
+    # logger.debug(f"Test file {test_file_path} does not import any target functions.")
     return False
 
 
@@ -540,7 +540,7 @@ def filter_test_files_by_imports(
     if not target_functions:
         return file_to_test_map
 
-    logger.debug(f"Target functions for import filtering: {target_functions}")
+    # logger.debug(f"Target functions for import filtering: {target_functions}")
 
     filtered_map = {}
     for test_file, test_functions in file_to_test_map.items():

codeflash/optimization/function_optimizer.py

@@ -2,6 +2,7 @@
 
 import ast
 import concurrent.futures
+import logging
 import os
 import queue
 import random
@@ -15,6 +16,7 @@
 from rich.console import Group
 from rich.panel import Panel
 from rich.syntax import Syntax
+from rich.text import Text
 from rich.tree import Tree
 
 from codeflash.api.aiservice import AiServiceClient, AIServiceRefinerRequest, LocalAiServiceClient
@@ -34,6 +36,7 @@
     create_rank_dictionary_compact,
     create_score_dictionary_from_metrics,
     diff_length,
+    encoded_tokens_len,
     extract_unique_errors,
     file_name_from_test_module_name,
     get_run_tmp_file,
@@ -46,6 +49,7 @@
     COVERAGE_THRESHOLD,
     INDIVIDUAL_TESTCASE_TIMEOUT,
     MIN_CORRECT_CANDIDATES,
+    OPTIMIZATION_CONTEXT_TOKEN_LIMIT,
     REFINED_CANDIDATE_RANKING_WEIGHTS,
     REPEAT_OPTIMIZATION_PROBABILITY,
     TOTAL_LOOPING_TIME_EFFECTIVE,
@@ -128,6 +132,69 @@
     from codeflash.verification.verification_utils import TestConfig
 
 
+def log_optimization_context(function_name: str, code_context: CodeOptimizationContext) -> None:
+    """Log optimization context details when in verbose mode using Rich formatting."""
+    if logger.getEffectiveLevel() > logging.DEBUG:
+        return
+
+    console.rule()
+    read_writable_tokens = encoded_tokens_len(code_context.read_writable_code.markdown)
+    read_only_tokens = (
+        encoded_tokens_len(code_context.read_only_context_code) if code_context.read_only_context_code else 0
+    )
+    total_tokens = read_writable_tokens + read_only_tokens
+    token_pct = min(total_tokens / OPTIMIZATION_CONTEXT_TOKEN_LIMIT, 1.0)
+
+    # Token bar color based on usage
+    bar_color = "green" if token_pct < 0.7 else "yellow" if token_pct < 0.9 else "red"
+
+    # Build compact info line
+    helper_names = [hf.qualified_name for hf in code_context.helper_functions]
+    helpers_str = f"[magenta]{', '.join(helper_names)}[/]" if helper_names else "[dim]none[/]"
+    read_writable_files = [str(cs.file_path) for cs in code_context.read_writable_code.code_strings]
+
+    # Create a tree view for the context
+    tree = Tree(f"[bold cyan]Context for {function_name}[/]")
+    tree.add(
+        Text.assemble(
+            ("Tokens: ", "dim"),
+            (f"{total_tokens:,}", "bold " + bar_color),
+            (f"/{OPTIMIZATION_CONTEXT_TOKEN_LIMIT:,} ", "dim"),
+            (f"({token_pct:.0%})", bar_color),
+            ("  [", "dim"),
+            (f"{read_writable_tokens:,}", "green"),
+            (" rw", "dim green"),
+            (" + ", "dim"),
+            (f"{read_only_tokens:,}", "yellow"),
+            (" ro", "dim yellow"),
+            ("]", "dim"),
+        )
+    )
+    tree.add(f"[dim]Helpers:[/] {helpers_str}")
+    files_branch = tree.add("[dim]Files:[/]")
+    for f in read_writable_files:
+        files_branch.add(f"[blue]{f}[/]")
+
+    console.print(tree)
+
+    console.print(
+        Panel(
+            Syntax(code_context.read_writable_code.markdown, "markdown", theme="monokai", word_wrap=True),
+            title="[green]Read-Writable Code[/]",
+            border_style="green",
+        )
+    )
+
+    if code_context.read_only_context_code:
+        console.print(
+            Panel(
+                Syntax(code_context.read_only_context_code, "markdown", theme="monokai", word_wrap=True),
+                title="[yellow]Read-Only Dependencies[/]",
+                border_style="yellow",
+            )
+        )
+
+
 class CandidateNode:
     __slots__ = ("candidate", "children", "parent")
 
@@ -409,6 +476,7 @@ def can_be_optimized(self) -> Result[tuple[bool, CodeOptimizationContext, dict[P
         if not is_successful(ctx_result):
             return Failure(ctx_result.failure())
         code_context: CodeOptimizationContext = ctx_result.unwrap()
+        log_optimization_context(self.function_to_optimize.function_name, code_context)
         original_helper_code: dict[Path, str] = {}
         helper_function_paths = {hf.file_path for hf in code_context.helper_functions}
         for helper_function_path in helper_function_paths: