Merge branch 'fiatrete:MVP' into MVP

.gitignore

@@ -2,3 +2,4 @@
 *.pyc
 rootfs/email/config.local.toml
 rootfs/data
+venv

doc/mvp/mvp plan 2.md

@@ -46,7 +46,7 @@ The previous plan, please see here: [MVP Plan](./mvp%20plan.md)
         - [ ] MPT-7B, S2
         - [ ] Vicuna, S2
     - [ ] Embeding,@photosssa,@lurenpluto , A4
-    - [ ] Txt2img,@glen0125,A4
+    - [x] Txt2img,@glen0125,A4
     - [ ] Img2txt(0.5.2),A3
     - [ ] Txt2voice,A3
     - [ ] Voice2txt, @wugren,A3
@@ -70,7 +70,7 @@ The previous plan, please see here: [MVP Plan](./mvp%20plan.md)
         - [ ] Telegram Tunnel,S2
         - [ ] Discord Tunnel,S2
     - [ ] Home IoT Environment (0.5.2), A4
-        - [] Compatible Home Assistant (0.5.2), A4
+        - [ ] Compatible Home Assistant (0.5.2), A4
 - [ ] Build-in Agents/Apps
     - [ ] Agent: Personal Information Assistant,@photosssa,@lurenpluto , A2
     - [ ] Agent: Bulter Jarvis,@waterflier, A2

doc/mvp/workflow.drawio

@@ -1,6 +1,6 @@
 <mxfile host="65bd71144e" pages="3">
     <diagram id="C5RBs43oDa-KdzZeNtuy" name="Page-1">
-        <mxGraphModel dx="1881" dy="676" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="827" pageHeight="1169" math="0" shadow="0">
+        <mxGraphModel dx="2069" dy="1139" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="827" pageHeight="1169" math="0" shadow="0">
             <root>
                 <mxCell id="WIyWlLk6GJQsqaUBKTNV-0"/>
                 <mxCell id="WIyWlLk6GJQsqaUBKTNV-1" parent="WIyWlLk6GJQsqaUBKTNV-0"/>
@@ -365,4 +365,30 @@
             </root>
         </mxGraphModel>
     </diagram>
+    <diagram id="zNgk-d6xdACtSja1wnUq" name="Page-4">
+        <mxGraphModel dx="2069" dy="1139" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="1100" pageHeight="850" math="0" shadow="0">
+            <root>
+                <mxCell id="0"/>
+                <mxCell id="1" parent="0"/>
+                <mxCell id="PKOyaMforMDOFoUY4PA1-1" value="ChatSession" style="rounded=1;whiteSpace=wrap;html=1;" vertex="1" parent="1">
+                    <mxGeometry x="140" y="150" width="120" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="PKOyaMforMDOFoUY4PA1-2" value="SubSession" style="rounded=1;whiteSpace=wrap;html=1;" vertex="1" parent="1">
+                    <mxGeometry x="230" y="240" width="120" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="PKOyaMforMDOFoUY4PA1-3" value="SubSession" style="rounded=1;whiteSpace=wrap;html=1;" vertex="1" parent="1">
+                    <mxGeometry x="230" y="330" width="120" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="PKOyaMforMDOFoUY4PA1-4" value="Agent Message" style="shape=hexagon;perimeter=hexagonPerimeter2;whiteSpace=wrap;html=1;fixedSize=1;" vertex="1" parent="1">
+                    <mxGeometry x="600" y="150" width="130" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="PKOyaMforMDOFoUY4PA1-5" value="Agent Message" style="shape=hexagon;perimeter=hexagonPerimeter2;whiteSpace=wrap;html=1;fixedSize=1;" vertex="1" parent="1">
+                    <mxGeometry x="600" y="230" width="130" height="60" as="geometry"/>
+                </mxCell>
+                <mxCell id="PKOyaMforMDOFoUY4PA1-6" value="Agent Message" style="shape=hexagon;perimeter=hexagonPerimeter2;whiteSpace=wrap;html=1;fixedSize=1;" vertex="1" parent="1">
+                    <mxGeometry x="600" y="320" width="130" height="60" as="geometry"/>
+                </mxCell>
+            </root>
+        </mxGraphModel>
+    </diagram>
 </mxfile>

rootfs/agents/agent_1/agent.toml

@@ -1,5 +1,6 @@
 instance_id = "agent_1"
 fullname = "tracy wang"
+
 [[prompt]]
 role = "system"
 content = "你是我的私人英文老师，和我用地道的美式英语进行交流。你会在和我交流的同时，调整我的输入成为更地道的美式句子，并根据你对我英文水平的预测，对可能发错英的单词标上音标。如果我给你发中文，说明我不知道这句话用美式英语怎么说，你依旧按上述规则回应我。"

rootfs/workflows/math_school/workflow.toml

@@ -7,7 +7,7 @@ GOAL="成为最好的学校"
 [[connected_env]]
 env_id = "calender"
 [[connected_env.event2msg]]
-timer = "现在是{data}"
+timer = "现在是{now}"
 role = "教导处主任"
 
 [filter]

rootfs/workflows/plan_maker/workflow.toml

@@ -42,11 +42,13 @@ fullname = "经理"
 agent="manager"
 [[roles.manager.prompt]] 
 role="system"
-content="""你是一个活动策划公司的经理,与客户对接并向团队下达指令。你的团队分为下面几个小组:嘉宾对接组,酒店预定组,行程预订组,财务组,活动摄像组。活动策划分为四个阶段：方案讨论，活动前，活动中，活动后。你会根据客户的需求，对团队进行分工，分别完成各个阶段的工作。你的基本工作模式是：
+content="""
+你是一个活动策划公司的经理,与客户对接并向团队下达指令。你的团队分为下面几个小组:嘉宾对接组,酒店预定组,行程预订组,财务组,活动摄像组。活动策划分为四个阶段：方案讨论，活动前，活动中，活动后。你会根据客户的需求，对团队进行分工，分别完成各个阶段的工作。你的基本工作模式是：
 1. 收到客户的明确的指令后,基于客户的已有信息和客户商量活动方案，和活动策划公司无关的业务你会回答‘与我无关’。当和客户完成活动方案的确认后,你会将拆解后的任务分配给各个小组 
 2. 根据目前已经确认的活动方案，你要根据时间适时的检查不同小组的工作情况。当收到小组的工作情况反馈后，你会站在全局的角度判断是否需要调整活动方案，如果需要调整，你会和客户商量重新确定方案，然后再将调整后的方案分配给各个小组。
 3. 有时工作小组会主动与你沟通，反馈一些问题。你会站在全局的角度给与指导，适当的调整工作小组的工作目标。如果反馈的问题需要你和客户沟通，你会和客户沟通后重新确定方案。再将调整后的方案分配给受到影响各个小组。
-4. 当你决定要和工作小组通信时，请使用 sendmsg(小组名称,内容) 的形式。"""
+4. 当你决定要和工作小组通信时，请使用 sendmsg(小组名称,内容) 的形式。
+"""
 
 
 

src/aios_kernel/__init__.py

@@ -1,5 +1,5 @@
 from .environment import Environment,EnvironmentEvent
-from .agent_message import AgentMsg,AgentMsgState
+from .agent_message import AgentMsg,AgentMsgStatus
 from .chatsession import AIChatSession
 from .agent import AIAgent,AIAgentTemplete,AgentPrompt
 from .compute_kernel import ComputeKernel,ComputeTask

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 .agent_message import AgentMsg, AgentMsgStatus, AgentMsgType
 from .chatsession import AIChatSession
+from .compute_task import ComputeTaskResult
+from .ai_function import AIFunction
+from .environment import Environment
 
 logger = logging.getLogger(__name__)
 
@@ -69,7 +73,7 @@ class AIAgent:
         self.prompt:AgentPrompt = None
         self.llm_model_name:str = None
         self.max_token_size:int = 3600
-        self.instance_id:str = None
+        self.agent_id:str = None
         self.template_id:str = None
         self.fullname:str = None
         self.powerby = None  
