Jarvis can query user knowlege base!

src/aios_kernel/agent.py

@@ -109,6 +109,7 @@ class AIAgent:
         self.fullname:str = None
         self.powerby = None  
         self.enable = True
+        self.enable_kb = False
 
         self.chat_db = None
         self.unread_msg = Queue() # msg from other agent
@@ -154,6 +155,8 @@ class AIAgent:
             self.max_token_size = config["max_token_size"]
         if config.get("enable_function") is not None:
             self.enable_function_list = config["enable_function"]
+        if config.get("enable_kb") is not None:
+            self.enable_kb = bool(config["enable_kb"])
         return True
 
 
@@ -300,9 +303,18 @@ class AIAgent:
             old_content = msg.get("content")
             msg["content"] = old_content.format_map(self.owner_env)
 
+    async def _get_knowlege_prompt(self,input_msg:AgentPrompt) -> AgentPrompt:
+        if self.enable_kb is False:
+            return None
+        
+        from .knowledge_base import KnowledgeBase
+        return await KnowledgeBase().query_prompt(input_msg)
+
     async def _process_msg(self,msg:AgentMsg) -> AgentMsg:
             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)
@@ -311,22 +323,27 @@ class AIAgent:
                     chatsession.append(msg)
                     logger.info(f"agent {self.agent_id} recv a group chat message from {msg.sender},but is not mentioned,ignore!")
                     return None
-            
+                
+            msg_prompt = AgentPrompt()
+            msg_prompt.messages = [{"role":"user","content":msg.body}]
+
             prompt = AgentPrompt()
             prompt.append(await self._get_agent_prompt())
+            self._format_msg_by_env_value(prompt)
+
             inner_functions,function_token_len = self._get_inner_functions()
-            # prompt.append(self._get_knowlege_prompt(the_role.get_name()))
+       
             system_prompt_len = prompt.get_prompt_token_len()
             input_len = len(msg.body)
             
             history_prmpt,history_token_len = await self._get_prompt_from_session(chatsession,system_prompt_len + function_token_len,input_len)
             prompt.append(history_prmpt) # chat context
             
-            msg_prompt = AgentPrompt()
-            msg_prompt.messages = [{"role":"user","content":msg.body}]
+            kb_prompt = await self._get_knowlege_prompt(msg_prompt)
+            prompt.append(kb_prompt)
             prompt.append(msg_prompt)
 
-            self._format_msg_by_env_value(prompt)
+
             logger.debug(f"Agent {self.agent_id} do llm token static system:{system_prompt_len},function:{function_token_len},history:{history_token_len},input:{input_len}，totoal prompt:{system_prompt_len + function_token_len + history_token_len} ")
             task_result:ComputeTaskResult = await ComputeKernel.get_instance().do_llm_completion(prompt,self.llm_model_name,self.max_token_size,inner_functions)
             final_result = task_result.result_str

src/aios_kernel/knowledge_base.py

@@ -165,8 +165,8 @@ class KnowledgeBase:
         objects = await self.query_objects(prompt)
         knowledge_prompt = self.prompt_from_objects(objects)
         logging.info(f"prompt_from_objects result: {knowledge_prompt.as_str()}")
-        prompt.append(knowledge_prompt)
-        return prompt
+       
+        return knowledge_prompt
 
     async def query_objects(self, prompt: AgentPrompt) -> [ObjectID]:
         results = []
@@ -204,7 +204,7 @@ class KnowledgeBase:
             else:
                 results[str(root_object_id)] = [root_object_id, object_id]
 
-        content = "I found the following contents described with json format:\n"
+        content = "*** I have provided the following known information for your reference with json format:\n"
         result_desc = []
         for result in results.values():
             # first element in result is the root object
@@ -239,7 +239,7 @@ class KnowledgeBase:
         content += ".\n"
 
         prompt = AgentPrompt()
-        prompt.messages.append({"role": "knowledge", "content": content})    
+        prompt.messages.append({"role": "user", "content": content})    
 
         return prompt
 

src/aios_kernel/knowledge_pipeline.py

@@ -24,6 +24,9 @@ import mailparser
 import hashlib
 import json
 import base64
+import chardet
+import aiofiles
+
 from bs4 import BeautifulSoup
 import requests
 import os
@@ -254,7 +257,16 @@ class KnowledgeDirSource:
 
     def path(self):
         return self.config["path"]
-
+    
+    @staticmethod
+    async def read_txt_file(file_path:str)->str:
+        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,'r',encoding=cur_encode) as f:
+            return await f.read()
+        
     async def run_once(self):
         logging.debug(f"knowledge dir source {self.id()} run once")
         journal_client = KnowledgeJournalClient()
@@ -278,11 +290,11 @@ class KnowledgeDirSource:
                 journal_client.insert(KnowledgeJournal("dir", self.id(), rel_path, str(image.calculate_id()), timestamp))
             if ext in ['.txt']:
                 logging.info(f"knowledge dir source {self.id()} found text file {file_path}")
-                with open(file_path, "r", encoding="utf-8") as f:
-                    text = f.read()
-                    document = DocumentObjectBuilder({}, {}, text).build()
-                    await KnowledgeBase().insert_object(document)
-                    journal_client.insert(KnowledgeJournal("dir", self.id(), rel_path, str(document.calculate_id()), timestamp))
+                text = await self.read_txt_file(file_path)
+    
+                document = DocumentObjectBuilder({}, {}, text).build()
+                await KnowledgeBase().insert_object(document)
+                journal_client.insert(KnowledgeJournal("dir", self.id(), rel_path, str(document.calculate_id()), timestamp))
             
 
 

src/requirements.txt

@@ -80,4 +80,5 @@ starlette==0.27.0
 sympy==1.12
 telegram==0.0.1
 tokenizers==0.14.0
-toml==0.10.0
+toml==0.10.0
+chardet