Merge pull request #69 from photosssa/MVP

Add image insert/query in knowledge base

rootfs/agents/Mia/agent.toml

@@ -1,10 +1,10 @@
 instance_id = "Mia"
 fullname = "Mia"
-llm_model_name = "gpt-3.5-turbo-16k-0613"
+llm_model_name = "gpt-4"
 max_token_size = 16000
 #enable_function =["add_event"]
 #enable_kb = "true"
-enable_timestamp = "true"
+enable_timestamp = "false"
 owner_prompt = "我是你的主人{name}"
 contact_prompt = "我是你的朋友{name}"
 owner_env = "knowledge"
@@ -18,6 +18,7 @@ content = """
 你在收到我的信息后，按如下规则处理
 1. 在第一次接受到一条信息时，优先尝试用合适的关键字查询去查询知识库。
 2. 如果信息中包含一段知识库的查询结果，尝试用查询结果处理，如果还是不能处理，尝试递增index继续查询。
-3. 如果知识库返回不了结果了，请尽力返回。
+3. 如果要返回知识库结果条目，在消息开头附上他的json字符串。
+4. 如果知识库返回不了结果了，请尽力返回。
 """
 

src/aios_kernel/agent.py

@@ -358,13 +358,6 @@ 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
@@ -393,11 +386,6 @@ class AIAgent:
             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
 
-            kb_prompt = await self._get_knowlege_prompt(msg_prompt)
-            prompt.append(kb_prompt)
-            prompt.append(msg_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)
             if task_result.result_code != ComputeTaskResultCode.OK:

src/aios_kernel/compute_kernel.py

@@ -5,6 +5,7 @@ import logging
 import asyncio
 from asyncio import Queue
 
+from knowledge import ObjectID
 from .agent import AgentPrompt
 from .compute_node import ComputeNode
 from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskResult, ComputeTaskType,ComputeTaskResultCode
@@ -152,6 +153,21 @@ class ComputeKernel:
 
         return "error!"
 