@@ -77,6 +81,8 @@ class AIAgent:
 
         self.chat_db = None
         self.unread_msg = Queue() # msg from other agent
+        self.owner_env : Environment = None
+        self.owenr_bus : AIBus = None
         
     @classmethod
     def create_from_templete(cls,templete:AIAgentTemplete, fullname:str):
@@ -85,7 +91,7 @@ class AIAgent:
         result_agent.llm_model_name = templete.llm_model_name
         result_agent.max_token_size = templete.max_token_size
         result_agent.template_id = templete.template_id
-        result_agent.instance_id = "agent#" + uuid.uuid4().hex
+        result_agent.agent_id = "agent#" + uuid.uuid4().hex
         result_agent.fullname = fullname
         result_agent.powerby = templete.author
         result_agent.prompt = templete.prompt
@@ -95,10 +101,10 @@ class AIAgent:
         if config.get("instance_id") is None:
             logger.error("agent instance_id is None!")
             return False
-        self.instance_id = config["instance_id"]
+        self.agent_id = config["instance_id"]
 
         if config.get("fullname") is None:
-            logger.error(f"agent {self.instance_id} fullname is None!")
+            logger.error(f"agent {self.agent_id} fullname is None!")
             return False
         self.fullname = config["fullname"]
 
@@ -123,83 +129,101 @@ 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,prompt:AgentPrompt,org_msg:AgentMsg) -> 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"
+        
+        ineternal_call_record = AgentMsg.create_internal_call_msg(func_name,arguments,org_msg.get_msg_id(),org_msg.target)
+
+        result_str:str = await func_node.execute(**arguments)
+        
+        inner_functions = self._get_inner_functions()
+        prompt.messages.append({"role":"function","content":result_str,"name":func_name})
+        task_result:ComputeTaskResult = await ComputeKernel().do_llm_completion(prompt,self.llm_model_name,self.max_token_size,inner_functions)
+        
+        ineternal_call_record.result_str = task_result.result_str
+        ineternal_call_record.done_time = time.time()
+        org_msg.inner_call_chain.append(ineternal_call_record)
+
+        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,prompt,org_msg)      
+        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)
+            chatsession = AIChatSession.get_session(self.agent_id,session_topic,self.chat_db)
+            if msg.mentions is not None:
+                if not self.agent_id in msg.mentions:
+                    chatsession.append(msg)
+                    logger.info(f"agent {self.agent_id} recv a group chat message from {msg.sender},but is not mentioned,ignore!")
+                    return None
+            
             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)
+            inner_functions = self._get_inner_functions()
-            final_result = result
+
-            result_type : str = self._get_llm_result_type(result)
+            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:
+                #TODO to save more token ,can i use msg_prompt?
+                final_result = await self._execute_func(inner_func_call_node,prompt,msg)
+            
+            result_type : str = self._get_llm_result_type(final_result)
             is_ignore = False
 
             match result_type:
-                # case "function":
-                #    callchain:CallChain = self._parse_function_call_chain(result)
-                #    resp = await callchain.exec()
-                #    if callchain.have_result():
-                #        # generator proc resp prompt with WAITING state
-                #        proc_resp_prompt:AgentPrompt = self._get_resp_prompt(resp,msg,the_role,prompt,chatsession)
-                #        final_result = await ComputeKernel().do_llm_completion(proc_resp_prompt,the_role.agent.get_llm_model_name(),the_role.agent.get_max_token_size())
-                #        return final_result
-
-                
-                # case "send_message":
-                #    # send message to other / sub workflow
-                #    next_msg:AgentMsg = self._parse_to_msg(result)
-                #    if next_msg is not None:
-                #        # TODO: Next Target can be another role in workflow
-                #        next_workflow:Workflow = self.get_workflow(next_msg.get_target())
-                #        inner_chat_session = the_role.agent.get_chat_session(next_msg.get_target(),next_msg.get_session_id())
-
-                #        inner_chat_session.append_post(next_msg)
-                #        resp = await next_workflow.send_msg(next_msg)
-                #        inner_chat_session.append_recv(resp)
-                #        # generator proc resp prompt with WAITING state
-                #        proc_resp_prompt:AgentPrompt = self._get_resp_prompt(resp,msg,the_role,prompt,chatsession)
-                #        final_result = await ComputeKernel().do_llm_completion(proc_resp_prompt,the_role.agent.get_llm_model_name(),the_role.agent.get_max_token_size())
-                        
-                #        return final_result
-                    
-                #case "post_message":
-                #    # post message to other / sub workflow
-                #    next_msg:AgentMsg = self._parse_to_msg(result)
-                #    if next_msg is not None:
-                #        next_workflow:Workflow = self.get_workflow(next_msg.get_target())
-                #        inner_chat_session = the_role.agent.get_chat_session(next_msg.get_target(),next_msg.get_session_id())
-                #        inner_chat_session.append_post(next_msg)
-                #        next_workflow.post_msg(next_msg)
-
                 case "ignore":
                     is_ignore = True
                     
             if is_ignore is not True:
-                # TODO : how to get inner chat session?
+                resp_msg = msg.create_resp_msg(final_result)
-                resp_msg = AgentMsg()
+                chatsession.append(msg)
-                resp_msg.set(self.instance_id,msg.sender,final_result)
+                chatsession.append(resp_msg)
-                resp_msg.topic = msg.topic
-
-                if chatsession is not None:
-                    chatsession.append_recv(msg)
-                    chatsession.append_post(resp_msg)
                 
                 return resp_msg
             
             return None
         
     def get_id(self) -> str:
-        return self.instance_id
+        return self.agent_id
     
     def get_fullname(self) -> str:
         return self.fullname
@@ -213,7 +237,8 @@ class AIAgent:
     def get_max_token_size(self) -> int:
         return self.max_token_size
     
-    async def _get_prompt_from_session(self,chatsession:AIChatSession) -> AgentPrompt:
+    async def _get_prompt_from_session(self,chatsession:AIChatSession,is_groupchat=False) -> AgentPrompt:
+        # TODO: get prompt from group chat is different from single chat
         messages = chatsession.read_history() # read last 10 message
         result_prompt = AgentPrompt()
         for msg in reversed(messages):

src/aios_kernel/agent_message.py

@@ -3,26 +3,100 @@ import uuid
 import time 
 import re
 
-class AgentMsgState(Enum):
+class AgentMsgType(Enum):
+    TYPE_MSG = 0
+    TYPE_INTERNAL_CALL = 1
+    TYPE_ACTION = 2
+    TYPE_EVENT = 3
+
+class AgentMsgStatus(Enum):
     RESPONSED = 0
     INIT = 1
     SENDING = 2
     PROCESSING = 3
     ERROR = 4
+    RECVED = 5
+    EXECUTED = 6
+
+# msg is a msg / msg resp
+# msg body可以有内容类型（MIME标签），text, image, voice, video, file,以及富文本(html)
+# msg is a inner function call with result
+# msg is a Action with result
+
+# qutoe Msg
+# forword msg
+# reply msg
+
+# 逻辑上的同一个Message在同一个session中看到的msgid相同
+#       在不同的session中看到的msgid不同
 
 class AgentMsg:
-    def __init__(self) -> None:
+    def __init__(self,msg_type=AgentMsgType.TYPE_MSG) -> None:
+        self.msg_id = ""
+        self.msg_type:AgentMsgType = msg_type
+        
+        self.prev_msg_id:str = None
+        self.quote_msg_id:str = None    
+        self.rely_msg_id:str = None # if not none means this is a respone msg
+        self.session_id:str = None
+
         self.create_time = 0
-        self.sender:str = None
+        self.done_time = 0
+        self.topic:str = None # topic is use to find session, not store in db
+
+        self.sender:str = None # obj_id.sub_objid@tunnel_id
         self.target:str = None
+        self.mentions:[] = None #use in group chat only
+        #self.title:str = None
         self.body:str = None
-        self.topic:str = "T#" + uuid.uuid4().hex
+        self.body_mime:str = None #//default is "text/plain",encode is utf8
-        #self.msg_type = 0
+
-        self.state = AgentMsgState.INIT
+        #type is call / action
+        self.func_name = None 
+        self.args = None
+        self.result_str = None
+
+        #type is event
+        self.event_name = None
+        self.event_args = None
+
+        self.status = AgentMsgStatus.INIT
+        self.inner_call_chain = []
         self.resp_msg = None
 
+    @classmethod
+    def create_internal_call_msg(self,func_name:str,args:dict,prev_msg_id:str,caller:str):
+        msg = AgentMsg(AgentMsgType.TYPE_INTERNAL_CALL)
+        msg.func_name = func_name
+        msg.args = args
+        msg.prev_msg_id = prev_msg_id
+        msg.sender = caller
+        return msg
+    
+    def create_action_msg(self,action_name:str,args:dict,caller:str):
+        msg = AgentMsg(AgentMsgType.TYPE_ACTION)
+        msg.func_name = action_name
+        msg.args = args
+        msg.prev_msg_id = self.msg_id
+        msg.topic  = self.topic
+        msg.sender = caller
+        return msg
+    
+    def create_resp_msg(self,resp_body):
+        resp_msg = AgentMsg()
+        resp_msg.msg_id = "msg#" + uuid.uuid4().hex
+        self.create_time = time.time()
+
+        resp_msg.rely_msg_id = self.msg_id
+        resp_msg.sender = self.target
+        resp_msg.target = self.sender
+        resp_msg.body = resp_body
+        resp_msg.topic = self.topic
+
+        return resp_msg
+
     def set(self,sender:str,target:str,body:str,topic:str=None) -> None:
-        self.id = "msg#" + uuid.uuid4().hex
+        self.msg_id = "msg#" + uuid.uuid4().hex
         self.sender = sender
         self.target = target
         self.body = body
@@ -31,7 +105,7 @@ class AgentMsg:
             self.topic = topic
 
     def get_msg_id(self) -> str:
-        return self.id
+        return self.msg_id
 
     def get_sender(self) -> str:
         return self.sender
@@ -39,6 +113,12 @@ class AgentMsg:
     def get_target(self) -> str:
         return self.target
     
+    def get_prev_msg_id(self) -> str:
+        return self.prev_msg_id
+    
+    def get_quote_msg_id(self) -> str:
+        return self.quote_msg_id
+    
     @classmethod
     def parse_function_call(cls,func_string:str):
         match = re.search(r'\s*(\w+)\s*\(\s*(.*)\s*\)\s*', func_string)

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/bus.py

@@ -1,5 +1,5 @@
-from typing import Any
+from typing import Coroutine,Dict,Any
-from .agent_message import AgentMsg,AgentMsgState
+from .agent_message import AgentMsg,AgentMsgStatus
 import asyncio
 from asyncio import Queue
 
@@ -8,19 +8,32 @@ import logging
 logger = logging.getLogger(__name__)
 
 class AIBusHandler:
-    def __init__(self,handler:Any) -> None:
+    def __init__(self,handler:Coroutine,owner_bus,enable_defualt_proc=True) -> None:
         self.handler = handler
         self.working_task = None
-        self.results = {}
+        self.results = {} # recv resps
         self.queue:Queue = Queue()
+        self.enable_defualt_proc = enable_defualt_proc
+        self.owner_bus = owner_bus
 
     async def handle_message(self,msg:AgentMsg) -> Any:
         if self.handler is None:
             return None
         
-        return await self.handler(msg)
+        if self.enable_defualt_proc:
+            # do default process
+            if msg.rely_msg_id is not None:
+                self.results[msg.rely_msg_id] = msg
+                return None
+            
 
+        resp_msg = await self.handler(msg)
+        if self.enable_defualt_proc:
+            if resp_msg is not None:
+                await self.owner_bus.post_message(resp_msg,False)
 
+        return resp_msg
+            
 class AIBus:
     _instance = None
     @classmethod
@@ -30,11 +43,13 @@ class AIBus:
         return cls._instance
 
     def __init__(self) -> None:
-        self.handlers = {}
+        self.handlers:Dict[AIBusHandler] = {}
-        self.unhandle_handler = None
+        self.unhandle_handler:Coroutine = None
-        
+
-    async def post_message(self,target_id,msg:AgentMsg,use_unhandle=True) -> bool:
+
-        target_id = target_id.split(".")[0]
+    async def post_message(self,msg:AgentMsg,use_unhandle=True) -> bool:
+        target_id = msg.target.split(".")[0]
+
         handler = self.handlers.get(target_id)
         if handler:
             handler.queue.put_nowait(msg)
@@ -44,46 +59,39 @@ class AIBus:
         if use_unhandle:
             if self.unhandle_handler is not None:
                 if await self.unhandle_handler(self,msg):
-                    return await self.post_message(target_id,msg,False)
+                    return await self.post_message(msg,False)
       
         logger.warn(f"post message to {msg.target} failed!,target not found")
         return False
 
-    def resp_message(self,my_id:str,org_msg_id:str,resp:AgentMsg) -> None:
+    async def resp_message(self,org_msg_id:str,resp:AgentMsg) -> None:
-        handler = self.handlers.get(my_id)
+        assert resp.rely_msg_id == org_msg_id
-        if handler is None:
+        return await self.post_message(resp)
-            return None
-        handler.results[org_msg_id] = resp
 
-    async def get_message_resp(self,name:str,msg_id:str) -> AgentMsg:
+    async def send_message(self,msg:AgentMsg) -> AgentMsg:
-        handler = self.handlers.get(name)
+        sender_id = msg.sender.split(".")[0]
-        if handler is None:
+        sender_handler = self.handlers.get(sender_id) # sender already register on bus
+        if sender_handler is None:
+            logger.warn(f"sender {sender_id} not register on AI_BUS!")
             return None
         
-        return handler.results.get(msg_id)
+        post_result = await self.post_message(msg)
-
-    async def send_message(self,target_id:str,msg:AgentMsg) -> AgentMsg:
-        target_id = target_id.split(".")[0]
-        post_result = await self.post_message(target_id,msg)
         if post_result is False:
             return None
-        
+
-        handler = self.handlers.get(target_id)
-        if handler is None:
-            return None
-                
         retry_times = 0
         while True:
-            resp = handler.results.get(msg.id)
+            resp = sender_handler.results.get(msg.msg_id)
             if resp is not None:
                 msg.resp_msg = resp
-                msg.state = AgentMsgState.RESPONSED
+                msg.status = AgentMsgStatus.RESPONSED
+                del sender_handler.results[msg.msg_id]
                 return resp
             
             await asyncio.sleep(0.2)
             retry_times += 1
             if retry_times > 100:
-                msg.state = AgentMsgState.ERROR
+                msg.status = AgentMsgStatus.ERROR
                 return None
             
         return None
@@ -93,7 +101,7 @@ class AIBus:
 
     # means sub
     def register_message_handler(self,handler_name:str,handler:Any) -> Queue:        
-        handler_node =  AIBusHandler(handler) 
+        handler_node =  AIBusHandler(handler,self) 
         self.handlers[handler_name] = handler_node
         return handler_node.queue
     
@@ -102,13 +110,12 @@ class AIBus:
             # Wait for a message
             message = await handler.queue.get()
 
-            try:
+            #try:
                 # Try to handle the message
-                result = await handler.handle_message(message)
+            await handler.handle_message(message)
-                handler.results[message.id] = result
+            #except Exception as e:
-            except Exception as e:
                 # If an error occurs, put the message back into the queue
-                logger.error(f"handle message {message.id} failed! {e}")
+            #    logger.error(f"handle message {message.msg_id} failed! {e}")
                 #self.queues[name].put_nowait(message)
         
         return
@@ -126,14 +133,4 @@ class AIBus:
             logger.warn(f"handler {target_name} is already working!")
             return
         
