refactoring

code_to_optimize/code_directories/circular_deps/constants.py

@@ -1,8 +1,2 @@
 DEFAULT_API_URL = "https://api.galileo.ai/"
 DEFAULT_APP_URL = "https://app.galileo.ai/"
-
-
-# function_names: GalileoApiClient.get_console_url
-# module_abs_path : /home/mohammed/Work/galileo-python/src/galileo/api_client.py
-# preexisting_objects: {('GalileoApiClient', ()), ('_set_destination', ()), ('get_console_url', (FunctionParent(name='GalileoApiClient', type='ClassDef'),))}
-# project_root_path: /home/mohammed/Work/galileo-python/src

codeflash/models/models.py

@@ -169,22 +169,19 @@ class CodeStringsMarkdown(BaseModel):
             ]
         )
 
-    def path_to_code_string(self) -> dict[str, str]:
-        return {str(code_string.file_path): code_string.code for code_string in self.code_strings}
-
     @staticmethod
-    def from_str_with_markers(code_with_markers: str) -> CodeStringsMarkdown:
+    def parse_splitter_markers(code_with_markers: str) -> dict[str, str]:
         pattern = rf"{LINE_SPLITTER_MARKER_PREFIX}([^\n]+)\n"
         matches = list(re.finditer(pattern, code_with_markers))
 
-        results = []
+        results = {}
         for i, match in enumerate(matches):
             start = match.end()
             end = matches[i + 1].start() if i + 1 < len(matches) else len(code_with_markers)
             file_path = match.group(1).strip()
             code = code_with_markers[start:end].lstrip("\n")
-            results.append(CodeString(file_path=file_path, code=code))
-        return CodeStringsMarkdown(code_strings=results)
+            results[file_path] = code
+        return results
 
 
 class CodeOptimizationContext(BaseModel):

codeflash/optimization/function_optimizer.py

@@ -621,11 +621,13 @@ class FunctionOptimizer:
         read_writable_functions_by_file_path[self.function_to_optimize.file_path].add(
             self.function_to_optimize.qualified_name
         )
-        file_to_code_context = CodeStringsMarkdown.from_str_with_markers(optimized_code).path_to_code_string()
-        logger.debug(f"Optimized code: {file_to_code_context}")
+
+        file_to_code_context = CodeStringsMarkdown.parse_splitter_markers(optimized_code)
+
         for helper_function in code_context.helper_functions:
             if helper_function.jedi_definition.type != "class":
                 read_writable_functions_by_file_path[helper_function.file_path].add(helper_function.qualified_name)
+
         for module_abspath, qualified_names in read_writable_functions_by_file_path.items():
             relative_module_path = str(module_abspath.relative_to(self.project_root))
             logger.debug(f"applying optimized code to: {relative_module_path}")
@@ -633,7 +635,8 @@ class FunctionOptimizer:
             scoped_optimized_code = file_to_code_context.get(relative_module_path, None)
             if scoped_optimized_code is None:
                 logger.warning(
-                    f"Optimized code not found for {relative_module_path}, existing files in the context are: {list(file_to_code_context.keys())}, re-check your 'split markers'"
+                    f"Optimized code not found for {relative_module_path} In the context\n-------\n{optimized_code}\n-------\n"
+                    "Existing files in the context are: {list(file_to_code_context.keys())}, re-check your 'split markers'"
                 )
                 scoped_optimized_code = ""
 

tests/test_code_context_extractor.py

@@ -2449,8 +2449,8 @@ class SimpleClass:
         assert "return 42" in code_content
 
 
-
-def test_replace_functions_and_add_imports():
+# This shouldn't happen as we are now using a scoped optimization context, but keep it just in case
+def test_circular_deps():
     path_to_root = Path(__file__).resolve().parent.parent / "code_to_optimize" / "code_directories" / "circular_deps"
     file_abs_path = path_to_root / "api_client.py"
     optimized_code = Path(path_to_root / "optimized.py").read_text(encoding="utf-8")

tests/test_multi_file_code_replacement.py

@@ -56,12 +56,7 @@ from collections.abc import Sequence
 
 from pydantic_ai_slim.pydantic_ai.messages import BinaryContent, UserContent
 
-# Compile regex once, as in original
 _TOKEN_SPLIT_RE = re.compile(r'[\\s",.:]+')
-
-# Precompute translation table for fast token splitting for string input
-# This covers the chars: whitespace (\\x09-\\x0d, space), " (0x22), , (0x2c),
-# Map those codepoints to ' '
 _translate_table = {{ord(c): ord(' ') for c in ' \\t\\n\\r\\x0b\\x0c",.:'}}
 
 def _estimate_string_tokens(content: str | Sequence[UserContent]) -> int: