refactor: refactor llm generation code for localization

Author: dengyinlin

agentless/fl/FL.py

@@ -211,24 +211,24 @@ class LLMFL(FL):
 Return just the locations.
 """
 
-    def __init__(self, instance_id, structure, problem_statement, **kwargs):
+    def __init__(
+        self, instance_id, structure, problem_statement, model_name, backend, **kwargs
+    ):
         super().__init__(instance_id, structure, problem_statement)
         self.max_tokens = 300
+        self.model_name = model_name
+        self.backend = backend
 
     def _parse_model_return_lines(self, content: str) -> list[str]:
         return content.strip().split("\n")
 
     def localize(self, top_n=1, mock=False) -> tuple[list, list, list, any]:
+        # lazy import, not sure if this is actually better?
+        from agentless.util.api_requests import num_tokens_from_messages
+        from agentless.util.model import make_model
 
         found_files = []
 
-        # lazy import, not sure if this is actually better?
-        from agentless.util.api_requests import (
-            create_chatgpt_config,
-            num_tokens_from_messages,
-            request_chatgpt_engine,
-        )
-
         message = self.obtain_relevant_files_prompt.format(
             problem_statement=self.problem_statement,
             structure=show_project_structure(self.structure).strip(),
@@ -246,23 +246,16 @@ def localize(self, top_n=1, mock=False) -> tuple[list, list, list, any]:
             }
             return [], {"raw_output_loc": ""}, traj
 
-        config = create_chatgpt_config(
-            message=message,
+        model = make_model(
+            model=self.model_name,
+            backend=self.backend,
             max_tokens=self.max_tokens,
             temperature=0,
             batch_size=1,
-            model="gpt-4o-2024-05-13",  # use gpt-4o for now.
         )
-        ret = request_chatgpt_engine(config)
-        raw_output = ret.choices[0].message.content
-        traj = {
-            "prompt": message,
-            "response": raw_output,
-            "usage": {
-                "prompt_tokens": ret.usage.prompt_tokens,
-                "completion_tokens": ret.usage.completion_tokens,
-            },
-        }
+        traj = model.codegen(message, num_samples=1)[0]
+        traj["prompt"] = message
+        raw_output = traj["response"]
         model_found_files = self._parse_model_return_lines(raw_output)
 
         files, classes, functions = get_full_file_paths_and_classes_and_functions(
@@ -288,11 +281,8 @@ def localize(self, top_n=1, mock=False) -> tuple[list, list, list, any]:
     def localize_function_for_files(
         self, file_names, mock=False
     ) -> tuple[list, dict, dict]:
-        from agentless.util.api_requests import (
-            create_chatgpt_config,
-            num_tokens_from_messages,
-            request_chatgpt_engine,
-        )
+        from agentless.util.api_requests import num_tokens_from_messages
+        from agentless.util.model import make_model
 
         files, classes, functions = get_full_file_paths_and_classes_and_functions(
             self.structure
@@ -339,23 +329,16 @@ def localize_function_for_files(
             }
             return [], {"raw_output_loc": ""}, traj
 
-        config = create_chatgpt_config(
-            message=message,
+        model = make_model(
+            model=self.model_name,
+            backend=self.backend,
             max_tokens=self.max_tokens,
             temperature=0,
             batch_size=1,
-            model="gpt-4o-2024-05-13",  # use gpt-4o for now.
         )
-        ret = request_chatgpt_engine(config)
-        raw_output = ret.choices[0].message.content
-        traj = {
-            "prompt": message,
-            "response": raw_output,
-            "usage": {
-                "prompt_tokens": ret.usage.prompt_tokens,
-                "completion_tokens": ret.usage.completion_tokens,
-            },
-        }
+        traj = model.codegen(message, num_samples=1)[0]
+        traj["prompt"] = message
+        raw_output = traj["response"]
 
         model_found_locs = extract_code_blocks(raw_output)
         model_found_locs_separated = extract_locs_for_files(
@@ -367,11 +350,8 @@ def localize_function_for_files(
         return model_found_locs_separated, {"raw_output_loc": raw_output}, traj
 
     def localize_function_from_compressed_files(self, file_names, mock=False):
-        from agentless.util.api_requests import (
-            create_chatgpt_config,
-            num_tokens_from_messages,
-            request_chatgpt_engine,
-        )
+        from agentless.util.api_requests import num_tokens_from_messages
+        from agentless.util.model import make_model
 
         file_contents = get_repo_files(self.structure, file_names)
         compressed_file_contents = {
@@ -397,30 +377,23 @@ def localize_function_from_compressed_files(self, file_names, mock=False):
                 "prompt": message,
                 "usage": {
                     "prompt_tokens": num_tokens_from_messages(
-                        message, "gpt-4o-2024-05-13"
+                        message,
+                        self.model_name,
                     ),
                 },
             }
             return [], {"raw_output_loc": ""}, traj
 
-        config = create_chatgpt_config(
-            message=message,
+        model = make_model(
+            model=self.model_name,
+            backend=self.backend,
             max_tokens=self.max_tokens,
             temperature=0,
             batch_size=1,
-            model="gpt-4o-2024-05-13",  # use gpt-4o for now.
         )
-        ret = request_chatgpt_engine(config)
-        raw_output = ret.choices[0].message.content
-        traj = {
-            "prompt": message,
-            "response": raw_output,
-            "usage": {
-                "prompt_tokens": ret.usage.prompt_tokens,
-                "completion_tokens": ret.usage.completion_tokens,
-            },
-        }
-
+        traj = model.codegen(message, num_samples=1)[0]
+        traj["prompt"] = message
+        raw_output = traj["response"]
         model_found_locs = extract_code_blocks(raw_output)
         model_found_locs_separated = extract_locs_for_files(
             model_found_locs, file_names
@@ -450,11 +423,8 @@ def localize_line_from_coarse_function_locs(
         num_samples: int = 1,
         mock=False,
     ):
-        from agentless.util.api_requests import (
-            create_chatgpt_config,
-            num_tokens_from_messages,
-            request_chatgpt_engine,
-        )
+        from agentless.util.api_requests import num_tokens_from_messages
+        from agentless.util.model import make_model
 
         file_contents = get_repo_files(self.structure, file_names)
         topn_content, file_loc_intervals = construct_topn_file_context(
@@ -488,21 +458,28 @@ def localize_line_from_coarse_function_locs(
                 },
             }
             return [], {"raw_output_loc": ""}, traj
-        config = create_chatgpt_config(
-            message=message,
+
+        model = make_model(
+            model=self.model_name,
+            backend=self.backend,
             max_tokens=self.max_tokens,
             temperature=temperature,
             batch_size=num_samples,
-            model="gpt-4o-2024-05-13",  # use gpt-4o for now.
         )
-        ret = request_chatgpt_engine(config)
-        raw_outputs = [choice.message.content for choice in ret.choices]
+        raw_trajs = model.codegen(message, num_samples=num_samples)
+
+        # Merge trajectories
+        raw_outputs = [raw_traj["response"] for raw_traj in raw_trajs]
         traj = {
             "prompt": message,
             "response": raw_outputs,
-            "usage": {
-                "prompt_tokens": ret.usage.prompt_tokens,
-                "completion_tokens": ret.usage.completion_tokens,
+            "usage": {  # merge token usage
+                "completion_tokens": sum(
+                    raw_traj["usage"]["completion_tokens"] for raw_traj in raw_trajs
+                ),
+                "prompt_tokens": sum(
+                    raw_traj["usage"]["prompt_tokens"] for raw_traj in raw_trajs
+                ),
             },
         }
         model_found_locs_separated_in_samples = []

agentless/fl/localize.py

@@ -73,6 +73,8 @@ def localize(args):
                 d["instance_id"],
                 structure,
                 problem_statement,
+                args.model,
+                args.backend,
             )
             found_files, additional_artifact_loc_file, file_traj = fl.localize(
                 mock=args.mock
@@ -101,6 +103,8 @@ def localize(args):
                     d["instance_id"],
                     structure,
                     problem_statement,
+                    args.model,
+                    args.backend,
                 )
 
                 additional_artifact_loc_related = []
@@ -128,6 +132,8 @@ def localize(args):
                 instance_id,
                 structure,
                 problem_statement,
+                args.model,
+                args.backend,
             )
             coarse_found_locs = {}
             for i, pred_file in enumerate(pred_files):
@@ -264,6 +270,10 @@ def main():
     parser.add_argument(
         "--mock", action="store_true", help="Mock run to compute prompt tokens."
     )
+    parser.add_argument(
+        "--model", type=str, default="gpt-4o-2024-05-13", choices=["gpt-4o-2024-05-13"]
+    )
+    parser.add_argument("--backend", type=str, default="openai", choices=["openai"])
 
     args = parser.parse_args()
 

agentless/util/model.py

@@ -19,7 +19,7 @@ def __init__(
         self.max_new_tokens = max_new_tokens
 
     @abstractmethod
-    def codegen(self, message: str, num_samples: int = 1) -> List[str]:
+    def codegen(self, message: str, num_samples: int = 1) -> List[dict]:
         pass
 
     @abstractmethod
@@ -37,7 +37,9 @@ class OpenAIChatDecoder(DecoderBase):
     def __init__(self, name: str, **kwargs) -> None:
         super().__init__(name, **kwargs)
 
-    def codegen(self, message: str, num_samples: int = 1) -> List[str]:
+    def codegen(self, message: str, num_samples: int = 1) -> List[dict]:
+        if self.temperature == 0:
+            assert num_samples == 1
         batch_size = min(self.batch_size, num_samples)
 
         config = create_chatgpt_config(