wip e2e

Author: misrasaurabh1

codeflash/api/aiservice.py

@@ -120,24 +120,28 @@ def _get_valid_candidates(
             )
         return candidates
 
-    def optimize_python_code(  # noqa: D417
+    def optimize_code(  # noqa: D417
         self,
         source_code: str,
         dependency_code: str,
         trace_id: str,
         experiment_metadata: ExperimentMetadata | None = None,
         *,
+        language: str = "python",
+        language_version: str | None = None,
         is_async: bool = False,
         n_candidates: int = 5,
     ) -> list[OptimizedCandidate]:
-        """Optimize the given python code for performance by making a request to the Django endpoint.
+        """Optimize the given code for performance by making a request to the Django endpoint.
 
         Parameters
         ----------
-        - source_code (str): The python code to optimize.
+        - source_code (str): The code to optimize.
         - dependency_code (str): The dependency code used as read-only context for the optimization
         - trace_id (str): Trace id of optimization run
         - experiment_metadata (Optional[ExperimentalMetadata, None]): Any available experiment metadata for this optimization
+        - language (str): Programming language ("python", "javascript", "typescript")
+        - language_version (str | None): Language version (e.g., "3.11.0" for Python, "ES2022" for JS)
         - is_async (bool): Whether the function being optimized is async
         - n_candidates (int): Number of candidates to generate
 
@@ -151,11 +155,12 @@ def optimize_python_code(  # noqa: D417
         start_time = time.perf_counter()
         git_repo_owner, git_repo_name = safe_get_repo_owner_and_name()
 
-        payload = {
+        # Build payload with language-specific fields
+        payload: dict[str, Any] = {
             "source_code": source_code,
             "dependency_code": dependency_code,
             "trace_id": trace_id,
-            "python_version": platform.python_version(),
+            "language": language,
             "experiment_metadata": experiment_metadata,
             "codeflash_version": codeflash_version,
             "current_username": get_last_commit_author_if_pr_exists(None),
@@ -165,6 +170,14 @@ def optimize_python_code(  # noqa: D417
             "call_sequence": self.get_next_sequence(),
             "n_candidates": n_candidates,
         }
+
+        # Add language-specific version fields
+        # Always include python_version for backward compatibility with older backend
+        payload["python_version"] = platform.python_version()
+        if language == "python":
+            pass  # python_version already set
+        else:
+            payload["language_version"] = language_version or "ES2022"
         logger.debug(f"Sending optimize request: trace_id={trace_id}, n_candidates={payload['n_candidates']}")
 
         try:
@@ -191,6 +204,28 @@ def optimize_python_code(  # noqa: D417
         console.rule()
         return []
 
+    # Backward-compatible alias
+    def optimize_python_code(
+        self,
+        source_code: str,
+        dependency_code: str,
+        trace_id: str,
+        experiment_metadata: ExperimentMetadata | None = None,
+        *,
+        is_async: bool = False,
+        n_candidates: int = 5,
+    ) -> list[OptimizedCandidate]:
+        """Backward-compatible alias for optimize_code() with language='python'."""
+        return self.optimize_code(
+            source_code=source_code,
+            dependency_code=dependency_code,
+            trace_id=trace_id,
+            experiment_metadata=experiment_metadata,
+            language="python",
+            is_async=is_async,
+            n_candidates=n_candidates,
+        )
+
     def optimize_python_code_line_profiler(  # noqa: D417
         self,
         source_code: str,
@@ -578,6 +613,9 @@ def generate_regression_tests(  # noqa: D417
         test_timeout: int,
         trace_id: str,
         test_index: int,
+        *,
+        language: str = "python",
+        language_version: str | None = None,
     ) -> tuple[str, str, str] | None:
         """Generate regression tests for the given function by making a request to the Django endpoint.
 
@@ -588,19 +626,30 @@ def generate_regression_tests(  # noqa: D417
         - helper_function_names (list[Source]): List of helper function names.
         - module_path (Path): The module path where the function is located.
         - test_module_path (Path): The module path for the test code.
-        - test_framework (str): The test framework to use, e.g., "pytest".
+        - test_framework (str): The test framework to use, e.g., "pytest", "jest".
         - test_timeout (int): The timeout for each test in seconds.
         - test_index (int): The index from 0-(n-1) if n tests are generated for a single trace_id
+        - language (str): Programming language ("python", "javascript", "typescript")
+        - language_version (str | None): Language version (e.g., "3.11.0" for Python, "ES2022" for JS)
 
         Returns
         -------
         - Dict[str, str] | None: The generated regression tests and instrumented tests, or None if an error occurred.
 
         """
-        assert test_framework in ["pytest", "unittest"], (
-            f"Invalid test framework, got {test_framework} but expected 'pytest' or 'unittest'"
-        )
-        payload = {
+        # Validate test framework based on language
+        python_frameworks = ["pytest", "unittest"]
+        javascript_frameworks = ["jest", "mocha", "vitest"]
+        if language == "python":
+            assert test_framework in python_frameworks, (
+                f"Invalid test framework for Python, got {test_framework} but expected one of {python_frameworks}"
+            )
+        elif language in ("javascript", "typescript"):
+            assert test_framework in javascript_frameworks, (
+                f"Invalid test framework for JavaScript, got {test_framework} but expected one of {javascript_frameworks}"
+            )
+
+        payload: dict[str, Any] = {
             "source_code_being_tested": source_code_being_tested,
             "function_to_optimize": function_to_optimize,
             "helper_function_names": helper_function_names,
@@ -610,11 +659,19 @@ def generate_regression_tests(  # noqa: D417
             "test_timeout": test_timeout,
             "trace_id": trace_id,
             "test_index": test_index,
-            "python_version": platform.python_version(),
+            "language": language,
             "codeflash_version": codeflash_version,
             "is_async": function_to_optimize.is_async,
             "call_sequence": self.get_next_sequence(),
         }
+
+        # Add language-specific version fields
+        # Always include python_version for backward compatibility with older backend
+        payload["python_version"] = platform.python_version()
+        if language == "python":
+            pass  # python_version already set
+        else:
+            payload["language_version"] = language_version or "ES2022"
         try:
             response = self.make_ai_service_request("/testgen", payload=payload, timeout=self.timeout)
         except requests.exceptions.RequestException as e:

codeflash/code_utils/code_extractor.py

@@ -1149,9 +1149,12 @@ def get_fn_references_jedi(
     source_code: str, file_path: Path, project_root: Path, target_function: str, target_class: str | None
 ) -> list[Path]:
     start_time = time.perf_counter()
-    function_position: CodePosition = find_specific_function_in_file(
+    function_position: CodePosition | None = find_specific_function_in_file(
         source_code, file_path, target_function, target_class
     )
+    if function_position is None:
+        # Function not found (may be non-Python code)
+        return []
     try:
         script = jedi.Script(code=source_code, path=file_path, project=jedi.Project(path=project_root))
         # Get references to the function

codeflash/code_utils/code_replacer.py

@@ -601,7 +601,8 @@ def replace_optimized_code(
                     [
                         callee.qualified_name
                         for callee in code_context.helper_functions
-                        if callee.file_path == module_path and callee.jedi_definition.type != "class"
+                        if callee.file_path == module_path
+                        and (callee.jedi_definition is None or callee.jedi_definition.type != "class")
                     ]
                 ),
                 candidate.source_code,

codeflash/context/unused_definition_remover.py

@@ -632,8 +632,8 @@ def _analyze_imports_in_optimized_code(
     helpers_by_file_and_func = defaultdict(dict)
     helpers_by_file = defaultdict(list)  # preserved for "import module"
     for helper in code_context.helper_functions:
-        jedi_type = helper.jedi_definition.type
-        if jedi_type != "class":
+        jedi_type = helper.jedi_definition.type if helper.jedi_definition else None
+        if jedi_type != "class":  # Include when jedi_definition is None (non-Python)
             func_name = helper.only_function_name
             module_name = helper.file_path.stem
             # Cache function lookup for this (module, func)
@@ -784,7 +784,8 @@ def detect_unused_helper_functions(
         # Find helper functions that are no longer called
         unused_helpers = []
         for helper_function in code_context.helper_functions:
-            if helper_function.jedi_definition.type != "class":
+            jedi_type = helper_function.jedi_definition.type if helper_function.jedi_definition else None
+            if jedi_type != "class":  # Include when jedi_definition is None (non-Python)
                 # Check if the helper function is called using multiple name variants
                 helper_qualified_name = helper_function.qualified_name
                 helper_simple_name = helper_function.only_function_name

codeflash/discovery/discover_unit_tests.py

@@ -616,7 +616,6 @@ def discover_tests_for_language(
                         end_col=func.ending_col,
                         is_async=func.is_async,
                         is_method=bool(func.parents and any(p.type == "ClassDef" for p in func.parents)),
-                        class_name=func.parents[0].name if func.parents and func.parents[0].type == "ClassDef" else None,
                         parents=parents,
                         language=Language(language),
                     )

codeflash/models/models.py

@@ -129,7 +129,7 @@ class FunctionSource:
     fully_qualified_name: str
     only_function_name: str
     source_code: str
-    jedi_definition: Name
+    jedi_definition: Name | None = None  # None for non-Python languages
 
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, FunctionSource):
@@ -323,10 +323,10 @@ def parse_markdown_code(markdown_code: str, expected_language: str = "python") -
                     detected_language = language
                 if file_path:
                     path = file_path.strip()
-                    code_string_list.append(CodeString(code=code, file_path=Path(path)))
+                    code_string_list.append(CodeString(code=code, file_path=Path(path), language=detected_language))
                 else:
                     # No file path specified - skip this block or create with None
-                    code_string_list.append(CodeString(code=code, file_path=None))
+                    code_string_list.append(CodeString(code=code, file_path=None, language=detected_language))
             return CodeStringsMarkdown(code_strings=code_string_list, language=detected_language)
         except ValidationError:
             # if any file is invalid, return an empty CodeStringsMarkdown for the entire context

codeflash/optimization/function_optimizer.py

@@ -434,10 +434,14 @@ def __init__(
             else function_to_optimize.file_path.read_text(encoding="utf8")
         )
         if not function_to_optimize_ast:
-            original_module_ast = ast.parse(function_to_optimize_source_code)
-            self.function_to_optimize_ast = get_first_top_level_function_or_method_ast(
-                function_to_optimize.function_name, function_to_optimize.parents, original_module_ast
-            )
+            # Skip Python AST parsing for JavaScript/TypeScript
+            if function_to_optimize.language in ("javascript", "typescript"):
+                self.function_to_optimize_ast = None
+            else:
+                original_module_ast = ast.parse(function_to_optimize_source_code)
+                self.function_to_optimize_ast = get_first_top_level_function_or_method_ast(
+                    function_to_optimize.function_name, function_to_optimize.parents, original_module_ast
+                )
         else:
             self.function_to_optimize_ast = function_to_optimize_ast
         self.function_to_tests = function_to_tests if function_to_tests else {}
@@ -450,7 +454,26 @@ def __init__(
 
         self.args = args  # Check defaults for these
         self.function_trace_id: str = str(uuid.uuid4())
-        self.original_module_path = module_name_from_file_path(self.function_to_optimize.file_path, self.project_root)
+        # For JavaScript/TypeScript, we need a relative path from the test file to the source file
+        # For Python, we use dot-separated module paths
+        if self.function_to_optimize.language in ("javascript", "typescript"):
+            # Compute relative path from tests directory to source file
+            # e.g., for source at /project/fibonacci.js and tests at /project/tests/
+            # the relative path should be ../fibonacci
+            try:
+                # Use os.path.relpath to compute relative path from tests_root to source file
+                import os
+                rel_path = os.path.relpath(
+                    str(self.function_to_optimize.file_path.with_suffix("")),
+                    str(test_cfg.tests_root)
+                )
+                self.original_module_path = rel_path
+            except ValueError:
+                # Fallback if paths are on different drives (Windows)
+                rel_path = self.function_to_optimize.file_path.relative_to(self.project_root)
+                self.original_module_path = "../" + rel_path.with_suffix("").as_posix()
+        else:
+            self.original_module_path = module_name_from_file_path(self.function_to_optimize.file_path, self.project_root)
 
         self.function_benchmark_timings = function_benchmark_timings if function_benchmark_timings else {}
         self.total_benchmark_timings = total_benchmark_timings if total_benchmark_timings else {}
@@ -510,15 +533,24 @@ def generate_and_instrument_tests(
     ]:
         """Generate and instrument tests for the function."""
         n_tests = get_effort_value(EffortKeys.N_GENERATED_TESTS, self.effort)
+        language = self.function_to_optimize.language
         generated_test_paths = [
             get_test_file_path(
-                self.test_cfg.tests_root, self.function_to_optimize.function_name, test_index, test_type="unit"
+                self.test_cfg.tests_root,
+                self.function_to_optimize.function_name,
+                test_index,
+                test_type="unit",
+                language=language,
             )
             for test_index in range(n_tests)
         ]
         generated_perf_test_paths = [
             get_test_file_path(
-                self.test_cfg.tests_root, self.function_to_optimize.function_name, test_index, test_type="perf"
+                self.test_cfg.tests_root,
+                self.function_to_optimize.function_name,
+                test_index,
+                test_type="perf",
+                language=language,
             )
             for test_index in range(n_tests)
         ]
@@ -1370,7 +1402,8 @@ def replace_function_and_helpers_with_optimized_code(
             self.function_to_optimize.qualified_name
         )
         for helper_function in code_context.helper_functions:
-            if helper_function.jedi_definition.type != "class":
+            # Skip class definitions (jedi_definition may be None for non-Python languages)
+            if helper_function.jedi_definition is None or 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():
             did_update |= replace_function_definitions_in_module(
@@ -1423,6 +1456,12 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, set[Functio
         func_qualname = self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root)
         if func_qualname not in function_to_all_tests:
             logger.info(f"Did not find any pre-existing tests for '{func_qualname}', will only use generated tests.")
+        # Skip existing test instrumentation for JavaScript/TypeScript - use generated tests only
+        elif self.function_to_optimize.language in ("javascript", "typescript"):
+            logger.info(
+                f"JavaScript/TypeScript detected - using generated tests only for '{func_qualname}'. "
+                "Existing test instrumentation not yet supported."
+            )
         else:
             test_file_invocation_positions = defaultdict(list)
             for tests_in_file in function_to_all_tests.get(func_qualname):
@@ -1577,11 +1616,12 @@ def generate_optimizations(
         """Generate optimization candidates for the function. Backend handles multi-model diversity."""
         n_candidates = get_effort_value(EffortKeys.N_OPTIMIZER_CANDIDATES, self.effort)
         future_optimization_candidates = self.executor.submit(
-            self.aiservice_client.optimize_python_code,
+            self.aiservice_client.optimize_code,
             read_writable_code.markdown,
             read_only_context_code,
             self.function_trace_id[:-4] + "EXP0" if run_experiment else self.function_trace_id,
             ExperimentMetadata(id=self.experiment_id, group="control") if run_experiment else None,
+            language=self.function_to_optimize.language,
             is_async=self.function_to_optimize.is_async,
             n_candidates=n_candidates,
         )
@@ -1600,11 +1640,12 @@ def generate_optimizations(
 
         if run_experiment:
             future_candidates_exp = self.executor.submit(
-                self.local_aiservice_client.optimize_python_code,
+                self.local_aiservice_client.optimize_code,
                 read_writable_code.markdown,
                 read_only_context_code,
                 self.function_trace_id[:-4] + "EXP1",
                 ExperimentMetadata(id=self.experiment_id, group="experiment"),
+                language=self.function_to_optimize.language,
                 is_async=self.function_to_optimize.is_async,
                 n_candidates=n_candidates,
             )