Add feature: Contact manager

rootfs/agents/Jarvis/agent.toml

@@ -3,7 +3,11 @@ fullname = "Jarvis"
 llm_model_name = "gpt-3.5-turbo-16k-0613"
 max_token_size = 16000
 #enable_function =["add_event"]
-enable_kb = "true"
+#enable_kb = "true"
+enable_timestamp = "true"
+owner_prompt = "我是你的主人{name}"
+contact_prompt = "我是你的朋友{name}"
+guest_prompt = "xxx"
 
 [[prompt]]
 role = "system"
@@ -15,6 +19,7 @@ Tracy,私人英语老师
 David,私人画家
 
 ***
+你看到的信息里有的有时会带上时间标签，这是为了让你更好的理解时间。你回复的信息不用创建这个时间标签。
 你在收到我的信息后，按如下规则处理
 1. 如果你认为团队里有人更适合处理该信息，用下面方法转发消息给他们处理
 ```

rootfs/agents/linux_admin/agent.toml

@@ -0,0 +1,12 @@
+instance_id = "linux_admin"
+fullname = "linux_admin"
+
+[[prompt]]
+role = "system"
+content = """你的名字是linux_admin,是非常资深的linux系统管理员，我理解一些linux,但对linux的sh命令记得不太清楚了。
+我给你的输入有两种
+1. 标准的linux bash命令，如果你认为这些命令在当前系统上是正确的，那么直接返回这些命令，我会去执行
+2. 如果你为这些命令需要调整后才能执行，先把你修正后的命令告诉我，我确认后再执行
+3. 如果我给你的信息不是linux bash操作，而是一些需求。你可以尝试理解后，给出一组实现这些需求的命令，我确认后再执行
+4. 其它信息，请站在你专业角度尽力执行。
+"""

src/aios_kernel/__init__.py

@@ -18,6 +18,7 @@ from .tunnel import AgentTunnel
 from .tg_tunnel import TelegramTunnel
 from .email_tunnel import EmailTunnel
 from .storage import ResourceLocation,AIStorage,UserConfig,UserConfigItem
+from .contact_manager import ContactManager,Contact,FamilyMember  
 
 AIOS_Version = "0.5.1, build 2023-9-17"
 

src/aios_kernel/agent.py

@@ -7,18 +7,22 @@ import uuid
 import time
 import json
 import shlex
+import datetime
 
 from .agent_message import AgentMsg, AgentMsgStatus, AgentMsgType,FunctionItem,LLMResult
 from .chatsession import AIChatSession
 from .compute_task import ComputeTaskResult
 from .ai_function import AIFunction
 from .environment import Environment
+from .contact_manager import ContactManager,Contact,FamilyMember
 
 logger = logging.getLogger(__name__)
 
 class AgentPrompt:
-    def __init__(self) -> None:
+    def __init__(self,prompt_str = None) -> None:
         self.messages = []
+        if prompt_str:
+            self.messages.append({"role":"user","content":prompt_str})
         self.system_message = None
 
     def as_str(self)->str:
@@ -110,6 +114,10 @@ class AIAgent:
         self.powerby = None  
         self.enable = True
         self.enable_kb = False
+        self.enable_timestamp = False
+        self.guest_prompt_str = None 
+        self.owner_promp_str = None
+        self.contact_prompt_str = None
 
         self.chat_db = None
         self.unread_msg = Queue() # msg from other agent
@@ -145,6 +153,16 @@ class AIAgent:
             self.prompt = AgentPrompt()
             self.prompt.load_from_config(config["prompt"])
 
+        if config.get("guest_prompt") is not None:
+            self.guest_prompt_str = config["guest_prompt"]
+
+        if config.get("owner_prompt") is not None:
+            self.owner_promp_str = config["owner_prompt"]
+        
+        if config.get("contact_prompt") is not None:
+            self.contact_prompt_str = config["contact_prompt"]
+
+
         if config.get("powerby") is not None:
             self.powerby = config["powerby"]
         if config.get("template_id") is not None:
@@ -157,6 +175,8 @@ class AIAgent:
             self.enable_function_list = config["enable_function"]
         if config.get("enable_kb") is not None:
             self.enable_kb = bool(config["enable_kb"])
+        if config.get("enable_timestamp") is not None:
+            self.enable_timestamp = bool(config["enable_timestamp"])
         return True
 
 
@@ -230,6 +250,32 @@ class AIAgent:
 
         return r
     
+    def _get_remote_user_prompt(self,remote_user:str) -> AgentPrompt:
+        cm = ContactManager.get_instance()
+        contact = cm.find_contact_by_name(remote_user)
+        if contact is None:
+            #create guest prompt
+            if self.guest_prompt_str is not None:
+                prompt = AgentPrompt()
+                prompt.system_message = {"role":"system","content":self.guest_prompt_str}
+                return prompt
+            return None
+        else:
+            if contact.is_family_member:
+                if self.owner_promp_str is not None:
+                    real_str = self.owner_promp_str.format_map(contact.to_dict())
+                    prompt = AgentPrompt()
+                    prompt.system_message = {"role":"system","content":real_str}
+                    return prompt
+            else:
+                if self.contact_prompt_str is not None:
+                    real_str = self.contact_prompt_str.format_map(contact.to_dict())
+                    prompt = AgentPrompt()
+                    prompt.system_message = {"role":"system","content":real_str}
+                    return prompt
+                
+        return None
+
     def _get_inner_functions(self) -> dict:
         if self.owner_env is None:
             return None
@@ -314,8 +360,6 @@ class AIAgent:
             from .compute_kernel import ComputeKernel
             from .bus import AIBus
             
-
-
             session_topic = msg.get_sender() + "#" + msg.topic
             chatsession = AIChatSession.get_session(self.agent_id,session_topic,self.chat_db)
             if msg.mentions is not None:
@@ -330,6 +374,7 @@ class AIAgent:
             prompt = AgentPrompt()
             prompt.append(await self._get_agent_prompt())
             self._format_msg_by_env_value(prompt)
+            prompt.append(self._get_remote_user_prompt(msg.sender))
 
             inner_functions,function_token_len = self._get_inner_functions()
        
@@ -406,9 +451,19 @@ class AIAgent:
         read_history_msg = 0
         for msg in reversed(messages):
             read_history_msg += 1
+            dt = datetime.datetime.fromtimestamp(float(msg.create_time))
+            formatted_time = dt.strftime('%m-%d %H:%M:%S')
+
             if msg.sender == self.agent_id:
+
+                if self.enable_timestamp:
+                    result_prompt.messages.append({"role":"assistant","content":f"(create on {formatted_time}) {msg.body} "})
+                else:
                     result_prompt.messages.append({"role":"assistant","content":msg.body})
                 
+            else:
+                if self.enable_timestamp:
+                    result_prompt.messages.append({"role":"user","content":f"(create on {formatted_time}) {msg.body} "})
                 else:
                     result_prompt.messages.append({"role":"user","content":msg.body})
 

src/aios_kernel/contact_manager.py

@@ -1,34 +1,137 @@
 from typing import List
+import toml
 
 class Contact:
-    def __init__(self,name:str) -> None:
+    def __init__(self, name, phone=None, email=None, telegram=None,added_by=None, tags=[], notes=""):
         self.name = name
-        self.tags = []
+        self.phone = phone
+        self.email = email
+        self.telegram = telegram
+        self.added_by = added_by
+        self.tags = tags
+        self.notes = notes
 
-    def is_zone_owner(self,zone_id=None) -> bool:
-        return True
+    def to_dict(self):
+        return {
+            "name": self.name,
+            "phone": self.phone,
+            "email": self.email,
+            "telegram" : self.telegram,
 
-    def get_tags(self)->List[str]:
-        return self.tags
+            "added_by": self.added_by,
+            "tags": self.tags,
+            "notes": self.notes
+        }
 
-    def get_name(self)->str:
-        return self.name
+    @classmethod
+    def from_dict(cls, data):
+        return Contact(data.get("name"), data.get("phone"), data.get("email"), data.get("telegram"),data.get("added_by"), data.get("tags"), data.get("notes"))
 
+class FamilyMember(Contact):
+    def __init__(self, name, relationship,phone=None, email=None,telegram=None):
+        super().__init__(name, phone, email, telegram)
+        self.name = name
+        self.relationship = relationship  
+        self.is_family_member = True
+
+    def to_dict(self):
+        result = super().to_dict()
+        result["relationship"] = self.relationship
+        return result
+
+    @classmethod
+    def from_dict(cls, data):
+        return FamilyMember(data.get("name"),data.get("relationship"),data.get("phone"), data.get("email"),data.get("telegram"))
 
 class ContactManager:
     _instance = None
     @classmethod
-    def get_instance(cls):
+    def get_instance(cls,filename=None) -> "ContactManager":
         if cls._instance is None:
-            cls._instance = ContactManager()
+            cls._instance = ContactManager(filename)
         return cls._instance
 
-    def __init__(self) -> None:
-        self.contacts = {}
-        self.contacts["liuzhicong"] = Contact("liuzhicong")
+    def __init__(self, filename="contacts.toml"):
+        self.filename = filename
+        self.contacts = []
+        self.family_members = []    
 
-    #def get_by_addr(self,addr:str) -> Contact:
-    #    pass
+        self.is_auto_create_contact_from_telegram = True
     
-    def get_by_name(self,name:str) -> Contact:
-        return self.contacts.get(name)
+    def load_data(self):
+        try:
+            with open(self.filename, "r") as f:
+                config =  toml.load(f)
+                return self.load_from_config(config)
+        except FileNotFoundError:
+            return {}
+        
+    def load_from_config(self,config_data:dict):
+        self.contacts = [Contact.from_dict(item) for item in config_data.get("contacts", [])]
+        self.family_members = [FamilyMember.from_dict(item) for item in config_data.get("family_members", [])]
+    
+    def save_data(self):
+        data = {
+            "contacts": [contact.to_dict() for contact in self.contacts],
+            "family_members": [member.to_dict() for member in self.family_members]
+        }
+        with open(self.filename, "w") as f:
+            toml.dump(data, f)
+
+    def add_contact(self, name:str, new_contact:Contact):
+        assert name == new_contact.name
+        self.contacts.append(new_contact)
+        self.save_data()
+
+    def remove_contact(self, name:str):
+        self.contacts = [contact for contact in self.contacts if contact.name != name]
+        self.save_data()
+
+    def find_contact_by_name(self, name:str):
+        for contact in self.contacts:
+            if contact.name == name:
+                return contact
+            
+        for member in self.family_members:
+            if member.name == name:
+                return member
+        return None
+    
+    def find_contact_by_telegram(self, telegram:str):
+        for contact in self.contacts:
+            if contact.telegram == telegram:
+                return contact
+        for member in self.family_members:
+            if member.telegram == telegram:
+                return member            
+        return None
+
+    def find_contact_by_email(self, email:str):
+        for contact in self.contacts:
+            if contact.email == email:
+                return contact
+        for member in self.family_members:
+            if member.email == email:
+                return member   
+        return None
+
+    def find_contact_by_phone(self, phone:str):
+        for contact in self.contacts:
+            if contact.phone == phone:
+                return contact
+        for member in self.family_members:
+            if member.phone == phone:
+                return member   
+        return None
+
+
+    def add_family_member(self, name, new_member:FamilyMember):
+        assert name == new_member.name
+        self.family_members.append(new_member)
+        self.save_data()
+
+    def list_contacts(self):
+        return self.contacts
+
+    def list_family_members(self):
+        return self.family_members

src/aios_kernel/tg_tunnel.py

@@ -14,9 +14,10 @@ from telegram.ext import Updater
 from telegram.error import Forbidden, NetworkError
 
 from .tunnel import AgentTunnel
-from .contact_manager import ContactManager
+from .contact_manager import ContactManager,Contact,FamilyMember
 from .agent_message import AgentMsg
 
+
 logger = logging.getLogger(__name__)
 
 class TelegramTunnel(AgentTunnel):
@@ -116,16 +117,31 @@ class TelegramTunnel(AgentTunnel):
 
 
     async def on_message(self, bot:Bot, update: Update) -> None:
-        cm = ContactManager.get_instance()
-        reomte_user_name = f"{update.effective_user.id}@telegram"
-        #contact = cm.get_by_name(update.effective_user.username)
-        #if contact is not None:
-        #    reomte_user_name = contact.get_name()
-        #if contact is None:
-        #    update.message.reply_text(f"{self.target_id} process message error, unknown user!")
-        #if not contact.is_zone_owner():
-        #    update.message.reply_text(f"{self.target_id} process message error, you are not my owner!")
+        if update.effective_user.is_bot:
+            logger.warning(f"ignore message from telegram bot {update.effective_user.id}")
+            return None
         
+        cm : ContactManager = ContactManager.get_instance()
+        reomte_user_name = f"{update.effective_user.id}@telegram"
+        contact : Contact = cm.find_contact_by_telegram(update.effective_user.username)
+        if contact is None:
+            contact = cm.find_contact_by_telegram(str(update.effective_user.id))
+
+        if contact is not None:
+            reomte_user_name = contact.name
+        else:
+            if cm.is_auto_create_contact_from_telegram:
+                contact_name = update.effective_user.first_name
+                if update.effective_user.last_name is not None:
+                    contact_name += " " + update.effective_user.last_name
+
+                contact = Contact(contact_name)
+                contact.telegram = update.effective_user.username if update.effective_user.username is not None else str(update.effective_user.id)
+                contact.added_by = self.target_id
+                cm.add_contact(contact.name, contact)
+                reomte_user_name = contact.name
+
+        # create gust account?
         agent_msg = await self.conver_tg_msg_to_agent_msg(update)
         agent_msg.sender = reomte_user_name
         self.ai_bus.register_message_handler(reomte_user_name, self._process_message)

src/aios_kernel/workflow.py

@@ -390,7 +390,7 @@ class Workflow:
             return result_func
         return None 
     
-    async def _role_execute_func(self,the_role:AIRole,inenr_func_call_node:dict,prompt:AgentPrompt,org_msg:AgentMsg) -> str:
+    async def _role_execute_func(self,the_role:AIRole,inenr_func_call_node:dict,prompt:AgentPrompt,org_msg:AgentMsg,stack_limit = 5) -> str:
         from .compute_kernel import ComputeKernel
       
         func_name = inenr_func_call_node.get("name")
@@ -413,10 +413,10 @@ class Workflow:
         ineternal_call_record.result_str = task_result.result_str
         ineternal_call_record.done_time = time.time()
         org_msg.inner_call_chain.append(ineternal_call_record)
-
+        if stack_limit > 0:
             inner_func_call_node = task_result.result_message.get("function_call")
         if inner_func_call_node:
-            return await self._role_execute_func(the_role,inner_func_call_node,prompt,org_msg)      
+            return await self._role_execute_func(the_role,inner_func_call_node,prompt,org_msg,stack_limit-1)      
         else:
             return task_result.result_str
     

src/service/aios_shell/aios_shell.py

@@ -27,7 +27,6 @@ import proxy
 from aios_kernel import *
 
 
-
 sys.path.append(directory + '/../../component/')
 from agent_manager import AgentManager
 from workflow_manager import WorkflowManager
@@ -88,6 +87,7 @@ class AIOS_Shell:
         await cal_env.start()
         Environment.set_env_by_id("calender",cal_env)
 
+
         await AgentManager.get_instance().initial()
         await WorkflowManager.get_instance().initial()
 
@@ -110,6 +110,10 @@ class AIOS_Shell:
         EmailTunnel.register_to_loader()
 
         user_data_dir = AIStorage.get_instance().get_myai_dir()
+        contact_config_path =os.path.abspath(f"{user_data_dir}/contacts.toml")
+        cm = ContactManager.get_instance(contact_config_path)
+        cm.load_data()
+        
         tunnels_config_path = os.path.abspath(f"{user_data_dir}/etc/tunnels.cfg.toml")
         tunnel_config = None
         try: 
@@ -139,8 +143,6 @@ class AIOS_Shell:
     async def _user_process_msg(self,msg:AgentMsg) -> AgentMsg:
         pass
 
-
-
     async def get_tunnel_config_from_input(self,tunnel_target,tunnel_type):
         tunnel_config = {}
         tunnel_config["tunnel_id"] = f"{tunnel_type}_2_{tunnel_target}"

test/test_contac_manager.py

@@ -0,0 +1,64 @@
+import unittest
+import os
+import toml
+import sys
+
+directory = os.path.dirname(__file__)
+sys.path.append(directory + '/../src')
+from aios_kernel import ContactManager, Contact, FamilyMember
+
+class TestContactManager(unittest.TestCase):
+
+    def setUp(self):
+        self.manager = ContactManager(filename="test_contacts.toml")
+        self.manager.load_data()
+
+    def tearDown(self):
+        if os.path.exists("test_contacts.toml"):
+            os.remove("test_contacts.toml")
+
+    def test_add_family_member(self):
+        new_member = FamilyMember("Alice", "123-456-7890", "sdfsd","alice@example.com")
+        self.manager.add_family_member("Alice", new_member)
+        members = self.manager.list_family_members()
+        self.assertEqual(len(members), 1)
+        self.assertEqual(members[0].name, "Alice")
+
+    def test_add_contact(self):
+        new_contact = Contact("Bob", "987-654-3210", "bob@example.com", "32323","Alice", ["Friend"], "Bob is Alice's friend.")
+        self.manager.add_contact("Bob", new_contact)
+        contacts = self.manager.list_contacts()
+        self.assertEqual(len(contacts), 1)
+        self.assertEqual(contacts[0].name, "Bob")
+        self.assertEqual(contacts[0].added_by, "Alice")
+
+    def test_remove_contact(self):
+        new_contact = Contact("Bob", "987-654-3210", "bob@example.com", "32323","Alice", ["Friend"], "Bob is Alice's friend.")
+        self.manager.add_contact("Bob", new_contact)
+        self.manager.remove_contact("Bob")
+        contacts = self.manager.list_contacts()
+        self.assertEqual(len(contacts), 0)
+
+    def test_find_contact_by_name(self):
+        new_contact = Contact("Bob", "987-654-3210", "bob@example.com", "32323","Alice", ["Friend"], "Bob is Alice's friend.")
+        self.manager.add_contact("Bob", new_contact)
+        contact = self.manager.find_contact_by_name("Bob")
+        self.assertIsNotNone(contact)
+        self.assertEqual(contact.name, "Bob")
+
+    def test_find_contact_by_email(self):
+        new_contact = Contact("Bob", "987-654-3210", "bob@example.com", "32323","Alice", ["Friend"], "Bob is Alice's friend.")
+        self.manager.add_contact("Bob", new_contact)
+        contact = self.manager.find_contact_by_email("bob@example.com")
+        self.assertIsNotNone(contact)
+        self.assertEqual(contact.email, "bob@example.com")
+
+    def test_find_contact_by_phone(self):
+        new_contact = Contact("Bob", "987-654-3210", "bob@example.com", "32323","Alice", ["Friend"], "Bob is Alice's friend.")
+        self.manager.add_contact("Bob", new_contact)
+        contact = self.manager.find_contact_by_phone("987-654-3210")
+        self.assertIsNotNone(contact)
+        self.assertEqual(contact.phone, "987-654-3210")
+
+if __name__ == '__main__':
+    unittest.main()