Update code_context_extractor.py

Author: KRRT7

codeflash/context/code_context_extractor.py

@@ -3,19 +3,16 @@
 import os
 from collections import defaultdict
 from itertools import chain
-from pathlib import Path
+from typing import TYPE_CHECKING
 
 import jedi
 import libcst as cst
 import tiktoken
-from jedi.api.classes import Name
-from libcst import CSTNode
 
 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 get_qualified_name, path_belongs_to_site_packages
 from codeflash.context.unused_definition_remover import remove_unused_definitions_by_function_names
-from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import (
     CodeContextType,
     CodeOptimizationContext,
@@ -25,6 +22,14 @@
 )
 from codeflash.optimization.function_context import belongs_to_function_qualified
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from jedi.api.classes import Name
+    from libcst import CSTNode
+
+    from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+
 
 def get_code_optimization_context(
     function_to_optimize: FunctionToOptimize,