restore optimizeFunction

Author: KRRT7

codeflash/discovery/functions_to_optimize.py

@@ -146,6 +146,15 @@ def qualified_name(self) -> str:
     def qualified_name_with_modules_from_root(self, project_root_path: Path) -> str:
         return f"{module_name_from_file_path(self.file_path, project_root_path)}.{self.qualified_name}"
 
+    @property
+    def server_info(self) -> dict[str, str | int]:
+        return {
+            "file_path": str(self.file_path),
+            "function_name": self.function_name,
+            "starting_line": self.starting_line,
+            "ending_line": self.ending_line,
+        }
+
 
 def get_functions_to_optimize(
     optimize_all: str | None,

codeflash/lsp/beta.py

@@ -40,6 +40,29 @@ def get_optimizable_functions(
     return path_to_qualified_names
 
 
+@server.feature("optimizeFunction")
+def optimize_function(server: CodeflashLanguageServer, params: OptimizeFunctionParams) -> dict[str, str]:
+    file_path = Path(uris.to_fs_path(params.textDocument.uri))
+    server.optimizer.args.function = params.functionName
+    server.optimizer.args.file = file_path
+    optimizable_funcs, _ = server.optimizer.get_optimizable_functions()
+    if not optimizable_funcs:
+        return {"functionName": params.functionName, "status": "not found", "args": None}
+    fto = optimizable_funcs.popitem()[1][0]
+    server.optimizer.current_function_being_optimized = fto
+    return {"functionName": params.functionName, "status": "success", "info": fto.server_info}
+
+
+@server.feature("second_step_in_optimize_function")
+def second_step_in_optimize_function(server: CodeflashLanguageServer, params: OptimizeFunctionParams) -> dict[str, str]:  # noqa: ARG001
+    return {
+        "functionName": params.functionName,
+        "status": "success",
+        "generated_tests": "5",
+        "generated_optimizations": "3",
+    }
+
+
 if __name__ == "__main__":
     from codeflash.cli_cmds.console import console