From b7b968d5f7d9e8ab88bac1d091dda2fd657bcab9 Mon Sep 17 00:00:00 2001 From: tsukasa Date: Fri, 1 Dec 2023 14:29:10 +0800 Subject: [PATCH] define bas environment --- rootfs/agents/JarvisPlus/agent.toml | 2 +- .../Mail/Spider/pipeline.toml | 4 - .../Mail/{Spider => Sync}/input.py | 0 .../Mail/Sync/pipeline.toml | 14 + rootfs/knowledge_pipelines/Mia/parser.py | 2 +- rootfs/knowledge_pipelines/pipelines.toml | 2 +- src/aios/agent/agent.py | 350 +++---- src/aios/agent/agent_base.py | 45 +- src/aios/agent/ai_function.py | 67 +- src/aios/environment/environment.py | 206 ++-- src/aios/environment/workspace_env.py | 927 ++++++------------ src/aios/knowledge/pipeline.py | 34 +- .../common_environment/local_document.py | 671 +++++++++++++ .../common_environment/to_learn_parser.py | 214 ++++ src/component/knowledge_manager/pipeline.py | 19 +- src/component/mail_environment/local.py | 2 +- src/component/mail_environment/mail.py | 129 ++- src/component/mail_environment/spider.py | 35 +- 18 files changed, 1700 insertions(+), 1023 deletions(-) delete mode 100644 rootfs/knowledge_pipelines/Mail/Spider/pipeline.toml rename rootfs/knowledge_pipelines/Mail/{Spider => Sync}/input.py (100%) create mode 100644 rootfs/knowledge_pipelines/Mail/Sync/pipeline.toml create mode 100644 src/component/common_environment/local_document.py create mode 100644 src/component/common_environment/to_learn_parser.py diff --git a/rootfs/agents/JarvisPlus/agent.toml b/rootfs/agents/JarvisPlus/agent.toml index dc249a2..b7cdac7 100644 --- a/rootfs/agents/JarvisPlus/agent.toml +++ b/rootfs/agents/JarvisPlus/agent.toml @@ -6,7 +6,7 @@ enable_timestamp = "true" owner_prompt = "I am your master {name} , now is {now}" contact_prompt = "I am your master's friend {name}" -[[do_prompt]] +[work.do] role = "system" content = """ My name is JarvisPlus, I am the master's super personal assistant. I think hard and try my best to complete TODOs. diff --git a/rootfs/knowledge_pipelines/Mail/Spider/pipeline.toml b/rootfs/knowledge_pipelines/Mail/Spider/pipeline.toml deleted file mode 100644 index df9f306..0000000 --- a/rootfs/knowledge_pipelines/Mail/Spider/pipeline.toml +++ /dev/null @@ -1,4 +0,0 @@ -name = "Mail.Issue" -input.module = "input.py" -input.params.path = "${myai_dir}/data" - diff --git a/rootfs/knowledge_pipelines/Mail/Spider/input.py b/rootfs/knowledge_pipelines/Mail/Sync/input.py similarity index 100% rename from rootfs/knowledge_pipelines/Mail/Spider/input.py rename to rootfs/knowledge_pipelines/Mail/Sync/input.py diff --git a/rootfs/knowledge_pipelines/Mail/Sync/pipeline.toml b/rootfs/knowledge_pipelines/Mail/Sync/pipeline.toml new file mode 100644 index 0000000..17fa1e2 --- /dev/null +++ b/rootfs/knowledge_pipelines/Mail/Sync/pipeline.toml @@ -0,0 +1,14 @@ +name = "Mail.Sync" +input.module = "input.py" +[input.params] +path = "${myai_dir}/mail" +imap_server = "imap.qq.com" +imap_port = 993 +address = "115620204@qq.com" +password = "zbbjpbukeonqbjja" +[input.params.fields] +from = "from" +to = "to" +subject = "subject" + + diff --git a/rootfs/knowledge_pipelines/Mia/parser.py b/rootfs/knowledge_pipelines/Mia/parser.py index cb8b5f8..6a3ca22 100644 --- a/rootfs/knowledge_pipelines/Mia/parser.py +++ b/rootfs/knowledge_pipelines/Mia/parser.py @@ -96,7 +96,7 @@ class EmbeddingParser: async def parse(self, object: ObjectID) -> str: obj = self.env.get_knowledge_store().load_object(object) await self.__do_embedding(obj) - return "insert into vector store" + return str(object) def init(env: KnowledgePipelineEnvironment, params: dict) -> EmbeddingParser: return EmbeddingParser(env, params) \ No newline at end of file diff --git a/rootfs/knowledge_pipelines/pipelines.toml b/rootfs/knowledge_pipelines/pipelines.toml index 7e69436..2008066 100644 --- a/rootfs/knowledge_pipelines/pipelines.toml +++ b/rootfs/knowledge_pipelines/pipelines.toml @@ -1,3 +1,3 @@ pipelines = [ - "Mail/Issue" + "Mail/Sync" ] \ No newline at end of file diff --git a/src/aios/agent/agent.py b/src/aios/agent/agent.py index 2b97326..4614216 100644 --- a/src/aios/agent/agent.py +++ b/src/aios/agent/agent.py @@ -18,6 +18,7 @@ from ..proto.compute_task import ComputeTaskResult,ComputeTaskResultCode from .agent_base import * from .chatsession import * from .ai_function import * +from ..environment.workspace_env import WorkspaceEnvironment, TodoListType from ..frame.contact_manager import ContactManager,Contact,FamilyMember from ..frame.compute_kernel import ComputeKernel @@ -145,17 +146,22 @@ class AIAgent(BaseAIAgent): self.contact_prompt_str = None self.history_len = 10 - self.review_todo_prompt = None - self.read_report_prompt = None - self.do_prompt = None - self.check_prompt = None - - self.goal_to_todo_prompt = None + todo_prompts = {} + todo_prompts[TodoListType.TO_WORK] = { + "do": DEFAULT_AGENT_DO_PROMPT, + "check": DEFAULT_AGENT_SELF_CHECK_PROMPT, + "review": None, + } + todo_prompts[TodoListType.TO_LEARN] = { + "do": DEFAULT_AGENT_LEARN_PROMPT, + "check": None, + "review": None, + } + self.todo_prompts = todo_prompts self.learn_token_limit = 4000 - self.learn_prompt = AgentPrompt(DEFAULT_AGENT_LEARN_PROMPT) self.chat_db = None self.unread_msg = Queue() # msg from other agent @@ -163,19 +169,18 @@ class AIAgent(BaseAIAgent): self.owenr_bus = None self.enable_function_list = None - - @classmethod - def create_from_templete(cls,templete:AIAgentTemplete, fullname:str): - # Agent just inherit from templete on craete,if template changed,agent will not change - result_agent = 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.agent_id = "agent#" + uuid.uuid4().hex - result_agent.fullname = fullname - result_agent.powerby = templete.author - result_agent.agent_prompt = templete.prompt - return result_agent + # @classmethod + # def create_from_templete(cls,templete:AIAgentTemplete, fullname:str): + # # Agent just inherit from templete on craete,if template changed,agent will not change + # result_agent = 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.agent_id = "agent#" + uuid.uuid4().hex + # result_agent.fullname = fullname + # result_agent.powerby = templete.author + # result_agent.agent_prompt = templete.prompt + # return result_agent def load_from_config(self,config:dict) -> bool: if config.get("instance_id") is None: @@ -200,11 +205,25 @@ class AIAgent(BaseAIAgent): self.agent_think_prompt = AgentPrompt() self.agent_think_prompt.load_from_config(config["think_prompt"]) - if config.get("do_prompt") is not None: - self.do_prompt = AgentPrompt() - self.do_prompt.load_from_config(config["do_prompt"]) - self.wake_up() - + def load_todo_config(todo_type:str) -> bool: + todo_config = config.get(todo_type) + if todo_config is not None: + if todo_config.get("do") is not None: + prompt = AgentPrompt() + prompt.load_from_config(todo_config["do"]) + self.todo_prompts[todo_type]["do"] = prompt + if todo_config.get("check") is not None: + prompt = AgentPrompt() + prompt.load_from_config(todo_config["check"]) + self.todo_prompts[todo_type]["check"] = prompt + if todo_config.get("review_prompt") is not None: + prompt = AgentPrompt() + prompt.load_from_config(todo_config["review_prompt"]) + self.todo_prompts[todo_type]["review"] = prompt + + load_todo_config(TodoListType.TO_WORK) + load_todo_config(TodoListType.TO_LEARN) + if config.get("guest_prompt") is not None: self.guest_prompt_str = config["guest_prompt"] @@ -234,6 +253,9 @@ class AIAgent(BaseAIAgent): self.enable_timestamp = bool(config["enable_timestamp"]) if config.get("history_len"): self.history_len = int(config.get("history_len")) + + self.wake_up() + return True def get_id(self) -> str: @@ -372,7 +394,7 @@ class AIAgent(BaseAIAgent): # self._format_msg_by_env_value(prompt) # inner_functions,function_token_len = self._get_inner_functions() - # system_prompt_len = prompt.get_prompt_token_len() + # system_prompt_len = self.token_len(prompt=prompt) # input_len = len(msg.body) # history_prmpt,history_token_len = await self._get_prompt_from_session_for_groupchat(chatsession,system_prompt_len + function_token_len,input_len) @@ -411,10 +433,10 @@ class AIAgent(BaseAIAgent): # resp_msg = msg.create_group_resp_msg(self.agent_id,final_result) # chatsession.append(msg) # chatsession.append(resp_msg) - # return resp_msg - # return None + + def get_workspace_by_msg(self,msg:AgentMsg) -> WorkspaceEnvironment: return self.agent_workspace @@ -528,7 +550,7 @@ class AIAgent(BaseAIAgent): have_known_info = True known_info_str += f"## todo\n{todos_str}\n" inner_functions,function_token_len = BaseAIAgent.get_inner_functions(self.owner_env) - system_prompt_len = prompt.get_prompt_token_len() + system_prompt_len = self.token_len(prompt=prompt) input_len = len(msg.body) if msg.msg_type == AgentMsgType.TYPE_GROUPMSG: history_str,history_token_len = await self._get_prompt_from_session_for_groupchat(chatsession,system_prompt_len + function_token_len,input_len) @@ -600,9 +622,7 @@ class AIAgent(BaseAIAgent): return None - async def _get_history_prompt_for_think(self,chatsession:AIChatSession,summary:str,system_token_len:int,pos:int)->(AgentPrompt,int): - history_len = (self.max_token_size * 0.7) - system_token_len messages = chatsession.read_history(self.history_len,pos,"natural") # read @@ -716,9 +736,7 @@ class AIAgent(BaseAIAgent): worksapce.set_work_summary(self.agent_id,task_result.result_str) - - # 尝试完成自己的TOOD （不依赖任何其他Agnet） - async def do_my_work(self) -> None: + async def _llm_run_todo_list(self, todo_list_type: TodoListType): workspace : WorkspaceEnvironment = self.get_workspace_by_msg(None) logger.info(f"agent {self.agent_id} do my work start!") @@ -726,134 +744,154 @@ class AIAgent(BaseAIAgent): #if await self.need_review_todolist(): # await self._llm_review_todolist(workspace) - todo_list = await workspace.get_todo_list(self.agent_id) + todo_list = workspace.todo_list[todo_list_type] + need_todo = todo_list.get_todo_list(self.agent_id) + check_count = 0 do_count = 0 + review_count = 0 - for todo in todo_list: + for todo in need_todo: if self.agent_energy <= 0: break + + do_prompts = self._can_do_todo(todo_list_type, todo) + if do_prompts: + prompt : AgentPrompt = AgentPrompt() + prompt.append(self.agent_prompt) + prompt.append(workspace.get_role_prompt(self.agent_id)) + prompt.append(do_prompts) + prompt.append(todo.to_prompt()) + + do_result : AgentTodoResult = await self._llm_do_todo(todo, prompt, workspace) + todo.last_do_time = datetime.datetime.now().timestamp() + todo.retry_count += 1 + + match do_result.result_code: + case AgentTodoResult.TODO_RESULT_CODE_LLM_ERROR: + continue + case AgentTodoResult.TODO_RESULT_CODE_OK: + await todo_list.update_todo(todo.todo_id,AgentTodo.TODO_STATE_WAITING_CHECK) + case AgentTodoResult.TODO_RESULT_CODE_EXEC_OP_ERROR: + await todo_list.update_todo(todo.todo_id,AgentTodo.TODO_STATE_EXEC_FAILED) - if await self.need_review_todo(todo,workspace): - review_result = await self._llm_review_todo(todo,workspace) - todo.last_review_time = datetime.datetime.now().timestamp() + await todo_list.append_worklog(todo,do_result) + self.agent_energy -= 2 + do_count += 1 + + # review_result = await self._llm_review_todo(todo,workspace) + # todo.last_review_time = datetime.datetime.now().timestamp() + continue - elif await self.can_check(todo,workspace): - check_result : AgentTodoResult = await self._llm_check_todo(todo,workspace) + check_prompts = self._can_check_todo(todo_list_type, todo) + if check_prompts: + prompt : AgentPrompt = AgentPrompt() + prompt.append(self.agent_prompt) + prompt.append(workspace.get_role_prompt(self.agent_id)) + prompt.append(check_prompts) + + if todo.last_check_result: + prompt.append(AgentPrompt(todo.last_check_result)) + + prompt.append(todo.detail) + prompt.append(todo.result) + + check_result: AgentTodoResult = await self._llm_check_todo(todo, prompt, workspace) todo.last_check_time = datetime.datetime.now().timestamp() match check_result.result_code: case AgentTodoResult.TODO_RESULT_CODE_LLM_ERROR: continue case AgentTodoResult.TODO_RESULT_CODE_OK: - await workspace.update_todo(todo.todo_id,AgentTodo.TODO_STATE_DONE) + await todo_list.update_todo(todo.todo_id,AgentTodo.TODO_STATE_DONE) case AgentTodoResult.TODO_RESULT_CODE_EXEC_OP_ERROR: - await workspace.update_todo(todo.todo_id,AgentTodo.TDDO_STATE_CHECKFAILED) + await todo_list.update_todo(todo.todo_id,AgentTodo.TDDO_STATE_CHECKFAILED) - await workspace.append_worklog(todo,check_result) + await todo_list.append_worklog(todo, check_result) self.agent_energy -= 1 check_count += 1 - elif await self.can_do(todo,workspace): - do_result : AgentTodoResult = await self._llm_do(todo,workspace) - todo.last_do_time = datetime.datetime.now().timestamp() - todo.retry_count += 1 + continue + + review_prompts = self._can_review_todo(todo_list_type, todo) + if review_prompts: + prompt.append(workspace.get_prompt()) + prompt.append(workspace.get_role_prompt(self.agent_id)) + prompt.append(review_prompts) + + todo_tree = todo_list.get_todo_tree("/") + prompt.append(AgentPrompt(todo_tree)) + + do_result : AgentTodoResult = await self._llm_review_todo(todo, prompt, workspace) + todo.last_review_time = datetime.datetime.now().timestamp() match do_result.result_code: case AgentTodoResult.TODO_RESULT_CODE_LLM_ERROR: continue - case AgentTodoResult.TODO_RESULT_CODE_OK: - await workspace.update_todo(todo.todo_id,AgentTodo.TODO_STATE_WAITING_CHECK) case AgentTodoResult.TODO_RESULT_CODE_EXEC_OP_ERROR: - await workspace.update_todo(todo.todo_id,AgentTodo.TODO_STATE_EXEC_FAILED) + continue + case AgentTodoResult.TODO_RESULT_CODE_OK: + await todo_list.update_todo(todo.todo_id,AgentTodo.TODO_STATE_REVIEWED) - await workspace.append_worklog(todo,do_result) - self.agent_energy -= 2 - do_count += 1 + await todo_list.append_worklog(todo,do_result) + self.agent_energy -= 1 + review_count += 1 + continue logger.info(f"agent {self.agent_id} ,check:{check_count} todo,do:{do_count} todo.") + + + def _can_review_todo(self, todo_list_type: TodoListType, todo:AgentTodo) -> AgentPrompt: + do_prompts = self.todo_prompts[todo_list_type].get("review") + if not do_prompts: + return None - def get_review_todo_prompt(self,todo:AgentTodo) -> AgentPrompt: - return self.review_todo_prompt + if todo.can_review() is False: + return None - async def _llm_review_todo(self,todo:AgentTodo,workspace:WorkspaceEnvironment): - prompt = AgentPrompt() + return do_prompts + - prompt.append(workspace.get_prompt()) - prompt.append(workspace.get_role_prompt(self.agent_id)) - prompt.append(self.get_review_todo_prompt(todo)) - - todo_tree = workspace.get_todo_tree("/") - prompt.append(AgentPrompt(todo_tree)) - inner_functions,_ = BaseAIAgent.get_inner_functions(self.owner_env) - - task_result:ComputeTaskResult = await self.do_llm_complection(prompt,inner_functions=inner_functions) - if task_result.result_code != ComputeTaskResultCode.OK: - logger.error(f"_llm_review_todos compute error:{task_result.error_str}") - return - - return - - def get_do_prompt(self,todo:AgentTodo) -> AgentPrompt: - return self.do_prompt - - def get_prompt_from_todo(self,todo:AgentTodo) -> AgentPrompt: - json_str = json.dumps(todo.raw_obj) - return AgentPrompt(json_str) - - async def need_review_todo(self,todo:AgentTodo,workspace:WorkspaceEnvironment) -> bool: - return False - - async def can_check(self,todo:AgentTodo,workspace:WorkspaceEnvironment) -> bool: - if self.get_check_prompt(todo) is None: - return False + def _can_check_todo(self, todo_list_type: TodoListType, todo:AgentTodo) -> AgentPrompt: + do_prompts = self.todo_prompts[todo_list_type].get("check") + if not do_prompts: + return None if todo.can_check() is False: - return False + return None if todo.checker is not None: if todo.checker != self.agent_id: - return False + return None else: if self.can_do_unassigned_task is False: - return False + return None else: todo.checker = self.agent_id - return True + return do_prompts - async def can_do(self,todo:AgentTodo,workspace:WorkspaceEnvironment) -> bool: + async def _can_do_todo(self, todo_list_type: TodoListType, todo:AgentTodo) -> AgentPrompt: + do_prompts = self.todo_prompts[todo_list_type].get("do") + if not do_prompts: + return None + if todo.can_do() is False: - return False + return None if todo.worker is not None: if todo.worker != self.agent_id: - return False + return None else: if self.can_do_unassigned_task is False: - return False + return None else: todo.worker = self.agent_id - return True + return do_prompts - async def _llm_do(self,todo:AgentTodo,workspace:WorkspaceEnvironment) -> AgentTodoResult: + async def _llm_do_todo(self, todo: AgentTodo, prompt: AgentPrompt, workspace: WorkspaceEnvironment) -> AgentTodoResult: result = AgentTodoResult() - prompt : AgentPrompt = AgentPrompt() - #prompt.append(self.agent_prompt) - prompt.append(workspace.get_role_prompt(self.agent_id)) - - do_prompt = workspace.get_do_prompt(todo) - if do_prompt is None: - do_prompt = self.get_do_prompt(todo) - - prompt.append(do_prompt) - - # There are general methods for executing todos, as well as customized ones that are more efficient for specific types of TODOS. - # Based on experience, an Agent can autonomously master/organize execution methods for a greater variety of TODO types. - - #prompt.append(work_log_prompt) - prompt.append(self.get_prompt_from_todo(todo)) - + task_result:ComputeTaskResult = await self.do_llm_complection(prompt) if task_result.error_str is not None: logger.error(f"_llm_do compute error:{task_result.error_str}") @@ -875,7 +913,7 @@ class AIAgent(BaseAIAgent): resp = await AIBus.get_default_bus().post_message(msg) logging.info(f"agent {self.agent_id} send msg to {msg.target} result:{resp}") - op_errors,have_error = await workspace.exec_op_list(llm_result.op_list,self.agent_id) + op_errors, have_error = await workspace.exec_op_list(llm_result.op_list, self.agent_id) if have_error: result.result_code = AgentTodoResult.TODO_RESULT_CODE_EXEC_OP_ERROR #result.error_str = error_str @@ -883,37 +921,30 @@ class AIAgent(BaseAIAgent): return result - async def append_toddo_result(self,todo,worksapce,llm_result,result_str): - pass - - def get_check_prompt(self,todo:AgentTodo) -> AgentPrompt: - return self.check_prompt - - async def _llm_check_todo(self, todo:AgentTodo,workspace:WorkspaceEnvironment) : - if self.get_check_prompt(todo) is None: - return None - - prompt : AgentPrompt = AgentPrompt() - prompt.append(self.agent_prompt) - prompt.append(workspace.get_role_prompt(self.agent_id)) - prompt.append(self.get_check_prompt(todo)) - if todo.last_check_result: - prompt.append(AgentPrompt(todo.last_check_result)) - - prompt.append(todo.detail) - prompt.append(todo.result) - + async def _llm_check_todo(self, todo: AgentTodo, prompt: AgentPrompt, workspace: WorkspaceEnvironment) -> AgentTodoResult: + result = AgentTodoResult() + inner_functions,_ = BaseAIAgent.get_inner_functions(workspace) task_result:ComputeTaskResult = await self.do_llm_complection(prompt,inner_functions=inner_functions,is_json_resp=True) - if task_result.result_code != ComputeTaskResultCode.OK: - logger.error(f"_llm_check_todo compute error:{task_result.error_str}") - return False - - if task_result.result_str == "OK": - return True + if task_result.error_str is not None: + logger.error(f"_llm_do compute error:{task_result.error_str}") + result.result_code = AgentTodoResult.TODO_RESULT_CODE_LLM_ERROR + result.error_str = task_result.error_str + return result + result.result_str = task_result.result_str todo.last_check_result = task_result.result_str - return False + return result + + async def _llm_review_todo(self, todo:AgentTodo, prompt: AgentPrompt, workspace: WorkspaceEnvironment): + inner_functions,_ = BaseAIAgent.get_inner_functions(self.owner_env) + + task_result:ComputeTaskResult = await self.do_llm_complection(prompt,inner_functions=inner_functions) + if task_result.result_code != ComputeTaskResultCode.OK: + logger.error(f"_llm_review_todos compute error:{task_result.error_str}") + return + + return # 尝试自我学习，会主动获取、读取资料并进行整理 # LLM的本质能力是处理海量知识，应该让LLM能基于知识把自己的工作处理的更好 @@ -1121,16 +1152,15 @@ class AIAgent(BaseAIAgent): used_energy = await self.think_chatsession(session_id) self.agent_energy -= used_energy - todo_logs = await self.get_todo_logs() - for todo_log in todo_logs: - if self.agent_energy <= 0: - break - used_energy = await self.think_todo_log(todo_log) - self.agent_energy -= used_energy + # todo_logs = await self.get_todo_logs() + # for todo_log in todo_logs: + # if self.agent_energy <= 0: + # break + # used_energy = await self.think_todo_log(todo_log) + # self.agent_energy -= used_energy return - async def think_todo_log(self,todo_log:AgentWorkLog): pass @@ -1146,7 +1176,7 @@ class AIAgent(BaseAIAgent): prompt:AgentPrompt = AgentPrompt() #prompt.append(self._get_agent_prompt()) prompt.append(await self._get_agent_think_prompt()) - system_prompt_len = prompt.get_prompt_token_len() + system_prompt_len = self.token_len(prompt=prompt) #think env? history_prompt,next_pos = await self._get_history_prompt_for_think(chatsession,summary,system_prompt_len,cur_pos) prompt.append(history_prompt) @@ -1220,11 +1250,7 @@ class AIAgent(BaseAIAgent): def need_self_think(self) -> bool: return False - def need_self_learn(self) -> bool: - if self.learn_prompt is not None: - return True - return False - + def wake_up(self) -> None: if self.agent_task is None: self.agent_task = asyncio.create_task(self._on_timer()) @@ -1248,26 +1274,20 @@ class AIAgent(BaseAIAgent): continue # complete & check todo - if self.need_work(): - await self.do_my_work() - - # review other's todo - # self.review_other_works() + await self._llm_run_todo_list(TodoListType.TO_WORK) + await self._llm_run_todo_list(TodoListType.TO_LEARN) + if self.need_self_think(): await self.do_self_think() - - if self.need_self_learn(): - await self.do_self_learn() - + + # review other's todo + # self.review_other_works() except Exception as e: tb_str = traceback.format_exc() logger.error(f"agent {self.agent_id} on timer error:{e},{tb_str}") continue - def token_len(self,text:str) -> int: - return ComputeKernel.llm_num_tokens_from_text(text,self.get_llm_model_name()) - diff --git a/src/aios/agent/agent_base.py b/src/aios/agent/agent_base.py index 4094c36..b706618 100644 --- a/src/aios/agent/agent_base.py +++ b/src/aios/agent/agent_base.py @@ -56,16 +56,6 @@ class AgentPrompt: self.messages.extend(prompt.messages) - def get_prompt_token_len(self): - result = 0 - - if self.system_message: - result += len(self.system_message.get("content")) - for msg in self.messages: - result += len(msg.get("content")) - - return result - def load_from_config(self,config:list) -> bool: if isinstance(config,list) is not True: logger.error("prompt is not list!") @@ -245,8 +235,9 @@ class AgentTodo: TODO_STATE_EXEC_FAILED = "exec_failed" TDDO_STATE_CHECKFAILED = "check_failed" - TODO_STATE_CASNCEL = "cancel" + TODO_STATE_CANCEL = "cancel" TODO_STATE_DONE = "done" + TODO_STATE_REVIEWED = "reviewed" TODO_STATE_EXPIRED = "expired" def __init__(self): @@ -341,6 +332,23 @@ class AgentTodo: result["retry_count"] = self.retry_count return result + + def to_prompt(self) -> AgentPrompt: + json_str = json.dumps(self.raw_obj) + return AgentPrompt(json_str) + + def can_review(self) -> bool: + if self.state != AgentTodo.TODO_STATE_DONE: + return False + + now = datetime.now().timestamp() + if self.last_review_time: + time_diff = now - self.last_review_time + if time_diff < 60*15: + logger.info(f"todo {self.title} is already reviewed, ignore") + return False + + return True def can_check(self)->bool: if self.state != AgentTodo.TODO_STATE_WAITING_CHECK: @@ -410,9 +418,18 @@ class BaseAIAgent(abc.ABC): def get_max_token_size(self) -> int: pass - @abstractmethod - async def _process_msg(self,msg:AgentMsg,workspace = None) -> AgentMsg: - pass + def token_len(self, text:str=None, prompt:AgentPrompt=None) -> int: + from .compute_kernel import ComputeKernel + if text: + return ComputeKernel.llm_num_tokens_from_text(text,self.get_llm_model_name()) + elif prompt: + result = 0 + if prompt.system_message: + result += ComputeKernel.llm_num_tokens_from_text(prompt.system_message.get("content"),self.get_llm_model_name()) + for msg in prompt.messages: + result += ComputeKernel.llm_num_tokens_from_text(msg.get("content"),self.get_llm_model_name()) + else: + return 0 @classmethod def get_inner_functions(cls, env:Environment) -> (dict,int): diff --git a/src/aios/agent/ai_function.py b/src/aios/agent/ai_function.py index bca41d2..92eb448 100644 --- a/src/aios/agent/ai_function.py +++ b/src/aios/agent/ai_function.py @@ -9,9 +9,6 @@ class ParameterDefine: class AIFunction: - def __init__(self) -> None: - self.description : str = None - @abstractmethod def get_name(self) -> str: """ @@ -24,7 +21,7 @@ class AIFunction: """ return a detailed description of what the function does """ - return self.description + pass @abstractmethod def get_parameters(self) -> Dict: @@ -112,6 +109,9 @@ class SimpleAIFunction(AIFunction): def get_name(self) -> str: return self.func_id + def get_description(self) -> str: + return self.description + def get_parameters(self) -> Dict: if self.parameters is not None: result = {} @@ -142,3 +142,62 @@ class SimpleAIFunction(AIFunction): def is_ready_only(self) -> bool: return False +class AIOperation: + @abstractmethod + def get_name(self) -> str: + """ + return the name of the operation (should be snake case) + """ + pass + + @abstractmethod + def get_description(self) -> str: + """ + return a detailed description of what the operation does + """ + pass + + + @abstractmethod + async def execute(self, params: dict) -> str: + """ + Execute the function and return a JSON serializable dict. + The parameters are passed in the form of kwargs + """ + pass + +class SimpleAIOperation(AIOperation): + def __init__(self,op:str,description:str,func_handler:Coroutine) -> None: + self.op = op + self.description = description + self.func_handler = func_handler + + def get_name(self) -> str: + return self.op + + def get_description(self) -> str: + return self.description + + async def execute(self, params: Dict) -> str: + if self.func_handler is None: + return "error: function not implemented" + + return await self.func_handler(**params) + + +class AIFunctionOperation(AIOperation): + def __init__(self, func: AIFunction) -> None: + self.func = func + super().__init__() + + @abstractmethod + def get_name(self) -> str: + return self.func.get_name() + + @abstractmethod + def get_description(self) -> str: + return self.func.get_description() + + @abstractmethod + async def execute(self, params: dict) -> str: + self.func.execute(**params) \ No newline at end of file diff --git a/src/aios/environment/environment.py b/src/aios/environment/environment.py index eafe74f..d4b426b 100644 --- a/src/aios/environment/environment.py +++ b/src/aios/environment/environment.py @@ -4,145 +4,127 @@ from abc import ABC, abstractmethod from typing import Any, Callable, Optional,Dict,Awaitable,List import logging +from ..agent.ai_function import AIFunction, AIOperation -from ..agent.ai_function import AIFunction logger = logging.getLogger(__name__) -class EnvironmentEvent(ABC): + +class BaseEnvironment: @abstractmethod - def display(self) -> str: + def get_id(self) -> str: + pass + + # @abstractmethod + # #TODO: how to use env? different env has different prompt + # def get_env_prompt(self) -> str: + # pass + + @abstractmethod + def get_ai_function(self,func_name:str) -> AIFunction: pass -EnvironmentEventHandler = Callable[[str,EnvironmentEvent],Awaitable[Any]] + @abstractmethod + def get_all_ai_functions(self) -> List[AIFunction]: + pass -class Environment: - _all_env = {} - @classmethod - def get_env_by_id(cls,env_id:str): - return cls._all_env.get(env_id) - @classmethod - def set_env_by_id(cls,id,env): - assert id == env.get_id() - cls._all_env[env.get_id()] = env + @abstractmethod + def get_ai_operation(self,op_name:str) -> AIOperation: + pass - def __init__(self,env_id:str) -> None: + @abstractmethod + def get_all_ai_operations(self) -> List[AIOperation]: + pass + + + + # _all_env = {} + # @classmethod + # def get_env_by_id(cls,env_id:str): + # return cls._all_env.get(env_id) + + # @classmethod + # def set_env_by_id(cls,id,env): + # assert id == env.get_id() + # cls._all_env[env.get_id()] = env + +class SimpleEnvironment(BaseEnvironment): + def __init__(self, env_id: str) -> None: self.env_id = env_id - self.values:Dict[str,str] = {} - self.get_handlers:Dict[str,Callable] = {} - self.owner_env:Dict[str,Environment] = {} - # self.valid_keys:Dict[str,bool] = None - self.event_handlers:Dict[str,List[EnvironmentEventHandler]]= {} - - self.functions : Dict[str,AIFunction] = {} + self.functions: Dict[str,AIFunction] = {} + self.operations: Dict[str,AIOperation] = {} def get_id(self) -> str: return self.env_id - - def add_owner_env(self,env) -> None: - self.owner_env[env.get_id()] = env - - #@abstractmethod - #TODO: how to use env? different env has different prompt - 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: func = self.functions.get(func_name) if func is not None: return func - - for owner_env in self.owner_env.values(): - func = owner_env.get_ai_function(func_name) - if func is not None: - return func - return None - #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]: func_list = [] func_list.extend(self.functions.values()) - for owner_env in self.owner_env.values(): - func_list.extend(owner_env.get_all_ai_functions()) return func_list - - @abstractmethod - def _do_get_value(self,key:str) -> Optional[str]: - pass - - def register_get_handler(self,key:str,handler:Callable) -> None: - h = self.get_handlers.get(key) - if h is not None: - logger.warn(f"register get_handler {key} in env {self.env_id}:handler already exist") - - self.get_handlers[key] = handler - - - def attach_event_handler(self,event_id:str,handler:Callable) -> None: - handler_list = self.event_handlers.get(event_id) - if handler_list is None: - handler_list = [] - self.event_handlers[event_id] = handler_list - - handler_list.append(handler) - - def remove_event_handler(self,event_id:str,handler:Callable) -> None: - handler_list = self.event_handlers.get(event_id) - if handler is not None: - handler_list.remove(handler) - return - - logger.warn(f"remove event_handler {event_id} in env {self.env_id}:handler not found") - - async def fire_event(self,event_id:str,event:EnvironmentEvent) -> None: - handler_list = self.event_handlers.get(event_id) - if handler_list is not None: - for handler in handler_list: - await handler(self.env_id,event) - else: - logger.debug(f"fire event {event_id} in env {self.env_id}:handler not found") - return - - def __getitem__(self, key): - return self.get_value(key) - - def get_value(self,key:str) -> Optional[str]: - handler = self.get_handlers.get(key) - if handler is not None: - return handler() - - s = self.values.get(key) - if isinstance(s,str): - return s - else: - logger.warn(f"get value {key} in env {self.env_id} failed!,type is not str") - - s = self._do_get_value(key) - if s is not None: - return s - if self.owner_env is not None: - for env in self.owner_env.values(): - s = env.get_value(key) - if s is not None: - return s - - logger.warn(f"get value {key} in env {self.env_id} failed!,not found") + + def add_ai_operation(self,op:AIOperation) -> None: + self.operations[op.get_name()] = op + + def get_ai_operation(self,op_name:str) -> AIOperation: + op = self.operations.get(op_name) + if op is not None: + return op return None - def set_value(self, key: str, str_value: str,is_storage:bool = True): - logger.info(f"set value {key} in env {self.env_id} to {str_value}") - self.values[key] = str_value + def get_all_ai_operations(self) -> List[AIOperation]: + op_list = [] + op_list.extend(self.operations.values()) + return op_list + + +class CompositeEnvironment(BaseEnvironment): + def __init__(self, env_id: str) -> None: + self.env_id = env_id + self.envs:Dict[str,BaseEnvironment] = {} + self.functions: Dict[str,AIFunction] = {} + self.operations: Dict[str,AIOperation] = {} + + def get_id(self) -> str: + return self.env_id + + def add_env(self, env: BaseEnvironment) -> None: + self.envs[env.get_id()] = env + functions = env.get_all_ai_functions() + for func in functions: + self.functions[func.get_name()] = func + operations = env.get_all_ai_operations() + for op in operations: + self.operations[op.get_name()] = op + + def get_ai_function(self,func_name:str) -> AIFunction: + func = self.functions.get(func_name) + if func is not None: + return func + return None + + def get_all_ai_functions(self) -> List[AIFunction]: + func_list = [] + func_list.extend(self.functions.values()) + return func_list + + def get_ai_operation(self,op_name:str) -> AIOperation: + op = self.operations.get(op_name) + if op is not None: + return op + return None + + def get_all_ai_operations(self) -> List[AIOperation]: + op_list = [] + op_list.extend(self.operations.values()) + return op_list \ No newline at end of file diff --git a/src/aios/environment/workspace_env.py b/src/aios/environment/workspace_env.py index d54c7b8..115e347 100644 --- a/src/aios/environment/workspace_env.py +++ b/src/aios/environment/workspace_env.py @@ -1,109 +1,252 @@ # this env is designed for workflow owner filesystem, support file/directory operations -import hashlib import json -import subprocess import logging -import tempfile -import threading -import traceback -import time -import ast -import sys import os -import re -import asyncio import aiofiles from typing import Any,List -import os import chardet -from markdown import Markdown -import PyPDF2 - -from ..proto.agent_msg import * -from ..agent.agent_base import AgentTodo,AgentPrompt,AgentTodoResult -from ..agent.ai_function import AIFunction,SimpleAIFunction +from ..agent.agent_base import AgentMsg,AgentTodo,AgentPrompt,AgentTodoResult +from ..agent.ai_function import AIFunction,SimpleAIFunction, SimpleAIOperation from ..storage.storage import AIStorage,ResourceLocation -from .simple_kb_db import SimpleKnowledgeDB -from .environment import Environment,EnvironmentEvent +from .environment import SimpleEnvironment, CompositeEnvironment + + logger = logging.getLogger(__name__) -class WorkspaceEnvironment(Environment): - def __init__(self, env_id: str) -> None: - super().__init__(env_id) - myai_path = AIStorage.get_instance().get_myai_dir() - self.root_path = f"{myai_path}/workspace/{env_id}" +class TodoListType: + TO_WORK = "work" + TO_LEARN = "learn" + +class TodoListEnvironment(SimpleEnvironment): + def __init__(self, root_path, list_type) -> None: + super.__init__(list_type) + self.root_path = os.path.join(root_path, list_type) if not os.path.exists(self.root_path): - os.makedirs(self.root_path+"/todos") - + os.makedirs(self.root_path) self.known_todo = {} - self.kb_db = SimpleKnowledgeDB(f"{self.root_path}/kb.db") - self.doc_dirs = {} - self._scan_thread = None - self._scan_dirthread = None + + async def create_todo(params): + todoObj = AgentTodo.from_dict(params["todo"]) + parent_id = params.get("parent") + return await self.create_todo(parent_id,todoObj) + self.add_ai_operation(SimpleAIOperation( + op="create_todo", + description="create todo", + func_handler=create_todo, + )) + + + async def update_todo(params): + todo_id = params["id"] + new_stat = params["state"] + return await self.update_todo(todo_id,new_stat) + self.add_ai_operation(SimpleAIOperation( + op="update_todo", + description="update todo", + func_handler=update_todo, + )) + + + # Task/todo system , create,update,delete,query + async def get_todo_tree(self,path:str = None,deep:int = 4): + if path: + directory_path = os.path.join(self.root_path, path) + else: + directory_path = self.root_path + + str_result:str = "/todos\n" + todo_count:int = 0 - def set_root_path(self,path:str): - self.root_path = path - - def get_prompt(self) -> AgentMsg: - return None - - def get_role_prompt(self,role_id:str) -> AgentPrompt: - return None - - def get_knowledge_base(self,root_dir=None,indent=0) -> str: - pass - - - def get_do_prompt(self,todo:AgentTodo=None)->AgentPrompt: - return None - - # result mean: list[op_error_str],have_error - async def exec_op_list(self,oplist:List,agent_id:str)->tuple[List[str],bool]: - result_str = "op list is none" - if oplist is None: - return None,False - - result_str = [] - have_error = False - for op in oplist: - if op["op"] == "create": - await self.create(op["path"],op["content"]) - elif op["op"] == "write_file": - is_append = op.get("is_append") - if is_append is None: - is_append = False - error_str = await self.write(op["path"],op["content"],is_append) - elif op["op"] == "delete": - error_str = await self.delete(op["path"]) - elif op["op"] == "rename": - error_str = await self.rename(op["path"],op["new_name"]) - elif op["op"] == "mkdir": - error_str = await self.mkdir(op["path"]) - elif op["op"] == "create_todo": - todoObj = AgentTodo.from_dict(op["todo"]) - todoObj.worker = agent_id - todoObj.createor = agent_id - parent_id = op.get("parent") - error_str = await self.create_todo(parent_id,todoObj) - elif op["op"] == "update_todo": - todo_id = op["id"] - new_stat = op["state"] - error_str = await self.update_todo(todo_id,new_stat) - else: - logger.error(f"execute op list failed: unknown op:{op['op']}") - error_str = f"execute op list failed: unknown op:{op['op']}" + async def scan_dir(directory_path:str,deep:int): + nonlocal str_result + nonlocal todo_count + if deep <= 0: + return - if error_str: - have_error = True - result_str.append(error_str) - else: - result_str.append(f"execute success!") - + if os.path.exists(directory_path) is False: + return + + for entry in os.scandir(directory_path): + is_dir = entry.is_dir() + if not is_dir: + continue + + if entry.name.startswith("."): + continue + + todo_count = todo_count + 1 + str_result = str_result + f"{' '*(4-deep)}{entry.name}\n" + await scan_dir(entry.path,deep-1) + + await scan_dir(directory_path,deep) + return str_result,todo_count + + async def get_todo_list(self,agent_id:str,path:str = None)->List[AgentTodo]: + logger.info("get_todo_list:%s,%s",agent_id,path) + if path: + directory_path = os.path.join(self.root_path, path) + else: + directory_path = self.root_path + + result_list:List[AgentTodo] = [] + + async def scan_dir(directory_path:str,deep:int,parent:AgentTodo=None): + nonlocal result_list + if os.path.exists(directory_path) is False: + return + + for entry in os.scandir(directory_path): + is_dir = entry.is_dir() + if not is_dir: + continue + + if entry.name.startswith("."): + continue + + todo = await self.get_todo_by_fullpath(entry.path) + if todo: + if todo.worker: + if todo.worker != agent_id: + continue + + if parent: + parent.sub_todos[todo.todo_id] = todo + + result_list.append(todo) + todo.rank = int(todo.create_time)>>deep + await scan_dir(entry.path,deep + 1,todo) + + return + + await scan_dir(directory_path,0) + #sort by rank + result_list.sort(key=lambda x:(x.rank,x.title)) + logger.info("get_todo_list return,todolist.length() is %d",len(result_list)) + return result_list + + async def get_todo_by_fullpath(self,path:str) -> AgentTodo: + logger.info("get_todo_by_fullpath:%s",path) + + detail_path = path + "/detail" + try: + async with aiofiles.open(detail_path, mode='r', encoding="utf-8") as f: + content = await f.read(4096) + logger.debug("get_todo_by_fullpath:%s,content:%s",path,content) + todo_dict = json.loads(content) + result_todo = AgentTodo.from_dict(todo_dict) + if result_todo: + relative_path = os.path.relpath(path, self.root_path) + if not relative_path.startswith('/'): + relative_path = '/' + relative_path + result_todo.todo_path = relative_path + self.known_todo[result_todo.todo_id] = result_todo + else: + logger.error("get_todo_by_path:%s,parse failed!",path) + + return result_todo + except Exception as e: + logger.error("get_todo_by_path:%s,failed:%s",path,e) + return None - return result_str,have_error + async def get_todo(self,id:str) -> AgentTodo: + return self.known_todo.get(id) + + async def create_todo(self,parent_id:str,todo:AgentTodo) -> str: + try: + if parent_id: + if parent_id not in self.known_todo: + logger.error("create_todo failed: parent_id not found!") + return False + + parent_path = self.known_todo.get(parent_id).todo_path + todo_path = f"{parent_path}/{todo.title}" + else: + todo_path = todo.title + + dir_path = f"{self.root_path}/{todo_path}" + + os.makedirs(dir_path) + detail_path = f"{dir_path}/detail" + if todo.todo_path is None: + todo.todo_path = todo_path + logger.info("create_todo %s",detail_path) + async with aiofiles.open(detail_path, mode='w', encoding="utf-8") as f: + await f.write(json.dumps(todo.to_dict())) + self.known_todo[todo.todo_id] = todo + except Exception as e: + logger.error("create_todo failed:%s",e) + return str(e) + + return None + + async def update_todo(self,todo_id:str,new_stat:str)->str: + try: + todo : AgentTodo = self.known_todo.get(todo_id) + if todo: + todo.state = new_stat + detail_path = f"{self.root_path}/{todo.todo_path}/detail" + async with aiofiles.open(detail_path, mode='w', encoding="utf-8") as f: + await f.write(json.dumps(todo.to_dict())) + return None + else: + return "todo not found." + except Exception as e: + return str(e) + + async def append_worklog(self, todo:AgentTodo, result:AgentTodoResult): + worklog = f"{self.root_path}/{todo.todo_path}/.worklog" + + async with aiofiles.open(worklog, mode='w+', encoding="utf-8") as f: + content = await f.read() + if len(content) > 0: + json_obj = json.loads(content) + else: + json_obj = {} + logs = json_obj.get("logs") + if logs is None: + logs = [] + logs.append(result.to_dict()) + json_obj["logs"] = logs + await f.write(json.dumps(json_obj)) + +class FilesystemEnvironment(SimpleEnvironment): + def __init__(self, root_path: str, env_id: str) -> None: + super().__init__(env_id) + self.root_path = root_path + + + # if op["op"] == "create": + # await self.create(op["path"],op["content"]) + + async def write(op): + is_append = op.get("is_append") + if is_append is None: + is_append = False + return await self.write(op["path"],op["content"],is_append) + self.add_ai_operation(SimpleAIOperation( + op="write", + description="write file", + func_handler=write, + )) + + async def delete(op): + return await self.delete(op["path"]) + self.add_ai_operation(SimpleAIOperation( + op="delete", + description="delete path", + func_handler=delete, + )) + + async def rename(op): + return await self.move(op["path"],op["new_name"]) + self.add_ai_operation(SimpleAIOperation( + op="rename", + description="rename path", + func_handler=rename, + )) # file system operation: list,read,write,delete,move,stat # inner_function @@ -201,560 +344,8 @@ class WorkspaceEnvironment(Environment): return str(e) return None - - # TODO use diff to update large file content - async def update_by_diff(self,path:str,diff): - pass - - # doc system （read_only,agent cann't modify doc） - - # inner_function - async def list_db(self) -> str: - pass - # inner_function - async def get_db_desc(self,db_name:str) -> str: - pass - # inner_function - async def query(self,db_name:str,sql:str) -> str: - pass - - # search (web) - # inner_function - async def google_search(self,keyword:str,opt=None) -> str: - pass - - # inner_function - async def local_search(self,keyword:str,root_path=None ,opt=None) -> str: - pass - - # inner_function, might be return a image is better - async def web_get(self,url:str) -> str: - pass - - # inner_function - async def blockchain_get(self,chainid:str,query:dict) -> str: - pass - - # code interpreter - # inner_function or operation - async def eval_code(self,pycode:str) -> str: - pass - - # operation or inner_function - async def improve_code(self,path:str): - pass - - # operation or inner_function - async def run(self,file_path:str)->str: - pass - - # operation or inner_function - async def pub_service(self,project_path:str): - pass - - # operation or inner_function - async def exec_tx(self,chain_id:str,tx:dict) -> str: - pass - - # social ability - # operation or inner_function - async def post_message(self,target:str,msg:AgentMsg,wait_time) -> AgentMsg: - pass - - # operation or inner_function - async def add_contact(self,name:str,contact_info) -> str: - pass - - # inner_function , include contact realtime info - async def get_contact(self,name_list:List[str],opt:dict) -> List: - pass - - - # Task/todo system , create,update,delete,query - async def get_todo_tree(self,path:str = None,deep:int = 4): - if path: - directory_path = self.root_path + "/todos/" + path - else: - directory_path = self.root_path + "/todos" - - - str_result:str = "/todos\n" - todo_count:int = 0 - - async def scan_dir(directory_path:str,deep:int): - nonlocal str_result - nonlocal todo_count - if deep <= 0: - return - - if os.path.exists(directory_path) is False: - return - - for entry in os.scandir(directory_path): - is_dir = entry.is_dir() - if not is_dir: - continue - - if entry.name.startswith("."): - continue - - todo_count = todo_count + 1 - str_result = str_result + f"{' '*(4-deep)}{entry.name}\n" - await scan_dir(entry.path,deep-1) - - await scan_dir(directory_path,deep) - return str_result,todo_count - - async def get_todo_list(self,agent_id:str,path:str = None)->List[AgentTodo]: - logger.info("get_todo_list:%s,%s",agent_id,path) - if path: - directory_path = self.root_path + "/todos/" + path - else: - directory_path = self.root_path + "/todos" - - result_list:List[AgentTodo] = [] - - async def scan_dir(directory_path:str,deep:int,parent:AgentTodo=None): - nonlocal result_list - if os.path.exists(directory_path) is False: - return - - for entry in os.scandir(directory_path): - is_dir = entry.is_dir() - if not is_dir: - continue - - if entry.name.startswith("."): - continue - - todo = await self.get_todo_by_fullpath(entry.path) - if todo: - if todo.worker: - if todo.worker != agent_id: - continue - - if parent: - parent.sub_todos[todo.todo_id] = todo - - result_list.append(todo) - todo.rank = int(todo.create_time)>>deep - await scan_dir(entry.path,deep + 1,todo) - - return - - await scan_dir(directory_path,0) - #sort by rank - result_list.sort(key=lambda x:(x.rank,x.title)) - logger.info("get_todo_list return,todolist.length() is %d",len(result_list)) - return result_list - - async def get_todo_by_fullpath(self,path:str) -> AgentTodo: - logger.info("get_todo_by_fullpath:%s",path) - - detail_path = path + "/detail" - try: - async with aiofiles.open(detail_path, mode='r', encoding="utf-8") as f: - content = await f.read(4096) - logger.debug("get_todo_by_fullpath:%s,content:%s",path,content) - todo_dict = json.loads(content) - result_todo = AgentTodo.from_dict(todo_dict) - if result_todo: - relative_path = os.path.relpath(path, self.root_path + "/todos/") - if not relative_path.startswith('/'): - relative_path = '/' + relative_path - result_todo.todo_path = relative_path - self.known_todo[result_todo.todo_id] = result_todo - else: - logger.error("get_todo_by_path:%s,parse failed!",path) - - return result_todo - except Exception as e: - logger.error("get_todo_by_path:%s,failed:%s",path,e) - return None - - async def get_todo(self,id:str) -> AgentTodo: - return self.known_todo.get(id) - - async def create_todo(self,parent_id:str,todo:AgentTodo) -> str: - try: - if parent_id: - if parent_id not in self.known_todo: - logger.error("create_todo failed: parent_id not found!") - return False - - parent_path = self.known_todo.get(parent_id).todo_path - todo_path = f"{parent_path}/{todo.title}" - else: - todo_path = todo.title - - dir_path = f"{self.root_path}/todos/{todo_path}" - - os.makedirs(dir_path) - detail_path = f"{dir_path}/detail" - if todo.todo_path is None: - todo.todo_path = todo_path - logger.info("create_todo %s",detail_path) - async with aiofiles.open(detail_path, mode='w', encoding="utf-8") as f: - await f.write(json.dumps(todo.to_dict())) - self.known_todo[todo.todo_id] = todo - except Exception as e: - logger.error("create_todo failed:%s",e) - return str(e) - - return None - - async def update_todo(self,todo_id:str,new_stat:str)->str: - try: - todo : AgentTodo = self.known_todo.get(todo_id) - if todo: - todo.state = new_stat - detail_path = f"{self.root_path}/todos/{todo.todo_path}/detail" - async with aiofiles.open(detail_path, mode='w', encoding="utf-8") as f: - await f.write(json.dumps(todo.to_dict())) - return None - else: - return "todo not found." - except Exception as e: - return str(e) - - async def append_worklog(self,todo:AgentTodo,result:AgentTodoResult): - worklog = f"{self.root_path}/todos/{todo.todo_path}/.worklog" - - async with aiofiles.open(worklog, mode='w+', encoding="utf-8") as f: - content = await f.read() - if len(content) > 0: - json_obj = json.loads(content) - else: - json_obj = {} - logs = json_obj.get("logs") - if logs is None: - logs = [] - logs.append(result.to_dict()) - json_obj["logs"] = logs - await f.write(json.dumps(json_obj)) - - async def set_wakeup_timer(self,todo_id:str,timestamp:int) -> str: - pass - - # knowledge base system - def get_knowledge_base_ai_functions(self): - all_inner_function = [] - - all_inner_function.append(SimpleAIFunction("get_knowledge_catalog","get knowledge catalog in tree format", - self.get_knowledege_catalog, - {"path":f"catalog path,none is /","depth":"max depth of catalog tree,default is 4"})) - all_inner_function.append(SimpleAIFunction("get_knowledge","get knowledge metadata", - self.get_knowledge, - {"path":f"knowledge path"})) - all_inner_function.append(SimpleAIFunction("load_knowledge_content","load knowledge content", - self.load_knowledge_content, - {"path":f"knowledge path","pos":"start position of content","length":"length of content"})) - result_func = [] - result_len = 0 - for inner_func in all_inner_function: - func_name = inner_func.get_name() - - this_func = {} - this_func["name"] = func_name - this_func["description"] = inner_func.get_description() - this_func["parameters"] = inner_func.get_parameters() - result_len += len(json.dumps(this_func)) / 4 - result_func.append(this_func) - - return result_func,result_len - - async def get_knowledege_catalog(self,path:str=None,only_dir =True,max_depth:int=5)->str: - if path: - full_path = f"{self.root_path}/knowledge/{path}" - else: - full_path = f"{self.root_path}/knowledge" - - catlogs,file_count = await self.get_directory_structure(full_path,max_depth,only_dir) - return catlogs - - async def get_directory_structure(self,root_dir, max_depth:int=4, only_dir=True, indent=1): - file_count = 0 - structure_str = '' - if os.path.isdir(root_dir): - sub_files = [] - with os.scandir(root_dir) as it: - for entry in it: - if entry.is_dir(): - sub_structure, sub_count = await self.get_directory_structure(entry.path, max_depth, only_dir, indent + 1) - if sub_structure: - structure_str += sub_structure - file_count += sub_count - else: - file_count += 1 - sub_files.append(entry.name) - - if only_dir is False: - for file_name in sub_files: - structure_str = structure_str + ' ' * (indent+1) + file_name + '\n' - - dir_name = os.path.basename(root_dir) - dir_info = f"{dir_name} " - - - structure_str = ' ' * indent + dir_info + '\n' + structure_str - - if indent - 1 >= max_depth: - return None, file_count - else: - return structure_str, file_count - - # inner_function - async def get_knowledge(self,path:str) -> str: - full_path = f"{self.root_path}/knowledge/{path}" - if os.islink(full_path): - org_path = os.readlink(full_path) - hash = self.kb_db.get_hash_by_doc_path(org_path) - if hash: - return self.kb_db.get_knowledge(org_path) - - return "not found" - - async def load_knowledge_content(self,path:str,pos:int=0,length:int=None) -> str: - if path.endswith("pdf"): - logger.info("load_knowledge_content:pdf") - dir_path = os.path.dirname(path) - base_name = os.path.basename(path) - text_content_path = f"{dir_path}/.{base_name}.txt" - if os.path.exists(text_content_path) is False: - return None - async with aiofiles.open(path, mode='r', encoding=cur_encode) as f: - await f.seek(pos) - content = await f.read(length) - return content - else: - async with aiofiles.open(path,'rb') as f: - cur_encode = chardet.detect(await f.read())['encoding'] - - async with aiofiles.open(path, mode='r', encoding=cur_encode) as f: - await f.seek(pos) - content = await f.read(length) - return content - - return "load content failed." - - def _add_document_dir(self,path:str): - self.doc_dirs[path] = 0 - - def _start_scan_document(self): - if self._scan_thread is None: - self._scan_thread = threading.Thread(target=self._scan_document) - self._scan_thread.start() - if self._scan_dirthread is None: - self._scan_dirthread = threading.Thread(target=self._scan_dir) - self._scan_dirthread.start() - - def _parse_pdf_bookmarks(self,bookmarks, parent:list): - - for item in bookmarks: - if isinstance(item,list): - self._parse_pdf_bookmarks(item,parent) - else: - if item.title: - new_item = {} - new_item["page"] = item.page.idnum - new_item["title"] = item.title - my_childs = [] - if item.childs: - if len(item.childs) > 0: - self._parse_pdf_bookmarks(item.childs, my_childs) - new_item["childs"] = my_childs - parent.append(new_item) - else: - logger.warning("parse pdf bookmarks failed: item.title is None!") - - return - - def _parse_pdf(self,doc_path:str): - metadata = {} - with open(doc_path, 'rb') as file: - reader = PyPDF2.PdfReader(file) - try: - doc_info = reader.metadata - if doc_info: - if doc_info.title: - metadata["title"] = doc_info.title - if doc_info.author: - metadata["authors"] = doc_info.author - except Exception as e: - logger.warn("parse pdf metadata failed:%s",e) - - dir_path = os.path.dirname(doc_path) - base_name = os.path.basename(doc_path) - text_content_path = f"{dir_path}/.{base_name}.txt" - full_text = "" - - for page in reader.pages: - text = page.extract_text() - full_text += text - with open(text_content_path, 'w', encoding='utf-8') as f: - f.write(full_text) - - try: - bookmarks = reader.outline - if bookmarks: - catalogs = [] - self._parse_pdf_bookmarks(bookmarks,catalogs) - metadata["catalogs"] = json.dumps(catalogs) - except Exception as e: - logger.warn("parse pdf bookmarks failed:%s",e) - - return metadata - - def _parse_txt(self,doc_path:str): - return {} - - def _parse_md(self,doc_path:str): - metadata = {} - cur_encode = "utf-8" - with open(doc_path,'rb') as f: - cur_encode = chardet.detect(f.read(1024))['encoding'] - - with open(doc_path, mode='r', encoding=cur_encode) as f: - content = f.read() - match = re.search(r'^# (.*)', content, re.MULTILINE) - if match: - metadata['title'] = match.group(1).strip() - md = Markdown(extensions=['toc']) - html_str = md.convert(content) - toc = md.toc - if toc: - metadata['catalogs'] = toc - - return metadata - - def _parse_document(self,doc_path:str): - hash_result = None - title = os.path.basename(doc_path) - meta_data = {} - - with open(doc_path, "rb") as f: - hash_md5 = hashlib.md5() - for chunk in iter(lambda: f.read(1024*1024), b""): - hash_md5.update(chunk) - hash_result = hash_md5.hexdigest() - try: - if doc_path.endswith(".md"): - meta_data = self._parse_md(doc_path) - elif doc_path.endswith(".pdf"): - meta_data = self._parse_pdf(doc_path) - except Exception as e: - logger.error("parse document %s failed:%s",doc_path,e) - traceback.print_exc() - - if meta_data.get("title"): - title = meta_data["title"] - logger.info("parse document %s!",doc_path) - return hash_result,title,meta_data - - - def _support_file(self,file_name:str) -> bool: - if file_name.startswith("."): - return False - - if file_name.endswith(".pdf"): - return True - if file_name.endswith(".md"): - return True - if file_name.endswith(".txt"): - return True - return False - - def _scan_dir(self): - while True: - time.sleep(10) - for directory in self.doc_dirs.keys(): - now = time.time() - if now - self.doc_dirs[directory] > 60*15: - self.doc_dirs[directory] = time.time() - else: - continue - - for root, dirs, files in os.walk(directory): - for file in files: - if self._support_file(file): - full_path = os.path.join(root, file) - full_path = os.path.normpath(full_path) - if self.kb_db.is_doc_exist(full_path): - continue - - file_stat = os.stat(full_path) - if file_stat.st_size < 1: - continue - - if file_stat.st_size < 1024*1024*8: - #parse and insert - hash,title,meta_data = self._parse_document(full_path) - self.kb_db.add_doc(full_path,file_stat.st_size,file_stat.st_mtime,hash) - self.kb_db.add_knowledge(hash,title,meta_data) - - else: - self.kb_db.add_doc(full_path,file_stat.st_size,file_stat.st_mtime) - - def _scan_document(self): - while True: - time.sleep(10) - parse_queue = self.kb_db.get_docs_without_hash() - for doc_path in parse_queue: - hash,title,meta_data = self._parse_document(doc_path) - self.kb_db.set_doc_hash(doc_path,hash) - self.kb_db.add_knowledge(hash,title,meta_data) - - - - -# merge to standard workspace env, **ABANDON this!** -class KnowledgeBaseFileSystemEnvironment(Environment): - def __init__(self, env_id: str) -> None: - super().__init__(env_id) - self.root_path = "." - - operator_param = { - "path": "full path of target directory", - } - self.add_ai_function(SimpleAIFunction("list", - "list the files and sub directory in target directory,result is a json array", - self.list,operator_param)) - - operator_param = { - "path": "full path of target file", - } - self.add_ai_function(SimpleAIFunction("cat", - "cat the file content in target path,result is a string", - self.cat,operator_param)) - - def set_root_path(self,path:str): - self.root_path = path - - - async def list(self,path:str) -> str: - directory_path = self.root_path + path - items = [] - - with await aiofiles.os.scandir(directory_path) as entries: - async for entry in entries: - item_type = "directory" if entry.is_dir() else "file" - items.append({"name": entry.name, "type": item_type}) - - return json.dumps(items) - - async def cat(self,path:str) -> str: - file_path = self.root_path + path - cur_encode = "utf-8" - async with aiofiles.open(file_path,'rb') as f: - cur_encode = chardet.detect(await f.read())['encoding'] - - async with aiofiles.open(file_path, mode='r', encoding=cur_encode) as f: - content = await f.read(2048) - return content - - -class ShellEnvironment(Environment): +class ShellEnvironment(SimpleEnvironment): def __init__(self, env_id: str) -> None: super().__init__(env_id) @@ -787,3 +378,57 @@ class ShellEnvironment(Environment): else: return f"Execute failed! stderr is:\n{stderr}\n" + +class WorkspaceEnvironment(CompositeEnvironment): + def __init__(self, env_id: str) -> None: + super().__init__(env_id) + myai_path = AIStorage.get_instance().get_myai_dir() + self.root_path = f"{myai_path}/workspace/{env_id}" + if not os.path.exists(self.root_path): + os.makedirs() + + self.todo_list = {} + self.todo_list[TodoListType.TO_WORK] = TodoListEnvironment(self.root_path,TodoListType.TO_WORK) + self.todo_list[TodoListType.TO_LEARN] = TodoListEnvironment(self.root_path,TodoListType.TO_LEARN) + + # default environments in workspace + self.add_env(self.todo_list[TodoListType.TO_WORK]) + self.add_env(ShellEnvironment("shell")) + self.add_env(FilesystemEnvironment(self.root_path, "fs")) + + def set_root_path(self,path:str): + self.root_path = path + + def get_prompt(self) -> AgentMsg: + return None + + def get_role_prompt(self,role_id:str) -> AgentPrompt: + return None + + def get_do_prompt(self,todo:AgentTodo=None)->AgentPrompt: + return None + + # result mean: list[op_error_str],have_error + async def exec_op_list(self,oplist:List,agent_id:str)->tuple[List[str],bool]: + result_str = "op list is none" + if oplist is None: + return None,False + + result_str = [] + have_error = False + for op in oplist: + operation = self.get_ai_operation(op["op"]) + if operation: + error_str = await operation.execute(op) + else: + logger.error(f"execute op list failed: unknown op:{op['op']}") + error_str = f"execute op list failed: unknown op:{op['op']}" + + if error_str: + have_error = True + result_str.append(error_str) + else: + result_str.append(f"execute success!") + + return result_str,have_error + diff --git a/src/aios/knowledge/pipeline.py b/src/aios/knowledge/pipeline.py index b902225..db2fb3e 100644 --- a/src/aios/knowledge/pipeline.py +++ b/src/aios/knowledge/pipeline.py @@ -6,17 +6,13 @@ from . import ObjectID, KnowledgeStore from enum import Enum class KnowledgePipelineJournal: - def __init__(self, time: datetime.datetime, object_id: str, input: str, parser: str): + def __init__(self, time: datetime.datetime, input: str, parser: str): self.time = time - self.object_id = None if object_id is None else ObjectID.from_base58(object_id) self.input = input self.parser = parser def is_finish(self) -> bool: - return self.object_id is None - - def get_object_id(self) -> ObjectID: - return self.object_id + return self.input is None def get_input(self) -> str: return self.input @@ -28,7 +24,7 @@ class KnowledgePipelineJournal: if self.is_finish(): return f"{self.time}: finished)" else: - return f"{self.time}: object:{self.object_id} input:{self.input}, parser:{self.parser})" + return f"{self.time}: input:{self.input}, parser:{self.parser})" # init sqlite3 client class KnowledgePipelineJournalClient: @@ -42,18 +38,17 @@ class KnowledgePipelineJournalClient: '''CREATE TABLE IF NOT EXISTS journal ( id INTEGER PRIMARY KEY AUTOINCREMENT, time DATETIME DEFAULT CURRENT_TIMESTAMP, - object_id TEXT, input TEXT, parser TEXT)''' ) conn.commit() - def insert(self, object_id: ObjectID, input: str, parser: str, timestamp: datetime.datetime = None): + def insert(self, input: str, parser: str, timestamp: datetime.datetime = None): timestamp = datetime.datetime.now() if timestamp is None else timestamp conn = sqlite3.connect(self.journal_path) conn.execute( - "INSERT INTO journal (time, object_id, input, parser) VALUES (?, ?, ?, ?)", - (timestamp, str(object_id), input, parser), + "INSERT INTO journal (time, input, parser) VALUES (?, ?, ?, ?)", + (timestamp, input, parser), ) conn.commit() @@ -61,7 +56,7 @@ class KnowledgePipelineJournalClient: conn = sqlite3.connect(self.journal_path) cursor = conn.cursor() cursor.execute("SELECT * FROM journal ORDER BY id DESC LIMIT ?", (topn,)) - return [KnowledgePipelineJournal(time, object_id, input, parser) for (_, time, object_id, input, parser) in cursor.fetchall()] + return [KnowledgePipelineJournal(time, input, parser) for (_, time, input, parser) in cursor.fetchall()] class KnowledgePipelineEnvironment: def __init__(self, pipeline_path: str): @@ -87,8 +82,12 @@ class KnowledgePipelineState(Enum): STOPPED = 2 FINISHED = 3 +class NullParser: + async def parse(self, object_id): + return "" + class KnowledgePipeline: - def __init__(self, name: str, env: KnowledgePipelineEnvironment, input_init, input_params, parser_init, parser_params): + def __init__(self, name: str, env: KnowledgePipelineEnvironment, input_init, input_params=None, parser_init=None, parser_params=None): self.name = name self.state = KnowledgePipelineState.INIT self.input_init = input_init @@ -108,18 +107,21 @@ class KnowledgePipeline: async def run(self): if self.state == KnowledgePipelineState.INIT: self.input = self.input_init(self.env, self.input_params) - self.parser = self.parser_init(self.env, self.parser_params) + if self.parser_init is None: + self.parser = NullParser() + else: + self.parser = self.parser_init(self.env, self.parser_params) self.state = KnowledgePipelineState.RUNNING if self.state == KnowledgePipelineState.RUNNING: async for input in self.input.next(): if input is None: self.state = KnowledgePipelineState.FINISHED - self.env.journal.insert(None, "finished", "finished") + self.env.journal.insert(None, None) return (object_id, input_journal) = input if object_id is not None: parser_journal = await self.parser.parse(object_id) - self.env.journal.insert(object_id, input_journal, parser_journal) + self.env.journal.insert(input_journal, parser_journal) else: return if self.state == KnowledgePipelineState.STOPPED: diff --git a/src/component/common_environment/local_document.py b/src/component/common_environment/local_document.py new file mode 100644 index 0000000..c003b43 --- /dev/null +++ b/src/component/common_environment/local_document.py @@ -0,0 +1,671 @@ +# import os +# import aiofiles +# import chardet +# import logging +# import string +# from knowledge import ImageObjectBuilder, DocumentObjectBuilder, KnowledgePipelineEnvironment, KnowledgePipelineJournal +# from aios_kernel.storage import AIStorage + + +import os +import aiofiles +import chardet +import logging +import string +import sqlite3 +import json +import threading +import logging +from datetime import datetime +from typing import Optional, List +from knowledge import ImageObjectBuilder, DocumentObjectBuilder, KnowledgePipelineEnvironment, KnowledgePipelineJournal +from aios_kernel import AIStorage, SimpleEnvironment + + +class ScanLocalDocument: + def __init__(self, env: KnowledgePipelineEnvironment, config): + self.env = env + path = string.Template(config["path"]).substitute(myai_dir=AIStorage.get_instance().get_myai_dir()) + config["path"] = path + self.config = config + + def path(self): + return self.config["path"] + + async def next(self): + while True: + journals = self.env.journal.latest_journals(1) + from_time = 0 + if len(journals) == 1: + latest_journal = journals[0] + if latest_journal.is_finish(): + yield None + continue + from_time = os.path.getctime(latest_journal.get_input()) + if os.path.getmtime(self.path()) <= from_time: + yield (None, None) + continue + + file_pathes = sorted(os.listdir(self.path()), key=lambda x: os.path.getctime(os.path.join(self.path(), x))) + for rel_path in file_pathes: + file_path = os.path.join(self.path(), rel_path) + timestamp = os.path.getctime(file_path) + if timestamp <= from_time: + continue + ext = os.path.splitext(file_path)[1].lower() + if ext in ['.pdf', '.md', '.txt']: + logging.info(f"knowledge dir source found document file {file_path}") + yield (file_path, file_path) + yield (None, None) + +class MetaDatabase: + def __init__(self,db_path:str): + self.db_path = db_path + self._get_conn() + + def _get_conn(self): + """ get db connection """ + local = threading.local() + if not hasattr(local, 'conn'): + local.conn = self._create_connection(self.db_path) + return local.conn + + + def _create_connection(self, db_file): + """ create a database connection to a SQLite database """ + conn = None + try: + conn = sqlite3.connect(db_file) + except Exception as e: + logger.error("Error occurred while connecting to database: %s", e) + return None + + if conn: + self._create_tables(conn) + + return conn + + def _create_tables(self,conn): + cursor = conn.cursor() + cursor.execute(''' + CREATE TABLE IF NOT EXISTS documents ( + doc_path TEXT PRIMARY KEY, + length INTEGER, + last_modify TEXT, + doc_hash TEXT, + create_time TEXT + ) + ''') + cursor.execute(''' + CREATE TABLE IF NOT EXISTS knowledge ( + doc_hash TEXT PRIMARY KEY, + title TEXT, + summary TEXT, + content TEXT, + catalogs TEXT, + tags TEXT, + llm_title TEXT, + llm_summary TEXT, + create_time TEXT + ) + ''') + + cursor.execute(''' + CREATE INDEX IF NOT EXISTS idx_documents_doc_hash + ON documents (doc_hash) + ''') + + cursor.execute(''' + CREATE INDEX IF NOT EXISTS idx_knowledge_tags + ON knowledge (tags) + ''') + + conn.commit() + + def add_doc(self, doc_path: str, length: int, last_modify: str, doc_hash: Optional[str] = None): + conn = self._get_conn() + cursor = conn.cursor() + create_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + cursor.execute(''' + INSERT INTO documents (doc_path, length, last_modify, doc_hash,create_time) + VALUES (?, ?, ?, ?,?) + ''', (doc_path, length, last_modify, doc_hash,create_time)) + conn.commit() + + def is_doc_exist(self, doc_path: str) -> bool: + conn = self._get_conn() + cursor = conn.cursor() + cursor.execute(''' + SELECT doc_path + FROM documents + WHERE doc_path = ? + ''', (doc_path,)) + return len(cursor.fetchall()) > 0 + + def set_doc_hash(self, doc_path: str, doc_hash: str): + conn = self._get_conn() + cursor = conn.cursor() + cursor.execute(''' + UPDATE documents + SET doc_hash = ? + WHERE doc_path = ? + ''', (doc_hash, doc_path)) + conn.commit() + + def get_docs_without_hash(self,limit:int=1024) -> List[str]: + conn = self._get_conn() + cursor = conn.cursor() + cursor.execute(''' + SELECT doc_path + FROM documents + WHERE doc_hash IS NULL OR doc_hash = '' + ORDER BY create_time DESC + LIMIT ? + ''',(limit,)) + return [row[0] for row in cursor.fetchall()] + + #metadata["summary"] + #metadata["catelogs"] + #metadata["tags"] + def add_knowledge(self, doc_hash: str, title: str, metadata: dict,content:str = None,): + conn = self._get_conn() + cursor = conn.cursor() + + create_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + summary = metadata.get("summary", "") + catalogs = metadata.get("catalogs","") + tags = ','.join(metadata.get("tags", [])) + + cursor.execute(''' + INSERT INTO knowledge (doc_hash, title , summary , catalogs , tags,create_time) + VALUES (?, ?, ?, ?, ?,?) + ''', (doc_hash, title, summary, catalogs, tags,create_time)) + conn.commit() + + #llm_result["summary"] + #llm_result["tags"] + #llm_result["catelog"] + def set_knowledge_llm_result(self, doc_hash: str, llm_result: dict): + conn = self._get_conn() + cursor = conn.cursor() + + title = llm_result.get("title", "") + summary = llm_result.get("summary", "") + catalogs = json.dumps(llm_result.get("catalogs", {})) + tags = ','.join(llm_result.get("tags", [])) + + cursor.execute(''' + UPDATE knowledge + SET llm_title = ?,llm_summary = ?, catalogs = ?, tags = ? + WHERE doc_hash = ? + ''', (title,summary, catalogs, tags, doc_hash)) + conn.commit() + + def get_hash_by_doc_path(self, doc_path: str) -> Optional[str]: + conn = self._get_conn() + cursor = conn.cursor() + cursor.execute(''' + SELECT doc_hash + FROM documents + WHERE doc_path = ? + ''', (doc_path,)) + row = cursor.fetchone() + if row is None: + return None + return row[0] + + def get_knowledge(self, doc_hash: str) -> Optional[dict]: + conn = self._get_conn() + cursor = conn.cursor() + cursor.execute(''' + SELECT title, summary, catalogs, tags, llm_title, llm_summary + FROM knowledge + WHERE doc_hash = ? + ''', (doc_hash,)) + row = cursor.fetchone() + if row is None: + return None + + # get doc path + cursor.execute(''' + SELECT doc_path + FROM documents + WHERE doc_hash = ? + ''', (doc_hash,)) + row2 = cursor.fetchone() + if row2 is None: + return None + doc_path = row2[0] + + + return { + "full_path": doc_path, + "title": row[0], + "summary": row[1], + "catalogs": row[2], + "tags": row[3], + "llm_title" : row[4], + "llm_summary" : row[5], + } + + def get_knowledge_without_llm_title(self,limit:int=16) -> List[str]: + conn = self._get_conn() + cursor = conn.cursor() + cursor.execute(''' + SELECT doc_hash + FROM knowledge + WHERE llm_title IS NULL OR llm_title = '' + ORDER BY create_time DESC + LIMIT ? + ''',(limit,)) + return [row[0] for row in cursor.fetchall()] + + def query_docs_by_tag(self, tag: str) -> List[str]: + conn = self._get_conn() + cursor = conn.cursor() + tag_json = json.dumps(tag) # 将标签转换为 JSON 字符串 + cursor.execute(''' + SELECT documents.doc_path + FROM documents + JOIN knowledge ON documents.doc_hash = knowledge.doc_hash + WHERE json_extract(knowledge.tags, '$') LIKE ? + ''', (tag)) + return [row[0] for row in cursor.fetchall()] + + +class DocumentKnowledgeBase(SimpleEnvironment): + async def get_knowledege_catalog(self,path:str=None,only_dir =True,max_depth:int=5)->str: + if path: + full_path = f"{self.root_path}/knowledge/{path}" + else: + full_path = f"{self.root_path}/knowledge" + + catlogs,file_count = await self.get_directory_structure(full_path,max_depth,only_dir) + return catlogs + + async def get_directory_structure(self,root_dir, max_depth:int=4, only_dir=True, indent=1): + file_count = 0 + structure_str = '' + if os.path.isdir(root_dir): + sub_files = [] + with os.scandir(root_dir) as it: + for entry in it: + if entry.is_dir(): + sub_structure, sub_count = await self.get_directory_structure(entry.path, max_depth, only_dir, indent + 1) + if sub_structure: + structure_str += sub_structure + file_count += sub_count + else: + file_count += 1 + sub_files.append(entry.name) + + if only_dir is False: + for file_name in sub_files: + structure_str = structure_str + ' ' * (indent+1) + file_name + '\n' + + dir_name = os.path.basename(root_dir) + dir_info = f"{dir_name} " + + + structure_str = ' ' * indent + dir_info + '\n' + structure_str + + if indent - 1 >= max_depth: + return None, file_count + else: + return structure_str, file_count + + # inner_function + async def get_knowledge(self,path:str) -> str: + full_path = f"{self.root_path}/knowledge/{path}" + if os.islink(full_path): + org_path = os.readlink(full_path) + hash = self.kb_db.get_hash_by_doc_path(org_path) + if hash: + return self.kb_db.get_knowledge(org_path) + + return "not found" + + +class ParseLocalDocument: + def _parse_pdf_bookmarks(self,bookmarks, parent:list): + + for item in bookmarks: + if isinstance(item,list): + self._parse_pdf_bookmarks(item,parent) + else: + if item.title: + new_item = {} + new_item["page"] = item.page.idnum + new_item["title"] = item.title + my_childs = [] + if item.childs: + if len(item.childs) > 0: + self._parse_pdf_bookmarks(item.childs, my_childs) + new_item["childs"] = my_childs + parent.append(new_item) + else: + logger.warning("parse pdf bookmarks failed: item.title is None!") + + return + + def _parse_pdf(self,doc_path:str): + metadata = {} + with open(doc_path, 'rb') as file: + reader = PyPDF2.PdfReader(file) + try: + doc_info = reader.metadata + if doc_info: + if doc_info.title: + metadata["title"] = doc_info.title + if doc_info.author: + metadata["authors"] = doc_info.author + except Exception as e: + logger.warn("parse pdf metadata failed:%s",e) + + dir_path = os.path.dirname(doc_path) + base_name = os.path.basename(doc_path) + text_content_path = f"{dir_path}/.{base_name}.txt" + full_text = "" + + for page in reader.pages: + text = page.extract_text() + full_text += text + with open(text_content_path, 'w', encoding='utf-8') as f: + f.write(full_text) + + try: + bookmarks = reader.outline + if bookmarks: + catalogs = [] + self._parse_pdf_bookmarks(bookmarks,catalogs) + metadata["catalogs"] = json.dumps(catalogs) + except Exception as e: + logger.warn("parse pdf bookmarks failed:%s",e) + + return metadata + + def _parse_txt(self,doc_path:str): + return {} + + def _parse_md(self,doc_path:str): + metadata = {} + cur_encode = "utf-8" + with open(doc_path,'rb') as f: + cur_encode = chardet.detect(f.read(1024))['encoding'] + + with open(doc_path, mode='r', encoding=cur_encode) as f: + content = f.read() + match = re.search(r'^# (.*)', content, re.MULTILINE) + if match: + metadata['title'] = match.group(1).strip() + md = Markdown(extensions=['toc']) + html_str = md.convert(content) + toc = md.toc + if toc: + metadata['catalogs'] = toc + + return metadata + + def _parse_document(self,doc_path:str): + hash_result = None + title = os.path.basename(doc_path) + meta_data = {} + + with open(doc_path, "rb") as f: + hash_md5 = hashlib.md5() + for chunk in iter(lambda: f.read(1024*1024), b""): + hash_md5.update(chunk) + hash_result = hash_md5.hexdigest() + try: + if doc_path.endswith(".md"): + meta_data = self._parse_md(doc_path) + elif doc_path.endswith(".pdf"): + meta_data = self._parse_pdf(doc_path) + except Exception as e: + logger.error("parse document %s failed:%s",doc_path,e) + traceback.print_exc() + + if meta_data.get("title"): + title = meta_data["title"] + logger.info("parse document %s!",doc_path) + return hash_result,title,meta_data + + async def parse(self, file_path: str) -> str: + + + +# async def get_knowledege_catalog(self,path:str=None,only_dir =True,max_depth:int=5)->str: +# if path: +# full_path = f"{self.root_path}/knowledge/{path}" +# else: +# full_path = f"{self.root_path}/knowledge" + +# catlogs,file_count = await self.get_directory_structure(full_path,max_depth,only_dir) +# return catlogs + +# async def get_directory_structure(self,root_dir, max_depth:int=4, only_dir=True, indent=1): +# file_count = 0 +# structure_str = '' +# if os.path.isdir(root_dir): +# sub_files = [] +# with os.scandir(root_dir) as it: +# for entry in it: +# if entry.is_dir(): +# sub_structure, sub_count = await self.get_directory_structure(entry.path, max_depth, only_dir, indent + 1) +# if sub_structure: +# structure_str += sub_structure +# file_count += sub_count +# else: +# file_count += 1 +# sub_files.append(entry.name) + +# if only_dir is False: +# for file_name in sub_files: +# structure_str = structure_str + ' ' * (indent+1) + file_name + '\n' + +# dir_name = os.path.basename(root_dir) +# dir_info = f"{dir_name} " + + +# structure_str = ' ' * indent + dir_info + '\n' + structure_str + +# if indent - 1 >= max_depth: +# return None, file_count +# else: +# return structure_str, file_count + +# # inner_function +# async def get_knowledge(self,path:str) -> str: +# full_path = f"{self.root_path}/knowledge/{path}" +# if os.islink(full_path): +# org_path = os.readlink(full_path) +# hash = self.kb_db.get_hash_by_doc_path(org_path) +# if hash: +# return self.kb_db.get_knowledge(org_path) + +# return "not found" + + async def load_knowledge_content(self,path:str,pos:int=0,length:int=None) -> str: + if path.endswith("pdf"): + logger.info("load_knowledge_content:pdf") + dir_path = os.path.dirname(path) + base_name = os.path.basename(path) + text_content_path = f"{dir_path}/.{base_name}.txt" + if os.path.exists(text_content_path) is False: + return None + async with aiofiles.open(path, mode='r', encoding=cur_encode) as f: + await f.seek(pos) + content = await f.read(length) + return content + else: + async with aiofiles.open(path,'rb') as f: + cur_encode = chardet.detect(await f.read())['encoding'] + + async with aiofiles.open(path, mode='r', encoding=cur_encode) as f: + await f.seek(pos) + content = await f.read(length) + return content + + return "load content failed." + + def _add_document_dir(self,path:str): + self.doc_dirs[path] = 0 + + + def _parse_pdf_bookmarks(self,bookmarks, parent:list): + + for item in bookmarks: + if isinstance(item,list): + self._parse_pdf_bookmarks(item,parent) + else: + if item.title: + new_item = {} + new_item["page"] = item.page.idnum + new_item["title"] = item.title + my_childs = [] + if item.childs: + if len(item.childs) > 0: + self._parse_pdf_bookmarks(item.childs, my_childs) + new_item["childs"] = my_childs + parent.append(new_item) + else: + logger.warning("parse pdf bookmarks failed: item.title is None!") + + return + + def _parse_pdf(self,doc_path:str): + metadata = {} + with open(doc_path, 'rb') as file: + reader = PyPDF2.PdfReader(file) + try: + doc_info = reader.metadata + if doc_info: + if doc_info.title: + metadata["title"] = doc_info.title + if doc_info.author: + metadata["authors"] = doc_info.author + except Exception as e: + logger.warn("parse pdf metadata failed:%s",e) + + dir_path = os.path.dirname(doc_path) + base_name = os.path.basename(doc_path) + text_content_path = f"{dir_path}/.{base_name}.txt" + full_text = "" + + for page in reader.pages: + text = page.extract_text() + full_text += text + with open(text_content_path, 'w', encoding='utf-8') as f: + f.write(full_text) + + try: + bookmarks = reader.outline + if bookmarks: + catalogs = [] + self._parse_pdf_bookmarks(bookmarks,catalogs) + metadata["catalogs"] = json.dumps(catalogs) + except Exception as e: + logger.warn("parse pdf bookmarks failed:%s",e) + + return metadata + + def _parse_txt(self,doc_path:str): + return {} + + def _parse_md(self,doc_path:str): + metadata = {} + cur_encode = "utf-8" + with open(doc_path,'rb') as f: + cur_encode = chardet.detect(f.read(1024))['encoding'] + + with open(doc_path, mode='r', encoding=cur_encode) as f: + content = f.read() + match = re.search(r'^# (.*)', content, re.MULTILINE) + if match: + metadata['title'] = match.group(1).strip() + md = Markdown(extensions=['toc']) + html_str = md.convert(content) + toc = md.toc + if toc: + metadata['catalogs'] = toc + + return metadata + + def _parse_document(self,doc_path:str): + hash_result = None + title = os.path.basename(doc_path) + meta_data = {} + + with open(doc_path, "rb") as f: + hash_md5 = hashlib.md5() + for chunk in iter(lambda: f.read(1024*1024), b""): + hash_md5.update(chunk) + hash_result = hash_md5.hexdigest() + try: + if doc_path.endswith(".md"): + meta_data = self._parse_md(doc_path) + elif doc_path.endswith(".pdf"): + meta_data = self._parse_pdf(doc_path) + except Exception as e: + logger.error("parse document %s failed:%s",doc_path,e) + traceback.print_exc() + + if meta_data.get("title"): + title = meta_data["title"] + logger.info("parse document %s!",doc_path) + return hash_result,title,meta_data + + + def _support_file(self,file_name:str) -> bool: + if file_name.startswith("."): + return False + + if file_name.endswith(".pdf"): + return True + if file_name.endswith(".md"): + return True + if file_name.endswith(".txt"): + return True + return False + + def _scan_dir(self): + while True: + time.sleep(10) + for directory in self.doc_dirs.keys(): + now = time.time() + if now - self.doc_dirs[directory] > 60*15: + self.doc_dirs[directory] = time.time() + else: + continue + + for root, dirs, files in os.walk(directory): + for file in files: + if self._support_file(file): + full_path = os.path.join(root, file) + full_path = os.path.normpath(full_path) + if self.kb_db.is_doc_exist(full_path): + continue + + file_stat = os.stat(full_path) + if file_stat.st_size < 1: + continue + + if file_stat.st_size < 1024*1024*8: + #parse and insert + hash,title,meta_data = self._parse_document(full_path) + self.kb_db.add_doc(full_path,file_stat.st_size,file_stat.st_mtime,hash) + self.kb_db.add_knowledge(hash,title,meta_data) + + else: + self.kb_db.add_doc(full_path,file_stat.st_size,file_stat.st_mtime) + + def _scan_document(self): + while True: + time.sleep(10) + parse_queue = self.kb_db.get_docs_without_hash() + for doc_path in parse_queue: + hash,title,meta_data = self._parse_document(doc_path) + self.kb_db.set_doc_hash(doc_path,hash) + self.kb_db.add_knowledge(hash,title,meta_data) + + \ No newline at end of file diff --git a/src/component/common_environment/to_learn_parser.py b/src/component/common_environment/to_learn_parser.py new file mode 100644 index 0000000..f90c8ac --- /dev/null +++ b/src/component/common_environment/to_learn_parser.py @@ -0,0 +1,214 @@ + # 尝试自我学习，会主动获取、读取资料并进行整理 + # LLM的本质能力是处理海量知识，应该让LLM能基于知识把自己的工作处理的更好 + async def do_self_learn(self) -> None: + # 不同的workspace是否应该有不同的学习方法？ + workspace = self.get_workspace_by_msg(None) + hash_list = workspace.kb_db.get_knowledge_without_llm_title() + for hash in hash_list: + if self.agent_energy <= 0: + break + + knowledge = workspace.kb_db.get_knowledge(hash) + if knowledge is None: + continue + + full_path = knowledge.get("full_path") + if full_path is None: + continue + + if os.path.exists(full_path) is False: + logger.warning(f"do_self_learn: knowledge {full_path} is not exists!") + continue + + #TODO 可以用v-db 对不同目录的名字进行选择后，先进行一次快速的插入。有时间再慢慢用LLM整理 + result_obj = await self._llm_read_article(knowledge,full_path) + + #根据结果更新knowledge + if result_obj is not None: + workspace.kb_db.set_knowledge_llm_result(hash,result_obj) + # 在知识库中创建软链接 + path_list = result_obj.get("path") + new_title = result_obj.get("title") + if path_list: + for new_path in path_list: + full_new_path = f"/knowledge{new_path}/{new_title}" + await workspace.symlink(full_path,full_new_path) + logger.info(f"create soft link {full_path} -> {full_new_path}") + + + self.agent_energy -= 1 + + # match item.type(): + # case "book": + # self.llm_read_book(kb,item) + # learn_power -= 1 + # case "article": + # + # self.llm_read_article(kb,item) + # learn_power -= 1 + # case "video": + # self.llm_watch_video(kb,item) + # learn_power -= 1 + # case "audio": + # self.llm_listen_audio(kb,item) + # learn_power -= 1 + # case "code_project": + # self.llm_read_code_project(kb,item) + # learn_power -= 1 + # case "image": + # self.llm_view_image(kb,item) + # learn_power -= 1 + # case "other": + # self.llm_read_other(kb,item) + # learn_power -= 1 + # case _: + # self.llm_learn_any(kb,item) + # pass + + + async def do_blance_knowledge_base(selft): + # 整理自己的知识库(让分类更平衡，更由于自己以后的工作)，并尝试更新学习目标 + current_path = "/" + current_list = kb.get_list(current_path) + self_assessment_with_goal = self.get_self_assessment_with_goal() + learn_goal = {} + + + llm_blance_knowledge_base(current_path,current_list,self_assessment_with_goal,learn_goal,learn_power) + + # 主动学习 + # 方法目前只有使用搜索引擎一种？ + for goal in learn_goal.items(): + self.llm_learn_with_search_engine(kb,goal,learn_power) + if learn_power <= 0: + break + + + def parser_learn_llm_result(self,llm_result:LLMResult): + pass + + async def gen_known_info_for_knowledge_prompt(self,knowledge_item:dict,temp_meta = None,need_catalogs = False) -> AgentPrompt: + workspace =self.get_workspace_by_msg(None) + kb_tree = await workspace.get_knowledege_catalog() + + + known_obj = {} + title = knowledge_item.get("title") + if title: + known_obj["title"] = title + summary = knowledge_item.get("summary") + if summary: + known_obj["summary"] = summary + tags = knowledge_item.get("tags") + if tags: + known_obj["tags"] = tags + if need_catalogs: + catalogs = knowledge_item.get("catalogs") + if catalogs: + known_obj["catalogs"] = catalogs + + if temp_meta: + for key in temp_meta.keys(): + known_obj[key] = temp_meta[key] + + org_path = knowledge_item.get("full_path") + known_obj["orginal_path"] = org_path + know_info_str = f"# Known information:\n## Current directory structure:\n{kb_tree}\n## Knowlege Metadata:\n{json.dumps(known_obj)}\n" + return AgentPrompt(know_info_str) + + async def _llm_read_article(self,knowledge_item:dict,full_path:str) -> ComputeTaskResult: + # Objectives: + # Obtain better titles, abstracts, table of contents (if necessary), tags + # Determine the appropriate place to put it (in line with the organization's goals) + # Known information: + # The reason why the target service's learn_prompt is being sorted + # Summary of the organization's work (if any) + # The current structure of the knowledge base (note the size control) gen_kb_tree_prompt (when empty, LLM should generate an appropriate initial directory structure) + # Original path, current title, abstract, table of contents + + # Sorting long files (general tricks) + # Indicate that the input is part of the content, let LLM generate intermediate results for the task + # Enter the content in sequence, when the last content block is input, LLM gets the result + + + #full_content = item.get_article_full_content() + workspace = self.get_workspace_by_msg(None) + full_content_len = self.token_len(full_content) + + if full_content_len < self.get_llm_learn_token_limit(): + + # 短文章不用总结catelog + #path_list,summary = llm_get_summary(summary,full_content) + #prompt = self.get_agent_role_prompt() + prompt = AgentPrompt() + prompt.append(self.get_learn_prompt()) + known_info_prompt = await self.gen_known_info_for_knowledge_prompt(knowledge_item) + prompt.append(known_info_prompt) + content_prompt = AgentPrompt(full_content) + prompt.append(content_prompt) + env_functions = None + #env_functions,function_len = workspace.get_knowledge_base_ai_functions() + task_result:ComputeTaskResult = await self.do_llm_complection(prompt,is_json_resp=True) + if task_result.result_code != ComputeTaskResultCode.OK: + result_obj = {} + result_obj["error_str"] = task_result.error_str + return result_obj + + result_obj = json.loads(task_result.result_str) + return result_obj + + else: + logger.warning(f"llm_read_article: article {full_path} use LLM loop learn!") + pos = 0 + read_len = int(self.get_llm_learn_token_limit() * 1.2) + + temp_meta_data = {} + is_final = False + while pos < str_len: + _content = full_content[pos:pos+read_len] + part_cotent_len = len(_content) + if part_cotent_len < read_len: + # last chunk + is_final = True + part_content = f"<>\n{_content}" + else: + part_content = f"<>\n{_content}" + + pos = pos + read_len + prompt = AgentPrompt() + prompt.append(self.get_learn_prompt()) + known_info_prompt = await self.gen_known_info_for_knowledge_prompt(knowledge_item,temp_meta_data) + prompt.append(known_info_prompt) + content_prompt = AgentPrompt(part_content) + prompt.append(content_prompt) + #env_functions,function_len = workspace.get_knowledge_base_ai_functions() + task_result:ComputeTaskResult = await self.do_llm_complection(prompt,is_json_resp=True) + if task_result.result_code != ComputeTaskResultCode.OK: + result_obj = {} + result_obj["error_str"] = task_result.error_str + return result_obj + + result_obj = json.loads(task_result.result_str) + temp_meta_data = result_obj + if is_final: + return result_obj + + return None + + + async def do_self_think(self): + session_id_list = AIChatSession.list_session(self.agent_id,self.chat_db) + for session_id in session_id_list: + if self.agent_energy <= 0: + break + used_energy = await self.think_chatsession(session_id) + self.agent_energy -= used_energy + + todo_logs = await self.get_todo_logs() + for todo_log in todo_logs: + if self.agent_energy <= 0: + break + used_energy = await self.think_todo_log(todo_log) + self.agent_energy -= used_energy + + return \ No newline at end of file diff --git a/src/component/knowledge_manager/pipeline.py b/src/component/knowledge_manager/pipeline.py index 31b105d..c8a1ded 100644 --- a/src/component/knowledge_manager/pipeline.py +++ b/src/component/knowledge_manager/pipeline.py @@ -44,14 +44,19 @@ class KnowledgePipelineManager: input_init = self.input_modules.get(input_module) input_params = config["input"].get("params") - parser_module = config["parser"]["module"] - _, ext = os.path.splitext(parser_module) - if ext == ".py": - parser_module = os.path.join(path, parser_module) - parser_init = runpy.run_path(parser_module)["init"] + parser_config = config.get("parser") + if parser_config is None: + parser_init = None + parser_params = None else: - parser_init = self.parser_modules.get(parser_module) - parser_params = config["parser"].get("params") + parser_module = parser_config["module"] + _, ext = os.path.splitext(parser_module) + if ext == ".py": + parser_module = os.path.join(path, parser_module) + parser_init = runpy.run_path(parser_module)["init"] + else: + parser_init = self.parser_modules.get(parser_module) + parser_params = parser_config.get("params") data_path = os.path.join(self.root_dir, name) diff --git a/src/component/mail_environment/local.py b/src/component/mail_environment/local.py index 21f3a94..af9726c 100644 --- a/src/component/mail_environment/local.py +++ b/src/component/mail_environment/local.py @@ -22,7 +22,7 @@ class LocalEmail: if latest_journal.is_finish(): yield None continue - parsed = str(latest_journal.get_object_id()) + parsed = latest_journal.get_input() mail_id = self.mail_storage.next_mail_id(parsed) if mail_id is None: diff --git a/src/component/mail_environment/mail.py b/src/component/mail_environment/mail.py index ad9ba2e..b49b0c7 100644 --- a/src/component/mail_environment/mail.py +++ b/src/component/mail_environment/mail.py @@ -7,17 +7,20 @@ import datetime from bs4 import BeautifulSoup import sqlite3 import html2text +from urllib.parse import urlparse from aios import * + + class Mail: def __init__(self, **kwargs) -> None: - self.from_addr = kwargs.get("From") - self.to_addr = kwargs.get("To") - self.subject = kwargs.get("Subject") - self.date = kwargs.get("Date") - self.bcc = kwargs.get("BCC") - self.cc = kwargs.get("CC") - self.reply_to = None + self.from_addr = kwargs.get("from") + self.to_addr = kwargs.get("to") + self.subject = kwargs.get("subject") + self.date = kwargs.get("date") + self.bcc = kwargs.get("bcc") + self.cc = kwargs.get("cc") + self.reply_to = kwargs.get("reply_to") self.id: str = None self.content: str = None @@ -192,20 +195,36 @@ class MailStorage: self.conn.commit() await asyncio.sleep(10) - def download(self, uid, mail: mailparser.MailParser): + def download(self, uid, parser: mailparser.MailParser, + save_image=True, + from_field="From", + to_field="To", + subject_field="Subject", + date_field="Date", + reply_to_field="In-Reply-To", + cc_field="CC", + bcc_field="BCC"): mail_dir = self.mail_dir(uid) - os.makedirs(dir) + if not os.path.exists(mail_dir): + os.makedirs(mail_dir) - meta = json.loads(mail.mail_json) - mail = Mail(**meta) - reply_to = meta.get("In-Reply-To") + src_meta = json.loads(parser.mail_json) + meta = {} + meta["from"] = src_meta.get(from_field) + meta["to"] = src_meta.get(to_field) + meta["subject"] = src_meta.get(subject_field) + meta["date"] = src_meta.get(date_field) + meta["bcc"] = src_meta.get(bcc_field) + meta["cc"] = src_meta.get(cc_field) + reply_to = src_meta.get(reply_to_field) if reply_to: - mail.reply_to = self.uid_to_object_id(reply_to) + meta["reply_to"] = self.uid_to_object_id(reply_to) + mail = Mail(**meta) h = html2text.HTML2Text() h.ignore_links = True h.ignore_images = True - mail_content = h.handle(mail.body) + mail_content = h.handle(parser.body) mail.content = mail_content mail.calculate_id() @@ -216,41 +235,52 @@ class MailStorage: with open(f"{mail_dir}/mail.txt", "w", encoding='utf-8') as f: f.write(mail_content) - for attachment in mail.attachments: - if attachment['mail_content_type'] in ['image/png', 'image/jpeg', 'image/gif']: - filename = attachment['filename'] - filefullname = f"{mail_dir}/{filename}" - image_data = attachment['payload'] + if save_image: + for attachment in parser.attachments: + if attachment['mail_content_type'] in ['image/png', 'image/jpg', 'image/jpeg', 'image/gif', 'image/svg']: + filename = attachment['filename'] + filefullname = f"{mail_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) + logging.info(f"save email image {filename} success") + + # get all image urls + soup = BeautifulSoup(parser.body, 'html.parser') + img_tags = soup.find_all('img') + img_urls = [img['src'] for img in img_tags if 'src' in img.attrs] + logging.info(f'Found {len(img_urls)} images in email body') + + name_count = 0 + + for img_url in img_urls: + # keep the original image filename(last of url) + url_result = urlparse(img_url) + if url_result.scheme not in ['http', 'https']: + continue + ext = url_result.path.split('/')[-1].split('.')[-1] + if ext in ['png', 'jpg', 'jpeg', 'gif', 'svg']: + img_filename = os.path.join(mail_dir, f"{name_count}.{ext}") + else : + img_filename = os.path.join(mail_dir, f"{name_count}") + name_count += 1 + # download image 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) - logging.info(f"save email image {filename} success") - - # get all image urls - soup = BeautifulSoup(mail.body, 'html.parser') - img_tags = soup.find_all('img') - img_urls = [img['src'] for img in img_tags if 'src' in img.attrs] - logging.info(f'Found {len(img_urls)} images in email body') - - name_count = 0 - - for img_url in img_urls: - # keep the original image filename(last of url) - ext = img_url.split('/')[-1].split('.')[-1] - img_filename = os.path.join(mail_dir, f"{name_count}.{ext}") - name_count += 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) - logging.info(f'Downloaded {img_url} to {img_filename}') - else: - logging.info(f'Failed to download {img_url}') + response = requests.get(img_url, stream=True) + except requests.exceptions.RequestException as e: + logging.error(f'Failed to download {img_url}: {e}') + continue + if response.status_code == 200: + with open(img_filename, 'wb') as img_file: + for chunk in response.iter_content(1024): + img_file.write(chunk) + logging.info(f'Downloaded {img_url} to {img_filename}') + else: + logging.error(f'Failed to download {img_url}') cursor = self.conn.cursor() cursor.execute( @@ -260,5 +290,8 @@ class MailStorage: """, (uid, mail.id, mail.date, mail.from_addr), ) + self.conn.commit() + + return mail.id \ No newline at end of file diff --git a/src/component/mail_environment/spider.py b/src/component/mail_environment/spider.py index cc6586b..0119919 100644 --- a/src/component/mail_environment/spider.py +++ b/src/component/mail_environment/spider.py @@ -1,9 +1,13 @@ import os import logging import json +import string import imaplib import mailparser -from aios import * + +from knowledge import * +from aios_kernel.storage import AIStorage +from .mail import Mail, MailStorage class EmailSpider: @@ -16,14 +20,22 @@ class EmailSpider: port=self.config.get('imap_port') ) self.client.login(self.config.get('address'), self.config.get('password')) - self.mail_local_root = os.path.join(self.env.pipeline_path, self.config.get("address")) - os.makedirs(self.mail_local_root) + self.client.select("INBOX") + local_path = string.Template(config["path"]).substitute(myai_dir=AIStorage.get_instance().get_myai_dir()) + local_path = os.path.join(local_path, self.config.get('address')) + self.mail_storage = MailStorage(local_path) + async def next(self): while True: - _, data = self.client.uid('search', None, "ALL") + try: + _, data = self.client.uid('search', None, "ALL") + except Exception as e: + self.env.get_logger().error(f"email spider error: {e}") + yield (None, None) + continue uid_list = data[0].split() - if uid_list.len() == 0: + if len(uid_list) == 0: yield (None, None) continue @@ -43,9 +55,16 @@ class EmailSpider: _uid = int.from_bytes(uid) if _uid > from_uid: message_parts = "(BODY.PEEK[])" - _, email_data = self.client.uid('fetch', uid, message_parts) - mail = mailparser.parse_from_bytes(email_data[0][1]) - self.save_email(_uid, mail) + try: + _, email_data = self.client.uid('fetch', uid, message_parts) + mail = mailparser.parse_from_bytes(email_data[0][1]) + id = self.mail_storage.download(_uid, mail) + except Exception as e: + self.env.get_logger().error(f"email spider error: {e}") + yield (None, None) + break + yield (ObjectID.from_base58(id), str(_uid)) + yield (None, None)