Complete AI Function support for Agent.

src/aios_kernel/agent.py

@@ -5,9 +5,13 @@ import asyncio
 import logging
 import uuid
 import time
+import json
 
 from .agent_message import AgentMsg
 from .chatsession import AIChatSession
+from .compute_task import ComputeTaskResult
+from .ai_function import AIFunction
+from .environment import Environment
 
 logger = logging.getLogger(__name__)
 
@@ -77,6 +81,7 @@ class AIAgent:
 
         self.chat_db = None
         self.unread_msg = Queue() # msg from other agent
+        self.owner_env : Environment = None
         
     @classmethod
     def create_from_templete(cls,templete:AIAgentTemplete, fullname:str):
@@ -123,25 +128,70 @@ class AIAgent:
             return "ignore"
         
         return "text"
+    
+    def _get_inner_functions(self) -> dict:
+        if self.owner_env is None:
+            return None
+        
+        all_inner_function = self.owner_env.get_all_ai_functions()
+        if all_inner_function is None:
+            return None
+        
+        result_func = []
+        for inner_func in all_inner_function:
+            this_func = {}
+            this_func["name"] = inner_func.get_name()
+            this_func["description"] = inner_func.get_description()
+            this_func["parameters"] = inner_func.get_parameters()
+            result_func.append(this_func)
+
+        return result_func 
+
+    async def _execute_func(self,inenr_func_call_node:dict,msg_prompt:AgentPrompt) -> str:
+        from .compute_kernel import ComputeKernel
+
+        func_name = inenr_func_call_node.get("name")
+        arguments = json.loads(inenr_func_call_node.get("arguments"))
+
+        func_node : AIFunction = self.owner_env.get_ai_function(func_name)
+        if func_node is None:
+            return "execute failed,function not found"
+        
+        result_str:str = await func_node.execute(**arguments)
+        inner_functions = self._get_inner_functions()
+        msg_prompt.messages.append({"role":"function","content":result_str,"name":func_name})
+        task_result:ComputeTaskResult = await ComputeKernel().do_llm_completion(msg_prompt,self.llm_model_name,self.max_token_size,inner_functions)
+        
+        inner_func_call_node = task_result.result_message.get("function_call")
+        if inner_func_call_node:
+            return await self._execute_func(inner_func_call_node,msg_prompt)       
+        else:
+            return task_result.result_str
 
     async def _process_msg(self,msg:AgentMsg) -> AgentMsg:
             from .compute_kernel import ComputeKernel
+
             session_topic = msg.get_sender() + "#" + msg.topic
             chatsession = AIChatSession.get_session(self.instance_id,session_topic,self.chat_db)
             prompt = AgentPrompt()
             prompt.append(self.prompt)
- 
-            # prompt.append(self._get_function_prompt(the_role.get_name()))
             # prompt.append(self._get_knowlege_prompt(the_role.get_name()))
             prompt.append(await self._get_prompt_from_session(chatsession)) # chat context
             
             msg_prompt = AgentPrompt()
             msg_prompt.messages = [{"role":"user","content":msg.body}]
             prompt.append(msg_prompt)
-    
-            result = await ComputeKernel().do_llm_completion(prompt,self.llm_model_name,self.max_token_size)
-            final_result = result
-            result_type : str = self._get_llm_result_type(result)
+
+            inner_functions = self._get_inner_functions()
+
+            task_result:ComputeTaskResult = await ComputeKernel().do_llm_completion(prompt,self.llm_model_name,self.max_token_size,inner_functions)
+            final_result = task_result.result_str
+
+            inner_func_call_node = task_result.result_message.get("function_call")
+            if inner_func_call_node:
+                final_result = await self._execute_func(inner_func_call_node,msg_prompt)
+            
+            result_type : str = self._get_llm_result_type(final_result)
             is_ignore = False
 
             match result_type:

src/aios_kernel/ai_function.py

@@ -1,9 +1,9 @@
 from abc import ABC, abstractmethod
-from typing import Dict
+from typing import Dict,Coroutine,Callable
 
 class AIFunction:
     def __init__(self) -> None:
-        self.intro : str = None
+        self.description : str = None
     
     @abstractmethod
     def get_name(self) -> str:
@@ -17,7 +17,7 @@ class AIFunction:
         """
         return a detailed description of what the function does
         """
-        pass
+        return self.description
 
     @abstractmethod
     def get_parameters(self) -> Dict:
@@ -25,11 +25,23 @@ class AIFunction:
         Return the list of parameters to execute this function in the form of
         JSON schema as specified in the OpenAI documentation:
         https://platform.openai.com/docs/api-reference/chat/create#chat/create-parameters
+
+        str = run_code(code:str)
+        parameters = {
+            "type": "object",
+            "properties": {
+                "code": {
+                    "type": "string",
+                    "description": "Python code which needs to be executed"
+                }
+            }
+        }
+
         """
         pass
     
     @abstractmethod
-    def execute(self, **kwargs) -> Dict:
+    async def execute(self, **kwargs) -> str:
         """
         Execute the function and return a JSON serializable dict.
         The parameters are passed in the form of kwargs
@@ -66,4 +78,35 @@ class CallChain:
         pass
 
     async def execute(self):
-        pass
+        pass
+
+class SimpleAIFunction(AIFunction):
+    def __init__(self,func_id:str,description:str,func_handler:Coroutine,parameters:Dict = None) -> None:
+        self.func_id = func_id
+        self.description = description
+        self.func_handler = func_handler
+        self.parameters = parameters
+
+    def get_name(self) -> str: 
+        return self.func_id
+    
+    def get_parameters(self) -> Dict:
+        if self.parameters is not None:
+            return self.parameters
+        return {"type": "object", "properties": {}}
+
+    async def execute(self,**kwargs) -> str:
+        if self.func_handler is None:
+            return "error: function not implemented"
+        
+        return await self.func_handler(**kwargs)
+
+    def is_local(self) -> bool:
+        return True
+
+    def is_in_zone(self) -> bool:
+        return True
+    
+    def is_ready_only(self) -> bool:
+        return False
+

src/aios_kernel/compute_kernel.py

@@ -54,7 +54,6 @@ class ComputeKernel:
 
         async def _run_task_loop():
             while True:
-                logger.info("compute_kernel is waiting for task...")
                 task = await self.task_queue.get()
                 logger.info(f"compute_kernel get task: {task.display()}")
                 c_node: ComputeNode = self._schedule(task)
@@ -91,16 +90,16 @@ class ComputeKernel:
         return True
 
     # friendly interface for use:
-    def llm_completion(self, prompt: AgentPrompt, mode_name: Optional[str] = None, max_token: int = 0):
+    def llm_completion(self, prompt: AgentPrompt, mode_name: Optional[str] = None, max_token: int = 0,inner_functions = None):
         # craete a llm_work_task ,push on queue's end
         # then task_schedule would run this task.(might schedule some work_task to another host)
         task_req = ComputeTask()
-        task_req.set_llm_params(prompt, mode_name, max_token)
+        task_req.set_llm_params(prompt, mode_name, max_token,inner_functions)
         self.run(task_req)
         return task_req
 
-    async def do_llm_completion(self, prompt: AgentPrompt, mode_name: Optional[str] = None, max_token: int = 0) -> str:
-        task_req = self.llm_completion(prompt, mode_name, max_token)
+    async def do_llm_completion(self, prompt: AgentPrompt, mode_name: Optional[str] = None, max_token: int = 0, inner_functions = None) -> str:
+        task_req = self.llm_completion(prompt, mode_name, max_token,inner_functions)
 
         async def check_timer():
             check_times = 0
@@ -120,6 +119,6 @@ class ComputeKernel:
 
         await asyncio.create_task(check_timer())
         if task_req.state == ComputeTaskState.DONE:
-            return task_req.result.result_str
+            return task_req.result
 
         return "error!"

src/aios_kernel/compute_task.py

@@ -11,7 +11,6 @@ class ComputeTaskState(Enum):
     ERROR = 3
     PENDING = 4
 
-
 class ComputeTaskType(Enum):
     NONE = -1
     LLM_COMPLETION = 0
@@ -36,7 +35,7 @@ class ComputeTask:
         self.result = None
         self.error_str = None
 
-    def set_llm_params(self, prompts, model_name, max_token_size, callchain_id=None):
+    def set_llm_params(self, prompts, model_name, max_token_size, inner_functions = None, callchain_id=None):
         self.task_type = ComputeTaskType.LLM_COMPLETION
         self.create_time = time.time()
         self.task_id = uuid.uuid4().hex
@@ -46,7 +45,13 @@ class ComputeTask:
             self.params["model_name"] = model_name
         else:
             self.params["model_name"] = "gpt-4-0613"
-        self.params["max_token_size"] = max_token_size
+        if max_token_size is None:
+            self.params["max_token_size"] = 4000
+        else:
+            self.params["max_token_size"] = max_token_size
+
+        if inner_functions is not None:
+            self.params["inner_functions"] = inner_functions
 
     def display(self) -> str:
         return f"ComputeTask: {self.task_id} {self.task_type} {self.state}"
@@ -59,9 +64,8 @@ class ComputeTaskResult:
         self.callchain_id: str = None
         self.worker_id: str = None
         self.result_code: int = 0
-        self.result_str: str = None
-
-        self.result: dict = {}
+        self.result_str: str = None # easy to use,can read from result
+        self.result_message: dict = {}
         self.result_refers: dict = None
         self.pading_data: bytearray = None
 

src/aios_kernel/environment.py

@@ -5,6 +5,8 @@ from abc import ABC, abstractmethod
 from typing import Any, Callable, Optional,Dict,Awaitable,List
 import logging
 
+from .ai_function import AIFunction
+
 logger = logging.getLogger(__name__)
 
 class EnvironmentEvent(ABC):
@@ -33,6 +35,8 @@ class Environment:
         # self.valid_keys:Dict[str,bool] = None
         self.event_handlers:Dict[str,List[EnvironmentEventHandler]]= {}
 
+        self.functions : Dict[str,AIFunction] = {}
+
     def get_id(self) -> str:
         return self.env_id
     
@@ -44,6 +48,24 @@ class Environment:
     #def get_env_prompt(self) -> str:
     #    pass
 
+    def add_ai_function(self,func:AIFunction) -> None:
+        if self.functions.get(func.get_name()) is not None:
+            logger.warn(f"add ai_function {func.get_name()} in env {self.env_id}:function already exist")
+            
+        self.functions[func.get_name()] = func
+
+    def get_ai_function(self,func_name:str) -> AIFunction:
+        return self.functions.get(func_name)
+
+    #def enable_ai_function(self,func_name:str) -> None:
+    #    pass
+
+    #def disable_ai_function(self,func_name:str) -> None:
+    #    pass
+
+    def get_all_ai_functions(self) -> List[AIFunction]:
+        return self.functions.values()
+
     @abstractmethod
     def _do_get_value(self,key:str) -> Optional[str]:
         pass

src/aios_kernel/open_ai_node.py

@@ -56,28 +56,35 @@ class OpenAI_ComputeNode(ComputeNode):
         logger.info(f"call openai {mode_name} prompts: {prompts}")
         resp = openai.ChatCompletion.create(model=mode_name,
                                             messages=prompts,
-                                            max_tokens=4000,
-                                            temperature=1.2)
+                                            functions=task.params["inner_functions"],
+                                            max_tokens=task.params["max_token_size"],
+                                            temperature=0.7) # TODO: add temperature to task params?
+        
         logger.info(f"openai response: {resp}")
