auto-coder 0.1.206__py3-none-any.whl → 0.1.208__py3-none-any.whl

autocoder/common/code_auto_generate.py

@@ -4,6 +4,8 @@ from autocoder.common import AutoCoderArgs
 import byzerllm
 from autocoder.utils.queue_communicate import queue_communicate, CommunicateEvent, CommunicateEventType
 from autocoder.common import sys_prompt
+from concurrent.futures import ThreadPoolExecutor
+from autocoder.common.types import CodeGenerateResult
 
 
 class CodeAutoGenerate:
@@ -12,13 +14,15 @@ class CodeAutoGenerate:
     ) -> None:
         self.llm = llm
         self.args = args
-        self.action = action
+        self.action = action        
+        self.generate_times_same_model = args.generate_times_same_model
         if not self.llm:
             raise ValueError(
                 "Please provide a valid model instance to use for code generation."
             )
-        if self.llm.get_sub_client("code_model"):
-            self.llm = self.llm.get_sub_client("code_model")
+        self.llms = self.llm.get_sub_client("code_model") or [self.llm]
+        if not isinstance(self.llms, list):
+            self.llms = [self.llms]
 
     @byzerllm.prompt(llm=lambda self: self.llm)
     def auto_implement_function(self, instruction: str, content: str) -> str:
@@ -145,7 +149,7 @@ class CodeAutoGenerate:
 
     def single_round_run(
         self, query: str, source_content: str
-    ) -> Tuple[str, Dict[str, str]]:
+    ) -> Tuple[List[str], Dict[str, str]]:
         llm_config = {"human_as_model": self.args.human_as_model}
 
         if self.args.request_id and not self.args.skip_events:
@@ -178,9 +182,27 @@ class CodeAutoGenerate:
         
         conversations.append({"role": "user", "content": init_prompt})
 
-
-        t = self.llm.chat_oai(conversations=conversations, llm_config=llm_config)
-        conversations.append({"role": "assistant", "content": t[0].output})
+        conversations_list = []
+        results = []
+        if not self.args.human_as_model:
+            with ThreadPoolExecutor(max_workers=len(self.llms) * self.generate_times_same_model) as executor:
+                futures = []
+                for llm in self.llms:
+                    for _ in range(self.generate_times_same_model):
+                        futures.append(executor.submit(
+                            llm.chat_oai, conversations=conversations, llm_config=llm_config))
+                results = [future.result()[0].output for future in futures]
+            for result in results:
+                conversations_list.append(
+                    conversations + [{"role": "assistant", "content": result}])
+        else:
+            results = []
+            conversations_list = []
+            for _ in range(self.args.human_model_num):
+                v = self.llms[0].chat_oai(
+                    conversations=conversations, llm_config=llm_config)
+                results.append(v[0].output)
+                conversations_list.append(conversations + [{"role": "assistant", "content": v[0].output}])
 
         if self.args.request_id and not self.args.skip_events:
             queue_communicate.send_event_no_wait(
@@ -191,7 +213,7 @@ class CodeAutoGenerate:
                 ),
             )
 
-        return [t[0].output], conversations
+        return CodeGenerateResult(contents=results, conversations=conversations_list)
 
     def multi_round_run(
         self, query: str, source_content: str, max_steps: int = 10
@@ -246,6 +268,6 @@ class CodeAutoGenerate:
                 or "/done" in t[0].output
                 or "__EOF__" in t[0].output
             ):
-                return result, conversations
+                return CodeGenerateResult(contents=["\n\n".join(result)], conversations=[conversations])
 
-        return result, conversations
+        return CodeGenerateResult(contents=["\n\n".join(result)], conversations=[conversations])

autocoder/common/code_auto_generate_diff.py

@@ -1,9 +1,12 @@
 from typing import List, Dict, Tuple
-from autocoder.common.types import Mode
+from autocoder.common.types import Mode, CodeGenerateResult
 from autocoder.common import AutoCoderArgs
 import byzerllm
 from autocoder.utils.queue_communicate import queue_communicate, CommunicateEvent, CommunicateEventType
 from autocoder.common import sys_prompt
+from concurrent.futures import ThreadPoolExecutor
+import json
+
 
 class CodeAutoGenerateDiff:
     def __init__(
@@ -12,12 +15,15 @@ class CodeAutoGenerateDiff:
         self.llm = llm
         self.args = args
         self.action = action
+        self.llms = []
+        self.generate_times_same_model = args.generate_times_same_model
         if not self.llm:
             raise ValueError(
                 "Please provide a valid model instance to use for code diff generation."
             )
-        if self.llm.get_sub_client("code_model"):
-            self.llm = self.llm.get_sub_client("code_model")
+        self.llms = self.llm.get_sub_client("code_model") or [self.llm]
+        if not isinstance(self.llms, list):
+            self.llms = [self.llms]
 
     @byzerllm.prompt(llm=lambda self: self.llm)
     def multi_round_instruction(
@@ -139,12 +145,12 @@ class CodeAutoGenerateDiff:
 
         每次生成一个文件的diff，然后询问我是否继续，当我回复继续，继续生成下一个文件的diff。当没有后续任务时，请回复 "__完成__" 或者 "__EOF__"。
         """
-        
+
         if not self.args.include_project_structure:
             return {
                 "structure": ""
             }
-        
+
         return {
             "structure": (
                 self.action.pp.get_tree_like_directory_structure()
@@ -271,12 +277,12 @@ class CodeAutoGenerateDiff:
 
         {{ instruction }}
         """
-        
+
         if not self.args.include_project_structure:
             return {
                 "structure": ""
             }
-        
+
         return {
             "structure": (
                 self.action.pp.get_tree_like_directory_structure()
@@ -287,72 +293,107 @@ class CodeAutoGenerateDiff:
 
     def single_round_run(
         self, query: str, source_content: str
-    ) -> Tuple[str, Dict[str, str]]:
+    ) -> CodeGenerateResult:
         llm_config = {"human_as_model": self.args.human_as_model}
 
-        if self.args.request_id and not self.args.skip_events:
-            queue_communicate.send_event_no_wait(
-                request_id=self.args.request_id,
-                event=CommunicateEvent(
-                    event_type=CommunicateEventType.CODE_GENERATE_START.value,
-                    data=query,
-                ),
+        if self.args.template == "common":
+            init_prompt = self.single_round_instruction.prompt(
+                instruction=query, content=source_content, context=self.args.context
+            )
+        elif self.args.template == "auto_implement":
+            init_prompt = self.auto_implement_function.prompt(
+                instruction=query, content=source_content
             )
-
-        init_prompt = self.single_round_instruction.prompt(
-            instruction=query, content=source_content, context=self.args.context
-        )
 
         with open(self.args.target_file, "w") as file:
             file.write(init_prompt)
 
         conversations = []
+
         if self.args.system_prompt and self.args.system_prompt.strip() == "claude":
-            conversations.append({"role": "system", "content": sys_prompt.claude_sys_prompt.prompt()})
-        else:
-            conversations.append({"role": "user", "content": init_prompt})
+            conversations.append(
+                {"role": "system", "content": sys_prompt.claude_sys_prompt.prompt()})
+        elif self.args.system_prompt:
+            conversations.append(
+                {"role": "system", "content": self.args.system_prompt})
 
-        t = self.llm.chat_oai(conversations=conversations, llm_config=llm_config)
-        conversations.append({"role": "assistant", "content": t[0].output})
+        conversations.append({"role": "user", "content": init_prompt})
 
         if self.args.request_id and not self.args.skip_events:
-            queue_communicate.send_event_no_wait(
+            _ = queue_communicate.send_event(
+                request_id=self.args.request_id,
+                event=CommunicateEvent(
+                    event_type=CommunicateEventType.CODE_GENERATE_START.value,
+                    data=json.dumps({}, ensure_ascii=False),
+                ),
+            )
+
+        conversations_list = []
+        results = []
+        if not self.args.human_as_model:
+            with ThreadPoolExecutor(max_workers=len(self.llms) * self.generate_times_same_model) as executor:
+                futures = []
+                for llm in self.llms:
+                    for _ in range(self.generate_times_same_model):
+                        futures.append(executor.submit(
+                            llm.chat_oai, conversations=conversations, llm_config=llm_config))
+                results = [future.result()[0].output for future in futures]
+            for result in results:
+                conversations_list.append(
+                    conversations + [{"role": "assistant", "content": result}])
+        else:
+            results = []
+            conversations_list = []
+            for _ in range(self.args.human_model_num):
+                v = self.llms[0].chat_oai(
+                    conversations=conversations, llm_config=llm_config)
+                results.append(v[0].output)
+                conversations_list.append(
+                    conversations + [{"role": "assistant", "content": v[0].output}])
+
+        if self.args.request_id and not self.args.skip_events:
+            _ = queue_communicate.send_event(
                 request_id=self.args.request_id,
                 event=CommunicateEvent(
                     event_type=CommunicateEventType.CODE_GENERATE_END.value,
-                    data="",
+                    data=json.dumps({}, ensure_ascii=False),
                 ),
             )
 
-        return [t[0].output], conversations
+        return CodeGenerateResult(contents=results, conversations=conversations_list)
 
     def multi_round_run(
         self, query: str, source_content: str, max_steps: int = 10
-    ) -> Tuple[List[str], List[Dict[str, str]]]:
+    ) -> CodeGenerateResult:
         llm_config = {"human_as_model": self.args.human_as_model}
         result = []
 
-        init_prompt = self.multi_round_instruction.prompt(
-            instruction=query, content=source_content, context=self.args.context
-        )
+        if self.args.template == "common":
+            init_prompt = self.multi_round_instruction.prompt(
+                instruction=query, content=source_content, context=self.args.context
+            )
+        elif self.args.template == "auto_implement":
+            init_prompt = self.auto_implement_function.prompt(
+                instruction=query, content=source_content
+            )
 
-        conversations = [{"role": "user", "content": init_prompt}]
+        conversations = []
+        # conversations.append({"role": "system", "content": sys_prompt.prompt()})
+        conversations.append({"role": "user", "content": init_prompt})
 
         with open(self.args.target_file, "w") as file:
             file.write(init_prompt)
 
-        t = self.llm.chat_oai(conversations=conversations, llm_config=llm_config)
+        code_llm = self.llms[0]
+        t = code_llm.chat_oai(conversations=conversations,
+                              llm_config=llm_config)
 
         result.append(t[0].output)
 
         conversations.append({"role": "assistant", "content": t[0].output})
 
-        if (
-            "__完成__" in t[0].output
-            or "/done" in t[0].output
-            or "__EOF__" in t[0].output
-        ):
-            return result, conversations
+        if "__完成__" in t[0].output or "/done" in t[0].output or "__EOF__" in t[0].output:
+            return CodeGenerateResult(contents=["\n\n".join(result)], conversations=[conversations])
 
         current_step = 0
 
@@ -363,17 +404,14 @@ class CodeAutoGenerateDiff:
             with open(self.args.target_file, "w") as file:
                 file.write("继续")
 
-            t = self.llm.chat_oai(conversations=conversations, llm_config=llm_config)
+            t = code_llm.chat_oai(
+                conversations=conversations, llm_config=llm_config)
 
             result.append(t[0].output)
             conversations.append({"role": "assistant", "content": t[0].output})
             current_step += 1
 
-            if (
-                "__完成__" in t[0].output
-                or "/done" in t[0].output
-                or "__EOF__" in t[0].output
-            ):
-                return result, conversations
+            if "__完成__" in t[0].output or "/done" in t[0].output or "__EOF__" in t[0].output:
+                return CodeGenerateResult(contents=["\n\n".join(result)], conversations=[conversations])
 
-        return result, conversations
+        return CodeGenerateResult(contents=["\n\n".join(result)], conversations=[conversations])

autocoder/common/code_auto_generate_editblock.py

@@ -1,5 +1,5 @@
 from typing import List, Dict, Tuple
-from autocoder.common.types import Mode
+from autocoder.common.types import Mode, CodeGenerateResult
 from autocoder.common import AutoCoderArgs
 import byzerllm
 from autocoder.common import sys_prompt
@@ -9,6 +9,7 @@ from autocoder.utils.queue_communicate import (
     CommunicateEventType,
 )
 import json
+from concurrent.futures import ThreadPoolExecutor
 
 
 class CodeAutoGenerateEditBlock:
@@ -25,14 +26,16 @@ class CodeAutoGenerateEditBlock:
         self.action = action
         self.fence_0 = fence_0
         self.fence_1 = fence_1
+        self.generate_times_same_model = args.generate_times_same_model
         if not self.llm:
             raise ValueError(
                 "Please provide a valid model instance to use for code generation."
             )
-        if self.llm.get_sub_client("code_model"):
-            self.llm = self.llm.get_sub_client("code_model")
+        self.llms = self.llm.get_sub_client("code_model") or [self.llm]
+        if not isinstance(self.llms, list):
+            self.llms = [self.llms]
 
-    @byzerllm.prompt(llm=lambda self: self.llm)
+    @byzerllm.prompt()
     def auto_implement_function(self, instruction: str, content: str) -> str:
         """
         下面是一些文件路径以及每个文件对应的源码：
@@ -49,7 +52,7 @@ class CodeAutoGenerateEditBlock:
 
         """
 
-    @byzerllm.prompt(llm=lambda self: self.llm)
+    @byzerllm.prompt()
     def multi_round_instruction(self, instruction: str, content: str, context: str = "") -> str:
         """
         如果你需要生成代码，对于每个需要更改的文件,你需要按 *SEARCH/REPLACE block* 的格式进行生成。
@@ -211,7 +214,7 @@ class CodeAutoGenerateEditBlock:
             "fence_1": self.fence_1,
         }
 
-    @byzerllm.prompt(llm=lambda self: self.llm)
+    @byzerllm.prompt()
     def single_round_instruction(self, instruction: str, content: str, context: str = "") -> str:
         """
         如果你需要生成代码，对于每个需要更改的文件,你需要按 *SEARCH/REPLACE block* 的格式进行生成。
@@ -372,7 +375,7 @@ class CodeAutoGenerateEditBlock:
 
     def single_round_run(
         self, query: str, source_content: str
-    ) -> Tuple[str, Dict[str, str]]:
+    ) -> CodeGenerateResult:
         llm_config = {"human_as_model": self.args.human_as_model}
 
         if self.args.template == "common":
@@ -388,12 +391,14 @@ class CodeAutoGenerateEditBlock:
             file.write(init_prompt)
 
         conversations = []
-        
+
         if self.args.system_prompt and self.args.system_prompt.strip() == "claude":
-            conversations.append({"role": "system", "content": sys_prompt.claude_sys_prompt.prompt()})
+            conversations.append(
+                {"role": "system", "content": sys_prompt.claude_sys_prompt.prompt()})
         elif self.args.system_prompt:
-            conversations.append({"role": "system", "content": self.args.system_prompt})
-        
+            conversations.append(
+                {"role": "system", "content": self.args.system_prompt})
+
         conversations.append({"role": "user", "content": init_prompt})
 
         if self.args.request_id and not self.args.skip_events:
@@ -404,11 +409,30 @@ class CodeAutoGenerateEditBlock:
                     data=json.dumps({}, ensure_ascii=False),
                 ),
             )
-            
-            t = self.llm.chat_oai(
-                conversations=conversations, llm_config=llm_config)
-            conversations.append({"role": "assistant", "content": t[0].output})
-            
+
+        conversations_list = []
+        results = []
+        if not self.args.human_as_model:
+            with ThreadPoolExecutor(max_workers=len(self.llms) * self.generate_times_same_model) as executor:
+                futures = []
+                for llm in self.llms:
+                    for _ in range(self.generate_times_same_model):
+                        futures.append(executor.submit(
+                            llm.chat_oai, conversations=conversations, llm_config=llm_config))
+                results = [future.result()[0].output for future in futures]
+            for result in results:
+                conversations_list.append(
+                    conversations + [{"role": "assistant", "content": result}])
+        else:
+            results = []
+            conversations_list = []
+            for _ in range(self.args.human_model_num):
+                v = self.llms[0].chat_oai(
+                    conversations=conversations, llm_config=llm_config)
+                results.append(v[0].output)
+                conversations_list.append(conversations + [{"role": "assistant", "content": v[0].output}])
+
+        if self.args.request_id and not self.args.skip_events:
             _ = queue_communicate.send_event(
                 request_id=self.args.request_id,
                 event=CommunicateEvent(
@@ -416,16 +440,12 @@ class CodeAutoGenerateEditBlock:
                     data=json.dumps({}, ensure_ascii=False),
                 ),
             )
-            return [t[0].output], conversations
-        else:
-            t = self.llm.chat_oai(
-                conversations=conversations, llm_config=llm_config)
-            conversations.append({"role": "assistant", "content": t[0].output})
-            return [t[0].output], conversations
+
+        return CodeGenerateResult(contents=results, conversations=conversations_list)
 
     def multi_round_run(
         self, query: str, source_content: str, max_steps: int = 10
-    ) -> Tuple[List[str], List[Dict[str, str]]]:
+    ) -> CodeGenerateResult:
         llm_config = {"human_as_model": self.args.human_as_model}
         result = []
 
