Update code_context_extractor.py

Author: KRRT7

codeflash/context/code_context_extractor.py

@@ -28,6 +28,7 @@
     from libcst import CSTNode
 
     from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+from typing import Callable
 
 
 def get_code_optimization_context(
@@ -462,7 +463,9 @@ def parse_code_and_prune_cst(
         helpers_of_helper_functions = set()
     module = cst.parse_module(code)
 
-    dispatch = {
+    dispatch: dict[
+        CodeContextType, Callable[[cst.CSTNode, set[str], set[str], str, bool], tuple[Optional[cst.CSTNode], bool]]
+    ] = {
         CodeContextType.READ_WRITABLE: prune_cst_for_read_writable_code,
         CodeContextType.READ_ONLY: prune_cst_for_read_only_code,
         CodeContextType.TESTGEN: prune_cst_for_testgen_code,
@@ -474,7 +477,7 @@ def parse_code_and_prune_cst(
         raise ValueError(msg)
 
     filtered_node, found_target = prune_func(
-        module, target_functions, helpers_of_helper_functions
+        module, target_functions, helpers_of_helper_functions, "", remove_docstrings
     )
 
     if not found_target: