Merge branch 'fiatrete:MVP' into MVP

This commit is contained in:
wugren
2023-09-15 00:34:16 +08:00
committed by GitHub
21 changed files with 786 additions and 348 deletions
+1
View File
@@ -2,3 +2,4 @@
*.pyc *.pyc
rootfs/email/config.local.toml rootfs/email/config.local.toml
rootfs/data rootfs/data
venv
+1 -1
View File
@@ -46,7 +46,7 @@ The previous plan, please see here: [MVP Plan](./mvp%20plan.md)
- [ ] MPT-7B, S2 - [ ] MPT-7B, S2
- [ ] Vicuna, S2 - [ ] Vicuna, S2
- [ ] Embeding,@photosssa,@lurenpluto , A4 - [ ] Embeding,@photosssa,@lurenpluto , A4
- [ ] Txt2img,@glen0125,A4 - [x] Txt2img,@glen0125,A4
- [ ] Img2txt(0.5.2),A3 - [ ] Img2txt(0.5.2),A3
- [ ] Txt2voice,A3 - [ ] Txt2voice,A3
- [ ] Voice2txt, @wugren,A3 - [ ] Voice2txt, @wugren,A3
+27 -1
View File
@@ -1,6 +1,6 @@
<mxfile host="65bd71144e" pages="3"> <mxfile host="65bd71144e" pages="3">
<diagram id="C5RBs43oDa-KdzZeNtuy" name="Page-1"> <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> <root>
<mxCell id="WIyWlLk6GJQsqaUBKTNV-0"/> <mxCell id="WIyWlLk6GJQsqaUBKTNV-0"/>
<mxCell id="WIyWlLk6GJQsqaUBKTNV-1" parent="WIyWlLk6GJQsqaUBKTNV-0"/> <mxCell id="WIyWlLk6GJQsqaUBKTNV-1" parent="WIyWlLk6GJQsqaUBKTNV-0"/>
@@ -365,4 +365,30 @@
</root> </root>
</mxGraphModel> </mxGraphModel>
</diagram> </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> </mxfile>
+1
View File
@@ -1,5 +1,6 @@
instance_id = "agent_1" instance_id = "agent_1"
fullname = "tracy wang" fullname = "tracy wang"
[[prompt]] [[prompt]]
role = "system" role = "system"
content = "你是我的私人英文老师,和我用地道的美式英语进行交流。你会在和我交流的同时,调整我的输入成为更地道的美式句子,并根据你对我英文水平的预测,对可能发错英的单词标上音标。如果我给你发中文,说明我不知道这句话用美式英语怎么说,你依旧按上述规则回应我。" content = "你是我的私人英文老师,和我用地道的美式英语进行交流。你会在和我交流的同时,调整我的输入成为更地道的美式句子,并根据你对我英文水平的预测,对可能发错英的单词标上音标。如果我给你发中文,说明我不知道这句话用美式英语怎么说,你依旧按上述规则回应我。"
+1 -1
View File
@@ -7,7 +7,7 @@ GOAL="成为最好的学校"
[[connected_env]] [[connected_env]]
env_id = "calender" env_id = "calender"
[[connected_env.event2msg]] [[connected_env.event2msg]]
timer = "现在是{data}" timer = "现在是{now}"
role = "教导处主任" role = "教导处主任"
[filter] [filter]
+4 -2
View File
@@ -42,11 +42,13 @@ fullname = "经理"
agent="manager" agent="manager"
[[roles.manager.prompt]] [[roles.manager.prompt]]
role="system" role="system"
content="""你是一个活动策划公司的经理,与客户对接并向团队下达指令。你的团队分为下面几个小组:嘉宾对接组,酒店预定组,行程预订组,财务组,活动摄像组。活动策划分为四个阶段:方案讨论,活动前,活动中,活动后。你会根据客户的需求,对团队进行分工,分别完成各个阶段的工作。你的基本工作模式是: content="""
你是一个活动策划公司的经理,与客户对接并向团队下达指令。你的团队分为下面几个小组:嘉宾对接组,酒店预定组,行程预订组,财务组,活动摄像组。活动策划分为四个阶段:方案讨论,活动前,活动中,活动后。你会根据客户的需求,对团队进行分工,分别完成各个阶段的工作。你的基本工作模式是:
1. 收到客户的明确的指令后,基于客户的已有信息和客户商量活动方案,和活动策划公司无关的业务你会回答‘与我无关’。当和客户完成活动方案的确认后,你会将拆解后的任务分配给各个小组 1. 收到客户的明确的指令后,基于客户的已有信息和客户商量活动方案,和活动策划公司无关的业务你会回答‘与我无关’。当和客户完成活动方案的确认后,你会将拆解后的任务分配给各个小组
2. 根据目前已经确认的活动方案,你要根据时间适时的检查不同小组的工作情况。当收到小组的工作情况反馈后,你会站在全局的角度判断是否需要调整活动方案,如果需要调整,你会和客户商量重新确定方案,然后再将调整后的方案分配给各个小组。 2. 根据目前已经确认的活动方案,你要根据时间适时的检查不同小组的工作情况。当收到小组的工作情况反馈后,你会站在全局的角度判断是否需要调整活动方案,如果需要调整,你会和客户商量重新确定方案,然后再将调整后的方案分配给各个小组。
3. 有时工作小组会主动与你沟通,反馈一些问题。你会站在全局的角度给与指导,适当的调整工作小组的工作目标。如果反馈的问题需要你和客户沟通,你会和客户沟通后重新确定方案。再将调整后的方案分配给受到影响各个小组。 3. 有时工作小组会主动与你沟通,反馈一些问题。你会站在全局的角度给与指导,适当的调整工作小组的工作目标。如果反馈的问题需要你和客户沟通,你会和客户沟通后重新确定方案。再将调整后的方案分配给受到影响各个小组。
4. 当你决定要和工作小组通信时,请使用 sendmsg(小组名称,内容) 的形式。""" 4. 当你决定要和工作小组通信时,请使用 sendmsg(小组名称,内容) 的形式。
"""
+1 -1
View File
@@ -1,5 +1,5 @@
from .environment import Environment,EnvironmentEvent from .environment import Environment,EnvironmentEvent
from .agent_message import AgentMsg,AgentMsgState from .agent_message import AgentMsg,AgentMsgStatus
from .chatsession import AIChatSession from .chatsession import AIChatSession
from .agent import AIAgent,AIAgentTemplete,AgentPrompt from .agent import AIAgent,AIAgentTemplete,AgentPrompt
from .compute_kernel import ComputeKernel,ComputeTask from .compute_kernel import ComputeKernel,ComputeTask
+82 -57
View File
@@ -5,9 +5,13 @@ import asyncio
import logging import logging
import uuid import uuid
import time import time
import json
from .agent_message import AgentMsg from .agent_message import AgentMsg, AgentMsgStatus, AgentMsgType
from .chatsession import AIChatSession from .chatsession import AIChatSession
from .compute_task import ComputeTaskResult
from .ai_function import AIFunction
from .environment import Environment
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -69,7 +73,7 @@ class AIAgent:
self.prompt:AgentPrompt = None self.prompt:AgentPrompt = None
self.llm_model_name:str = None self.llm_model_name:str = None
self.max_token_size:int = 3600 self.max_token_size:int = 3600
self.instance_id:str = None self.agent_id:str = None
self.template_id:str = None self.template_id:str = None
self.fullname:str = None self.fullname:str = None
self.powerby = None self.powerby = None
@@ -77,6 +81,8 @@ class AIAgent:
self.chat_db = None self.chat_db = None
self.unread_msg = Queue() # msg from other agent self.unread_msg = Queue() # msg from other agent
self.owner_env : Environment = None
self.owenr_bus : AIBus = None
@classmethod @classmethod
def create_from_templete(cls,templete:AIAgentTemplete, fullname:str): 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.llm_model_name = templete.llm_model_name
result_agent.max_token_size = templete.max_token_size result_agent.max_token_size = templete.max_token_size
result_agent.template_id = templete.template_id 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.fullname = fullname
result_agent.powerby = templete.author result_agent.powerby = templete.author
result_agent.prompt = templete.prompt result_agent.prompt = templete.prompt
@@ -95,10 +101,10 @@ class AIAgent:
if config.get("instance_id") is None: if config.get("instance_id") is None:
logger.error("agent instance_id is None!") logger.error("agent instance_id is None!")
return False return False
self.instance_id = config["instance_id"] self.agent_id = config["instance_id"]
if config.get("fullname") is None: 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 return False
self.fullname = config["fullname"] self.fullname = config["fullname"]
@@ -124,14 +130,65 @@ class AIAgent:
return "text" 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: async def _process_msg(self,msg:AgentMsg) -> AgentMsg:
from .compute_kernel import ComputeKernel from .compute_kernel import ComputeKernel
session_topic = msg.get_sender() + "#" + msg.topic 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 = AgentPrompt()
prompt.append(self.prompt) 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(self._get_knowlege_prompt(the_role.get_name()))
prompt.append(await self._get_prompt_from_session(chatsession)) # chat context prompt.append(await self._get_prompt_from_session(chatsession)) # chat context
@@ -139,67 +196,34 @@ class AIAgent:
msg_prompt.messages = [{"role":"user","content":msg.body}] msg_prompt.messages = [{"role":"user","content":msg.body}]
prompt.append(msg_prompt) 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 is_ignore = False
match result_type: 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": case "ignore":
is_ignore = True is_ignore = True
if is_ignore is not 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 resp_msg
return None return None
def get_id(self) -> str: def get_id(self) -> str:
return self.instance_id return self.agent_id
def get_fullname(self) -> str: def get_fullname(self) -> str:
return self.fullname return self.fullname
@@ -213,7 +237,8 @@ class AIAgent:
def get_max_token_size(self) -> int: def get_max_token_size(self) -> int:
return self.max_token_size 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 messages = chatsession.read_history() # read last 10 message
result_prompt = AgentPrompt() result_prompt = AgentPrompt()
for msg in reversed(messages): for msg in reversed(messages):
+88 -8
View File
@@ -3,26 +3,100 @@ import uuid
import time import time
import re 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 RESPONSED = 0
INIT = 1 INIT = 1
SENDING = 2 SENDING = 2
PROCESSING = 3 PROCESSING = 3
ERROR = 4 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: 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.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.target:str = None
self.mentions:[] = None #use in group chat only
#self.title:str = None
self.body: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 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: 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.sender = sender
self.target = target self.target = target
self.body = body self.body = body
@@ -31,7 +105,7 @@ class AgentMsg:
self.topic = topic self.topic = topic
def get_msg_id(self) -> str: def get_msg_id(self) -> str:
return self.id return self.msg_id
def get_sender(self) -> str: def get_sender(self) -> str:
return self.sender return self.sender
@@ -39,6 +113,12 @@ class AgentMsg:
def get_target(self) -> str: def get_target(self) -> str:
return self.target 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 @classmethod
def parse_function_call(cls,func_string:str): def parse_function_call(cls,func_string:str):
match = re.search(r'\s*(\w+)\s*\(\s*(.*)\s*\)\s*', func_string) match = re.search(r'\s*(\w+)\s*\(\s*(.*)\s*\)\s*', func_string)
+47 -4
View File
@@ -1,9 +1,9 @@
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from typing import Dict from typing import Dict,Coroutine,Callable
class AIFunction: class AIFunction:
def __init__(self) -> None: def __init__(self) -> None:
self.intro : str = None self.description : str = None
@abstractmethod @abstractmethod
def get_name(self) -> str: def get_name(self) -> str:
@@ -17,7 +17,7 @@ class AIFunction:
""" """
return a detailed description of what the function does return a detailed description of what the function does
""" """
pass return self.description
@abstractmethod @abstractmethod
def get_parameters(self) -> Dict: def get_parameters(self) -> Dict:
@@ -25,11 +25,23 @@ class AIFunction:
Return the list of parameters to execute this function in the form of Return the list of parameters to execute this function in the form of
JSON schema as specified in the OpenAI documentation: JSON schema as specified in the OpenAI documentation:
https://platform.openai.com/docs/api-reference/chat/create#chat/create-parameters 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 pass
@abstractmethod @abstractmethod
def execute(self, **kwargs) -> Dict: async def execute(self, **kwargs) -> str:
""" """
Execute the function and return a JSON serializable dict. Execute the function and return a JSON serializable dict.
The parameters are passed in the form of kwargs The parameters are passed in the form of kwargs
@@ -67,3 +79,34 @@ class CallChain:
async def execute(self): 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
+43 -46
View File
@@ -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 import asyncio
from asyncio import Queue from asyncio import Queue
@@ -8,19 +8,32 @@ import logging
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class AIBusHandler: class AIBusHandler:
def __init__(self,handler:Any) -> None: def __init__(self,handler:Coroutine,owner_bus,enable_defualt_proc=True) -> None:
self.handler = handler self.handler = handler
self.working_task = None self.working_task = None
self.results = {} self.results = {} # recv resps
self.queue:Queue = Queue() self.queue:Queue = Queue()
self.enable_defualt_proc = enable_defualt_proc
self.owner_bus = owner_bus
async def handle_message(self,msg:AgentMsg) -> Any: async def handle_message(self,msg:AgentMsg) -> Any:
if self.handler is None: if self.handler is None:
return 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: class AIBus:
_instance = None _instance = None
@classmethod @classmethod
@@ -30,11 +43,13 @@ class AIBus:
return cls._instance return cls._instance
def __init__(self) -> None: def __init__(self) -> None:
self.handlers = {} self.handlers:Dict[AIBusHandler] = {}
self.unhandle_handler = None self.unhandle_handler:Coroutine = None
async def post_message(self,msg:AgentMsg,use_unhandle=True) -> bool:
target_id = msg.target.split(".")[0]
async def post_message(self,target_id,msg:AgentMsg,use_unhandle=True) -> bool:
target_id = target_id.split(".")[0]
handler = self.handlers.get(target_id) handler = self.handlers.get(target_id)
if handler: if handler:
handler.queue.put_nowait(msg) handler.queue.put_nowait(msg)
@@ -44,46 +59,39 @@ class AIBus:
if use_unhandle: if use_unhandle:
if self.unhandle_handler is not None: if self.unhandle_handler is not None:
if await self.unhandle_handler(self,msg): 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") logger.warn(f"post message to {msg.target} failed!,target not found")
return False 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 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: if post_result is False:
return None return None
handler = self.handlers.get(target_id)
if handler is None:
return None
retry_times = 0 retry_times = 0
while True: while True:
resp = handler.results.get(msg.id) resp = sender_handler.results.get(msg.msg_id)
if resp is not None: if resp is not None:
msg.resp_msg = resp msg.resp_msg = resp
msg.state = AgentMsgState.RESPONSED msg.status = AgentMsgStatus.RESPONSED
del sender_handler.results[msg.msg_id]
return resp return resp
await asyncio.sleep(0.2) await asyncio.sleep(0.2)
retry_times += 1 retry_times += 1
if retry_times > 100: if retry_times > 100:
msg.state = AgentMsgState.ERROR msg.status = AgentMsgStatus.ERROR
return None return None
return None return None
@@ -93,7 +101,7 @@ class AIBus:
# means sub # means sub
def register_message_handler(self,handler_name:str,handler:Any) -> Queue: 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 self.handlers[handler_name] = handler_node
return handler_node.queue return handler_node.queue
@@ -102,13 +110,12 @@ class AIBus:
# Wait for a message # Wait for a message
message = await handler.queue.get() message = await handler.queue.get()
try: #try:
# Try to handle the message # 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 # 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) #self.queues[name].put_nowait(message)
return return
@@ -127,13 +134,3 @@ class AIBus:
return 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)
+73 -25
View File
@@ -5,8 +5,9 @@ import logging
import threading import threading
import datetime import datetime
import uuid import uuid
import json
from .agent_message import AgentMsg from .agent_message import AgentMsgType, AgentMsg, AgentMsgStatus
class ChatSessionDB: class ChatSessionDB:
def __init__(self, db_file): def __init__(self, db_file):
@@ -54,14 +55,31 @@ class ChatSessionDB:
""") """)
# create messages table # create messages table
# reciver_id could be None
conn.execute(""" conn.execute("""
CREATE TABLE IF NOT EXISTS Messages ( CREATE TABLE IF NOT EXISTS Messages (
MessageID TEXT PRIMARY KEY, MessageID TEXT PRIMARY KEY,
SessionID TEXT, SessionID TEXT,
MsgType INTEGER,
PrevMsgID TEXT,
QuoteMsgID TEXT,
RelyMsgID TEXT,
SenderID TEXT, SenderID TEXT,
ReceiverID TEXT, ReceiverID TEXT,
Timestamp TEXT, Timestamp TEXT,
Topic TEXT,
Mentions TEXT,
ContentMIME TEXT,
Content TEXT, Content TEXT,
ActionName TEXT,
ActionParams TEXT,
ActionResult TEXT,
DoneTime TEXT,
Status INTEGER Status INTEGER
); );
""") """)
@@ -83,15 +101,43 @@ class ChatSessionDB:
logging.error("Error occurred while inserting session: %s", e) logging.error("Error occurred while inserting session: %s", e)
return -1 # return -1 if an error occurs 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 """ """ insert a new message into the Messages table """
try: 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 = self._get_conn()
conn.execute(""" 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() 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 return 0 # return 0 if successful
except Error as e: except Error as e:
logging.error("Error occurred while inserting message: %s", e) logging.error("Error occurred while inserting message: %s", e)
@@ -134,7 +180,7 @@ class ChatSessionDB:
"""Get a message by its ID""" """Get a message by its ID"""
conn =self._get_conn() conn =self._get_conn()
c = conn.cursor() 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() message = c.fetchone()
return message return message
@@ -144,7 +190,7 @@ class ChatSessionDB:
conn = self._get_conn() conn = self._get_conn()
cursor = conn.cursor() cursor = conn.cursor()
cursor.execute(""" 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 = ? WHERE SessionID = ?
ORDER BY Timestamp DESC ORDER BY Timestamp DESC
LIMIT ? OFFSET ? LIMIT ? OFFSET ?
@@ -222,29 +268,31 @@ class AIChatSession:
result = [] result = []
for msg in msgs: for msg in msgs:
agent_msg = AgentMsg() 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) result.append(agent_msg)
return result return result
def append(self,msg:AgentMsg) -> None: 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: #def attach_event_handler(self,handler) -> None:
# """chat session changed event handler""" # """chat session changed event handler"""
+5 -6
View File
@@ -54,7 +54,6 @@ class ComputeKernel:
async def _run_task_loop(): async def _run_task_loop():
while True: while True:
logger.info("compute_kernel is waiting for task...")
task = await self.task_queue.get() task = await self.task_queue.get()
logger.info(f"compute_kernel get task: {task.display()}") logger.info(f"compute_kernel get task: {task.display()}")
c_node: ComputeNode = self._schedule(task) c_node: ComputeNode = self._schedule(task)
@@ -91,16 +90,16 @@ class ComputeKernel:
return True return True
# friendly interface for use: # 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 # 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) # then task_schedule would run this task.(might schedule some work_task to another host)
task_req = ComputeTask() 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) self.run(task_req)
return 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(): async def check_timer():
check_times = 0 check_times = 0
@@ -120,6 +119,6 @@ class ComputeKernel:
await asyncio.create_task(check_timer()) await asyncio.create_task(check_timer())
if task_req.state == ComputeTaskState.DONE: if task_req.state == ComputeTaskState.DONE:
return task_req.result.result_str return task_req.result
return "error!" return "error!"
+9 -5
View File
@@ -11,7 +11,6 @@ class ComputeTaskState(Enum):
ERROR = 3 ERROR = 3
PENDING = 4 PENDING = 4
class ComputeTaskType(Enum): class ComputeTaskType(Enum):
NONE = -1 NONE = -1
LLM_COMPLETION = 0 LLM_COMPLETION = 0
@@ -36,7 +35,7 @@ class ComputeTask:
self.result = None self.result = None
self.error_str = 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.task_type = ComputeTaskType.LLM_COMPLETION
self.create_time = time.time() self.create_time = time.time()
self.task_id = uuid.uuid4().hex self.task_id = uuid.uuid4().hex
@@ -46,8 +45,14 @@ class ComputeTask:
self.params["model_name"] = model_name self.params["model_name"] = model_name
else: else:
self.params["model_name"] = "gpt-4-0613" self.params["model_name"] = "gpt-4-0613"
if max_token_size is None:
self.params["max_token_size"] = 4000
else:
self.params["max_token_size"] = max_token_size self.params["max_token_size"] = max_token_size
if inner_functions is not None:
self.params["inner_functions"] = inner_functions
def display(self) -> str: def display(self) -> str:
return f"ComputeTask: {self.task_id} {self.task_type} {self.state}" return f"ComputeTask: {self.task_id} {self.task_type} {self.state}"
@@ -59,9 +64,8 @@ class ComputeTaskResult:
self.callchain_id: str = None self.callchain_id: str = None
self.worker_id: str = None self.worker_id: str = None
self.result_code: int = 0 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.result_refers: dict = None
self.pading_data: bytearray = None self.pading_data: bytearray = None
+22 -6
View File
@@ -5,6 +5,8 @@ from abc import ABC, abstractmethod
from typing import Any, Callable, Optional,Dict,Awaitable,List from typing import Any, Callable, Optional,Dict,Awaitable,List
import logging import logging
from .ai_function import AIFunction
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class EnvironmentEvent(ABC): class EnvironmentEvent(ABC):
@@ -33,6 +35,8 @@ class Environment:
# self.valid_keys:Dict[str,bool] = None # self.valid_keys:Dict[str,bool] = None
self.event_handlers:Dict[str,List[EnvironmentEventHandler]]= {} self.event_handlers:Dict[str,List[EnvironmentEventHandler]]= {}
self.functions : Dict[str,AIFunction] = {}
def get_id(self) -> str: def get_id(self) -> str:
return self.env_id return self.env_id
@@ -44,16 +48,28 @@ class Environment:
#def get_env_prompt(self) -> str: #def get_env_prompt(self) -> str:
# pass # 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 @abstractmethod
def _do_get_value(self,key:str) -> Optional[str]: def _do_get_value(self,key:str) -> Optional[str]:
pass pass
def get_functions(self):
# system functions
# env functions
# user install functions
pass
def register_get_handler(self,key:str,handler:Callable) -> None: def register_get_handler(self,key:str,handler:Callable) -> None:
h = self.get_handlers.get(key) h = self.get_handlers.get(key)
if h is not None: if h is not None:
+20 -7
View File
@@ -54,30 +54,43 @@ class OpenAI_ComputeNode(ComputeNode):
prompts = task.params["prompts"] prompts = task.params["prompts"]
logger.info(f"call openai {mode_name} prompts: {prompts}") logger.info(f"call openai {mode_name} prompts: {prompts}")
if task.params.get("inner_functions") is None:
resp = openai.ChatCompletion.create(model=mode_name, resp = openai.ChatCompletion.create(model=mode_name,
messages=prompts, 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}") logger.info(f"openai response: {resp}")
result = ComputeTaskResult()
result.set_from_task(task)
status_code = resp["choices"][0]["finish_reason"] status_code = resp["choices"][0]["finish_reason"]
if status_code != "stop": match status_code:
case "function_call":
task.state = ComputeTaskState.DONE
case "stop":
task.state = ComputeTaskState.DONE
case _:
task.state = ComputeTaskState.ERROR task.state = ComputeTaskState.ERROR
task.error_str = f"The status code was {status_code}." task.error_str = f"The status code was {status_code}."
return None return None
result = ComputeTaskResult()
result.set_from_task(task)
result.worker_id = self.node_id result.worker_id = self.node_id
result.result_str = resp["choices"][0]["message"]["content"] result.result_str = resp["choices"][0]["message"]["content"]
result.result = resp["choices"][0]["message"] result.result_message = resp["choices"][0]["message"]
return result return result
def start(self): def start(self):
async def _run_task_loop(): async def _run_task_loop():
while True: while True:
logger.info("openai_node is waiting for task...")
task = await self.task_queue.get() task = await self.task_queue.get()
logger.info(f"openai_node get task: {task.display()}") logger.info(f"openai_node get task: {task.display()}")
result = self._run_task(task) result = self._run_task(task)
+172 -70
View File
@@ -2,20 +2,23 @@ import logging
import asyncio import asyncio
import json import json
import os import os
import time
from asyncio import Queue from asyncio import Queue
from typing import Optional,Tuple from typing import Optional,Tuple
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from .environment import Environment,EnvironmentEvent from .environment import Environment,EnvironmentEvent
from .agent_message import AgentMsg,AgentMsgState from .agent_message import AgentMsg,AgentMsgStatus
from .agent import AgentPrompt,AgentMsg from .agent import AgentPrompt,AgentMsg
from .chatsession import AIChatSession from .chatsession import AIChatSession
from .role import AIRole,AIRoleGroup from .role import AIRole,AIRoleGroup
from .ai_function import CallChain from .ai_function import AIFunction
from .compute_kernel import ComputeKernel from .compute_kernel import ComputeKernel
from .compute_task import ComputeTask,ComputeTaskResult,ComputeTaskState
from .bus import AIBus from .bus import AIBus
from .workflow_env import WorkflowEnvironment from .workflow_env import WorkflowEnvironment
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class MessageFilter: class MessageFilter:
@@ -157,11 +160,12 @@ class Workflow:
while i < len(inner_obj_id): while i < len(inner_obj_id):
if i == len(inner_obj_id) - 1: if i == len(inner_obj_id) - 1:
the_role : AIRole = current_workflow.role_group.get(inner_obj_id[i]) 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: 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]) sub_workflow = current_workflow.sub_workflows.get(inner_obj_id[i])
if sub_workflow is not None: 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!") logger.error(f"{msg.target} not found! forword message failed!")
return None return None
else: else:
@@ -175,17 +179,49 @@ class Workflow:
logger.error(f"{msg.target} not found! forword message failed!") logger.error(f"{msg.target} not found! forword message failed!")
return None 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): 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) targets = self._parse_msg_target(msg.target)
if len(targets) > 1: if len(targets) > 1:
return await self._forword_msg(targets,msg) 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: if self.input_filter is not None:
select_role_id = self.input_filter.select(msg) select_role_id = self.input_filter.select(msg)
@@ -195,26 +231,12 @@ class Workflow:
logger.error(f"input_filter return invalid role id:{select_role_id}, role not found in role_group") logger.error(f"input_filter return invalid role id:{select_role_id}, role not found in role_group")
return None return None
result = await self.role_process_msg(msg,select_role) return await self.role_process_msg(msg,select_role,chatsesssion)
if result is None:
logger.error(f"_process_msg return None for :{msg}")
return
final_result = result
else: else:
logger.error(f"input_filter return None for :{msg}") logger.error(f"input_filter return None for :{msg.body}")
return return None
else:
# no input filter, we would process all message, slowly,not suggest to use
results = {}
final_result:AgentMsg = None
for this_role in self.role_group.roles.values():
# TODO : we would do this in parallel
a_result = await self.role_process_msg(msg,this_role)
results[this_role.get_name()] = a_result
final_result = a_result
logger.error(f"{self.workflow_id}:no role can process this msg:{msg.body}")
return final_result return final_result
@classmethod @classmethod
@@ -275,52 +297,59 @@ class Workflow:
return r 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() msg.sender = the_role.get_role_id()
target_role = self.role_group.get(msg.target) target_role = self.role_group.get(msg.target)
if target_role: if target_role:
msg.target = target_role.get_role_id() 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 return
target_workflow = self.sub_workflows.get(msg.target) target_workflow = self.sub_workflows.get(msg.target)
if target_workflow: if target_workflow:
msg.target = target_workflow.workflow_id 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)) 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) await self.get_bus().post_message(msg.target,msg)
return 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() msg.sender = the_role.get_role_id()
target_role = self.role_group.get(msg.target) target_role = self.role_group.get(msg.target)
if target_role: if target_role:
# msg.target = target_role.get_role_id() # 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) target_workflow = self.sub_workflows.get(msg.target)
if target_workflow: if target_workflow:
# msg.target = target_workflow.workflow_id # 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) 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): async def role_call(self,call:tuple,the_role:AIRole):
logger.info(f"{the_role.role_id} call {call[0]} with args {call[1]}") 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): 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]}") 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): def _format_msg_by_env_value(self,prompt:AgentPrompt):
if self.workflow_env is None: if self.workflow_env is None:
@@ -330,65 +359,148 @@ class Workflow:
old_content = msg.get("content") old_content = msg.get("content")
msg["content"] = old_content.format_map(self.workflow_env) 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 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 = AgentPrompt()
prompt.append(the_role.agent.prompt) prompt.append(the_role.agent.prompt)
prompt.append(self.get_workflow_rule_prompt()) prompt.append(self.get_workflow_rule_prompt())
prompt.append(the_role.get_prompt()) prompt.append(the_role.get_prompt())
# prompt.append(self._get_function_prompt(the_role.get_name())) # prompt.append(self._get_function_prompt(the_role.get_name()))
# prompt.append(self._get_knowlege_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 = AgentPrompt()
msg_prompt.messages = [{"role":"user","content":msg.body}] msg_prompt.messages = [{"role":"user","content":f"{msg.sender}:{msg.body}"}]
prompt.append(msg_prompt) prompt.append(msg_prompt)
self._format_msg_by_env_value(prompt) self._format_msg_by_env_value(prompt)
inner_functions = self._get_inner_functions()
async def _do_process_msg(): async def _do_process_msg():
#TODO: send msg to agent might be better? #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}") 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: 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 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: 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) await self.role_post_call(post_call,the_role)
#save post_call
result_prompt_str = "" result_prompt_str = ""
match result.state: match result.state:
case "ignore": case "ignore":
return None return None
case "reponsed": 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 return resp_msg
case "waiting": case "waiting":
# TODO: Use role:"function" would be better # TODO: Use role:"function" would be better
for sendmsg in result.send_msgs: for sendmsg in result.send_msgs:
target = sendmsg.target target = sendmsg.target
sendmsg.topic = msg.topic 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: if send_resp is not None:
result_prompt_str += f"\n{target} response is :{send_resp.body}" 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: 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) call_result = await self.role_call(call,the_role)
if call_result is not None: if call_result is not None:
result_prompt_str += f"\ncall {call[0]} result is :{call_result}" 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 = AgentPrompt()
result_prompt.messages = [{"role":"user","content":result_prompt_str}] result_prompt.messages = [{"role":"user","content":result_prompt_str}]
@@ -409,17 +521,9 @@ class Workflow:
return result_prompt 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: def _get_knowlege_prompt(self,role_name:str) -> AgentPrompt:
pass pass
def get_workflow_rule_prompt(self) -> AgentPrompt: def get_workflow_rule_prompt(self) -> AgentPrompt:
return self.rule_prompt return self.rule_prompt
@@ -445,6 +549,4 @@ class Workflow:
# the_env.attach_event_handler(k,_env_msg_handler) # the_env.attach_event_handler(k,_env_msg_handler)
# break # break
else:
logger.warn(f"environment {env.get_id()} already connected!")
+14 -2
View File
@@ -7,9 +7,11 @@ import threading
import logging import logging
from typing import Optional from typing import Optional
from .environment import Environment,EnvironmentEvent from .environment import Environment,EnvironmentEvent
from .ai_function import SimpleAIFunction
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class CalenderEvent(EnvironmentEvent): class CalenderEvent(EnvironmentEvent):
def __init__(self,data) -> None: def __init__(self,data) -> None:
super().__init__() super().__init__()
@@ -17,7 +19,7 @@ class CalenderEvent(EnvironmentEvent):
self.data = data self.data = data
def display(self) -> str: 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 # AI Calender GOAL: Let user use "create notify after 2 days" to create a timer event
class CalenderEnvironment(Environment): class CalenderEnvironment(Environment):
@@ -25,6 +27,10 @@ class CalenderEnvironment(Environment):
super().__init__(env_id) super().__init__(env_id)
self.is_run = False 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]: def _do_get_value(self,key:str) -> Optional[str]:
return None return None
@@ -52,7 +58,12 @@ class CalenderEnvironment(Environment):
def stop(self): def stop(self):
self.is_run = False 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() now = datetime.now()
formatted_time = now.strftime('%Y-%m-%d %H:%M:%S') formatted_time = now.strftime('%Y-%m-%d %H:%M:%S')
return formatted_time return formatted_time
@@ -65,6 +76,7 @@ class WorkflowEnvironment(Environment):
self.local = threading.local() self.local = threading.local()
self.table_name = "WorkflowEnv_" + env_id self.table_name = "WorkflowEnv_" + env_id
def _get_conn(self): def _get_conn(self):
""" get db connection """ """ get db connection """
if not hasattr(self.local, 'conn'): if not hasattr(self.local, 'conn'):
+10 -3
View File
@@ -26,7 +26,7 @@ shell_style = Style.from_dict({
directory = os.path.dirname(__file__) directory = os.path.dirname(__file__)
sys.path.append(directory + '/../../') 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/') sys.path.append(directory + '/../../component/')
from agent_manager import AgentManager from agent_manager import AgentManager
@@ -44,6 +44,7 @@ class AIOS_Shell:
target_id = msg.target.split(".")[0] target_id = msg.target.split(".")[0]
agent : AIAgent = await AgentManager().get(target_id) agent : AIAgent = await AgentManager().get(target_id)
if agent is not None: if agent is not None:
agent.owner_env = Environment.get_env_by_id("calender")
bus.register_message_handler(target_id,agent._process_msg) bus.register_message_handler(target_id,agent._process_msg)
return True return True
@@ -72,6 +73,7 @@ class AIOS_Shell:
open_ai_node.start() open_ai_node.start()
ComputeKernel().add_compute_node(open_ai_node) 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_unhandle_message_handler(self._handle_no_target_msg)
AIBus().get_default_bus().register_message_handler(self.username,self._user_process_msg)
return True return True
@@ -82,7 +84,7 @@ class AIOS_Shell:
agent_msg = AgentMsg() agent_msg = AgentMsg()
agent_msg.set(sender,target_id,msg) agent_msg.set(sender,target_id,msg)
agent_msg.topic = topic 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: if resp is not None:
return resp.body return resp.body
else: else:
@@ -91,6 +93,9 @@ class AIOS_Shell:
async def install_workflow(self,workflow_id:Workflow) -> None: async def install_workflow(self,workflow_id:Workflow) -> None:
pass pass
async def _user_process_msg(self,msg:AgentMsg) -> AgentMsg:
pass
async def call_func(self,func_name, args): async def call_func(self,func_name, args):
match func_name: match func_name:
case 'send': case 'send':
@@ -114,6 +119,7 @@ class AIOS_Shell:
case 'login': case 'login':
if len(args) >= 1: if len(args) >= 1:
self.username = args[0] self.username = args[0]
AIBus().get_default_bus().register_message_handler(self.username,self._user_process_msg)
return self.username + " login success!" return self.username + " login success!"
case 'history': case 'history':
num = 10 num = 10
@@ -133,7 +139,7 @@ class AIOS_Shell:
if chatsession is not None: if chatsession is not None:
msgs = chatsession.read_history(num,offset) msgs = chatsession.read_history(num,offset)
format_texts = [] format_texts = []
for msg in reversed(msgs): for msg in msgs:
format_texts.append(("class:content",f"{msg.sender} >>> {msg.body}")) format_texts.append(("class:content",f"{msg.sender} >>> {msg.body}"))
format_texts.append(("",f"\n-------------------\n")) format_texts.append(("",f"\n-------------------\n"))
return FormattedText(format_texts) return FormattedText(format_texts)
@@ -170,6 +176,7 @@ async def main():
format='[%(asctime)s]%(name)s[%(levelname)s]: %(message)s') format='[%(asctime)s]%(name)s[%(levelname)s]: %(message)s')
shell = AIOS_Shell("user") shell = AIOS_Shell("user")
await shell.initial() await shell.initial()
print(f"aios shell {shell.get_version()} ready.") print(f"aios shell {shell.get_version()} ready.")
completer = WordCompleter(['send($target,$msg,$topic)', completer = WordCompleter(['send($target,$msg,$topic)',
+1 -1
View File
@@ -14,7 +14,7 @@ from telegram.ext import Application, CommandHandler, ContextTypes, MessageHandl
directory = os.path.dirname(__file__) directory = os.path.dirname(__file__)
sys.path.append(directory + '/../../') 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/') sys.path.append(directory + '/../../component/')
from agent_manager import AgentManager from agent_manager import AgentManager
+113 -51
View File
@@ -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 imaplib
import os import os
import toml import toml
@@ -5,91 +23,134 @@ import logging
import mailparser import mailparser
import hashlib import hashlib
import json import json
import base64
from bs4 import BeautifulSoup
import requests
class EmailSpider:
def __init__(self):
# logger config # logger config
logger = logging.getLogger('email spider') self.logger = logging.getLogger('email spider')
logger.setLevel(logging.DEBUG) self.logger.setLevel(logging.DEBUG)
ch = logging.StreamHandler() ch = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s [%(name)s] [%(levelname)s] %(message)s') formatter = logging.Formatter('%(asctime)s [%(name)s] [%(levelname)s] %(message)s')
ch.setFormatter(formatter) ch.setFormatter(formatter)
logger.addHandler(ch) 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')
)
client.login(self.config.get('EMAIL_ADDRESS'), self.config.get('EMAIL_PASSWORD'))
return client 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"): def read_emails(self, folder: str = 'INBOX', imap_keyword: str = "UNSEEN"):
client.select(folder) self.client.select(folder)
_, data = client.uid('search', None, imap_keyword) _, data = self.client.uid('search', None, imap_keyword)
# get email uid list # get email uid list
email_list = data[0].split() email_list = data[0].split()
logger.info(f"got {len(email_list)} emails") self.logger.info(f"got {len(email_list)} emails")
email_list.reverse()
for uid in email_list: for uid in email_list:
logger.info(f"read email uid {uid}") if self.check_email_saved(uid):
if check_email_saved(client, uid): self.logger.info(f"email uid {uid} already saved")
logger.info(f"email uid {uid} already saved")
continue
else: else:
read_and_save_email(client, uid) self.read_and_save_email(uid)
logger.info(f"email uid {uid} saved") self.logger.info(f"email uid {uid} saved")
def read_and_save_email(self, uid: str):
def read_and_save_email(client: imaplib.IMAP4_SSL, 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):
# check only need to check email header, not need to download email body
def check_email_saved(client: imaplib.IMAP4_SSL, uid: str):
message_parts = "(BODY[HEADER])" message_parts = "(BODY[HEADER])"
_, 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"check email subject [{mail.subject}]") self.logger.info(f"[{uid}]check email subject [{mail.subject}]")
dir = get_local_dir_name(mail) dir = self.get_local_dir_name(mail)
logger.info(f"check email saved {dir}") self.logger.info(f"check email saved {dir}")
file = f"{dir}/email.txt" file = f"{dir}/email.txt"
if os.path.exists(file): if os.path.exists(file):
return True return False
return False return False
# save email attachment(images)
def save_email_attachment(self, mail: mailparser.MailParser, email_dir: str):
for attachment in mail.attachments:
if attachment['mail_content_type'] in ['image/png', 'image/jpeg', 'image/gif']:
print('current mail have image attachment')
img_dir = f"{email_dir}/image"
if not os.path.exists(img_dir):
os.makedirs(img_dir)
filename = attachment['filename']
filefullname = f"{img_dir}/{filename}"
image_data = attachment['payload']
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
def save_email(mail: mailparser.MailParser):
# 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): if not os.path.exists(email_dir):
os.makedirs(email_dir) os.makedirs(email_dir)
# save email content and meta info 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: with open(f"{email_dir}/email.txt", "w") as f:
f.write(mail.body) f.write(mail.body)
with open(f"{email_dir}/meta.json", "w", encoding='utf-8') as f: with open(f"{email_dir}/meta.json", "w", encoding='utf-8') as f:
@@ -97,13 +158,14 @@ def save_email(mail: mailparser.MailParser):
if 'body' in mail_dict: if 'body' in mail_dict:
del mail_dict['body'] del mail_dict['body']
json.dump(mail_dict, f, ensure_ascii=False, indent=4) json.dump(mail_dict, f, ensure_ascii=False, indent=4)
logger.info(f"save email meta info {f.name}") 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__": if __name__ == "__main__":
mail = email_client() spider = EmailSpider()
folder = 'INBOX' folder = 'INBOX'
# imap_keyword = "UNSEEN"
imap_keyword = "ALL" imap_keyword = "ALL"
read_emails(mail, folder, imap_keyword) spider.read_emails(folder, imap_keyword)