From 97f403080e469d841ebdae7529504a2a6fd08d26 Mon Sep 17 00:00:00 2001
From: Liu Zhicong <waterflier@gmail.com>
Date: Thu, 21 Sep 2023 00:39:39 -0700
Subject: [PATCH] Add feature: Agent/Role enable_function config.

---
 rootfs/agents/Jarvis/agent.toml |  1 +
 src/aios_kernel/agent.py        | 14 ++++++++++++--
 src/aios_kernel/role.py         |  4 ++++
 src/aios_kernel/workflow.py     | 12 +++++++++---
 4 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/rootfs/agents/Jarvis/agent.toml b/rootfs/agents/Jarvis/agent.toml
index 3eb81ec..575dc87 100644
--- a/rootfs/agents/Jarvis/agent.toml
+++ b/rootfs/agents/Jarvis/agent.toml
@@ -2,6 +2,7 @@ instance_id = "Jarvis"
 fullname = "Jarvis"
 llm_model_name = "gpt-3.5-turbo-16k-0613"
 max_token_size = 16000
+#enable_function =["add_event"]
 
 [[prompt]]
 role = "system"
diff --git a/src/aios_kernel/agent.py b/src/aios_kernel/agent.py
index fb07ac5..18eeeec 100644
--- a/src/aios_kernel/agent.py
+++ b/src/aios_kernel/agent.py
@@ -94,6 +94,7 @@ def load_from_config(self,config:dict) -> bool:
             if self.prompt.load_from_config(config["prompt"]) is False:
                 logger.error("load prompt from config failed!")
                 return False
+
         
         return True
     
@@ -113,6 +114,7 @@ def __init__(self) -> None:
         self.unread_msg = Queue() # msg from other agent
         self.owner_env : Environment = None
         self.owenr_bus = None
+        self.enable_function_list = []
         
     @classmethod
     def create_from_templete(cls,templete:AIAgentTemplete, fullname:str):
@@ -150,7 +152,8 @@ def load_from_config(self,config:dict) -> bool:
             self.llm_model_name = config["llm_model_name"]
         if config.get("max_token_size") is not None:
             self.max_token_size = config["max_token_size"]
-
+        if config.get("enable_function") is not None:
+            self.enable_function_list = config["enable_function"]
         return True
 
 
@@ -235,8 +238,15 @@ def _get_inner_functions(self) -> dict:
         result_func = []
         result_len = 0
         for inner_func in all_inner_function:
+            func_name = inner_func.get_name()
+            if self.enable_function_list:
+                if len(self.enable_function_list) > 0:
+                    if func_name not in self.enable_function_list:
+                        logger.debug(f"ageint {self.agent_id} ignore inner func:{func_name}")
+                        continue
+
             this_func = {}
-            this_func["name"] = inner_func.get_name()
+            this_func["name"] = func_name
             this_func["description"] = inner_func.get_description()
             this_func["parameters"] = inner_func.get_parameters()
             result_len += len(json.dumps(this_func)) / 4
diff --git a/src/aios_kernel/role.py b/src/aios_kernel/role.py
index 1d63e61..2f58894 100644
--- a/src/aios_kernel/role.py
+++ b/src/aios_kernel/role.py
@@ -12,6 +12,7 @@ def __init__(self) -> None:
         self.prompt : AgentPrompt = None
         self.introduce : str = None
         self.agent = None
+        self.enable_function_list : list[str] = None
 
     def load_from_config(self,config:dict) -> bool:
         name_node = config.get("name")
@@ -37,6 +38,9 @@ def load_from_config(self,config:dict) -> bool:
         intro_node = config.get("intro")
         if intro_node is not None:
             self.introduce = intro_node
+
+        if config.get("enable_function") is not None:
+            self.enable_function_list = config["enable_function"]
     
     def get_role_id(self) -> str:
         return self.role_id
diff --git a/src/aios_kernel/workflow.py b/src/aios_kernel/workflow.py
index 48e8a4b..998787c 100644
--- a/src/aios_kernel/workflow.py
+++ b/src/aios_kernel/workflow.py
@@ -368,15 +368,21 @@ def _format_msg_by_env_value(self,prompt:AgentPrompt):
             old_content = msg.get("content")
             msg["content"] = old_content.format_map(self.workflow_env)
 
-    def _get_inner_functions(self) -> dict:
+    def _get_inner_functions(self,the_role:AIRole) -> dict:
         all_inner_function = self.workflow_env.get_all_ai_functions()
         if all_inner_function is None:
             return None
         
         result_func = []
         for inner_func in all_inner_function:
+            func_name = inner_func.get_name()
+            if the_role.enable_function_list:
+                if len(the_role.enable_function_list) > 0:
+                    if func_name not in the_role.enable_function_list:
+                        logger.debug(f"ageint {self.agent_id} ignore inner func:{func_name}")
+                        continue
             this_func = {}
-            this_func["name"] = inner_func.get_name()
+            this_func["name"] = func_name
             this_func["description"] = inner_func.get_description()
             this_func["parameters"] = inner_func.get_parameters()
             result_func.append(this_func)
@@ -436,7 +442,7 @@ async def role_process_msg(self,msg:AgentMsg,the_role:AIRole,workflow_chat_sessi
         prompt.append(msg_prompt)
 
         self._format_msg_by_env_value(prompt)
-        inner_functions = self._get_inner_functions()
+        inner_functions = self._get_inner_functions(the_role)
         
         async def _do_process_msg():
             #TODO: send msg to agent might be better?