-        handler.working_task = asyncio.create_task(self.process_queue(handler))
+        handler.working_task = asyncio.create_task(self.process_queue(handler))
-
-    
-
-#send message to target logic:
-# find target handler:
-#   process_message(msg):
-#       session = get_session(msg.sender,msg.target)
-
-
-# history: open(sender,target,topic)

src/aios_kernel/chatsession.py

@@ -5,8 +5,9 @@ import logging
 import threading
 import datetime
 import uuid
+import json
 
-from .agent_message import AgentMsg
+from .agent_message import AgentMsgType, AgentMsg, AgentMsgStatus
 
 class ChatSessionDB:
     def __init__(self, db_file):
@@ -54,14 +55,31 @@ class ChatSessionDB:
             """)
 
             # create messages table
+            # reciver_id could be None
+        
             conn.execute("""
                 CREATE TABLE IF NOT EXISTS Messages (
                     MessageID TEXT PRIMARY KEY,
                     SessionID TEXT,
-                    SenderID TEXT,
+                    MsgType INTEGER,
+                    PrevMsgID TEXT,
+                    QuoteMsgID TEXT,
+                    RelyMsgID TEXT,
+                    
+                    SenderID TEXT, 
                     ReceiverID TEXT,
                     Timestamp TEXT,
+                    
+                    Topic TEXT,
+                    Mentions TEXT,
+                    ContentMIME TEXT,
                     Content TEXT,
+                    
+                    ActionName TEXT,
+                    ActionParams TEXT,
+                    ActionResult TEXT,
+                    DoneTime TEXT,     
+                         
                     Status INTEGER
                 );
             """)
@@ -83,15 +101,43 @@ class ChatSessionDB:
             logging.error("Error occurred while inserting session: %s", e)
             return -1  # return -1 if an error occurs
 
-    def insert_message(self, message_id, session_id, sender_id, receiver_id, timestamp, content, status):
+    def insert_message(self, msg:AgentMsg):
         """ insert a new message into the Messages table """
         try:
+            action_name = None
+            action_params = None
+            action_result = None
+            mentions = None
+            if msg.mentions:
+                mentions = json.dumps(msg.mentions)
+
+            match msg.msg_type:
+                case AgentMsgType.TYPE_MSG:
+                    pass
+                case AgentMsgType.TYPE_ACTION:
+                    action_name = msg.func_name
+                    action_params = json.dumps(msg.args)
+                    action_result = msg.result_str
+                case AgentMsgType.TYPE_INTERNAL_CALL:
+                    action_name = msg.func_name
+                    action_params = json.dumps(msg.args)
+                    action_result = msg.result_str
+                case AgentMsgType.TYPE_EVENT:
+                    action_name = msg.event_name
+                    action_params = json.dumps(msg.event_args)
+
+
             conn = self._get_conn()
             conn.execute("""
-                INSERT INTO Messages (MessageID, SessionID, SenderID, ReceiverID, Timestamp, Content, Status)
+                INSERT INTO Messages (MessageID, SessionID, MsgType, PrevMsgID, SenderID, ReceiverID, Timestamp, Topic,Mentions,ContentMIME,Content,ActionName,ActionParams,ActionResult,DoneTime,Status)
-                VALUES (?, ?, ?, ?, ?, ?, ?)
+                VALUES (?, ?, ?, ?, ?, ?, ?,?, ?, ?, ?, ?, ?, ?, ?, ?)
-            """, (message_id, session_id, sender_id, receiver_id, timestamp, content, status))
+            """, (msg.msg_id, msg.session_id, msg.msg_type.value, msg.prev_msg_id, msg.sender, msg.target, msg.create_time, msg.topic,mentions,msg.body_mime,msg.body,action_name,action_params,action_result,msg.done_time,msg.status.value))
             conn.commit()
+
+            if msg.inner_call_chain:
+                for inner_call in msg.inner_call_chain:
+                    self.insert_message(inner_call)
+
             return 0  # return 0 if successful
         except Error as e:
             logging.error("Error occurred while inserting message: %s", e)
@@ -134,7 +180,7 @@ class ChatSessionDB:
         """Get a message by its ID"""
         conn =self._get_conn()
         c = conn.cursor()
-        c.execute("SELECT MessageID,SessionID,SenderID,ReceiverID,Timestamp,Content,Status FROM Messages WHERE MessageID = ?", (message_id,))
+        c.execute("SELECT MessageID, SessionID, MsgType, PrevMsgID, SenderID, ReceiverID, Timestamp, Topic,Mentions,ContentMIME,Content,ActionName,ActionParams,ActionResult,DoneTime,Status FROM Messages WHERE MessageID = ?", (message_id,))
         message = c.fetchone()
         return message
 
@@ -144,7 +190,7 @@ class ChatSessionDB:
             conn = self._get_conn()
             cursor = conn.cursor()
             cursor.execute("""
-                SELECT MessageID,SessionID,SenderID,ReceiverID,Timestamp,Content,Status FROM Messages
+                SELECT MessageID, SessionID, MsgType, PrevMsgID, SenderID, ReceiverID, Timestamp, Topic,Mentions,ContentMIME,Content,ActionName,ActionParams,ActionResult,DoneTime,Status FROM Messages
                 WHERE SessionID = ?
                 ORDER BY Timestamp DESC
                 LIMIT ? OFFSET ?
@@ -222,29 +268,31 @@ class AIChatSession:
         result = []
         for msg in msgs:
             agent_msg = AgentMsg()
-            agent_msg.id = msg[0]
+            agent_msg.msg_id = msg[0]
-            agent_msg.sender = msg[2]
+            agent_msg.session_id = msg[1]
-            agent_msg.target = msg[3]
+            agent_msg.msg_type = AgentMsgType(msg[2])
-            agent_msg.create_time = msg[4]
+            agent_msg.prev_msg_id = msg[3]
-            agent_msg.body = msg[5]
+            agent_msg.sender = msg[4]
-            # agent_msg.state = msg[6]
+            agent_msg.target = msg[5]
+            agent_msg.create_time = msg[6]
+            agent_msg.topic = msg[7]
+            if msg[8] is not None:
+                agent_msg.mentions = json.loads(msg[8])
+            agent_msg.body_mime = msg[9]
+            agent_msg.body = msg[10]
+            agent_msg.func_name = msg[11]
+            if msg[12] is not None:
+                agent_msg.args = json.loads(msg[12])
+            agent_msg.result_str = msg[13]
+            agent_msg.done_time = msg[14]
+            agent_msg.status = AgentMsgStatus(msg[15])
 
             result.append(agent_msg)
         return result
 
     def append(self,msg:AgentMsg) -> None:
-        self.db.insert_message(msg.id,self.session_id,msg.sender,msg.target,msg.create_time,msg.body,0)
+        msg.session_id = self.session_id
-
+        self.db.insert_message(msg)
-    def append_post(self,msg:AgentMsg) -> None:
-        """append msg to session, msg is post from session (owner => msg.target)"""
-        #assert msg.sender == self.owner_id,"post message means msg.sender == self.owner_id"
-        self.append(msg)
-        
-
-    def append_recv(self,msg:AgentMsg) -> None:
-        """append msg to session, msg is recv from msg'sender (msg.sender => owner)"""
-        #assert msg.target == self.owner_id,"recv message means msg.target == self.owner_id"
-        self.append(msg)        
 
     #def attach_event_handler(self,handler) -> None:
     #    """chat session changed event handler"""

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:
+    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)
+        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_str: str = None # easy to use,can read from result
-
+        self.result_message: dict = {}
-        self.result: 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,16 +48,28 @@ 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
 
-    def get_functions(self):
-        # system functions
-        # env functions
-        # user install functions
-        pass
-
     def register_get_handler(self,key:str,handler:Callable) -> None:
         h = self.get_handlers.get(key)
         if h is not None:

src/aios_kernel/open_ai_node.py

