Enhanced llm editor (#9174)

Co-authored-by: jianchuanli <jianchuanli@langcode.com>
Co-authored-by: Xingyao Wang <xingyao@all-hands.dev>

config.template.toml

@@ -200,6 +200,7 @@ model = "gpt-4o"
 # https://github.com/All-Hands-AI/OpenHands/pull/4711
 #native_tool_calling = None
 
+
 # Safety settings for models that support them (e.g., Mistral AI, Gemini)
 # Example for Mistral AI:
 # safety_settings = [
@@ -218,6 +219,9 @@ model = "gpt-4o"
 # ]
 #safety_settings = []
 
+[llm.draft_editor]
+# The number of times llm_editor tries to fix an error when editing.
+correct_num = 5
 
 [llm.gpt4o-mini]
 api_key = ""
@@ -335,6 +339,9 @@ classpath = "my_package.my_module.MyCustomAgent"
 # Enable GPU support in the runtime
 #enable_gpu = false
 
+# When there are multiple cards, you can specify the GPU by ID
+#cuda_visible_devices = ''
+
 # Additional Docker runtime kwargs
 #docker_runtime_kwargs = {}
 

openhands/core/config/sandbox_config.py

@@ -88,6 +88,7 @@ class SandboxConfig(BaseModel):
         description="Volume mounts in the format 'host_path:container_path[:mode]', e.g. '/my/host/dir:/workspace:rw'. Multiple mounts can be specified using commas, e.g. '/path1:/workspace/path1,/path2:/workspace/path2:ro'",
     )
 
+    cuda_visible_devices: str | None = Field(default=None)
     model_config = ConfigDict(extra='forbid')
 
     @classmethod

openhands/runtime/impl/docker/docker_runtime.py

@@ -360,7 +360,21 @@ class DockerRuntime(ActionExecutionClient):
         )
 
         command = self.get_action_execution_server_startup_command()
-
+        if self.config.sandbox.enable_gpu:
+            gpu_ids = self.config.sandbox.cuda_visible_devices
+            if gpu_ids is None:
+                device_requests = [
+                    docker.types.DeviceRequest(capabilities=[['gpu']], count=-1)
+                ]
+            else:
+                device_requests = [
+                    docker.types.DeviceRequest(
+                        capabilities=[['gpu']],
+                        device_ids=[str(i) for i in gpu_ids.split(',')],
+                    )
+                ]
+        else:
+            device_requests = None
         try:
             if self.runtime_container_image is None:
                 raise ValueError('Runtime container image is not set')
@@ -376,11 +390,7 @@ class DockerRuntime(ActionExecutionClient):
                 detach=True,
                 environment=environment,
                 volumes=volumes,  # type: ignore
-                device_requests=(
-                    [docker.types.DeviceRequest(capabilities=[['gpu']], count=-1)]
-                    if self.config.sandbox.enable_gpu
-                    else None
-                ),
+                device_requests=device_requests,
                 **(self.config.sandbox.docker_runtime_kwargs or {}),
             )
             self.log('debug', f'Container started. Server url: {self.api_url}')

openhands/runtime/utils/edit.py

