beswarm 0.2.82__py3-none-any.whl → 0.2.84__py3-none-any.whl

beswarm/agents/planact.py

@@ -5,17 +5,16 @@ import json
 import difflib
 import asyncio
 import tomllib
-import platform
 from pathlib import Path
-from datetime import datetime
 from typing import List, Dict, Union
 
 from ..broker import MessageBroker
 from ..aient.aient.models import chatgpt
 from ..aient.aient.plugins import get_function_call_list, registry
-from ..prompt import worker_system_prompt, instruction_system_prompt
+from ..prompt import worker_system_prompt, instruction_system_prompt, Goal
 from ..utils import extract_xml_content, get_current_screen_image_message, replace_xml_content, register_mcp_tools, setup_logger
 from ..aient.aient.models.chatgpt import ModelNotFoundError, TaskComplete, RetryFailedError, InputTokenCountExceededError, BadRequestError
+from ..aient.aient.architext.architext import Messages, SystemMessage, UserMessage, AssistantMessage, ToolCalls, ToolResults, Texts, RoleMessage, Images, Files
 
 try:
     from importlib import metadata
@@ -31,24 +30,22 @@ except metadata.PackageNotFoundError:
 
 class BaseAgent:
     """Base class for agents, handling common initialization and disposal."""
-    def __init__(self, goal: str, tools_json: List, agent_config: Dict, work_dir: str, cache_messages: Union[bool, List[Dict]], broker: MessageBroker, listen_topic: str, publish_topic: str, status_topic: str, graph_update_topic: str = None):
+    def __init__(self, goal: str, tools_json: List, agent_config: Dict, work_dir: str, cache_messages: Union[bool, List[Dict]], broker: MessageBroker, listen_topic: str, publish_topic: str, status_topic: str):
         self.goal = goal
         self.tools_json = tools_json
         self.work_dir = work_dir
+        self.pkl_file = Path(work_dir) / ".beswarm" / "history.pkl"
         self.cache_file = Path(work_dir) / ".beswarm" / "work_agent_conversation_history.json"
         self.config = agent_config
         self.logger = agent_config.get("logger", None)
         self.cache_messages = cache_messages
         if cache_messages and isinstance(cache_messages, bool) and cache_messages == True:
-            self.cache_messages = json.loads(self.cache_file.read_text(encoding="utf-8"))
+            self.cache_messages = Messages.load(self.pkl_file)
         self.broker = broker
         self.listen_topic = listen_topic
         self.error_topic = listen_topic + ".error"
         self.publish_topic = publish_topic
         self.status_topic = status_topic
-        self.graph_update_topic = graph_update_topic
-        if self.graph_update_topic:
-            self.graph_update_subscription = self.broker.subscribe(self.handle_graph_update, self.graph_update_topic)
 
         self._subscription = self.broker.subscribe(self.handle_message, [self.listen_topic, self.error_topic])
 
@@ -58,11 +55,6 @@ class BaseAgent:
         """Process incoming messages. Must be implemented by subclasses."""
         raise NotImplementedError
 
-    def handle_graph_update(self, message: Dict):
-        """Handle graph update messages."""
-        if message.get("message") == "graph_updated":
-            self.graph_tree = message.get("graph")
-
     def dispose(self):
         """Cancels the subscription and cleans up resources."""
         if self._subscription:
@@ -71,16 +63,16 @@ class BaseAgent:
 
 class InstructionAgent(BaseAgent):
     """Generates instructions and publishes them to a message broker."""
-    def __init__(self, goal: str, tools_json: List, agent_config: Dict, work_dir: str, cache_messages: Union[bool, List[Dict]], broker: MessageBroker, listen_topic: str, publish_topic: str, status_topic: str, graph_update_topic: str):
-        super().__init__(goal, tools_json, agent_config, work_dir, cache_messages, broker, listen_topic, publish_topic, status_topic, graph_update_topic)
+    def __init__(self, goal: str, tools_json: List, agent_config: Dict, work_dir: str, cache_messages: Union[bool, List[Dict]], broker: MessageBroker, listen_topic: str, publish_topic: str, status_topic: str):
+        super().__init__(goal, tools_json, agent_config, work_dir, cache_messages, broker, listen_topic, publish_topic, status_topic)
 
         self.last_instruction = None
         self.agent = chatgpt(**self.config)
 
         self.goal_diff = None
 
-        if self.cache_messages and isinstance(self.cache_messages, list) and len(self.cache_messages) > 1:
-            old_goal = extract_xml_content(self.cache_messages[1]["content"], "goal")
+        if self.cache_messages and self.cache_messages.provider("goal"):
+            old_goal = self.cache_messages.provider("goal").content
             if old_goal.strip() != goal.strip():
                 diff_generator = difflib.ndiff(old_goal.splitlines(), goal.splitlines())
                 changed_lines = []
