Adapt SWEAgentTranslator for local model use via Ollama and vLLM

.gitignore

@@ -11,6 +11,7 @@ __pycache__
 
 tmp*
 .env
+.envrc
 .venv
 
 .request_cache_*.pkl

src/translate/swe_agent/swe_agent_translator.py

@@ -5,6 +5,7 @@
 import shutil
 import subprocess
 import json
+import time
 from typing import List, Optional, Dict, Any
 
 # local imports
@@ -20,6 +21,7 @@ class SWEAgentTranslator(Translator):
     TRAJECTORIES_DIR = "trajectories"
     PATCH_FILENAME = "temp.patch"
     EXPERIMENT_METADATA_FILENAME = "experiment_metadata.json"
+    SERVE_CHECK_COOLDOWN = 10
 
     # File extensions to remove from output
     REMOVE_EXTENSIONS = (".cu", ".cuh")
@@ -32,6 +34,10 @@ class SWEAgentTranslator(Translator):
     # Instance variables
     _swe_agent_model_name: str
     _swe_agent_per_instance_cost_limit: float
+    _swe_agent_config: Optional[str]
+    _swe_agent_parser: Optional[str]
+    _swe_agent_max_input_token: Optional[int]
+
     _temp_repo_path: str
     _translation_task_path: str
     _output_path: str
@@ -47,7 +53,10 @@ def __init__(
         dry: bool = False,
         hide_progress: bool = False,
         swe_agent_model_name: Optional[str] = None,
-        swe_agent_per_instance_cost_limit: float = 0.06
+        swe_agent_per_instance_cost_limit: float = 0.06,
+        swe_agent_config: Optional[str] = None,
+        swe_agent_parser: Optional[str] = None,
+        swe_agent_max_input_token: Optional[int] = None,
     ) -> None:
         super().__init__(
             input_repo,
@@ -62,31 +71,82 @@ def __init__(
 
         self._swe_agent_model_name = swe_agent_model_name
         self._swe_agent_per_instance_cost_limit = swe_agent_per_instance_cost_limit
+        self._swe_agent_config = swe_agent_config
+        self._swe_agent_parser = swe_agent_parser
+        self._swe_agent_max_input_token = swe_agent_max_input_token
+
         self._temp_repo_path = self.TEMP_REPO_PATH
         self._translation_task_path = os.path.join(
             self._input_repo.path, self.TRANSLATION_TASK_FILENAME
         )
         self._output_path = os.path.join(self._output_paths[0], "repo")
 
+        if self._is_ollama_model(self._swe_agent_model_name):
+            if self._swe_agent_parser is None:
+                self._swe_agent_parser = "thought_action"
+            if self._swe_agent_max_input_token is None:
+                self._swe_agent_max_input_token = 4096
+            self._launch_ollama_server()
+
+    @staticmethod
+    def _is_ollama_model(name: str) -> bool:
+        name = (name or "").lower()
+        return name.startswith("ollama/")
+
+
+    def _launch_ollama_server(self) -> None:
+        """Launch an Ollama server in the background."""
+        # Check that ollama is installed
+        if not shutil.which("ollama"):
+            raise ValueError("Ollama is not in the path. Please install Ollama and add it to the path.")
+        # Early exit if ollama is already running
+        if subprocess.run(["ollama", "list"], capture_output=True, text=True).returncode == 0:
+            return
+        ollama_command = ["ollama", "serve"]
+        subprocess.Popen(ollama_command,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.STDOUT,       
+        stdin=subprocess.DEVNULL,   
+        start_new_session=True,
+        env=env)
+        # Check that the server is running
+        checking = True
+        while checking:
+            status = subprocess.run(["ollama", "list"], capture_output=True, text=True)
+            if status.returncode == 0:
+                checking = False
+            else:
+                print(f"Ollama server not ready, checking again after {self.SERVE_CHECK_COOLDOWN} seconds...")
+                time.sleep(self.SERVE_CHECK_COOLDOWN)
+        print(f"Ollama server ready.")
+        return
+
 
     @staticmethod
     def add_args(parser: Any) -> None:
         """Add command line arguments for SWE-agent configuration."""
         parser.add_argument("--swe-agent-model-name", type=str,
-                            help="Name of the agent model to use (e.g. 'gpt-4o').")
+                            help="Name of the agent model to use (e.g. 'gpt-4o', 'ollama/llama3.2:latest').")
         parser.add_argument("--swe-agent-per-instance-cost-limit", type=float,
-                            help="Per-instance cost limit for the agent model.")
-
+                            help="Per-instance cost limit for the agent model; set to 0 for local models.")
+        parser.add_argument("--swe-agent-config", type=str,
+                            help="Custom config file provided to SWE-Agent. Default config file is used if none is provided.")
+        parser.add_argument("--swe-agent-parser", type=str, choices=["thought_action", "function_call"],
+                            help="Parsing strategy. Use 'thought_action' for local/Ollama models.")
+        parser.add_argument("--swe-agent-max-input-token", type=int,
+                            help="Override max input tokens to avoid local-model warnings.")
 
     @staticmethod
     def parse_args(args: Any) -> Dict[str, Any]:
         """Parse command line arguments for SWE-agent configuration."""
         return {
             "swe_agent_model_name": args.swe_agent_model_name,
-            "swe_agent_per_instance_cost_limit": args.swe_agent_per_instance_cost_limit
+            "swe_agent_per_instance_cost_limit": args.swe_agent_per_instance_cost_limit,
+            "swe_agent_config": args.swe_agent_config,
+            "swe_agent_parser": args.swe_agent_parser,
+            "swe_agent_max_input_token": args.swe_agent_max_input_token,
         }
 
-
     def translate(self) -> None:
         """Execute the complete translation process using SWE-agent.
 
@@ -175,7 +235,6 @@ def _initialize_git_repo(self) -> None:
         subprocess.run(self.GIT_ADD_ALL, cwd=self._temp_repo_path, check=True)
         subprocess.run(self.GIT_COMMIT_INITIAL, cwd=self._temp_repo_path, check=True)
 
-
     def run_swe_agent(self) -> bool:
         """Run the SWE-agent command and apply the resulting patch."""
         command = self._build_swe_agent_command()
@@ -197,15 +256,23 @@ def run_swe_agent(self) -> bool:
 
     def _build_swe_agent_command(self) -> List[str]:
         """Build the SWE-agent command with all required parameters."""
-        return [
+        cmd = [
             "sweagent", "run",
             f"--agent.model.name={self._swe_agent_model_name}",
             f"--agent.model.per_instance_cost_limit={self._swe_agent_per_instance_cost_limit}",
             f"--env.repo.path={self._temp_repo_path}",
-            "--env.deployment.image=python",
             f"--problem_statement.path={self._translation_task_path}",
         ]
 
+        if self._swe_agent_parser:
+            cmd.append(f"--agent.tools.parse_function.type={self._swe_agent_parser}")
+        if self._swe_agent_max_input_token:
+            cmd.append(f"--agent.model.max_input_tokens={self._swe_agent_max_input_token}")
+        if self._swe_agent_config:
+            cmd.extend(["--config", self._swe_agent_config])
+
+        return cmd
+
     def _apply_swe_agent_patch(self) -> bool:
         """Find and apply the patch file generated by SWE-agent."""
         print("Applying patch...")