@@ -39,6 +39,40 @@ Within the `<updated_code>` tag, include only the final code after updation. Do
 <update_snippet>{draft_changes}</update_snippet>
     """
 
+CORRECT_SYS_MSG = """You are a code repair assistant. Now you have an original file content and error information from a static code checking tool (lint tool). Your task is to automatically modify and return the repaired complete code based on these error messages and refer to the current file content.
+
+The following are the specific task steps you need to complete:
+
+Carefully read the current file content to ensure that you fully understand its code structure.
+
+According to the lint error prompt, accurately locate and analyze the cause of the problem.
+
+Modify the original file content and fix all errors prompted by the lint tool.
+
+Return complete, runnable, and error-fixed code, paying attention to maintaining the overall style and specifications of the original code.
+
+Please note:
+
+Please strictly follow the lint error prompts to make modifications and do not miss any problems.
+
+The modified code must be complete and cannot introduce new errors or bugs.
+
+The modified code must maintain the original code function and logic, and no changes unrelated to error repair should be made."""
+
+CORRECT_USER_MSG = """
+THE FOLLOWING ARE THE ORIGINAL FILE CONTENTS AND THE ERROR INFORMATION REPORTED BY THE LINT TOOL
+
+# CURRENT FILE CONTENT:
+```
+{file_content}
+```
+
+# ERROR MESSAGE FROM STATIC CODE CHECKING TOOL:
+```
+{lint_error}
+```
+""".strip()
+
 
 def _extract_code(string: str) -> str | None:
     pattern = r'<updated_code>(.*?)</updated_code>'
@@ -196,7 +230,7 @@ class FileEditRuntimeMixin(FileEditRuntimeInterface):
                 return ErrorObservation(error_message)
         return None
 
-    def llm_based_edit(self, action: FileEditAction) -> Observation:
+    def llm_based_edit(self, action: FileEditAction, retry_num: int = 0) -> Observation:
         obs = self.read(FileReadAction(path=action.path))
         if (
             isinstance(obs, ErrorObservation)
@@ -253,7 +287,14 @@ class FileEditRuntimeMixin(FileEditRuntimeInterface):
                     diff,
                 )
                 if error_obs is not None:
-                    return error_obs
+                    self.write(
+                        FileWriteAction(path=action.path, content=updated_content)
+                    )
+                    return self.correct_edit(
+                        file_content=updated_content,
+                        error_obs=error_obs,
+                        retry_num=retry_num,
+                    )
 
             obs = self.write(FileWriteAction(path=action.path, content=updated_content))
             return FileEditObservation(
@@ -280,7 +321,8 @@ class FileEditRuntimeMixin(FileEditRuntimeInterface):
             error_msg = (
                 f'[Edit error: The range of lines to edit is too long.]\n'
                 f'[The maximum number of lines allowed to edit at once is {self.MAX_LINES_TO_EDIT}. '
-                f'Got (L{start_idx + 1}-L{end_idx}) {length_of_range} lines.]\n'  # [start_idx, end_idx), so no need to + 1
+                f'Got (L{start_idx + 1}-L{end_idx}) {length_of_range} lines.]\n'
+                # [start_idx, end_idx), so no need to + 1
             )
             # search for relevant ranges to hint the agent
             topk_chunks: list[Chunk] = get_top_k_chunk_matches(
@@ -333,7 +375,12 @@ class FileEditRuntimeMixin(FileEditRuntimeInterface):
             )
             if error_obs is not None:
                 error_obs.llm_metrics = self.draft_editor_llm.metrics
-                return error_obs
+                self.write(FileWriteAction(path=action.path, content=updated_content))
+                return self.correct_edit(
+                    file_content=updated_content,
+                    error_obs=error_obs,
+                    retry_num=retry_num,
+                )
 
         obs = self.write(FileWriteAction(path=action.path, content=updated_content))
         ret_obs = FileEditObservation(
@@ -345,3 +392,40 @@ class FileEditRuntimeMixin(FileEditRuntimeInterface):
         )
         ret_obs.llm_metrics = self.draft_editor_llm.metrics
         return ret_obs
+
+    def check_retry_num(self, retry_num):
+        correct_num = self.draft_editor_llm.config.correct_num
+        return correct_num < retry_num
+
+    def correct_edit(
+        self, file_content: str, error_obs: ErrorObservation, retry_num: int = 0
+    ) -> Observation:
+        import openhands.agenthub.codeact_agent.function_calling as codeact_function_calling
+        from openhands.agenthub.codeact_agent.tools import LLMBasedFileEditTool
+        from openhands.llm.llm_utils import check_tools
+
+        _retry_num = retry_num + 1
+        if self.check_retry_num(_retry_num):
+            return error_obs
+        tools = check_tools([LLMBasedFileEditTool], self.draft_editor_llm.config)
+        messages = [
+            {'role': 'system', 'content': CORRECT_SYS_MSG},
+            {
+                'role': 'user',
+                'content': CORRECT_USER_MSG.format(
+                    file_content=file_content, lint_error=error_obs.content
+                ),
+            },
+        ]
+        params: dict = {'messages': messages, 'tools': tools}
+        try:
+            response = self.draft_editor_llm.completion(**params)
+            actions = codeact_function_calling.response_to_actions(response)
+            if len(actions) != 1:
+                return error_obs
+            for action in actions:
+                if isinstance(action, FileEditAction):
+                    return self.llm_based_edit(action, _retry_num)
+        except Exception as e:
+            logger.error(f'correct lint error is failed: {e}')
+        return error_obs