@@ -54,30 +54,43 @@ class OpenAI_ComputeNode(ComputeNode):
         prompts = task.params["prompts"]
 
         logger.info(f"call openai {mode_name} prompts: {prompts}")
-        resp = openai.ChatCompletion.create(model=mode_name,
+        if task.params.get("inner_functions") is None:
+            resp = openai.ChatCompletion.create(model=mode_name,
                                             messages=prompts,
-                                            max_tokens=4000,
+                                            max_tokens=task.params["max_token_size"],
-                                            temperature=1.2)
+                                            temperature=0.7)
+        else:
+            resp = openai.ChatCompletion.create(model=mode_name,
+                                                messages=prompts,
+                                                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

@@ -2,20 +2,23 @@ import logging
 import asyncio
 import json
 import os
+import time 
 from asyncio import Queue
 from typing import Optional,Tuple
 from abc import ABC, abstractmethod
 
 from .environment import Environment,EnvironmentEvent
-from .agent_message import AgentMsg,AgentMsgState
+from .agent_message import AgentMsg,AgentMsgStatus
 from .agent import AgentPrompt,AgentMsg
 from .chatsession import AIChatSession
 from .role import AIRole,AIRoleGroup
-from .ai_function import CallChain
+from .ai_function import AIFunction
 from .compute_kernel import ComputeKernel
+from .compute_task import ComputeTask,ComputeTaskResult,ComputeTaskState
 from .bus import AIBus
 from .workflow_env import WorkflowEnvironment
 
+
 logger = logging.getLogger(__name__)
 
 class MessageFilter:
@@ -157,11 +160,12 @@ class Workflow:
         while i < len(inner_obj_id):
             if i == len(inner_obj_id) - 1:
                 the_role : AIRole = current_workflow.role_group.get(inner_obj_id[i])
+                current_workflow_chatsession = AIChatSession.get_session(current_workflow.workflow_id,msg.sender + "#" + msg.topic,current_workflow.db_file)
                 if the_role is not None:
-                    return await current_workflow.role_process_msg(msg,the_role)
+                    return await current_workflow.role_process_msg(msg,the_role,current_workflow_chatsession)
                 sub_workflow = current_workflow.sub_workflows.get(inner_obj_id[i])
                 if sub_workflow is not None:
-                    return await sub_workflow.process_msg(msg)
+                    return await sub_workflow._process_msg(msg)
                 logger.error(f"{msg.target} not found! forword message failed!")
                 return None
             else:
@@ -175,18 +179,50 @@ class Workflow:
         logger.error(f"{msg.target} not found! forword message failed!")
         return None
         
+    def get_workflow_id_from_target(self,target:str) -> str:
+        target_list = target.split(".")
+        if len(target_list) == 0:
+            return target
+        else:
+            result_str = ""
+            p = 0
+            for s in target_list:
+                p = p + 1
+                result_str += s
+                if p < len(target_list)-1:
+                    result_str += "."
+                else:
+                    return result_str
 
     async def _process_msg(self,msg:AgentMsg):
-        # workflow can be a message handler, but never be a message sender
+        real_target = msg.target.split(".")[0]
-        # all message forword to roles or sub workflow
-        # workflow no chatsession record, but role have
-        
-        final_result = None
-
         targets = self._parse_msg_target(msg.target)
         if len(targets) > 1:
             return await self._forword_msg(targets,msg)
-            
+        #0 we don't support workflow join a group right now, this cloud be a feture in future
+        if msg.mentions is not None:
+            logger.warn(f"workflow {self.workflow_id} recv a group chat message,not support ignore!")
+            return None
+
+        #1. workflow start process message
+        final_result = None
+        
+        # this is workflow's group_chat session
+        session_topic = msg.sender + "#" + msg.topic
+        chatsesssion = AIChatSession.get_session(self.workflow_id,session_topic,self.db_file)
+        
+        #2. find role by msg.mentions or workflow's selector logic
+        if msg.mentions is not None:
+            if not self.workflow_id in msg.mentions:
+                chatsesssion.append(msg)
+                logger.info(f"workflow {self.workflow_id} recv a group chat message from {msg.sender},but is not mentioned,ignore!")
+                return None
+
+            for mention in msg.mentions:
+                this_role = self.role_group.get(mention)
+                if this_role is not None:
+                    return await self.role_process_msg(msg,this_role,chatsesssion)
+                
         if self.input_filter is not None:
             select_role_id = self.input_filter.select(msg)
             if select_role_id is not None: 
@@ -194,27 +230,13 @@ class Workflow:
                 if select_role is None:
                     logger.error(f"input_filter return invalid role id:{select_role_id}, role not found in role_group")
                     return None
-
-                result = await self.role_process_msg(msg,select_role)
-                if result is None:
-                    logger.error(f"_process_msg return None for :{msg}")
-                    return
-
-                final_result = result
-            else:
-                logger.error(f"input_filter return None for :{msg}")
-                return
                 
-        else:
+                return await self.role_process_msg(msg,select_role,chatsesssion)
-            # no input filter, we would process all message, slowly,not suggest to use 
+            else:
-            results = {}
+                logger.error(f"input_filter return None for :{msg.body}")
-            final_result:AgentMsg  = None 
+                return None
-            for this_role in self.role_group.roles.values():
+                
-                # TODO : we would do this in parallel
+        logger.error(f"{self.workflow_id}:no role can process this msg:{msg.body}")
-                a_result = await self.role_process_msg(msg,this_role)
-                results[this_role.get_name()] = a_result
-                final_result = a_result
-        
         return final_result
 
     @classmethod
@@ -275,52 +297,59 @@ class Workflow:
 
         return r
 
-    async def role_post_msg(self,msg:AgentMsg,the_role:AIRole):
+    async def role_post_msg(self,msg:AgentMsg,the_role:AIRole,workflow_chat_session:AIChatSession):
         msg.sender = the_role.get_role_id()
         
         target_role = self.role_group.get(msg.target)
         if target_role:
             msg.target = target_role.get_role_id()
-            logger.info(f"{msg.sender} post message {msg.id} to inner role: {msg.target}")
+            logger.info(f"{msg.sender} post message {msg.msg_id} to inner role: {msg.target}")
-            asyncio.create_task(self.role_process_msg(msg,target_role))
+            asyncio.create_task(self.role_process_msg(msg,target_role,workflow_chat_session))
             return
         
         target_workflow = self.sub_workflows.get(msg.target)
         if target_workflow:
             msg.target = target_workflow.workflow_id
-            logger.info(f"{msg.sender} post message {msg.id} to sub workflow: {msg.target}")
+            logger.info(f"{msg.sender} post message {msg.msg_id} to sub workflow: {msg.target}")
             asyncio.create_task(target_workflow._process_msg(msg))
 
-        logger.info(f"{msg.sender} post message {msg.id} to AIBus: {msg.target}")
+        logger.info(f"{msg.sender} post message {msg.msg_id} to AIBus: {msg.target}")
         await self.get_bus().post_message(msg.target,msg)
         return
 
 
-    async def role_send_msg(self,msg:AgentMsg,the_role:AIRole):
+    async def role_send_msg(self,msg:AgentMsg,the_role:AIRole,workflow_chat_session:AIChatSession):
         msg.sender = the_role.get_role_id()
         target_role = self.role_group.get(msg.target)
         if target_role:
             # msg.target = target_role.get_role_id()
-            logger.info(f"{msg.sender} send message {msg.id} to inner role: {msg.target}")
+            logger.info(f"{msg.sender} send message {msg.msg_id} to inner role: {msg.target}")
-            return await self.role_process_msg(msg,target_role)
+            return await self.role_process_msg(msg,target_role,workflow_chat_session)
 
-        
         target_workflow = self.sub_workflows.get(msg.target)
         if target_workflow:
             # msg.target = target_workflow.workflow_id
-            logger.info(f"{msg.sender} send message {msg.id} to sub workflow: {msg.target}")
+            logger.info(f"{msg.sender} send message {msg.msg_id} to sub workflow: {msg.target}")
             return await target_workflow._process_msg(msg)
         
-        logger.info(f"{msg.sender} post message {msg.id} to AIBus: {msg.target}")
+        logger.info(f"{msg.sender} post message {msg.msg_id} to AIBus: {msg.target}")
-        return await self.get_bus().send_message(msg.target,msg)
+        return await self.get_bus().send_message(msg)
 
     async def role_call(self,call:tuple,the_role:AIRole):
         logger.info(f"{the_role.role_id} call {call[0]} with args {call[1]}")
-        return """{result:"timeout"}"""
+        func_name = call[0]
+        arguments = call[1]
+
+        func_node : AIFunction = self.workflow_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)        
+        return result_str
 
     async def role_post_call(self,call:tuple,the_role:AIRole):
         logger.info(f"{the_role.role_id} post call {call[0]} with args {call[1]}")