@@ -89,38 +81,14 @@ class InstructionAgent(BaseAgent):
                         changed_lines.append(line)
                 self.goal_diff = '\n'.join(changed_lines).strip()
 
-    def get_conversation_history(self, raw_conversation_history: List[Dict]):
-
-        for index, message in enumerate(raw_conversation_history):
-            if message.get("content") and isinstance(message["content"], str):
-                if "<knowledge_graph_tree>" in message["content"] and self.graph_tree:
-                    message["content"] = replace_xml_content(message["content"], "knowledge_graph_tree", self.graph_tree)
-                    raw_conversation_history[index] = message
-                if "\n\nYour message **must** end with [done] to signify the end of your output." in message["content"]:
-                    message["content"] = message["content"].replace("\n\nYour message **must** end with [done] to signify the end of your output.", "")
-
+    async def get_conversation_history(self, raw_conversation_history: List[Dict]):
         conversation_history = copy.deepcopy(raw_conversation_history)
-        self.cache_file.write_text(json.dumps(conversation_history, ensure_ascii=False, indent=4), encoding="utf-8")
-
-        work_agent_system_prompt = conversation_history.pop(0)
-        if conversation_history:
-            original_content = work_agent_system_prompt["content"]
-            regex = r"<latest_file_content>(.*?)</latest_file_content>"
-            match = re.search(regex, original_content, re.DOTALL)
-            if match:
-                extracted_content = f"<latest_file_content>{match.group(1)}</latest_file_content>\n\n"
-            else:
-                extracted_content = ""
-            if isinstance(conversation_history[0]["content"], str):
-                conversation_history[0]["content"] = extracted_content + conversation_history[0]["content"]
-            elif isinstance(conversation_history[0]["content"], list) and extracted_content:
-                conversation_history[0]["content"].append({"type": "text", "text": extracted_content})
-
-            for message in conversation_history:
-                if message.get("content") and isinstance(message["content"], str):
-                    if "<knowledge_graph_tree>" in message["content"] and self.graph_tree:
-                        message["content"] = replace_xml_content(message["content"], "knowledge_graph_tree", self.graph_tree)
-
+        conversation_history.save(self.pkl_file)
+        self.cache_file.write_text(json.dumps(await conversation_history.render_latest(), ensure_ascii=False, indent=4), encoding="utf-8")
+        latest_file_content = conversation_history.pop("files")
+        conversation_history.pop(0)
+        if conversation_history and latest_file_content:
+            conversation_history[0] = latest_file_content + conversation_history[0]
 
         return conversation_history
 
@@ -141,7 +109,7 @@ class InstructionAgent(BaseAgent):
                 f"这是你上次给assistant的错误格式的指令：\n{self.last_instruction}"
             )
 
-        self.agent.conversation["default"][1:] = self.get_conversation_history(message["conversation"])
+        self.agent.conversation["default"][1:] = await self.get_conversation_history(message["conversation"])
 
         if "find_and_click_element" in json.dumps(self.tools_json):
             instruction_prompt = await get_current_screen_image_message(instruction_prompt)
@@ -183,7 +151,7 @@ class InstructionAgent(BaseAgent):
                     "现在开始执行第一步：\n"
                     f"{instruction}"
                 )
-            self.broker.publish({"instruction": instruction + "\n\nYour message **must** end with [done] to signify the end of your output.", "conversation": message["conversation"]}, self.publish_topic)
+            self.broker.publish({"instruction": instruction, "conversation": message["conversation"]}, self.publish_topic)
             self.last_instruction = None
         else:
             self.logger.error("\n❌ 指令智能体生成的指令不符合要求，正在重新生成。")
@@ -193,8 +161,8 @@ class InstructionAgent(BaseAgent):
 
 class WorkerAgent(BaseAgent):
     """Executes instructions and publishes results to a message broker."""
-    def __init__(self, goal: str, tools_json: List, agent_config: Dict, work_dir: str, cache_messages: Union[bool, List[Dict]], broker: MessageBroker, listen_topic: str, publish_topic: str, status_topic: str, graph_update_topic: str):
-        super().__init__(goal, tools_json, agent_config, work_dir, cache_messages, broker, listen_topic, publish_topic, status_topic, graph_update_topic)
+    def __init__(self, goal: str, tools_json: List, agent_config: Dict, work_dir: str, cache_messages: Union[bool, List[Dict]], broker: MessageBroker, listen_topic: str, publish_topic: str, status_topic: str):
+        super().__init__(goal, tools_json, agent_config, work_dir, cache_messages, broker, listen_topic, publish_topic, status_topic)
 
         if self.cache_messages and isinstance(self.cache_messages, list) and len(self.cache_messages) > 1:
             first_user_message = replace_xml_content(self.cache_messages[1]["content"], "goal", goal)
