From 04b1e30321731c7ec2f2dd3902ab400f1bfd9df1 Mon Sep 17 00:00:00 2001 From: aseembits93 Date: Mon, 16 Mar 2026 16:48:57 -0700 Subject: [PATCH 1/2] Improve context extraction and test discovery performance --- codeflash/discovery/discover_unit_tests.py | 147 +++++++++++------- .../python/context/code_context_extractor.py | 120 ++++++++++++-- tests/test_code_context_extractor.py | 51 ++++++ tests/test_unit_test_discovery.py | 65 ++++++++ 4 files changed, 311 insertions(+), 72 deletions(-) diff --git a/codeflash/discovery/discover_unit_tests.py b/codeflash/discovery/discover_unit_tests.py index fa1ebb16e..a184fd302 100644 --- a/codeflash/discovery/discover_unit_tests.py +++ b/codeflash/discovery/discover_unit_tests.py @@ -120,6 +120,8 @@ def __init__(self, project_root_path: Path) -> None: ) self.memory_cache = {} + self.pending_rows: list[tuple[str, str, str, str, str, str, int | str, int, int]] = [] + self.writes_enabled = True def insert_test( self, @@ -134,10 +136,8 @@ def insert_test( col_number: int, ) -> None: test_type_value = test_type.value if hasattr(test_type, "value") else test_type - self.cur.execute( - "INSERT INTO discovered_tests VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + self.pending_rows.append( ( - self.project_root_path, file_path, file_hash, qualified_name_with_modules_from_root, @@ -147,9 +147,47 @@ def insert_test( test_type_value, line_number, col_number, - ), + ) ) - self.connection.commit() + + def flush(self) -> None: + if not self.writes_enabled or not self.pending_rows: + return + try: + self.cur.executemany( + "INSERT INTO discovered_tests VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + [ + ( + self.project_root_path, + file_path, + file_hash, + qualified_name_with_modules_from_root, + function_name, + test_class, + test_function, + test_type_value, + line_number, + col_number, + ) + for ( + file_path, + file_hash, + qualified_name_with_modules_from_root, + function_name, + test_class, + test_function, + test_type_value, + line_number, + col_number, + ) in self.pending_rows + ], + ) + self.connection.commit() + except sqlite3.OperationalError as e: + logger.debug(f"Failed to persist discovered test cache, disabling cache writes: {e}") + self.writes_enabled = False + finally: + self.pending_rows.clear() def get_function_to_test_map_for_file( self, file_path: str, file_hash: str @@ -196,6 +234,7 @@ def compute_file_hash(path: Path) -> str: return h.hexdigest() def close(self) -> None: + self.flush() self.cur.close() self.connection.close() @@ -833,6 +872,10 @@ def process_test_files( function_to_test_map = defaultdict(set) num_discovered_tests = 0 num_discovered_replay_tests = 0 + functions_to_optimize_by_name: dict[str, list[FunctionToOptimize]] = defaultdict(list) + if functions_to_optimize: + for function_to_optimize in functions_to_optimize: + functions_to_optimize_by_name[function_to_optimize.function_name].append(function_to_optimize) # Set up sys_path for Jedi to resolve imports correctly import sys @@ -875,8 +918,8 @@ def process_test_files( test_functions = set() all_names = script.get_names(all_scopes=True, references=True) - all_defs = script.get_names(all_scopes=True, definitions=True) all_names_top = script.get_names(all_scopes=True) + all_defs = [name for name in all_names if name.is_definition()] top_level_functions = {name.name: name for name in all_names_top if name.type == "function"} top_level_classes = {name.name: name for name in all_names_top if name.type == "class"} @@ -951,10 +994,9 @@ def process_test_files( test_function_names_set = set(test_functions_by_name.keys()) relevant_names = [] - - names_with_full_name = [name for name in all_names if name.full_name is not None] - - for name in names_with_full_name: + for name in all_names: + if name.full_name is None: + continue match = FUNCTION_NAME_REGEX.search(name.full_name) if match and match.group(1) in test_function_names_set: relevant_names.append((name, match.group(1))) @@ -969,56 +1011,50 @@ def process_test_files( if not definition or definition[0].type != "function": # Fallback: Try to match against functions_to_optimize when Jedi can't resolve # This handles cases where Jedi fails with pytest fixtures - if functions_to_optimize and name.name: - for func_to_opt in functions_to_optimize: + if functions_to_optimize_by_name and name.name: + for func_to_opt in functions_to_optimize_by_name.get(name.name, []): # Check if this unresolved name matches a function we're looking for - if func_to_opt.function_name == name.name: - # Check if the test file imports the class/module containing this function - qualified_name_with_modules = func_to_opt.qualified_name_with_modules_from_root( - project_root_path - ) + qualified_name_with_modules = func_to_opt.qualified_name_with_modules_from_root( + project_root_path + ) - # Only add if this test actually tests the function we're optimizing - for test_func in test_functions_by_name[scope]: - if test_func.parameters is not None: - if test_framework == "pytest": - scope_test_function = ( - f"{test_func.function_name}[{test_func.parameters}]" - ) - else: # unittest - scope_test_function = ( - f"{test_func.function_name}_{test_func.parameters}" - ) - else: - scope_test_function = test_func.function_name - - function_to_test_map[qualified_name_with_modules].add( - FunctionCalledInTest( - tests_in_file=TestsInFile( - test_file=test_file, - test_class=test_func.test_class, - test_function=scope_test_function, - test_type=test_func.test_type, - ), - position=CodePosition(line_no=name.line, col_no=name.column), - ) - ) - tests_cache.insert_test( - file_path=str(test_file), - file_hash=file_hash, - qualified_name_with_modules_from_root=qualified_name_with_modules, - function_name=scope, - test_class=test_func.test_class or "", - test_function=scope_test_function, - test_type=test_func.test_type, - line_number=name.line, - col_number=name.column, + # Only add if this test actually tests the function we're optimizing + for test_func in test_functions_by_name[scope]: + if test_func.parameters is not None: + if test_framework == "pytest": + scope_test_function = f"{test_func.function_name}[{test_func.parameters}]" + else: # unittest + scope_test_function = f"{test_func.function_name}_{test_func.parameters}" + else: + scope_test_function = test_func.function_name + + function_to_test_map[qualified_name_with_modules].add( + FunctionCalledInTest( + tests_in_file=TestsInFile( + test_file=test_file, + test_class=test_func.test_class, + test_function=scope_test_function, + test_type=test_func.test_type, + ), + position=CodePosition(line_no=name.line, col_no=name.column), ) + ) + tests_cache.insert_test( + file_path=str(test_file), + file_hash=file_hash, + qualified_name_with_modules_from_root=qualified_name_with_modules, + function_name=scope, + test_class=test_func.test_class or "", + test_function=scope_test_function, + test_type=test_func.test_type, + line_number=name.line, + col_number=name.column, + ) - if test_func.test_type == TestType.REPLAY_TEST: - num_discovered_replay_tests += 1 + if test_func.test_type == TestType.REPLAY_TEST: + num_discovered_replay_tests += 1 - num_discovered_tests += 1 + num_discovered_tests += 1 continue definition_obj = definition[0] definition_path = str(definition_obj.module_path) @@ -1074,6 +1110,7 @@ def process_test_files( logger.debug(str(e)) continue + tests_cache.flush() progress.advance(task_id) tests_cache.close() diff --git a/codeflash/languages/python/context/code_context_extractor.py b/codeflash/languages/python/context/code_context_extractor.py index 83eb49bea..d1d4b2399 100644 --- a/codeflash/languages/python/context/code_context_extractor.py +++ b/codeflash/languages/python/context/code_context_extractor.py @@ -4,6 +4,8 @@ import hashlib import os from collections import defaultdict, deque +from dataclasses import dataclass, field +from functools import cache from itertools import chain from pathlib import Path from typing import TYPE_CHECKING @@ -45,6 +47,22 @@ from codeflash.languages.python.context.unused_definition_remover import UsageInfo +@dataclass +class ContextProcessingCache: + source_by_file: dict[Path, str] = field(default_factory=dict) + pruned_source_by_key: dict[tuple[Path, frozenset[str]], str] = field(default_factory=dict) + code_string_by_key: dict[tuple[Path, frozenset[str], frozenset[str], CodeContextType, bool], CodeString | None] = ( + field(default_factory=dict) + ) + + +@cache +def _get_jedi_project(project_root_path: str): # type: ignore[no-untyped-def] + import jedi + + return jedi.Project(path=project_root_path) + + def build_testgen_context( helpers_of_fto_dict: dict[Path, set[FunctionSource]], helpers_of_helpers_dict: dict[Path, set[FunctionSource]], @@ -53,6 +71,7 @@ def build_testgen_context( remove_docstrings: bool = False, include_enrichment: bool = True, function_to_optimize: FunctionToOptimize | None = None, + processing_cache: ContextProcessingCache | None = None, ) -> CodeStringsMarkdown: testgen_context = extract_code_markdown_context_from_files( helpers_of_fto_dict, @@ -60,6 +79,7 @@ def build_testgen_context( project_root_path, remove_docstrings=remove_docstrings, code_context_type=CodeContextType.TESTGEN, + processing_cache=processing_cache, ) if include_enrichment: @@ -88,15 +108,23 @@ def get_code_optimization_context( testgen_token_limit: int = TESTGEN_CONTEXT_TOKEN_LIMIT, call_graph: DependencyResolver | None = None, ) -> CodeOptimizationContext: + processing_cache = ContextProcessingCache() + project_root_path = project_root_path.resolve() + jedi_project = _get_jedi_project(str(project_root_path)) + # Get FunctionSource representation of helpers of FTO fto_input = {function_to_optimize.file_path: {function_to_optimize.qualified_name}} if call_graph is not None: helpers_of_fto_dict, helpers_of_fto_list = call_graph.get_callees(fto_input) else: - helpers_of_fto_dict, helpers_of_fto_list = get_function_sources_from_jedi(fto_input, project_root_path) + helpers_of_fto_dict, helpers_of_fto_list = get_function_sources_from_jedi( + fto_input, project_root_path, jedi_project=jedi_project + ) # Add function to optimize into helpers of FTO dict, as they'll be processed together - fto_as_function_source = get_function_to_optimize_as_function_source(function_to_optimize, project_root_path) + fto_as_function_source = get_function_to_optimize_as_function_source( + function_to_optimize, project_root_path, jedi_project=jedi_project + ) helpers_of_fto_dict[function_to_optimize.file_path].add(fto_as_function_source) # Format data to search for helpers of helpers using get_function_sources_from_jedi @@ -110,7 +138,7 @@ def get_code_optimization_context( qualified_names.update({f"{qn.rsplit('.', 1)[0]}.__init__" for qn in qualified_names if "." in qn}) helpers_of_helpers_dict, helpers_of_helpers_list = get_function_sources_from_jedi( - helpers_of_fto_qualified_names_dict, project_root_path + helpers_of_fto_qualified_names_dict, project_root_path, jedi_project=jedi_project ) # Extract code context for optimization @@ -120,6 +148,7 @@ def get_code_optimization_context( project_root_path, remove_docstrings=False, code_context_type=CodeContextType.READ_WRITABLE, + processing_cache=processing_cache, ) # Ensure the target file is first in the code blocks so the LLM knows which file to optimize @@ -135,6 +164,7 @@ def get_code_optimization_context( project_root_path, remove_docstrings=False, code_context_type=CodeContextType.READ_ONLY, + processing_cache=processing_cache, ) hashing_code_context = extract_code_markdown_context_from_files( helpers_of_fto_dict, @@ -142,6 +172,7 @@ def get_code_optimization_context( project_root_path, remove_docstrings=True, code_context_type=CodeContextType.HASHING, + processing_cache=processing_cache, ) # Handle token limits @@ -163,7 +194,11 @@ def get_code_optimization_context( if final_read_writable_tokens + read_only_tokens > optim_token_limit: logger.debug("Code context has exceeded token limit, removing docstrings from read-only code") read_only_code_no_docstrings = extract_code_markdown_context_from_files( - helpers_of_fto_dict, helpers_of_helpers_dict, project_root_path, remove_docstrings=True + helpers_of_fto_dict, + helpers_of_helpers_dict, + project_root_path, + remove_docstrings=True, + processing_cache=processing_cache, ) read_only_context_code = read_only_code_no_docstrings.markdown if final_read_writable_tokens + encoded_tokens_len(read_only_context_code) > optim_token_limit: @@ -172,7 +207,11 @@ def get_code_optimization_context( # Progressive fallback for testgen context token limits testgen_context = build_testgen_context( - helpers_of_fto_dict, helpers_of_helpers_dict, project_root_path, function_to_optimize=function_to_optimize + helpers_of_fto_dict, + helpers_of_helpers_dict, + project_root_path, + function_to_optimize=function_to_optimize, + processing_cache=processing_cache, ) if encoded_tokens_len(testgen_context.markdown) > testgen_token_limit: @@ -183,6 +222,7 @@ def get_code_optimization_context( project_root_path, remove_docstrings=True, function_to_optimize=function_to_optimize, + processing_cache=processing_cache, ) if encoded_tokens_len(testgen_context.markdown) > testgen_token_limit: @@ -193,6 +233,7 @@ def get_code_optimization_context( project_root_path, remove_docstrings=True, include_enrichment=False, + processing_cache=processing_cache, ) if encoded_tokens_len(testgen_context.markdown) > testgen_token_limit: @@ -222,16 +263,38 @@ def process_file_context( remove_docstrings: bool, project_root_path: Path, helper_functions: list[FunctionSource], + processing_cache: ContextProcessingCache | None = None, ) -> CodeString | None: + cache_key = ( + file_path, + frozenset(primary_qualified_names), + frozenset(secondary_qualified_names), + code_context_type, + remove_docstrings, + ) + if processing_cache is not None and cache_key in processing_cache.code_string_by_key: + return processing_cache.code_string_by_key[cache_key] + try: - original_code = file_path.read_text("utf8") + if processing_cache is not None and file_path in processing_cache.source_by_file: + original_code = processing_cache.source_by_file[file_path] + else: + original_code = file_path.read_text("utf8") + if processing_cache is not None: + processing_cache.source_by_file[file_path] = original_code except Exception as e: logger.exception(f"Error while parsing {file_path}: {e}") return None try: all_names = primary_qualified_names | secondary_qualified_names - code_without_unused_defs = remove_unused_definitions_by_function_names(original_code, all_names) + pruned_source_key = (file_path, frozenset(all_names)) + if processing_cache is not None and pruned_source_key in processing_cache.pruned_source_by_key: + code_without_unused_defs = processing_cache.pruned_source_by_key[pruned_source_key] + else: + code_without_unused_defs = remove_unused_definitions_by_function_names(original_code, all_names) + if processing_cache is not None: + processing_cache.pruned_source_by_key[pruned_source_key] = code_without_unused_defs pruned_module = parse_code_and_prune_cst( code_without_unused_defs, code_context_type, @@ -241,6 +304,8 @@ def process_file_context( ) except ValueError as e: logger.debug(f"Error while getting read-only code: {e}") + if processing_cache is not None: + processing_cache.code_string_by_key[cache_key] = None return None if pruned_module.code.strip(): @@ -259,7 +324,12 @@ def process_file_context( relative_path = file_path.resolve().relative_to(project_root_path.resolve()) except ValueError: relative_path = file_path - return CodeString(code=code_context, file_path=relative_path) + result = CodeString(code=code_context, file_path=relative_path) + if processing_cache is not None: + processing_cache.code_string_by_key[cache_key] = result + return result + if processing_cache is not None: + processing_cache.code_string_by_key[cache_key] = None return None @@ -269,6 +339,7 @@ def extract_code_markdown_context_from_files( project_root_path: Path, remove_docstrings: bool = False, code_context_type: CodeContextType = CodeContextType.READ_ONLY, + processing_cache: ContextProcessingCache | None = None, ) -> CodeStringsMarkdown: """Extract code context from files containing target functions and their helpers, formatting them as markdown. @@ -316,6 +387,7 @@ def extract_code_markdown_context_from_files( remove_docstrings=remove_docstrings, project_root_path=project_root_path, helper_functions=helper_functions, + processing_cache=processing_cache, ) if result is not None: @@ -333,6 +405,7 @@ def extract_code_markdown_context_from_files( remove_docstrings=remove_docstrings, project_root_path=project_root_path, helper_functions=helper_functions, + processing_cache=processing_cache, ) if result is not None: @@ -341,12 +414,15 @@ def extract_code_markdown_context_from_files( def get_function_to_optimize_as_function_source( - function_to_optimize: FunctionToOptimize, project_root_path: Path + function_to_optimize: FunctionToOptimize, project_root_path: Path, *, jedi_project: object | None = None ) -> FunctionSource: import jedi # Use jedi to find function to optimize - script = jedi.Script(path=function_to_optimize.file_path, project=jedi.Project(path=project_root_path)) + script = jedi.Script( + path=function_to_optimize.file_path, + project=jedi_project if jedi_project is not None else _get_jedi_project(str(project_root_path.resolve())), + ) # Get all names in the file names = script.get_names(all_scopes=True, definitions=True, references=False) @@ -377,22 +453,32 @@ def get_function_to_optimize_as_function_source( def get_function_sources_from_jedi( - file_path_to_qualified_function_names: dict[Path, set[str]], project_root_path: Path + file_path_to_qualified_function_names: dict[Path, set[str]], + project_root_path: Path, + *, + jedi_project: object | None = None, ) -> tuple[dict[Path, set[FunctionSource]], list[FunctionSource]]: import jedi + project_root_path = project_root_path.resolve() + project = jedi_project if jedi_project is not None else _get_jedi_project(str(project_root_path)) file_path_to_function_source = defaultdict(set) function_source_list: list[FunctionSource] = [] for file_path, qualified_function_names in file_path_to_qualified_function_names.items(): - script = jedi.Script(path=file_path, project=jedi.Project(path=project_root_path)) + script = jedi.Script(path=file_path, project=project) file_refs = script.get_names(all_scopes=True, definitions=False, references=True) + refs_by_qualified_function: dict[str, list[Name]] = defaultdict(list) + for ref in file_refs: + try: + parent = ref.parent() + if parent is None or parent.type != "function": + continue + refs_by_qualified_function[get_qualified_name(parent.module_name, parent.full_name)].append(ref) + except (AttributeError, ValueError): + continue for qualified_function_name in qualified_function_names: - names = [ - ref - for ref in file_refs - if ref.full_name and belongs_to_function_qualified(ref, qualified_function_name) - ] + names = refs_by_qualified_function.get(qualified_function_name, []) for name in names: try: definitions: list[Name] = name.goto(follow_imports=True, follow_builtin_imports=False) diff --git a/tests/test_code_context_extractor.py b/tests/test_code_context_extractor.py index af466cd3a..9ca05ab22 100644 --- a/tests/test_code_context_extractor.py +++ b/tests/test_code_context_extractor.py @@ -9,6 +9,7 @@ import pytest +import codeflash.languages.python.context.code_context_extractor as code_context_extractor from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.languages.python.context.code_context_extractor import ( collect_type_names_from_annotation, @@ -2541,6 +2542,56 @@ def target_method(self): assert code_ctx1.hashing_code_context_hash == expected_hash +def test_code_context_reuses_target_file_reads(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + code_path = tmp_path / "target.py" + code_path.write_text( + "def helper(value):\n return value + 1\n\n\ndef target(value):\n return helper(value)\n", + encoding="utf-8", + ) + function_to_optimize = FunctionToOptimize(function_name="target", file_path=code_path, parents=[]) + + original_read_text = Path.read_text + target_read_count = 0 + + def counting_read_text(path: Path, *args, **kwargs): # type: ignore[no-untyped-def] + nonlocal target_read_count + if path == code_path: + target_read_count += 1 + return original_read_text(path, *args, **kwargs) + + monkeypatch.setattr(Path, "read_text", counting_read_text) + + get_code_optimization_context(function_to_optimize=function_to_optimize, project_root_path=tmp_path) + + assert target_read_count <= 2 + + +def test_code_context_reuses_cached_jedi_project(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + pytest.importorskip("jedi") + import jedi + + code_context_extractor._get_jedi_project.cache_clear() + + code_path = tmp_path / "target.py" + code_path.write_text("def target(value):\n return value + 1\n", encoding="utf-8") + function_to_optimize = FunctionToOptimize(function_name="target", file_path=code_path, parents=[]) + + real_project = jedi.Project + project_calls = 0 + + def counting_project(*args, **kwargs): # type: ignore[no-untyped-def] + nonlocal project_calls + project_calls += 1 + return real_project(*args, **kwargs) + + monkeypatch.setattr(jedi, "Project", counting_project) + + get_code_optimization_context(function_to_optimize=function_to_optimize, project_root_path=tmp_path) + get_code_optimization_context(function_to_optimize=function_to_optimize, project_root_path=tmp_path) + + assert project_calls == 1 + + def test_hashing_code_context_different_code_different_hash(tmp_path: Path) -> None: """Test that different code produces different hashes.""" code1 = """ diff --git a/tests/test_unit_test_discovery.py b/tests/test_unit_test_discovery.py index e232a5b71..fd636d959 100644 --- a/tests/test_unit_test_discovery.py +++ b/tests/test_unit_test_discovery.py @@ -1,8 +1,12 @@ import os +import sqlite3 import tempfile from pathlib import Path +import pytest + from codeflash.discovery.discover_unit_tests import ( + TestsCache, analyze_imports_in_test_file, discover_unit_tests, filter_test_files_by_imports, @@ -2045,3 +2049,64 @@ def test_discover_unit_tests_caching(): assert non_cached_num_discovered_tests == num_discovered_tests assert non_cached_function_to_tests == tests assert non_cached_num_discovered_replay_tests == num_discovered_replay_tests + + +def test_tests_cache_batches_commits(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + real_connect = sqlite3.connect + commit_count = 0 + + class CountingConnection: + def __init__(self, path: str) -> None: + self._connection = real_connect(path) + + def __getattr__(self, name: str): + return getattr(self._connection, name) + + def commit(self) -> None: + nonlocal commit_count + commit_count += 1 + self._connection.commit() + + db_path = tmp_path / "tests_cache.sqlite" + monkeypatch.setattr( + "codeflash.discovery.discover_unit_tests.codeflash_cache_db", + db_path, + ) + monkeypatch.setattr( + "codeflash.discovery.discover_unit_tests.sqlite3.connect", + lambda path: CountingConnection(str(path)), + ) + + cache = TestsCache(project_root_path=tmp_path) + initial_commit_count = commit_count + + cache.insert_test( + file_path=str(tmp_path / "test_a.py"), + file_hash="hash-a", + qualified_name_with_modules_from_root="pkg.fn_a", + function_name="fn_a", + test_class="", + test_function="test_fn_a", + test_type=TestType.EXISTING_UNIT_TEST, + line_number=1, + col_number=0, + ) + cache.insert_test( + file_path=str(tmp_path / "test_b.py"), + file_hash="hash-b", + qualified_name_with_modules_from_root="pkg.fn_b", + function_name="fn_b", + test_class="", + test_function="test_fn_b", + test_type=TestType.EXISTING_UNIT_TEST, + line_number=2, + col_number=0, + ) + + assert commit_count == initial_commit_count + + cache.flush() + assert commit_count == initial_commit_count + 1 + assert cache.pending_rows == [] + + cache.close() From 7944e43eb3239262fd9d7f9752a332df090d25a3 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Tue, 17 Mar 2026 00:28:25 +0000 Subject: [PATCH 2/2] fix: handle Windows 8.3 short-name path mismatch in target block reordering On Windows, Path.resolve() expands 8.3 short names for directories but not always for files. After project_root_path.resolve() converts RUNNER~1 to runneradmin, file_path.resolve() may still return the 8.3 form, causing relative_to() to raise ValueError and silently returning empty target_code. Wrap the cosmetic reordering block in try/except to match the existing pattern in process_file_context, preventing the crash without any correctness loss. Co-authored-by: Aseem Saxena --- .../python/context/code_context_extractor.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/codeflash/languages/python/context/code_context_extractor.py b/codeflash/languages/python/context/code_context_extractor.py index d1d4b2399..67229bbe9 100644 --- a/codeflash/languages/python/context/code_context_extractor.py +++ b/codeflash/languages/python/context/code_context_extractor.py @@ -152,11 +152,14 @@ def get_code_optimization_context( ) # Ensure the target file is first in the code blocks so the LLM knows which file to optimize - target_relative = function_to_optimize.file_path.resolve().relative_to(project_root_path.resolve()) - target_blocks = [cs for cs in final_read_writable_code.code_strings if cs.file_path == target_relative] - other_blocks = [cs for cs in final_read_writable_code.code_strings if cs.file_path != target_relative] - if target_blocks: - final_read_writable_code.code_strings = target_blocks + other_blocks + try: + target_relative = function_to_optimize.file_path.resolve().relative_to(project_root_path) + target_blocks = [cs for cs in final_read_writable_code.code_strings if cs.file_path == target_relative] + other_blocks = [cs for cs in final_read_writable_code.code_strings if cs.file_path != target_relative] + if target_blocks: + final_read_writable_code.code_strings = target_blocks + other_blocks + except ValueError: + pass read_only_code_markdown = extract_code_markdown_context_from_files( helpers_of_fto_dict,