AgentAssistant

rootfs/agents/AgentAssistant/environment.py

@@ -0,0 +1,59 @@
+import logging
+from typing import Optional
+
+import toml
+
+from aios_kernel import Environment, SimpleAIFunction
+import os
+
+
+local_path = os.path.split(os.path.realpath(__file__))[0]
+
+logger = logging.getLogger(__name__)
+
+
+class AgentAssistantEnvironment(Environment):
+    def __init__(self):
+        super().__init__("agent_assistant_env")
+        self.add_ai_function(SimpleAIFunction("read_agent",
+                                              "Read the agent with the specified name",
+                                              self.read_agent,
+                                              {"agent_id": "The id of the agent to be read"}))
+        self.add_ai_function(SimpleAIFunction("save_agent",
+                                              "save the agent with the specified name",
+                                              self.save_agent,
+                                              {"agent_id": "The id of the agent to be saved",
+                                                          "agent_data": "The toml data of the agent to be saved",
+                                                          "is_new": "Whether to create a new agent, The value is true if it is created, and it is false if it is modified."}))
+
+    def _do_get_value(self, key: str) -> Optional[str]:
+        pass
+
+    async def read_agent(self, agent_id: str) -> str:
+        agent_dir = os.path.join(local_path, "..", agent_id)
+        if not os.path.exists(agent_dir):
+            return "exec failed.agent is not exists."
+
+        with open(os.path.join(agent_dir, "agent.toml"), "r", encoding='utf-8') as f:
+            agent_data = f.read()
+
+        return "exec success. agent data:\n" + agent_data
+
+    async def save_agent(self, agent_id: str, agent_data: str, is_new: str) -> str:
+        logger.info(f"save_agent: {agent_id} {agent_data}")
+
+        agent_dir = os.path.join(local_path, "..", agent_id)
+        if os.path.exists(agent_dir) and is_new == "true":
+            return "exec failed.agent already exists, please change the agent id and try again."
+
+        if not os.path.exists(agent_dir):
+            os.mkdir(agent_dir)
+
+        with open(os.path.join(agent_dir, "agent.toml"), "w", encoding='utf-8') as f:
+            f.write(agent_data)
+
+        return "exec success."
+
+
+def init():
+    return AgentAssistantEnvironment()