-        return
+        return await self.role_call(call,the_role)
 
     def _format_msg_by_env_value(self,prompt:AgentPrompt):
         if self.workflow_env is None:
@@ -330,65 +359,148 @@ class Workflow:
             old_content = msg.get("content")
             msg["content"] = old_content.format_map(self.workflow_env)
 
-    async def role_process_msg(self,msg:AgentMsg,the_role:AIRole):
+    def _get_inner_functions(self) -> dict:
-        session_topic = f"{msg.sender}#{msg.topic}"
+        all_inner_function = self.workflow_env.get_all_ai_functions()
-        session_owner = the_role.get_role_id()
+        if all_inner_function is None:
-        chatsession : AIChatSession = AIChatSession.get_session(session_owner,session_topic,self.db_file)
-        if chatsession is None:
-            logger.error(f"get session {session_topic}@{session_owner} failed!")
             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)
+        if len(result_func) > 0:
+            return result_func
+        return None 
+    
+    async def _role_execute_func(self,the_role:AIRole,inenr_func_call_node:dict,prompt:AgentPrompt,org_msg:AgentMsg) -> 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.workflow_env.get_ai_function(func_name)
+        if func_node is None:
+            return "execute failed,function not found"
+        
+        ineternal_call_record = AgentMsg.create_internal_call_msg(func_name,arguments,org_msg.get_msg_id(),org_msg.target)
+
+        result_str:str = await func_node.execute(**arguments)
+        
+        inner_functions = self._get_inner_functions()
+        prompt.messages.append({"role":"function","content":result_str,"name":func_name})
+        task_result:ComputeTaskResult = await ComputeKernel().do_llm_completion(prompt,
+                                                                                the_role.agent.llm_model_name,the_role.agent.max_token_size,
+                                                                                inner_functions)
+        
+        ineternal_call_record.result_str = task_result.result_str
+        ineternal_call_record.done_time = time.time()
+        org_msg.inner_call_chain.append(ineternal_call_record)
+
+        inner_func_call_node = task_result.result_message.get("function_call")
+        if inner_func_call_node:
+            return await self._role_execute_func(the_role,inner_func_call_node,prompt,org_msg)      
+        else:
+            return task_result.result_str
+    
+    def _is_in_same_workflow(self,msg) -> bool:
+        pass
+
+    async def role_process_msg(self,msg:AgentMsg,the_role:AIRole,workflow_chat_session:AIChatSession):        
+        msg.target = the_role.get_role_id()
+  
+
         prompt = AgentPrompt()
         prompt.append(the_role.agent.prompt)
         prompt.append(self.get_workflow_rule_prompt())
         prompt.append(the_role.get_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))
+        
+        #support group chat, user content include sender name!
+        prompt.append(await self._get_prompt_from_session(workflow_chat_session))
 
         msg_prompt = AgentPrompt()
-        msg_prompt.messages = [{"role":"user","content":msg.body}]
+        msg_prompt.messages = [{"role":"user","content":f"{msg.sender}:{msg.body}"}]
         prompt.append(msg_prompt)
 
         self._format_msg_by_env_value(prompt)
-  
+        inner_functions = self._get_inner_functions()
+        
         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(),inner_functions)
-            result = Workflow.prase_llm_result(result_str)
+            result_str = task_result.result_str
             logger.info(f"{the_role.role_id} process {msg.sender}:{msg.body},llm str is :{result_str}")
+            
+            inner_func_call_node = task_result.result_message.get("function_call")
+            
+            if inner_func_call_node:
+                #TODO to save more token ,can i use msg_prompt?
+                result_str = await self._role_execute_func(the_role,inner_func_call_node,prompt,msg)
+
+            result = Workflow.prase_llm_result(result_str)
             for postmsg in result.post_msgs:
+                postmsg.prev_msg_id = msg.get_msg_id()
+                # might be craete a new msg.topic for this postmsg
                 postmsg.topic = msg.topic
-                await self.role_post_msg(postmsg,the_role)
+                
+                await self.role_post_msg(postmsg,the_role,workflow_chat_session)
+                if not self._is_in_same_workflow(postmsg):
+                    role_sesion = AIChatSession.get_session(the_role.get_role_id(),f"{postmsg.target}#{msg.topic}",self.db_file)
+                    role_sesion.append(postmsg)
+                else:
+                    # message will be saved in role.process_message
+                    pass
+  
                 
             for post_call in result.post_calls:
+                action_msg = msg.create_action_msg(post_call[0],post_call[1],the_role.get_role_id())
+                workflow_chat_session.append(action_msg)
                 await self.role_post_call(post_call,the_role)
+                #save post_call
                 
             result_prompt_str = ""
             match result.state:
                 case "ignore":
                     return None
                 case "reponsed":
-                    resp_msg = AgentMsg()
+                    resp_msg = msg.create_resp_msg(result.resp)
-                    resp_msg.topic = msg.topic
+                    resp_msg.sender = the_role.get_role_id()
-                    resp_msg.set(session_owner,msg.sender,result.resp)
+                    # It is always the person handling the messages who puts them into the session.
-            
+                    workflow_chat_session.append(msg)
-                    chatsession.append_recv(msg)
+                    workflow_chat_session.append(resp_msg)
-                    chatsession.append_post(resp_msg)
+                    #await self.get_bus().resp_message(resp_msg)
                     return resp_msg
                 case "waiting":
                     # TODO: Use role:"function" would be better
                     for sendmsg in result.send_msgs:
                         target = sendmsg.target
                         sendmsg.topic = msg.topic
-                        send_resp = await self.role_send_msg(sendmsg,the_role)
+                        sendmsg.prev_msg_id = msg.get_msg_id()
+                        send_resp = await self.role_send_msg(sendmsg,the_role,workflow_chat_session)
                         if send_resp is not None:
                             result_prompt_str += f"\n{target} response is :{send_resp.body}"
+
+                        if not self._is_in_same_workflow(sendmsg):
+                            role_sesion = AIChatSession.get_session(the_role.get_role_id(),f"{sendmsg.target}#{sendmsg.topic}",self.db_file)
+                            role_sesion.append(sendmsg)
+                            role_sesion.append(send_resp)
+                        else:
+                             # message will be saved in role.process_message
+                            pass
                        
                     for call in result.calls:
+                        action_msg = msg.create_action_msg(call[0],call[1],call_result,the_role.get_role_id)
                         call_result = await self.role_call(call,the_role)
+          
                         if call_result is not None:
                             result_prompt_str += f"\ncall {call[0]} result is :{call_result}"
+                            #save action
+                            action_msg.result_str = call_result
+                            workflow_chat_session.append(action_msg)
 
                     result_prompt = AgentPrompt()
                     result_prompt.messages = [{"role":"user","content":result_prompt_str}]
@@ -408,18 +520,10 @@ class Workflow:
                 result_prompt.messages.append({"role":"user","content":f"{msg.body}"})
         
         return result_prompt
-
-
-    def _get_function_prompt(self,role_name:str) -> AgentPrompt:
-        # system functions
-        # env functions
-        # user install functions
-        pass    
     
     def _get_knowlege_prompt(self,role_name:str) -> AgentPrompt:
         pass
 