@@ -445,7 +465,8 @@ class CodeAutoGenerateEditBlock:
         with open(self.args.target_file, "w") as file:
             file.write(init_prompt)
 
-        t = self.llm.chat_oai(conversations=conversations,
+        code_llm = self.llms[0]
+        t = code_llm.chat_oai(conversations=conversations,
                               llm_config=llm_config)
 
         result.append(t[0].output)
@@ -453,7 +474,7 @@ class CodeAutoGenerateEditBlock:
         conversations.append({"role": "assistant", "content": t[0].output})
 
         if "__完成__" in t[0].output or "/done" in t[0].output or "__EOF__" in t[0].output:
-            return result, conversations
+            return CodeGenerateResult(contents=["\n\n".join(result)], conversations=[conversations])
 
         current_step = 0
 
@@ -464,7 +485,7 @@ class CodeAutoGenerateEditBlock:
             with open(self.args.target_file, "w") as file:
                 file.write("继续")
 
-            t = self.llm.chat_oai(
+            t = code_llm.chat_oai(
                 conversations=conversations, llm_config=llm_config)
 
             result.append(t[0].output)
@@ -472,6 +493,7 @@ class CodeAutoGenerateEditBlock:
             current_step += 1
 
             if "__完成__" in t[0].output or "/done" in t[0].output or "__EOF__" in t[0].output:
-                return result, conversations
 
-        return result, conversations
+                return CodeGenerateResult(contents=["\n\n".join(result)], conversations=[conversations])
+
+        return CodeGenerateResult(contents=["\n\n".join(result)], conversations=[conversations])