-
-        status_code = resp["choices"][0]["finish_reason"]
-        if status_code != "stop":
-            task.state = ComputeTaskState.ERROR
-            task.error_str = f"The status code was {status_code}."
-            return None
-
+        
         result = ComputeTaskResult()
         result.set_from_task(task)
+
+        status_code = resp["choices"][0]["finish_reason"]
+        match status_code:
+            case "function_call":
+                task.state = ComputeTaskState.DONE
+            case "stop":
+                task.state = ComputeTaskState.DONE
+            case _:
+                task.state = ComputeTaskState.ERROR
+                task.error_str = f"The status code was {status_code}."
+                return None
+
         result.worker_id = self.node_id
         result.result_str = resp["choices"][0]["message"]["content"]
-        result.result = resp["choices"][0]["message"]
+        result.result_message = resp["choices"][0]["message"]
 
         return result
 
     def start(self):
         async def _run_task_loop():
             while True:
-                logger.info("openai_node is waiting for task...")
                 task = await self.task_queue.get()
                 logger.info(f"openai_node get task: {task.display()}")
                 result = self._run_task(task)

src/aios_kernel/workflow.py

@@ -13,6 +13,7 @@ from .chatsession import AIChatSession
 from .role import AIRole,AIRoleGroup
 from .ai_function import CallChain
 from .compute_kernel import ComputeKernel
+from .compute_task import ComputeTask,ComputeTaskResult,ComputeTaskState
 from .bus import AIBus
 from .workflow_env import WorkflowEnvironment
 
@@ -354,7 +355,8 @@ class Workflow:
   
         async def _do_process_msg():
             #TODO: send msg to agent might be better?
-            result_str = await ComputeKernel().do_llm_completion(prompt,the_role.agent.get_llm_model_name(),the_role.agent.get_max_token_size())
+            task_result:ComputeTaskResult = await ComputeKernel().do_llm_completion(prompt,the_role.agent.get_llm_model_name(),the_role.agent.get_max_token_size())
+            result_str = task_result.result_str
             result = Workflow.prase_llm_result(result_str)
             logger.info(f"{the_role.role_id} process {msg.sender}:{msg.body},llm str is :{result_str}")
             for postmsg in result.post_msgs:

src/aios_kernel/workflow_env.py

@@ -7,9 +7,11 @@ import threading
 import logging
 from typing import Optional
 from .environment import Environment,EnvironmentEvent
+from .ai_function import SimpleAIFunction
 
 logger = logging.getLogger(__name__)
 
+
 class CalenderEvent(EnvironmentEvent):
     def __init__(self,data) -> None:
         super().__init__()
@@ -17,7 +19,7 @@ class CalenderEvent(EnvironmentEvent):
         self.data = data
 
     def display(self) -> str:
-        return f"#event timer:{self.event_data}"    
+        return f"#event timer:{self.data}"    
     
 # AI Calender GOAL: Let user use "create notify after 2 days" to create a timer event
 class CalenderEnvironment(Environment):
@@ -25,6 +27,10 @@ class CalenderEnvironment(Environment):
         super().__init__(env_id)
         self.is_run = False
 
+        self.add_ai_function(SimpleAIFunction("get_time",
+                                        "get current time",
+                                        self.get_now))
+
     def _do_get_value(self,key:str) -> Optional[str]:
         return None
 
@@ -52,7 +58,7 @@ class CalenderEnvironment(Environment):
     def stop(self):
         self.is_run = False
 
-    def get_now(self,key:str) -> str:
+    async def get_now(self) -> str:
         now = datetime.now()
         formatted_time = now.strftime('%Y-%m-%d %H:%M:%S')
         return formatted_time

src/service/aios_shell/aios_shell.py

@@ -44,6 +44,7 @@ class AIOS_Shell:
         target_id = msg.target.split(".")[0]
         agent : AIAgent = await AgentManager().get(target_id)
         if agent is not None:
+            agent.owner_env = Environment.get_env_by_id("calender") 
             bus.register_message_handler(target_id,agent._process_msg)
             return True