-
     def get_workflow_rule_prompt(self) -> AgentPrompt:
         return self.rule_prompt
 
@@ -445,6 +549,4 @@ class Workflow:
 
             #            the_env.attach_event_handler(k,_env_msg_handler)
             #            break
-        else:
-            logger.warn(f"environment {env.get_id()} already connected!")
 

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,12 @@ class CalenderEnvironment(Environment):
     def stop(self):
         self.is_run = False
 
-    def get_now(self,key:str) -> str:
+    def get_now(self,key)->str:
+        now = datetime.now()
+        formatted_time = now.strftime('%Y-%m-%d %H:%M:%S')
+        return formatted_time    
+    
+    async def _get_now(self) -> str:
         now = datetime.now()
         formatted_time = now.strftime('%Y-%m-%d %H:%M:%S')
         return formatted_time
@@ -65,6 +76,7 @@ class WorkflowEnvironment(Environment):
         self.local = threading.local()
         self.table_name = "WorkflowEnv_" + env_id
 
+
     def _get_conn(self):
         """ get db connection """
         if not hasattr(self.local, 'conn'):

src/service/aios_shell/aios_shell.py

@@ -26,7 +26,7 @@ shell_style = Style.from_dict({
 
 directory = os.path.dirname(__file__)
 sys.path.append(directory + '/../../')
-from aios_kernel import Workflow,AIAgent,AgentMsg,AgentMsgState,ComputeKernel,OpenAI_ComputeNode,AIBus,AIChatSession
+from aios_kernel import Workflow,AIAgent,AgentMsg,AgentMsgStatus,ComputeKernel,OpenAI_ComputeNode,AIBus,AIChatSession
 
 sys.path.append(directory + '/../../component/')
 from agent_manager import AgentManager
@@ -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
         
@@ -72,6 +73,7 @@ class AIOS_Shell:
         open_ai_node.start()
         ComputeKernel().add_compute_node(open_ai_node)
         AIBus().get_default_bus().register_unhandle_message_handler(self._handle_no_target_msg)
+        AIBus().get_default_bus().register_message_handler(self.username,self._user_process_msg)
         return True 
         
 
@@ -82,7 +84,7 @@ class AIOS_Shell:
         agent_msg = AgentMsg()
         agent_msg.set(sender,target_id,msg)
         agent_msg.topic = topic
-        resp = await AIBus.get_default_bus().send_message(target_id,agent_msg)
+        resp = await AIBus.get_default_bus().send_message(agent_msg)
         if resp is not None:
             return resp.body
         else:
@@ -90,6 +92,9 @@ class AIOS_Shell:
 
     async def install_workflow(self,workflow_id:Workflow) -> None:
         pass
+    
+    async def _user_process_msg(self,msg:AgentMsg) -> AgentMsg:
+        pass
 
     async def call_func(self,func_name, args):
         match func_name:
@@ -114,6 +119,7 @@ class AIOS_Shell:
             case 'login':
                 if len(args) >= 1:
                     self.username = args[0]
+                AIBus().get_default_bus().register_message_handler(self.username,self._user_process_msg)
                 return self.username + " login success!"    
             case 'history':
                 num = 10
@@ -133,7 +139,7 @@ class AIOS_Shell:
                 if chatsession is not None:
                     msgs = chatsession.read_history(num,offset)
                     format_texts = []
-                    for msg in reversed(msgs):
+                    for msg in msgs:
                         format_texts.append(("class:content",f"{msg.sender} >>> {msg.body}"))
                         format_texts.append(("",f"\n-------------------\n"))
                     return FormattedText(format_texts)
@@ -170,6 +176,7 @@ async def main():
                         format='[%(asctime)s]%(name)s[%(levelname)s]: %(message)s')
     shell = AIOS_Shell("user")
     await shell.initial()
+
     print(f"aios shell {shell.get_version()} ready.")
 
     completer = WordCompleter(['send($target,$msg,$topic)', 

src/service/aios_shell/tg_aios_shell.py

@@ -14,7 +14,7 @@ from telegram.ext import Application, CommandHandler, ContextTypes, MessageHandl
 
 directory = os.path.dirname(__file__)
 sys.path.append(directory + '/../../')
-from aios_kernel import Workflow,AIAgent,AgentMsg,AgentMsgState,ComputeKernel,OpenAI_ComputeNode,AIBus,AIChatSession
+from aios_kernel import Workflow,AIAgent,AgentMsg,AgentMsgStatus,ComputeKernel,OpenAI_ComputeNode,AIBus,AIChatSession
 
 sys.path.append(directory + '/../../component/')
 from agent_manager import AgentManager

src/service/spider/email_spider.py

@@ -1,3 +1,21 @@
+"""
+Capture your email locally, and parse out the pictures in the email body and the pictures, videos and other files in the attachment. Subsequently, it supports vectorized analysis of your personal data and serves as a knowledge base to enable large language model answers. Better results.
+
+An example of a local file is as follows:
+├── data
+│ └── alex0072@gmail.com
+│   └── 5de3e52f3a6b90cabe6cbdd4ae3a5c5b
+│     ├── email.txt
+│     ├── meta.json
+│     ├── image
+│     │   ├── 0648B869@99C03070.DB94B354.jpg
+│     └── body_image
+│         ├── 11044884873.jpg
+│         ├── 282985198265470.gif
+│         └── dd-login-service-min.png
+
+"""
+
 import imaplib
 import os
 import toml
@@ -5,105 +23,149 @@ import logging
 import mailparser
 import hashlib
 import json
+import base64
+from bs4 import BeautifulSoup
+import requests
 
-# logger config
+class EmailSpider:
-logger = logging.getLogger('email spider')
+    def __init__(self):
-logger.setLevel(logging.DEBUG)
+        # logger config
-ch = logging.StreamHandler()
+        self.logger = logging.getLogger('email spider')
-formatter = logging.Formatter('%(asctime)s [%(name)s] [%(levelname)s] %(message)s')
+        self.logger.setLevel(logging.DEBUG)
-ch.setFormatter(formatter)
+        ch = logging.StreamHandler()
-logger.addHandler(ch)
+        formatter = logging.Formatter('%(asctime)s [%(name)s] [%(levelname)s] %(message)s')
+        ch.setFormatter(formatter)
+        self.logger.addHandler(ch)
 
-# read config from toml file
+        # read config from toml file
-# and read from config config.local.toml if exists (config.local.toml is ignored by git)
+        # and read from config config.local.toml if exists (config.local.toml is ignored by git)
-config = toml.load('./rootfs/email/config.toml')
+        self.config = toml.load('./rootfs/email/config.toml')
-if os.path.exists('./rootfs/email/config.local.toml'):
+        if os.path.exists('./rootfs/email/config.local.toml'):
-    config = toml.load('./rootfs/email/config.local.toml')
+            self.config = toml.load('./rootfs/email/config.local.toml')
 
+        self.client = self.email_client()
 
-# create email client
+    def email_client(self) -> imaplib.IMAP4_SSL:
-def email_client() -> imaplib.IMAP4_SSL:
+        self.logger.info(f"read email config from {self.config.get('EMAIL_IMAP_SERVER')}")
-    logger.info(f"read email config from {config.get('EMAIL_IMAP_SERVER')}")
+        client = imaplib.IMAP4_SSL(
-    client = imaplib.IMAP4_SSL(config.get('EMAIL_IMAP_SERVER'))
+            host=self.config.get('EMAIL_IMAP_SERVER'), 
-    client.login(config.get('EMAIL_ADDRESS'), config.get('EMAIL_PASSWORD'))
+            port=self.config.get('EMAIL_IMAP_PORT')
-    return client
+        )
+        client.login(self.config.get('EMAIL_ADDRESS'), self.config.get('EMAIL_PASSWORD'))
+        return client
 
-def list_box(mail:  imaplib.IMAP4_SSL):
+    def list_box(self):
-    _, mailbox_list = mail.list()
+        _, mailbox_list = self.client.list()
-    for mailbox in mailbox_list:
+        for mailbox in mailbox_list:
-        print(mailbox.decode())
+            print(mailbox.decode())
-    
-def read_emails(client: imaplib.IMAP4_SSL, folder: str = 'INBOX', imap_keyword: str = "UNSEEN"):
-    client.select(folder)
-    _, data = client.uid('search', None, imap_keyword)
-    # get email uid list
-    email_list = data[0].split()
-    logger.info(f"got {len(email_list)} emails")
-    for uid in email_list:
-        logger.info(f"read email uid {uid}")
-        if check_email_saved(client, uid):
-            logger.info(f"email uid {uid} already saved")
-            continue
-        else:
-            read_and_save_email(client, uid)
-            logger.info(f"email uid {uid} saved")
 
+    def read_emails(self, folder: str = 'INBOX', imap_keyword: str = "UNSEEN"):
+        self.client.select(folder)
+        _, data = self.client.uid('search', None, imap_keyword)
+        
+        # get email uid list
+        email_list = data[0].split()
+        self.logger.info(f"got {len(email_list)} emails")
+        email_list.reverse()
+        for uid in email_list:
+            if self.check_email_saved(uid):
+                self.logger.info(f"email uid {uid} already saved")
+            else:
+                self.read_and_save_email(uid)
+                self.logger.info(f"email uid {uid} saved")
 
-def read_and_save_email(client: imaplib.IMAP4_SSL, uid: str):
+    def read_and_save_email(self, uid: str):
-    message_parts = "(BODY.PEEK[])"
+        message_parts = "(BODY.PEEK[])"
-    _, email_data = client.uid('fetch', uid, message_parts)
+        _, email_data = self.client.uid('fetch', uid, message_parts)
-    mail = mailparser.parse_from_bytes(email_data[0][1])
+        mail = mailparser.parse_from_bytes(email_data[0][1])
-    logger.info(f"got email subject [{mail.subject}]")
+        self.logger.info(f"got email subject [{mail.subject}]")
-    save_email(mail)
+        self.save_email(mail)
 
-def get_local_dir_name(mail: mailparser.MailParser) -> str:
+    def get_local_dir_name(self, mail: mailparser.MailParser) -> str:
-    dir =  f"{config.get('LOCAL_DIR')}/{config.get('EMAIL_ADDRESS')}"
+        dir =  f"{self.config.get('LOCAL_DIR')}/{self.config.get('EMAIL_ADDRESS')}"
-    name = f"{mail.subject}__{mail.date}"
+        name = f"{mail.subject}__{mail.date}"
-    name = hashlib.md5(name.encode('utf-8')).hexdigest()
+        name = hashlib.md5(name.encode('utf-8')).hexdigest()
-    return f"{dir}/{name}"
+        return f"{dir}/{name}"
 
+    def check_email_saved(self, uid: str):
+        message_parts = "(BODY[HEADER])"
+        _, email_data = self.client.uid('fetch', uid, message_parts)
+        mail = mailparser.parse_from_bytes(email_data[0][1])
+        self.logger.info(f"[{uid}]check email subject [{mail.subject}]")
+        dir = self.get_local_dir_name(mail)
+        self.logger.info(f"check email saved {dir}")
+        file = f"{dir}/email.txt"
+        if os.path.exists(file):
+            return False
+        return False
 
-# check only need to check email header, not need to download email body
+    # save email attachment(images)
-def check_email_saved(client: imaplib.IMAP4_SSL, uid: str):
+    def save_email_attachment(self, mail: mailparser.MailParser, email_dir: str):
-    message_parts = "(BODY[HEADER])"
+        for attachment in mail.attachments:
-    _, email_data = client.uid('fetch', uid, message_parts)
+            if attachment['mail_content_type'] in ['image/png', 'image/jpeg', 'image/gif']:
-    mail = mailparser.parse_from_bytes(email_data[0][1])
+                print('current mail have image attachment')
-    logger.info(f"check email subject [{mail.subject}]")
+                img_dir = f"{email_dir}/image"
-    dir = get_local_dir_name(mail)
+                if not os.path.exists(img_dir):
-    logger.info(f"check email saved {dir}")
+                    os.makedirs(img_dir)
-    file = f"{dir}/email.txt"
+                filename = attachment['filename']
-    if os.path.exists(file):
+                filefullname = f"{img_dir}/{filename}"
-        return True
+                image_data = attachment['payload']
-    return False
+                try:
+                    image_data = base64.b64decode(image_data)
+                except base64.binascii.Error:
+                    image_data = image_data.encode()
+                with open(filefullname, 'wb') as f:
+                    f.write(image_data)
+                    self.logger.info(f"save email image {filename} success")
 
+    # save email body images(html content)
+    def save_body_images(self, html_content: str, email_dir: str):
+        # get all image urls
+        soup = BeautifulSoup(html_content, 'html.parser')
+        img_tags = soup.find_all('img')
+        img_urls = [img['src'] for img in img_tags if 'src' in img.attrs]
+        self.logger.info(f'Found {len(img_urls)} images in email body')
 
-# save email to local file by each folder
+        if not os.path.exists(email_dir):
-def save_email(mail: mailparser.MailParser):
+            os.makedirs(email_dir)
-    # create email account dir
-    dir = f"{config.get('LOCAL_DIR')}/{config.get('EMAIL_ADDRESS')}"
-    if not os.path.exists(dir):
-        os.makedirs(dir)
-    # create email local dir
-    email_dir = get_local_dir_name(mail)
-    logger.info(f"save email to {email_dir}")
-    if not os.path.exists(email_dir):
-        os.makedirs(email_dir)
-    
-    # save email content and meta info
-    with open(f"{email_dir}/email.txt", "w") as f:
-        f.write(mail.body)
-    with open(f"{email_dir}/meta.json", "w", encoding='utf-8') as f:
-        mail_dict = json.loads(mail.mail_json)
-        if 'body' in mail_dict:
-            del mail_dict['body']
-        json.dump(mail_dict, f, ensure_ascii=False, indent=4)
-        logger.info(f"save email meta info {f.name}")
 
+        for img_url in img_urls:
+            # keep the original image filename(last of url)
+            img_filename = os.path.join(email_dir, img_url.split('/')[-1])
+            # download image 
+            response = requests.get(img_url, stream=True)
+            if response.status_code == 200:
+                with open(img_filename, 'wb') as img_file:
+                    for chunk in response.iter_content(1024):
+                        img_file.write(chunk)
+                self.logger.info(f'Downloaded {img_url} to {img_filename}')
+            else:
+                self.logger.info(f'Failed to download {img_url}')
+
+    # save email content to local dir
+    def save_email(self, mail: mailparser.MailParser):
+        dir = f"{self.config.get('LOCAL_DIR')}/{self.config.get('EMAIL_ADDRESS')}"
+        if not os.path.exists(dir):
+            os.makedirs(dir)
+        email_dir = self.get_local_dir_name(mail)
+        self.logger.info(f"save email to {email_dir}")
+        if not os.path.exists(email_dir):
+            os.makedirs(email_dir)
+        with open(f"{email_dir}/email.txt", "w") as f:
+            f.write(mail.body)
+        with open(f"{email_dir}/meta.json", "w", encoding='utf-8') as f:
+            mail_dict = json.loads(mail.mail_json)
+            if 'body' in mail_dict:
+                del mail_dict['body']
+            json.dump(mail_dict, f, ensure_ascii=False, indent=4)
+            self.logger.info(f"save email meta info {f.name}")
+        
+        self.save_email_attachment(mail, email_dir)
+        self.save_body_images(mail.body, f"{email_dir}/body_image")
 
 
 if __name__ == "__main__":
-    mail = email_client()
+    spider = EmailSpider()
     folder = 'INBOX'
-    # imap_keyword = "UNSEEN"
     imap_keyword = "ALL"
-    read_emails(mail, folder, imap_keyword)
+    spider.read_emails(folder, imap_keyword)