+    def image_embedding(self,input:ObjectID,model_name:Optional[str] = None):
+        task_req = ComputeTask()
+        task_req.set_image_embedding_params(input,model_name)
+        self.run(task_req)
+        return task_req
+    
+    async def do_image_embedding(self,input:ObjectID,model_name:Optional[str] = None) -> [float]:
+        task_req = self.image_embedding(input,model_name)
+        task_result = await self._send_task(task_req)
+
+        if task_req.state == ComputeTaskState.DONE:
+            return task_result.result_str
+
+        return "error!"
+
     async def do_text_to_speech(self,
                        input:str,
                        language_code:Optional[str] = None,

src/aios_kernel/knowledge_base.py

@@ -23,6 +23,7 @@ class KnowledgeBase:
         self.store = KnowledgeStore()
         self.compute_kernel = ComputeKernel.get_instance()
         self._default_text_model = "all-MiniLM-L6-v2"
+        self._default_image_model = "clip-ViT-B-32"
 
     async def __embedding_document(self, document: DocumentObject):
         for chunk_id in document.get_chunk_list():
@@ -35,15 +36,16 @@ class KnowledgeBase:
             await self.store.get_vector_store(self._default_text_model).insert(vector, chunk_id)
 
     async def __embedding_image(self, image: ImageObject):
-        desc = {}
+        # desc = {}
-        if not not image.get_meta():
+        # if not not image.get_meta():
-            desc["meta"] = image.get_meta()
+        #     desc["meta"] = image.get_meta()
-        if not not image.get_exif():
+        # if not not image.get_exif():
-            desc["exif"] = image.get_exif()
+        #     desc["exif"] = image.get_exif()
-        if not not image.get_tags():
+        # if not not image.get_tags():
-            desc["tags"] = image.get_tags()
+        #     desc["tags"] = image.get_tags()
-        vector = await self.compute_kernel.do_text_embedding(json.dumps(desc), self._default_text_model)
+        # vector = await self.compute_kernel.do_text_embedding(json.dumps(desc), self._default_text_model)
-        await self.store.get_vector_store(self._default_text_model).insert(vector, image.calculate_id())
+        vector = await self.compute_kernel.do_image_embedding(image.calculate_id(), self._default_image_model)
+        await self.store.get_vector_store(self._default_image_model).insert(vector, image.calculate_id())
 
     async def __embedding_video(self, vedio: VideoObject):
         desc = {}
@@ -163,9 +165,16 @@ class KnowledgeBase:
         self.store.get_object_store().put_object(object.calculate_id(), object.encode())
         await self.__do_embedding(object)
     
-    async def query_objects(self, tokens: str, topk: int) -> [ObjectID]:
+    async def query_objects(self, tokens: str, types: list[str], topk: int) -> [ObjectID]:
+        texts = []
+        if "text" in types:
             vector = await self.compute_kernel.do_text_embedding(tokens, self._default_text_model)
-        return await self.store.get_vector_store(self._default_text_model).query(vector, topk)
+            texts = await self.store.get_vector_store(self._default_text_model).query(vector, topk)
+        images = []
+        if "image" in types:
+            vector = await self.compute_kernel.do_text_embedding(tokens, self._default_image_model)
+            images = await self.store.get_vector_store(self._default_image_model).query(vector, topk)
+        return texts + images
 
     def __load_object(self, object_id: ObjectID) -> KnowledgeObject:
         if object_id.get_object_type() == ObjectType.Document:
@@ -213,8 +222,9 @@ class KnowledgeBase:
                 if object_id.get_object_type() == ObjectType.Chunk:
                     upper_list.append({"type": "text", "content": self.store.get_chunk_reader().get_chunk(object_id).read().decode("utf-8")})
                 if object_id.get_object_type() == ObjectType.Image:
-                    image = self.__load_object(object_id)
+                    # image = self.__load_object(object_id)
-                    desc = image.get_desc()
+                    desc = dict()
+                    desc["id"] = str(object_id)
                     desc["type"] = "image"
                     upper_list.append(desc)
                 if object_id.get_object_type() == ObjectType.Video:
@@ -229,13 +239,39 @@ class KnowledgeBase:
 
         return content
     
+    def parse_object_in_message(self, message: str) -> KnowledgeObject:
+        # get message's first line 
+        lines = message.split("\n")
+        if len(lines) > 0:
+            message = lines[0]
+            try:
+                desc = json.loads(message)
+                object_id = desc["object_id"]
+            except:
+                return None
+            
+            if object_id is not None:
+                return self.__load_object(ObjectID(object_id))
+            
+            
+    def bytes_from_object(self, object: KnowledgeObject) -> bytes:
+        if object.get_object_type() == ObjectType.Image:
+            image_object = object
+            return self.store.get_chunk_reader().read_chunk_list_to_single_bytes(image_object.get_chunk_list())
+           
+
+
+           
+
+
 
 class KnowledgeEnvironment(Environment):
     def __init__(self, env_id: str) -> None:
         super().__init__(env_id)
 
         query_param = {
-            "tokens": "tokens to query", 
+            "tokens": "key words to query", 
+            "types": "prefered knowledge types, one or more of [text, image]",
             "index": "index of query result"
         }
         self.add_ai_function(SimpleAIFunction("query_knowledge", 
@@ -243,10 +279,10 @@ class KnowledgeEnvironment(Environment):
                                             self._query, 
                                             query_param))
         
-    async def _query(self, tokens: str, index: int=0):
+    async def _query(self, tokens: str, types: list[str] = ["text"], index: int=0):
-        object_ids = await KnowledgeBase().query_objects(tokens, 4)
+        object_ids = await KnowledgeBase().query_objects(tokens, types, 4)
         if len(object_ids) <= index:
             return "*** I have no more information for your reference.\n"
         else:
             content = "*** I have provided the following known information for your reference with json format:\n"
-            return content + KnowledgeBase().tokens_from_objects(object_ids[index:index + 1])
+            return content + KnowledgeBase().tokens_from_objects(object_ids[index:index+1])

src/aios_kernel/knowledge_pipeline.py

@@ -267,6 +267,7 @@ class KnowledgeEmailSource:
 class KnowledgeDirSource:
     def __init__(self, config):
         self.config = config
+        config["path"] = os.path.abspath(config["path"])
         self.config["type"] = "dir"
 
     @classmethod
@@ -342,6 +343,11 @@ class KnowledgePipline:
                     self.add_email_source(KnowledgeEmailSource(source_config))
                 if source_config['type'] == 'dir':
                     self.add_dir_source(KnowledgeDirSource(source_config))
+        user_data_dir = AIStorage.get_instance().get_myai_dir()
+        default_dir = os.path.abspath(f"{user_data_dir}/data")
+        if not os.path.exists(default_dir):
+            os.makedirs(default_dir)
+        self.add_dir_source(KnowledgeDirSource({"path": default_dir}))
 
         return True
 

src/aios_kernel/local_st_compute_node.py

@@ -146,7 +146,7 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
                 return None
 
             file_size = image_obj.get_file_size()
-            print(f"got image object: {source.to_base58()}, size: {file_size}")
+            # print(f"got image object: {source.to_base58()}, size: {file_size}")
 
             image_data = (
                 KnowledgeStore()
@@ -207,7 +207,7 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
                         "error": {"code": -1, "message": "load image failed"},
                     }
 
-                sentence_embeddings = self.model.encode(img)
+                sentence_embeddings = self.model.encode(img, show_progress_bar=False).tolist()
 
                 # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}")
                 return {

src/aios_kernel/tg_tunnel.py

@@ -9,6 +9,10 @@ from telegram import Bot
 from telegram.ext import Updater
 from telegram.error import Forbidden, NetworkError
 
+from knowledge.object.object_id import ObjectType
+
+from .knowledge_base import KnowledgeBase
+
 from .tunnel import AgentTunnel
 from .storage import AIStorage
 from .contact_manager import ContactManager,Contact,FamilyMember
@@ -171,6 +175,13 @@ class TelegramTunnel(AgentTunnel):
         else:
             if resp_msg.body_mime is None:
                 if resp_msg.body is not None:
+                    knowledge_object = KnowledgeBase().parse_object_in_message(resp_msg.body)
+                    if knowledge_object is not None:
+                        if knowledge_object.get_object_type() == ObjectType.Image:
+                            image = KnowledgeBase().bytes_from_object(knowledge_object)
+                            await update.message.reply_photo(image)
+                            return
+                    else:
                         pos = resp_msg.body.find("audio file")
                         if pos != -1:
                             audio_file = resp_msg.body[pos+11:].strip()

src/service/aios_shell/aios_shell.py

@@ -319,7 +319,7 @@ class AIOS_Shell:
     
     async def handle_knowledge_commands(self, args):
         show_text = FormattedText([("class:title", "sub command not support!\n" 
-                              "/knowledge add email | dir\n"
+                              "/knowledge add dir\n"
                               "/knowledge journal [$topn]\n")])
         if len(args) < 1:
             return show_text
@@ -585,7 +585,7 @@ def print_welcome_screen():
 \033[1;94m\tGive your Agent a Telegram account :\033[0m /connect $agent_name
 \033[1;94m\tAdd personal files to the AI Knowledge Base. \033[0m
 \t\t1) Copy your file to ~/myai/data 
-\t\t2) /knowlege add dir
+\t\t2) /knowledge add dir
 \033[1;94m\tSearch your knowledge base :\033[0m /open Mia
 \033[1;94m\tCheck the progress of AI reading personal data :\033[0m /knowledge journal
 \033[1;94m\tOpen AI Bash (For Developer Only):\033[0m /open ai_bash
@@ -665,7 +665,7 @@ async def main():
                                '/history $num $offset',
                                '/connect $target',
                                '/contact $name',
-                               '/knowledge add email | dir',
+                               '/knowledge add dir',
                                '/knowledge journal [$topn]', 
                                '/set_config $key',
                                '/enable $feature',