From 1bbdf2e7225662666a8efd159c4066bd9cb19acc Mon Sep 17 00:00:00 2001 From: Liu Zhicong Date: Mon, 4 Sep 2023 22:36:59 -0700 Subject: [PATCH] Add basic enviroment support. Can use env values in prompts right now. --- rootfs/workflows/math_school/workflow.toml | 17 ++- src/aios_kernel/__init__.py | 3 +- src/aios_kernel/environment.py | 100 ++++++++++++++- src/aios_kernel/workflow.py | 69 ++++++++--- src/aios_kernel/workflow_env.py | 136 +++++++++++++++++++++ src/aios_kernel/worksapce_env.py | 0 src/service/aios_shell/aios_shell.py | 6 +- test/env_test.py | 32 +++++ 8 files changed, 340 insertions(+), 23 deletions(-) create mode 100644 src/aios_kernel/workflow_env.py create mode 100644 src/aios_kernel/worksapce_env.py create mode 100644 test/env_test.py diff --git a/rootfs/workflows/math_school/workflow.toml b/rootfs/workflows/math_school/workflow.toml index 9f5af27..48c9831 100644 --- a/rootfs/workflows/math_school/workflow.toml +++ b/rootfs/workflows/math_school/workflow.toml @@ -1,5 +1,15 @@ name = "math_school" +[enviroment] +GOAL="成为最好的学校" + + +[[connected_env]] +env_id = "calender" +[[connected_env.event2msg]] +timer = "现在是{data}" +role = "教导处主任" + [filter] "*" = "小学老师" @@ -9,7 +19,7 @@ fullname = "Ada Zhang" agent="math_teacher" [[roles."小学老师".prompt]] role="system" -content="""你在学校任职，担任小学老师。学校由 小学老师、初中老师、高中老师、教导处主任 组成。 +content="""现在时间是:{now}你在学校任职，担任小学老师。学校由 小学老师、初中老师、高中老师、教导处主任 组成。 当你发现学生的水平不是小学生时，应使用 sendmsg(老师名称,问题) 的方法，把学生的问题转发给学校里合适的老师 当学生发来作业时，进行批改(满分5分)，并把批改结果以 postmsg(教导处主任,学生名_作业结果) 的方法，将一次作业情况汇报给教导处主任。 你会根据教导处主任的指示，定期调整教学方法""" @@ -45,6 +55,7 @@ agent="math_teacher" [[roles."教导处主任".prompt]] role="system" -content="""你在学校任职，担任教导处主任。 -你收到老师发来的信息时，如果是类似 学生名_作业分数 的结果，会在合适的情况下根据学生作业的整体情况，对老师的教学方法进行必要的调整。""" +content="""你在学校任职，担任教导处主任。你的目标是{GOAL} +你收到老师发来的信息时，如果是类似 学生名_作业分数 的结果，会在合适的情况下根据学生作业的整体情况，对老师的教学方法进行必要的调整。 +当收到非老师发来的时间信息时，回复那一天学生的平均分。""" diff --git a/src/aios_kernel/__init__.py b/src/aios_kernel/__init__.py index bd1434b..bef117f 100644 --- a/src/aios_kernel/__init__.py +++ b/src/aios_kernel/__init__.py @@ -7,4 +7,5 @@ from .compute_node import ComputeNode,LocalComputeNode from .open_ai_node import OpenAI_ComputeNode from .role import AIRole,AIRoleGroup from .workflow import Workflow -from .bus import AIBus \ No newline at end of file +from .bus import AIBus +from .workflow_env import WorkflowEnvironment,CalenderEnvironment,CalenderEvent \ No newline at end of file diff --git a/src/aios_kernel/environment.py b/src/aios_kernel/environment.py index 64c6481..9979e2c 100644 --- a/src/aios_kernel/environment.py +++ b/src/aios_kernel/environment.py @@ -2,20 +2,112 @@ # we have some built-in environment: Calender(include timer),Home(connect to IoT device in your home), ,KnwoledgeBase,FileSystem, from abc import ABC, abstractmethod -from typing import Callable +from typing import Any, Callable, Optional,Dict,Awaitable,List +import logging + +logger = logging.getLogger(__name__) class EnvironmentEvent(ABC): @abstractmethod def display(self) -> str: pass +EnvironmentEventHandler = Callable[[str,EnvironmentEvent],Awaitable[Any]] + class Environment: - def __init__(self) -> None: - 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 + + 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]]= {} 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 + + @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: - pass + 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.warn(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(key) + + 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") + 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 diff --git a/src/aios_kernel/workflow.py b/src/aios_kernel/workflow.py index f7c13e8..69f39c3 100644 --- a/src/aios_kernel/workflow.py +++ b/src/aios_kernel/workflow.py @@ -2,6 +2,7 @@ import logging import asyncio import json +import os from asyncio import Queue from typing import Optional,Tuple from abc import ABC, abstractmethod @@ -14,6 +15,7 @@ from .role import AIRole,AIRoleGroup from .ai_function import CallChain from .compute_kernel import ComputeKernel from .bus import AIBus +from .workflow_env import WorkflowEnvironment logger = logging.getLogger(__name__) @@ -56,6 +58,8 @@ class Workflow: self.sub_workflows = {} self.owner_workflow = None self.db_file = None + self.env_db_file = None + self.workflow_env:WorkflowEnvironment = None self.is_start = False self.msg_queue = Queue() @@ -100,6 +104,27 @@ class Workflow: if self.input_filter.load_from_config(config.get("filter")) is False: logger.error("Workflow load input_filter failed") return False + + if self.owner_workflow is None: + self.env_db_file = os.path.dirname(self.db_file) + "/" + self.workflow_id + "_env.db" + else: + self.env_db_file = self.owner_workflow.env_db_file + self.workflow_env = WorkflowEnvironment(self.workflow_id,self.env_db_file) + + env_ndoe = config.get("enviroment") + if env_ndoe is not None: + if self._load_env_from_config(env_ndoe) is False: + logger.error("Workflow load env failed") + return False + + connected_env_ndoe = config.get("connected_env") + if connected_env_ndoe is not None: + for _node in connected_env_ndoe: + remote_env = Environment.get_env_by_id(_node.get("env_id")) + if remote_env is None: + logger.error(f"Workflow load connected_env failed, env {env_id} not found!") + return False + self.connect_to_environment(remote_env,_node.get("event2msg")) sub_workflows = config.get("sub_workflows") if sub_workflows is not None: @@ -107,10 +132,12 @@ class Workflow: logger.error("Workflow load sub workflows failed") return False - #TODO: load env - return True + def _load_env_from_config(self,config:dict) -> bool: + for k,v in config.items(): + self.workflow_env.set_value(k,v,False) + def _load_sub_workflows(self,config:dict) -> bool: for k,v in config.items(): sub_workflow = Workflow() @@ -296,6 +323,13 @@ class Workflow: logger.info(f"{the_role.role_id} post call {call[0]} with args {call[1]}") return + def _format_msg_by_env_value(self,prompt:AgentPrompt): + if self.workflow_env is None: + return + + for msg in prompt.messages: + old_content = msg.get("content") + msg["content"] = old_content.format_map(self.workflow_env) async def role_process_msg(self,msg:AgentMsg,the_role:AIRole): session_topic = f"{msg.sender}#{msg.topic}" @@ -316,8 +350,9 @@ class Workflow: msg_prompt = AgentPrompt() msg_prompt.messages = [{"role":"user","content":msg.body}] prompt.append(msg_prompt) - - + + self._format_msg_by_env_value(prompt) + async def _do_process_msg(): #TODO: send msg to agent might be better? result_str = await ComputeKernel().do_llm_completion(prompt,the_role.agent.get_llm_model_name(),the_role.agent.get_max_token_size()) @@ -392,16 +427,22 @@ class Workflow: def get_inner_environment(self,env_id:str) -> Environment: pass - def connect_to_environment(self,env:Environment) -> None: - the_env = self.connected_environment.get(env.get_id()) - if the_env is None: - self.connected_environment[env.get_id()] = env - def _env_msg_handler(env_event:EnvironmentEvent) -> None: - the_msg:AgentMsg= self._env_event_to_msg(env_event) - self.post_msg(the_msg) - - # register all event handler - the_env.attach_event_handler(None,_env_msg_handler) + def connect_to_environment(self,the_env:Environment,conn_info:dict) -> None: + if the_env is not None: + self.workflow_env.add_owner_env(the_env) + + #for event2msg in conn_info: + # for k,v in event2msg: + # if k == "role": + # continue + # else: + # + # def _env_msg_handler(env_event:EnvironmentEvent) -> None: + # the_msg:AgentMsg= self._env_event_to_msg(env_event) + # self.role_post_msg + + # the_env.attach_event_handler(k,_env_msg_handler) + # break else: logger.warn(f"environment {env.get_id()} already connected!") diff --git a/src/aios_kernel/workflow_env.py b/src/aios_kernel/workflow_env.py new file mode 100644 index 0000000..9ef00cb --- /dev/null +++ b/src/aios_kernel/workflow_env.py @@ -0,0 +1,136 @@ + +from datetime import datetime +import asyncio +import sqlite3 # Because sqlite3 IO operation is small, so we can use sqlite3 directly.(so we don't need to use async sqlite3 now) +from sqlite3 import Error +import threading +import logging +from typing import Optional +from .environment import Environment,EnvironmentEvent + +logger = logging.getLogger(__name__) + +class CalenderEvent(EnvironmentEvent): + def __init__(self,data) -> None: + super().__init__() + self.event_name = "timer" + self.data = data + + def display(self) -> str: + return f"#event timer:{self.event_data}" + +# AI Calender GOAL: Let user use "create notify after 2 days" to create a timer event +class CalenderEnvironment(Environment): + def __init__(self, env_id: str) -> None: + super().__init__(env_id) + self.is_run = False + + def _do_get_value(self,key:str) -> Optional[str]: + return None + + def start(self) -> None: + if self.is_run: + return + self.is_run = True + + self.register_get_handler("now",self.get_now) + async def timer_loop(): + while True: + if self.is_run == False: + break + + await asyncio.sleep(1.0) + now = datetime.now() + formatted_time = now.strftime('%Y-%m-%d %H:%M:%S') + env_event:CalenderEvent = CalenderEvent(formatted_time) + await self.fire_event("timer",env_event) + + return + + asyncio.create_task(timer_loop()) + + def stop(self): + self.is_run = False + + def get_now(self,key:str) -> str: + now = datetime.now() + formatted_time = now.strftime('%Y-%m-%d %H:%M:%S') + return formatted_time + +# Default Workflow Environment(Context) +class WorkflowEnvironment(Environment): + def __init__(self, env_id: str,db_file:str) -> None: + super().__init__(env_id) + self.db_file = db_file + self.local = threading.local() + self.table_name = "WorkflowEnv_" + env_id + + def _get_conn(self): + """ get db connection """ + if not hasattr(self.local, 'conn'): + self.local.conn = self._create_connection() + return self.local.conn + + def _create_connection(self): + """ create a database connection to a SQLite database """ + conn = None + try: + conn = sqlite3.connect(self.db_file) + except Error as e: + logging.error("Error occurred while connecting to database: %s", e) + return None + + if conn: + self._create_table(conn) + + return conn + + def close(self): + if not hasattr(self.local, 'conn'): + return + self.local.conn.close() + + def _create_table(self, conn): + """ create table """ + try: + # create sessions table + conn.execute(f""" + CREATE TABLE IF NOT EXISTS """ + self.table_name + """ ( + EnvKey TEXT PRIMARY KEY, + EnvValue TEXT, + UpdateTime TEXT + ); + """) + conn.commit() + except Error as e: + logging.error("Error occurred while creating tables: %s", e) + + def _do_get_value(self, key: str) -> str | None: + try: + conn = self._get_conn() + c = conn.cursor() + c.execute("SELECT EnvValue FROM " + self.table_name +" WHERE EnvKey = ?", (key,)) + value = c.fetchone() + if value is None: + return None + return value[0] + except Error as e: + logging.error(f"Error occurred while _do_get_value{key}: {e}") + return None + + def set_value(self, key: str, str_value: str, is_storage:bool=True): + super().set_value(key,str_value) + if is_storage is False: + return + + try: + conn = self._get_conn() + conn.execute(""" + INSERT OR REPLACE INTO """ + self.table_name+ """ (EnvKey, EnvValue, UpdateTime) + VALUES (?, ?, ?) + """, (key, str_value, datetime.now())) + conn.commit() + return 0 # return 0 if successful + except Error as e: + logging.error(f"Error occurred while update env{self.env_id}.{key} ,error:{e}") + diff --git a/src/aios_kernel/worksapce_env.py b/src/aios_kernel/worksapce_env.py new file mode 100644 index 0000000..e69de29 diff --git a/src/service/aios_shell/aios_shell.py b/src/service/aios_shell/aios_shell.py index 5236e7e..053c5ec 100644 --- a/src/service/aios_shell/aios_shell.py +++ b/src/service/aios_shell/aios_shell.py @@ -31,7 +31,7 @@ from aios_kernel import Workflow,AIAgent,AgentMsg,AgentMsgState,ComputeKernel,Op sys.path.append(directory + '/../../component/') from agent_manager import AgentManager from workflow_manager import WorkflowManager - +from aios_kernel import CalenderEnvironment,Environment class AIOS_Shell: @@ -62,6 +62,10 @@ class AIOS_Shell: return False async def initial(self) -> bool: + cal_env = CalenderEnvironment("calender") + cal_env.start() + Environment.set_env_by_id("calender",cal_env) + AgentManager().initial(os.path.abspath(directory + "/../../../rootfs/")) WorkflowManager().initial(os.path.abspath(directory + "/../../../rootfs/workflows/")) open_ai_node = OpenAI_ComputeNode() diff --git a/test/env_test.py b/test/env_test.py new file mode 100644 index 0000000..694acbb --- /dev/null +++ b/test/env_test.py @@ -0,0 +1,32 @@ +import asyncio +import os +import sys + +directory = os.path.dirname(__file__) +sys.path.append(directory + '/../src') + +from aios_kernel import CalenderEnvironment,WorkflowEnvironment + + +async def test_buildin_envs(): + c_env = CalenderEnvironment("calender") + c_env.start() + print(c_env.get_value("now")) + async def show_event(eventid,event): + print(event.data) + c_env.attach_event_handler("timer",show_event) + + w_env = WorkflowEnvironment("workflow",os.path.abspath(directory + "/../rootfs/workflow_env.db")) + w_env.set_value("test","test_aaaa") + print(w_env.get_value("test")) + + await asyncio.sleep(10) + + +if __name__ == "__main__": + #test_rstr = "abc is {abc}" + #values = {"abc":"123"} + #new_str = test_rstr.format_map(values) + #print(new_str) + + asyncio.run(test_buildin_envs()) \ No newline at end of file