@@ -211,18 +179,12 @@ class WorkerAgent(BaseAgent):
             }, self.publish_topic)
             return
 
-        for index, raw_message in enumerate(self.agent.conversation["default"]):
-            if raw_message.get("content") and isinstance(raw_message["content"], str) \
-            and "<knowledge_graph_tree>" in raw_message["content"] and self.graph_tree:
-                raw_message["content"] = replace_xml_content(raw_message["content"], "knowledge_graph_tree", self.graph_tree)
-                self.agent.conversation["default"][index] = raw_message
-
         instruction = message["instruction"]
         if "find_and_click_element" in json.dumps(self.tools_json):
             instruction = await get_current_screen_image_message(instruction)
 
         try:
-            response = await self.agent.ask_async(instruction)
+            response = await self.agent.ask_async(UserMessage(instruction, Texts("\n\nYour message **must** end with [done] to signify the end of your output.", name="done")))
         except TaskComplete as e:
             self.broker.publish({"status": "finished", "result": e.completion_message}, self.status_topic)
             return
@@ -263,7 +225,6 @@ class BrokerWorker:
         self.INSTRUCTION_TOPIC = self.channel + ".instructions"
         self.WORKER_RESPONSE_TOPIC = self.channel + ".worker_responses"
         self.TASK_STATUS_TOPIC =self.channel +  ".task_status"
-        self.GRAPH_UPDATE_TOPIC = self.channel + ".knowledge_graph"
 
         self.setup()
 
@@ -272,7 +233,6 @@ class BrokerWorker:
         cache_dir.mkdir(parents=True, exist_ok=True)
         self.task_manager.set_root_path(self.work_dir)
         self.kgm.set_root_path(self.work_dir)
-        self.kgm.set_publish_topic(self.GRAPH_UPDATE_TOPIC)
         self.cache_file = cache_dir / "work_agent_conversation_history.json"
         if not self.cache_file.exists():
             self.cache_file.write_text("[]", encoding="utf-8")
@@ -313,25 +273,22 @@ class BrokerWorker:
             self.logger.info(message.get("result"))
 
     def _setup_agents(self):
+        instruction_system_prompt.provider("tools").update(self.tools_json)
+        instruction_system_prompt.provider("workspace_path").update(str(self.work_dir))
         instruction_agent_config = {
             "api_key": os.getenv("API_KEY"), "api_url": os.getenv("BASE_URL"),
             "engine": os.getenv("MODEL"),
-            "system_prompt": instruction_system_prompt.format(
-                os_version=platform.platform(), tools_list=self.tools_json,
-                workspace_path=self.work_dir, current_time=datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-            ),
+            "system_prompt": instruction_system_prompt,
             "print_log": os.getenv("DEBUG", "false").lower() in ("true", "1", "t", "yes"),
             "temperature": 0.7, "use_plugins": False, "logger": self.logger
         }
 
+        worker_system_prompt.provider("tools").update(self.tools_json)
+        worker_system_prompt.provider("workspace_path").update(str(self.work_dir))
         worker_agent_config = {
             "api_key": os.getenv("API_KEY"), "api_url": os.getenv("BASE_URL"),
             "engine": os.getenv("FAST_MODEL") or os.getenv("MODEL"),
-            "system_prompt": worker_system_prompt.format(
-                os_version=platform.platform(), workspace_path=self.work_dir,
-                shell=os.getenv('SHELL', 'Unknown'), current_time=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-                tools_list=self.tools_json
-            ),
+            "system_prompt": worker_system_prompt,
             "print_log": True, "temperature": 0.5, "function_call_max_loop": 100, "logger": self.logger,
             "check_done": True
         }
@@ -340,14 +297,12 @@ class BrokerWorker:
             goal=self.goal, tools_json=self.tools_json, agent_config=instruction_agent_config, work_dir=self.work_dir, cache_messages=self.cache_messages,
             broker=self.broker, listen_topic=self.WORKER_RESPONSE_TOPIC,
             publish_topic=self.INSTRUCTION_TOPIC, status_topic=self.TASK_STATUS_TOPIC,
-            graph_update_topic=self.GRAPH_UPDATE_TOPIC
         )
 
         worker_agent = WorkerAgent(
             goal=self.goal, tools_json=self.tools_json, agent_config=worker_agent_config, work_dir=self.work_dir, cache_messages=self.cache_messages,
             broker=self.broker, listen_topic=self.INSTRUCTION_TOPIC,
             publish_topic=self.WORKER_RESPONSE_TOPIC, status_topic=self.TASK_STATUS_TOPIC,
-            graph_update_topic=self.GRAPH_UPDATE_TOPIC
         )
         return instruction_agent, worker_agent