Merge pull request #68 from photosssa/MVP

Add local text/image embedding node; Add an agent Mia to use local knowledge

rootfs/agents/Mia/agent.toml

@@ -2,34 +2,22 @@ instance_id = "Mia"
 fullname = "Mia"
 llm_model_name = "gpt-3.5-turbo-16k-0613"
 max_token_size = 16000
-enable_kb = "true"
-#enable_timestamp = "true"
-#owner_prompt = "我是你的主人{name}"
-contact_prompt = "我是你的主人朋友{name}，请回答主人允许回答的问题"
-guest_prompt = "我是你的的主人的客人{name}，请不要回答我的任何问题"
-owner_env = "calender"
+#enable_function =["add_event"]
+#enable_kb = "true"
+enable_timestamp = "true"
+owner_prompt = "我是你的主人{name}"
+contact_prompt = "我是你的朋友{name}"
+owner_env = "knowledge"
 
 [[prompt]]
 role = "system"
 content = """
-你叫Jarvis,是我的超级私人助理。
-你领导一个团队为我服务，团队的成员有：
-Tracy,私人英语老师
-David,私人画家
+你叫Mia,你可以访问我的个人知识库。
 
 ***
-你看到的信息里有的有时会带上时间标签，这是为了让你更好的理解时间。你回复的信息不用创建这个时间标签。
 你在收到我的信息后，按如下规则处理
-1. 如果你认为团队里有人更适合处理该信息，用下面方法转发消息给他们处理
-```
-##/send_msg 成员名字
-消息内容
-```
-2.你可以访问我的Calender,查看我的日程安排。如果你在处理信息的过程中需要修改我的日程安排，请直接用合适的方法修改。
-3.不符合上述规则的信息，请尽力处理。
+1. 在第一次接受到一条信息时，优先尝试用合适的关键字查询去查询知识库。
+2. 如果信息中包含一段知识库的查询结果，尝试用查询结果处理，如果还是不能处理，尝试递增index继续查询。
+3. 如果知识库返回不了结果了，请尽力返回。
 """
 
-
-#3.你可以访问我的个人信息库，当你处理我的信息时，如果需要用到我的个人信息，请先用合适的方法进行查询，然后再基于查询的结果进行进一步处理后再将结果发给我。
-#4.你能根据我的需要对系统进行配置，但在修改任何配置前，请先和我确认。
-

src/aios_kernel/__init__.py

@@ -5,7 +5,7 @@ from .agent import AIAgent,AIAgentTemplete,AgentPrompt
 from .compute_kernel import ComputeKernel,ComputeTask,ComputeTaskResult,ComputeTaskState,ComputeTaskType
 from .compute_node import ComputeNode,LocalComputeNode
 from .open_ai_node import OpenAI_ComputeNode
-from .knowledge_base import KnowledgeBase 
+from .knowledge_base import KnowledgeBase, KnowledgeEnvironment
 from .knowledge_pipeline import KnowledgeEmailSource, KnowledgeDirSource, KnowledgePipline
 from .role import AIRole,AIRoleGroup
 from .workflow import Workflow
@@ -23,6 +23,7 @@ from .text_to_speech_function import TextToSpeechFunction
 from .workspace_env import WorkspaceEnvironment
 from .local_stability_node import Local_Stability_ComputeNode
 from .stability_node import Stability_ComputeNode
+from .local_st_compute_node import LocalSentenceTransformer_Text_ComputeNode,LocalSentenceTransformer_Image_ComputeNode
 
 AIOS_Version = "0.5.1, build 2023-9-27"
 

src/aios_kernel/compute_kernel.py

@@ -148,7 +148,7 @@ class ComputeKernel:
         task_result = await self._send_task(task_req)
 
         if task_req.state == ComputeTaskState.DONE:
-            return task_result.result
+            return task_result.result_str
 
         return "error!"
 

src/aios_kernel/compute_task.py

@@ -2,6 +2,8 @@
 from enum import Enum
 import uuid
 import time
+from typing import Union
+from knowledge import ObjectID
 
 class ComputeTaskResultCode(Enum):
     OK = 0
@@ -25,6 +27,7 @@ class ComputeTaskType(Enum):
     VOICE_2_TEXT = "voice_2_text"
     TEXT_2_VOICE = "text_2_voice"
     TEXT_EMBEDDING ="text_embedding"
+    IMAGE_EMBEDDING ="image_embedding"
 
 
 class ComputeTask:
@@ -60,7 +63,7 @@ class ComputeTask:
         if inner_functions is not None:
             self.params["inner_functions"] = inner_functions
 
-    def set_text_embedding_params(self, input, model_name=None, callchain_id = None):
+    def set_text_embedding_params(self, input: str, model_name=None, callchain_id = None):
         self.task_type = ComputeTaskType.TEXT_EMBEDDING
         self.create_time = time.time()
         self.task_id = uuid.uuid4().hex
@@ -70,6 +73,17 @@ class ComputeTask:
         else:
             self.params["model_name"] = "text-embedding-ada-002"
         self.params["input"] = input
+        
+    def set_image_embedding_params(self, input = Union[ObjectID, bytes], model_name=None, callchain_id = None):
+        self.task_type = ComputeTaskType.IMAGE_EMBEDDING
+        self.create_time = time.time()
+        self.task_id = uuid.uuid4().hex
+        self.callchain_id = callchain_id
+        if model_name is not None:
+            self.params["model_name"] = model_name
+        else:
+            self.params["model_name"] = None
+        self.params["input"] = input
     
     def set_text_2_image_params(self, prompt: str, model_name, negative_prompt="", callchain_id=None):
         self.task_type = ComputeTaskType.TEXT_2_IMAGE

src/aios_kernel/knowledge_base.py

@@ -4,6 +4,8 @@ import logging
 from .agent import AgentPrompt
 from .compute_kernel import ComputeKernel 
 from .storage import AIStorage
+from .environment import Environment
+from .ai_function import SimpleAIFunction
 from knowledge import *
 
 
@@ -20,6 +22,7 @@ class KnowledgeBase:
     def __singleton_init__(self) -> None:
         self.store = KnowledgeStore()
         self.compute_kernel = ComputeKernel.get_instance()
+        self._default_text_model = "all-MiniLM-L6-v2"
 
     async def __embedding_document(self, document: DocumentObject):
         for chunk_id in document.get_chunk_list():
@@ -28,8 +31,8 @@ class KnowledgeBase:
                 raise ValueError(f"text chunk not found: {chunk_id}")
         
             text = chunk.read().decode("utf-8")
-            vector = await self.compute_kernel.do_text_embedding(text)
-            await self.store.get_vector_store("default").insert(vector, chunk_id)
+            vector = await self.compute_kernel.do_text_embedding(text, self._default_text_model)
+            await self.store.get_vector_store(self._default_text_model).insert(vector, chunk_id)
 
     async def __embedding_image(self, image: ImageObject):
         desc = {}
@@ -39,8 +42,8 @@ class KnowledgeBase:
             desc["exif"] = image.get_exif()
         if not not image.get_tags():
             desc["tags"] = image.get_tags()
-        vector = await self.compute_kernel.do_text_embedding(json.dumps(desc))
-        await self.store.get_vector_store("default").insert(vector, image.calculate_id())
+        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())
 
     async def __embedding_video(self, vedio: VideoObject):
         desc = {}
@@ -50,8 +53,8 @@ class KnowledgeBase:
             desc["info"] = vedio.get_info()
         if not not vedio.get_tags():
             desc["tags"] = vedio.get_tags()
-        vector = await self.compute_kernel.do_text_embedding(json.dumps(desc))
-        await self.store.get_vector_store("default").insert(vector, vedio.calculate_id())
+        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, vedio.calculate_id())
 
     async def __embedding_rich_text(self, rich_text: RichTextObject):
         for document_id in rich_text.get_documents().values():
@@ -68,8 +71,8 @@ class KnowledgeBase:
             await self.__embedding_rich_text(rich_text)
 
     async def __embedding_email(self, email: EmailObject):
-        vector = await self.compute_kernel.do_text_embedding(json.dumps(email.get_desc()))
-        await self.store.get_vector_store("default").insert(vector, email.calculate_id())
+        vector = await self.compute_kernel.do_text_embedding(json.dumps(email.get_desc()), self._default_text_model)
+        await self.store.get_vector_store(self._default_text_model).insert(vector, email.calculate_id())
         await self.__embedding_rich_text(email.get_rich_text())
 
 
@@ -159,23 +162,10 @@ class KnowledgeBase:
     async def insert_object(self, object: KnowledgeObject):
         self.store.get_object_store().put_object(object.calculate_id(), object.encode())
         await self.__do_embedding(object)
-
-    async def query_prompt(self, prompt: AgentPrompt):
-        logging.info(f"query_prompt: {prompt}")
-        objects = await self.query_objects(prompt)
-        knowledge_prompt = self.prompt_from_objects(objects)
-        logging.info(f"prompt_from_objects result: {knowledge_prompt.as_str()}")
-       
-        return knowledge_prompt
-
-    async def query_objects(self, prompt: AgentPrompt) -> [ObjectID]:
-        results = []
-        for msg in prompt.messages:
-            if msg["role"] == "user":
-                vector = await self.compute_kernel.do_text_embedding(msg["content"])
-                object_ids = await self.store.get_vector_store("default").query(vector, 10)
-                results.extend(object_ids)
-        return results
+    
+    async def query_objects(self, tokens: str, topk: int) -> [ObjectID]:
+        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)
 
     def __load_object(self, object_id: ObjectID) -> KnowledgeObject:
         if object_id.get_object_type() == ObjectType.Document:
@@ -192,7 +182,7 @@ class KnowledgeBase:
             pass
         
 
-    def prompt_from_objects(self, object_ids: [ObjectID]) -> AgentPrompt:
+    def tokens_from_objects(self, object_ids: [ObjectID]) -> list[str]:
         results = dict()
         for object_id in object_ids:
             parents = self.store.get_relation_store().get_related_root_objects(object_id)
@@ -203,8 +193,7 @@ class KnowledgeBase:
                 results[str(root_object_id)].append(object_id)
             else:
                 results[str(root_object_id)] = [root_object_id, object_id]
-
-        content = "*** I have provided the following known information for your reference with json format:\n"
+        content = ""
         result_desc = []
         for result in results.values():
             # first element in result is the root object
@@ -236,12 +225,28 @@ class KnowledgeBase:
                 else:
                     pass
         content += json.dumps(result_desc)
-        content += ".\n"
+        content += ".\n"  
 
-        prompt = AgentPrompt()
-        prompt.messages.append({"role": "user", "content": content})    
-
-        return prompt
+        return content
 
 
-    
+class KnowledgeEnvironment(Environment):
+    def __init__(self, env_id: str) -> None:
+        super().__init__(env_id)
+
+        query_param = {
+            "tokens": "tokens to query", 
+            "index": "index of query result"
+        }
+        self.add_ai_function(SimpleAIFunction("query_knowledge", 
+                                            "vector query content from local knowledge base",
+                                            self._query, 
+                                            query_param))
+        
+    async def _query(self, tokens: str, index: int=0):
+        object_ids = await KnowledgeBase().query_objects(tokens, 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])

src/aios_kernel/knowledge_pipeline.py

@@ -32,7 +32,7 @@ import requests
 import os
 import toml
 from .storage import AIStorage, UserConfigItem
-from .knowledge_base import KnowledgeBase, ImageObjectBuilder, ObjectID, ObjectType, DocumentObjectBuilder
+from .knowledge_base import KnowledgeBase, ImageObjectBuilder, ObjectID, ObjectType, DocumentObjectBuilder, EmailObjectBuilder, EmailObject
 
 class KnowledgeJournal:
     def __init__(self, source_type: str, source_id: str, item_id: str, object_id: str, timestamp=None):
@@ -53,7 +53,10 @@ class KnowledgeJournal:
                 pass
             return f"Add {object_type} from {os.path.join(self.source_id, self.item_id)}"
         if self.source_type == "email":
-            pass
+            object_id = ObjectID.from_base58(self.object_id)
+            email = EmailObject.decode(KnowledgeBase().store.get_object_store().get_object(object_id))
+            meta = email.get_meta()
+            return f'Add email from {os.path.join(self.source_id)} subject {meta["subject"]}'
 
 
 # init sqlite3 client
@@ -108,7 +111,7 @@ class KnowledgeEmailSource:
         self.config["type"] = "email"
     
     def id(self):
-        "::".join([self.config["imap_server"], self.config["address"]])
+        return self.config["address"]
 
     @classmethod
     def user_config_items(cls):
@@ -121,13 +124,15 @@ class KnowledgeEmailSource:
     @classmethod
     def local_root(cls):
         user_data_dir = AIStorage.get_instance().get_myai_dir()
-        return os.path.abspath(f"{user_data_dir}/email")    
+        return os.path.abspath(f"{user_data_dir}/knowledge/email")    
 
     async def run_once(self):
         # read config from toml file
         # and read from config config.local.toml if exists (config.local.toml is ignored by git)
+        logging.debug(f"knowledge email source {self.id()} run once")
+        filter = "ALL"  
         self.client = self.email_client()
-        await self.read_emails()
+        await self.read_emails(imap_keyword=filter)
 
     def email_client(self) -> imaplib.IMAP4_SSL:
         logging.info(f"read email config from {self.config.get('imap_server')}")
@@ -139,26 +144,37 @@ class KnowledgeEmailSource:
         return client
 
     async def read_emails(self, folder: str = 'INBOX', imap_keyword: str = "UNSEEN"):
+        journal_client = KnowledgeJournalClient()
+        latest_journal = journal_client.latest_journal(self.id())
+        latest_uid = 0 if latest_journal is None else int(latest_journal.item_id)
         self.client.select(folder)
         _, data = self.client.uid('search', None, imap_keyword)
         
         # get email uid list
         email_list = data[0].split()
         logging.info(f"got {len(email_list)} emails")
-        email_list.reverse()
+        journal_client = KnowledgeJournalClient()
         for uid in email_list:
-            if self.check_email_saved(uid):
-                logging.info(f"email uid {uid} already saved")
-            else:
-                self.read_and_save_email(uid)
-                logging.info(f"email uid {uid} saved")
+            _uid = int.from_bytes(uid)
+            if _uid > latest_uid:
+                email_dir = self.check_email_saved(uid)
+                if email_dir is not None:
+                    logging.info(f"email uid {uid} already saved")
+                else:
+                    email_dir = self.read_and_save_email(uid)
+                    logging.info(f"email uid {uid} saved")
+                email_object = EmailObjectBuilder({}, email_dir).build()
+                await KnowledgeBase().insert_object(email_object)
+                journal_client.insert(KnowledgeJournal("email", self.id(), str(int.from_bytes(uid)), str(email_object.calculate_id())))
 
-    def read_and_save_email(self, uid: str):
+
+    def read_and_save_email(self, uid: str) -> str:
         message_parts = "(BODY.PEEK[])"
         _, email_data = self.client.uid('fetch', uid, message_parts)
         mail = mailparser.parse_from_bytes(email_data[0][1])
         logging.info(f"got email subject [{mail.subject}]")
         self.save_email(mail)
+        return self.get_local_dir_name(mail)
 
     def get_local_dir_name(self, mail: mailparser.MailParser) -> str:
         dir =  f"{self.local_root()}/{self.config.get('address')}"
@@ -166,7 +182,7 @@ class KnowledgeEmailSource:
         name = hashlib.md5(name.encode('utf-8')).hexdigest()
         return f"{dir}/{name}"
 
-    def check_email_saved(self, uid: str):
+    def check_email_saved(self, uid: str) -> str:
         message_parts = "(BODY[HEADER])"
         _, email_data = self.client.uid('fetch', uid, message_parts)
         mail = mailparser.parse_from_bytes(email_data[0][1])
@@ -175,8 +191,8 @@ class KnowledgeEmailSource:
         logging.info(f"check email saved {dir}")
         file = f"{dir}/email.txt"
         if os.path.exists(file):
-            return False
-        return False
+            return dir
+        return None
 
     # save email attachment(images)
     def save_email_attachment(self, mail: mailparser.MailParser, email_dir: str):
@@ -205,12 +221,16 @@ class KnowledgeEmailSource:
         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
+        
         if not os.path.exists(email_dir):
             os.makedirs(email_dir)
 
         for img_url in img_urls:
             # keep the original image filename(last of url)
-            img_filename = os.path.join(email_dir, img_url.split('/')[-1])
+            ext = img_url.split('/')[-1].split('.')[-1]
+            img_filename = os.path.join(email_dir, f"{name_count}.{ext}")
+            name_count += 1
             # download image 
             response = requests.get(img_url, stream=True)
             if response.status_code == 200:
@@ -230,7 +250,8 @@ class KnowledgeEmailSource:
         logging.info(f"save email to {email_dir}")
         if not os.path.exists(email_dir):
             os.makedirs(email_dir)
-        with open(f"{email_dir}/email.txt", "w") as f:
+        with open(f"{email_dir}/email.txt", "w", encoding='utf-8') as f:
+            # soup = BeautifulSoup(mail.body, 'html.parser')
             f.write(mail.body)
         with open(f"{email_dir}/meta.json", "w", encoding='utf-8') as f:
             mail_dict = json.loads(mail.mail_json)

src/aios_kernel/local_st_compute_node.py

@@ -0,0 +1,246 @@
+import logging
+import requests
+from typing import Optional, List
+from pydantic import BaseModel
+from typing import Union
+from PIL import Image
+import io
+from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskType
+from .queue_compute_node import Queue_ComputeNode
+from knowledge import ObjectID
+
+logger = logging.getLogger(__name__)
+
+
+
+class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode):
+    # For valid pretrained models, see https://www.sbert.net/docs/pretrained_models.html
+    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
+        super().__init__()
+
+        self.node_id = "local_sentence_transformer_text_embedding_node"
+        self.model_name = model_name
+        self.model = None
+
+    def initial(self) -> bool:
+        logger.info(
+            f"LocalSentenceTransformer_Text_ComputeNode init, model_name: {self.model_name}"
+        )
+
+        assert self.model_name is not None
+        assert self.model is None
+        try:
+            from sentence_transformers import SentenceTransformer
+
+            self.model = SentenceTransformer(self.model_name)
+        except Exception as err:
+            logger.error(f"load model {self.model} failed: {err}")
+            return False
+        self.start()
+        return True
+
+    async def execute_task(
+        self, task: ComputeTask
+    ) -> {
+        "task_type": str,
+        "content": str,
+        "message": str,
+        "state": ComputeTaskState,
+        "error": {
+            "code": int,
+            "message": str,
+        },
+    }:
+        try:
+            # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task: {task}")
+            if task.task_type == ComputeTaskType.TEXT_EMBEDDING:
+                input = task.params["input"]
+                logger.debug(
+                    f"LocalSentenceTransformer_Text_ComputeNode task input: {input}"
+                )
+                sentence_embeddings = self.model.encode(input, show_progress_bar=False).tolist()
+                # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}")
+                return {
+                    "state": ComputeTaskState.DONE,
+                    "content": sentence_embeddings,
+                    "message": None,
+                }
+            else:
+                return {
+                    "state": ComputeTaskState.ERROR,
+                    "error": {"code": -1, "message": "unsupport embedding task type"},
+                }
+        except Exception as err:
+            import traceback
+
+            logger.error(f"{traceback.format_exc()}, error: {err}")
+
+            return {
+                "state": ComputeTaskState.ERROR,
+                "error": {"code": -1, "message": "unknown exception: " + str(err)},
+            }
+
+    def display(self) -> str:
+        return f"LocalSentenceTransformer_Text_ComputeNode: {self.node_id}, {self.model_name}"
+
+    def get_capacity(self):
+        pass
+
+    def is_support(self, task: ComputeTask) -> bool:
+        return task.task_type == ComputeTaskType.TEXT_EMBEDDING and task.params["model_name"] == "all-MiniLM-L6-v2"
+
+    def is_local(self) -> bool:
+        return True
+
+
+class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
+    # For valid pretrained models, see https://www.sbert.net/docs/pretrained_models.html
+    def __init__(
+        self,
+        model_name: str = "clip-ViT-B-32",
+        multi_model_name: str = "clip-ViT-B-32-multilingual-v1",
+    ):
+        super().__init__()
+
+        self.node_id = "local_sentence_transformer_image_embedding_node"
+        self.model_name = model_name
+        self.multi_model_name = multi_model_name
+        self.model = None
+        self.multi_model = None
+
+    def initial(self) -> bool:
+        logger.info(
+            f"LocalSentenceTransformer_Image_ComputeNode init, model_name: {self.model_name} {self.multi_model_name}"
+        )
+
+        assert self.model_name is not None
+        assert self.multi_model_name is not None
+        assert self.model is None
+        assert self.multi_model is None
+
+        try:
+            from sentence_transformers import SentenceTransformer
+
+            self.model = SentenceTransformer(self.model_name)
+            self.multi_model = SentenceTransformer(self.multi_model_name)
+        except Exception as err:
+            logger.error(f"load model {self.model} failed: {err}")
+            return False
+        self.start()
+        return True
+
+    def _load_image(self, source: Union[ObjectID, bytes]) -> Optional[Image]:
+        image_data = None
+        if isinstance(source, ObjectID):
+            from knowledge import KnowledgeStore, ImageObject
+
+            buf = KnowledgeStore().get_object_store().get_object(source)
+            if buf is None:
+                logger.error(f"load image object but not found! {source}")
+                return None
+
+            try:
+                image_obj = ImageObject.decode(buf)
+            except Exception as err:
+                logger.error(f"decode ImageObject from buffer failed: {source}, {err}")
+                return None
+
+            file_size = image_obj.get_file_size()
+            print(f"got image object: {source.to_base58()}, size: {file_size}")
+
+            image_data = (
+                KnowledgeStore()
+                .get_chunk_reader()
+                .read_chunk_list_to_single_bytes(image_obj.get_chunk_list())
+            )
+
+        elif isinstance(source, bytes):
+            image_data = source
+        else:
+            logger.error(f"unsupport image source type: {type(source)}, {source}")
+            return None
+
+        try:
+            img = Image.open(io.BytesIO(image_data))
+
+            return img
+        except Exception as err:
+            logger.error(f"load image from buffer failed: {source}, {err}")
+            return None
+
+    async def execute_task(
+        self, task: ComputeTask
+    ) -> {
+        "task_type": str,
+        "content": str,
+        "message": str,
+        "state": ComputeTaskState,
+        "error": {
+            "code": int,
+            "message": str,
+        },
+    }:
+        try:
+            # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task: {task}")
+            if task.task_type == ComputeTaskType.TEXT_EMBEDDING:
+                input = task.params["input"]
+                logger.debug(
+                    f"LocalSentenceTransformer_Text_ComputeNode task text input: {input}"
+                )
+                sentence_embeddings = self.multi_model.encode(input, show_progress_bar=False).tolist()
+                # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}")
+                return {
+                    "state": ComputeTaskState.DONE,
+                    "content": sentence_embeddings,
+                    "message": None,
+                }
+            elif task.task_type == ComputeTaskType.IMAGE_EMBEDDING:
+                input = task.params["input"]
+                logger.debug(
+                    f"LocalSentenceTransformer_Image_ComputeNode task image input: {input}"
+                )
+
+                img = self._load_image(input)
+                if img is None:
+                    return {
+                        "state": ComputeTaskState.ERROR,
+                        "error": {"code": -1, "message": "load image failed"},
+                    }
+
+                sentence_embeddings = self.model.encode(img)
+
+                # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}")
+                return {
+                    "state": ComputeTaskState.DONE,
+                    "content": sentence_embeddings,
+                    "message": None,
+                }
+            else:
+                return {
+                    "state": ComputeTaskState.ERROR,
+                    "error": {"code": -1, "message": "unsupport embedding task type"},
+                }
+        except Exception as err:
+            import traceback
+
+            logger.error(f"{traceback.format_exc()}, error: {err}")
+
+            return {
+                "state": ComputeTaskState.ERROR,
+                "error": {"code": -1, "message": "unknown exception: " + str(err)},
+            }
+
+    def display(self) -> str:
+        return f"LocalSentenceTransformer_Image_ComputeNode: {self.node_id}, {self.model_name}"
+
+    def get_capacity(self):
+        pass
+
+    def is_support(self, task: ComputeTask) -> bool:
+        return (
+            (task.task_type == ComputeTaskType.TEXT_EMBEDDING and task.params["model_name"] == "clip-ViT-B-32")
+            or task.task_type == ComputeTaskType.IMAGE_EMBEDDING
+        )
+
+    def is_local(self) -> bool:
+        return True

src/aios_kernel/open_ai_node.py

@@ -98,7 +98,7 @@ class OpenAI_ComputeNode(ComputeNode):
                 task.state = ComputeTaskState.DONE
                 result.result_code = ComputeTaskResultCode.OK
                 result.worker_id = self.node_id
-                result.result = resp["data"][0]["embedding"]
+                result.result_str = resp["data"][0]["embedding"]
 
                 return result
             case ComputeTaskType.LLM_COMPLETION:

src/aios_kernel/queue_compute_node.py

@@ -4,7 +4,7 @@ from asyncio import Queue
 import logging
 from abc import abstractmethod
 
-from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
+from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskResultCode, ComputeTaskState, ComputeTaskType
 from .compute_node import ComputeNode
 
 logger = logging.getLogger(__name__)
@@ -13,6 +13,7 @@ class Queue_ComputeNode(ComputeNode):
     def __init__(self):
         super().__init__()
         self.task_queue = Queue()
+        self.is_start = False
 
     @abstractmethod
     async def execute_task(self, task: ComputeTask) -> {
@@ -39,28 +40,32 @@ class Queue_ComputeNode(ComputeNode):
         resp = await self.execute_task(task)
 
         result = ComputeTaskResult()
-        result.set_from_task(task)
-
-        task.state = resp["state"]
-
-        if task.state == ComputeTaskState.ERROR:
-            task.error_str = resp["error"]["message"]
-
 
         result.worker_id = self.node_id
-        result.result_str = resp["content"]
-        result.result_message = resp["message"]
+        task.state = resp["state"]
+        
+        if task.state == ComputeTaskState.ERROR:
+            result.result_code = ComputeTaskResultCode.ERROR
+            task.error_str = resp["error"]["message"]
+        else:
+            result.result_code = ComputeTaskResultCode.OK
+            result.result_str = resp["content"]
+            result.result_message = resp["message"]
+
+        result.set_from_task(task)
 
         return result
 
-    async def start(self):
+    def start(self):
+        if self.is_start is True:
+            return
+        self.is_start = True
+
         async def _run_task_loop():
             while True:
                 task = await self.task_queue.get()
-                logger.info(f"{self.display()} get task: {task.display()}")
-                result = await self._run_task(task)
-                if result is not None:
-                    task.result = result
+                logger.info(f"openai_node get task: {task.display()}")
+                await self._run_task(task)
 
         asyncio.create_task(_run_task_loop())
 

src/knowledge/core_object/document_object.py

@@ -50,11 +50,7 @@ class DocumentObjectBuilder:
         return self
 
     def build(self) -> DocumentObject:
-        chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_text(
-            self.text,
-            1024 * 4,
-            ".?!\n"
-        )
+        chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_text(self.text)
         doc = DocumentObject(self.meta, self.tags, chunk_list)
         doc_id = doc.calculate_id()
         

src/knowledge/core_object/email_object.py

@@ -21,13 +21,13 @@ class EmailObject(KnowledgeObject):
 
         super().__init__(ObjectType.Email, desc, body)
 
-    def get_meta(self):
+    def get_meta(self) -> dict:
         return self.desc["meta"]
 
-    def get_tags(self):
+    def get_tags(self) -> dict:
         return self.desc["tags"]
 
-    def get_rich_text(self):
+    def get_rich_text(self) -> RichTextObject:
         return self.body["content"]
 
 

src/knowledge/core_object/image_object.py

@@ -2,6 +2,7 @@ from ..object import KnowledgeObject
 from ..data import ChunkList, ChunkListWriter
 from ..object import ObjectType
 from .. import KnowledgeStore
+import os
 
 # desc
 #   meta
@@ -13,30 +14,34 @@ from .. import KnowledgeStore
 
 
 class ImageObject(KnowledgeObject):
-    def __init__(self, meta: dict, tags: dict, exif: dict, chunk_list: ChunkList):
+    def __init__(self, meta: dict, tags: dict, exif: dict, file_size: int, chunk_list: ChunkList):
         desc = dict()
         body = dict()
         desc["meta"] = meta
         desc["exif"] = exif
         desc["tags"] = tags
         desc["hash"] = chunk_list.hash.to_base58()
+        desc["file_size"] = file_size
         body["chunk_list"] = chunk_list.chunk_list
 
         super().__init__(ObjectType.Image, desc, body)
 
-    def get_meta(self):
+    def get_meta(self) -> dict:
         return self.desc["meta"]
 
-    def get_exif(self):
+    def get_exif(self) -> dict:
         return self.desc["exif"]
 
-    def get_tags(self):
+    def get_tags(self) -> dict:
         return self.desc["tags"]
     
-    def get_hash(self):
+    def get_hash(self) -> str:
         return self.desc["hash"]
 
-    def get_chunk_list(self):
+    def get_file_size(self) -> int:
+        return self.desc["file_size"]
+    
+    def get_chunk_list(self) -> ChunkList:
         return self.body["chunk_list"]
 
 
@@ -82,8 +87,10 @@ class ImageObjectBuilder:
         return self
 
     def build(self) -> ImageObject:
+        
+        file_size = os.path.getsize(self.image_file)
         chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_file(
             self.image_file, 1024 * 1024 * 4, self.restore_file
         )
         exif = get_exif_data(self.image_file)
-        return ImageObject(self.meta, self.tags, exif, chunk_list)
+        return ImageObject(self.meta, self.tags, exif, file_size, chunk_list)

src/knowledge/data/reader.py

@@ -12,7 +12,7 @@ class Chunk:
         self.range_start = range_start
         self.size = size
 
-    def read(self):
+    def read(self) -> bytes:
         with open(self.file_path, 'rb') as f:
             f.seek(self.range_start)
             return f.read(self.size)
@@ -26,6 +26,8 @@ class ChunkReader:
 
     def get_chunk(self, chunk_id: ChunkID) -> Chunk:
         positions = self.chunk_tracker.get_position(chunk_id)
+        logging.info(f"chunk positions: {chunk_id}, {positions}")
+        
         if positions is None:
             logging.warning(f"chunk not found: {chunk_id}")
             return None 
@@ -54,15 +56,23 @@ class ChunkReader:
     def get_chunk_list(self, chunk_list: List[ChunkID]) -> List[Chunk]:
         return [self.get_chunk(chunk_id) for chunk_id in chunk_list]
         
-    def read_chunk_list(self, chunk_ids: List[ChunkID]):
+    def read_chunk_list(self, chunk_ids: List[ChunkID]) -> bytes:
         for chunk_id in chunk_ids:
             chunk = self.get_chunk(chunk_id)
             if chunk is None:
                 raise ValueError(f"chunk not found: {chunk_id}")
             
-            yield from chunk.read()
+            yield chunk.read()
     
-    def read_text_chunk_list(self, chunk_ids: List[ChunkID]):
+    def read_chunk_list_to_single_bytes(self, chunk_ids: List[ChunkID]) -> bytes:
+        chunks = []
+        for chunk in self.read_chunk_list(chunk_ids):
+            chunks.append(chunk)
+            
+        image_data = b''.join(chunks)
+        return image_data
+    
+    def read_text_chunk_list(self, chunk_ids: List[ChunkID]) -> str:
         for chunk_id in chunk_ids:
             chunk = self.get_chunk(chunk_id)
             if chunk is None:

src/knowledge/data/writer.py

@@ -1,13 +1,139 @@
 import os
 import hashlib
 import re
-from typing import Tuple, List
+import tiktoken
+import logging
+from typing import Callable, Iterable, Optional, Tuple, List
 from .chunk_store import ChunkStore
 from .chunk import ChunkID, PositionFileRange, PositionType
 from ..object import HashValue
 from .tracker import ChunkTracker
 from .chunk_list import ChunkList
 
+def _join_docs(docs: List[str], separator: str) -> Optional[str]:
+    text = separator.join(docs)
+    text = text.strip()
+    if text == "":
+        return None
+    else:
+        return text
+
+def _merge_splits(
+        splits: Iterable[str], 
+        separator: str, 
+        chunk_size: int, 
+        chunk_overlap: int,  
+        length_function: Callable[[str], int]
+    ) -> List[str]:
+    # We now want to combine these smaller pieces into medium size
+    # chunks to send to the LLM.
+    separator_len = length_function(separator)
+
+    docs = []
+    current_doc: List[str] = []
+    total = 0
+    for d in splits:
+        _len = length_function(d)
+        if (
+            total + _len + (separator_len if len(current_doc) > 0 else 0)
+            > chunk_size
+        ):
+            if total > chunk_size:
+                logging.warning(
+                    f"Created a chunk of size {total}, "
+                    f"which is longer than the specified {self._chunk_size}"
+                )
+            if len(current_doc) > 0:
+                doc = _join_docs(current_doc, separator)
+                if doc is not None:
+                    docs.append(doc)
+                # Keep on popping if:
+                # - we have a larger chunk than in the chunk overlap
+                # - or if we still have any chunks and the length is long
+                while total > chunk_overlap or (
+                    total + _len + (separator_len if len(current_doc) > 0 else 0)
+                    > chunk_size
+                    and total > 0
+                ):
+                    total -= length_function(current_doc[0]) + (
+                        separator_len if len(current_doc) > 1 else 0
+                    )
+                    current_doc = current_doc[1:]
+        current_doc.append(d)
+        total += _len + (separator_len if len(current_doc) > 1 else 0)
+    doc = _join_docs(current_doc, separator)
+    if doc is not None:
+        docs.append(doc)
+    return docs
+
+
+def _split_text_with_regex(
+    text: str, separator: str, keep_separator: bool
+) -> List[str]:
+    # Now that we have the separator, split the text
+    if separator:
+        if keep_separator:
+            # The parentheses in the pattern keep the delimiters in the result.
+            _splits = re.split(f"({separator})", text)
+            splits = [_splits[i] + _splits[i + 1] for i in range(1, len(_splits), 2)]
+            if len(_splits) % 2 == 0:
+                splits += _splits[-1:]
+            splits = [_splits[0]] + splits
+        else:
+            splits = re.split(separator, text)
+    else:
+        splits = list(text)
+    return [s for s in splits if s != ""]
+
+
+def _split_text(
+        text: str, 
+        separators: List[str], 
+        chunk_size: int, 
+        chunk_overlap: int, 
+        length_function: Callable[[str], int]
+    ) -> List[str]:
+
+    """Split incoming text and return chunks."""
+    final_chunks = []
+    # Get appropriate separator to use
+    separator = separators[-1]
+    new_separators = []
+    for i, _s in enumerate(separators):
+        _separator = re.escape(_s)
+        if _s == "":
+            separator = _s
+            break
+        if re.search(_separator, text):
+            separator = _s
+            new_separators = separators[i + 1 :]
+            break
+
+    keep_separator = True
+    _separator = re.escape(separator)
+    splits = _split_text_with_regex(text, _separator, keep_separator)
+
+    # Now go merging things, recursively splitting longer texts.
+    _good_splits = []
+    _separator = "" if keep_separator else separator
+    for s in splits:
+        if length_function(s) < chunk_size:
+            _good_splits.append(s)
+        else:
+            if _good_splits:
+                merged_text = _merge_splits(_good_splits, _separator, chunk_size, chunk_overlap, length_function)
+                final_chunks.extend(merged_text)
+                _good_splits = []
+            if not new_separators:
+                final_chunks.append(s)
+            else:
+                other_info = _split_text(s, new_separators, chunk_size, chunk_overlap, length_function)
+                final_chunks.extend(other_info)
+    if _good_splits:
+        merged_text = _merge_splits(_good_splits, _separator, chunk_size, chunk_overlap, length_function)
+        final_chunks.extend(merged_text)
+    return final_chunks
+
 class ChunkListWriter:
     def __init__(self, chunk_store: ChunkStore, chunk_tracker: ChunkTracker):
         self.chunk_store = chunk_store
@@ -27,6 +153,8 @@ class ChunkListWriter:
                 chunk = file.read(chunk_size)
                 if not chunk:
                     break
+                
+                chunk_len = len(chunk)
                 chunk_id = ChunkID.hash_data(chunk)
                 chunk_list.append(chunk_id)
 
@@ -38,8 +166,9 @@ class ChunkListWriter:
                     )
                     self.chunk_store.put_chunk(chunk_id, chunk)
                 else:
+                    pos = file.tell()
                     file_range = PositionFileRange(
-                        file_path, file.tell() - chunk_size, chunk_size
+                        file_path, pos - chunk_len, pos
                     )
                     self.chunk_tracker.add_position(
                         chunk_id, str(file_range), PositionType.FileRange
@@ -51,9 +180,24 @@ class ChunkListWriter:
         return ChunkList(chunk_list, file_hash)
 
     def create_chunk_list_from_text(
-        self, text: str, chunk_max_words: int, separator_chars: str = ".,"
+        self, 
+        text: str, 
+        chunk_size: int = 4000, 
+        chunk_overlap: int = 200, 
+        separators: str = ["\n\n", "\n", " ", ""]
     ) -> ChunkList:
-        text_list = self._split_text_list(text, chunk_max_words, separator_chars)
+        enc = tiktoken.encoding_for_model("gpt-3.5-turbo")
+
+        def length_function(text: str) -> int:
+            return len(
+                enc.encode(
+                    text,
+                    allowed_special=set(),
+                    disallowed_special="all",
+                )
+            )
+        
+        text_list = _split_text(text, separators, chunk_size, chunk_overlap, length_function)
         chunk_list = []
         hash_obj = hashlib.sha256()
 
@@ -67,27 +211,4 @@ class ChunkListWriter:
             self.chunk_store.put_chunk(chunk_id, chunk_bytes)
 
         hash = HashValue(hash_obj.digest())
-        return ChunkList(chunk_list, hash)
-
-    @staticmethod
-    def _split_text_list(
-        text: str, chunk_max_words: int, separator_chars: str = ".,"
-    ) -> List[str]:
-        sentences = re.split(f"[{separator_chars}]", text)
-        # chunk_list = []
-        # chunk = []
-        # word_count = 0
-        # for sentence in sentences:
-        #     words = sentence.split()
-        #     for word in words:
-        #         if word_count < chunk_max_words:
-        #             chunk.append(word)
-        #             word_count += 1
-        #         else:
-        #             chunk_list.append(" ".join(chunk))
-        #             chunk = [word]
-        #             word_count = 1
-        # if chunk:
-        #     chunk_list.append(" ".join(chunk))
-        # return chunk_list
-        return sentences
+        return ChunkList(chunk_list, hash)

src/knowledge/object.py

@@ -1,65 +0,0 @@
-
-# define a object type enum
-from abc import ABC, abstractmethod
-from enum import Enum
-
-class ObjectType(Enum):
-    TextChunk = 1
-    Image = 2
-    Email = 101
-
-
-# define a object ID class to identify a object
-class ObjectID: # pylint: disable=too-few-public-methods
-    def __init__(self, object_type, digist):
-        self.object_type = object_type
-        self.digist = digist
-
-    def __str__(self):
-        return f"{self.object_type.name}:{self.digist}"
-    
-
-# define a object class
-class KnowledgeObject(ABC): # pylint: disable=too-few-public-methods
-    def __init__(self, object_type: ObjectType):
-        self.object_type = object_type
-
-    @abstractmethod
-    def get_id(self) -> ObjectID:
-        pass
-
-    # define a to binary method to convert object to binary
-    @abstractmethod
-    def to_binary(self) -> bytes:
-        pass
-
-    # define a from binary method to convert binary to object
-    @abstractmethod
-    def from_binary(self, binary: bytes):
-        pass
-
-
-# define a text chunk class
-class TextChunkObject(KnowledgeObject): # pylint: disable=too-few-public-methods
-    def __init__(self, text: str):
-        super().__init__(ObjectType.TextChunk)
-        self.text = text
-
-
-# define a image class
-class ImageObject(KnowledgeObject): # pylint: disable=too-few-public-methods
-    def __init__(self, meta, path):
-        super().__init__(ObjectType.Image)
-        self.meta = meta
-        self.path = path
-
-
-# define a email class
-class EmailObject(KnowledgeObject): # pylint: disable=too-few-public-methods
-    def __init__(self, meta):
-        super().__init__(ObjectType.Email)
-        self.meta = meta
-        self.text = [ObjectID]
-        self.images = [ObjectID]
-
-    

src/knowledge/object/blob.py

@@ -1,6 +1,8 @@
 import os
 import shutil
 from .object import ObjectID
+import logging
+logger = logging.getLogger(__name__)
 
 
 class FileBlobStorage:
@@ -38,16 +40,23 @@ class FileBlobStorage:
 
     def put(self, object_id: ObjectID, contents: bytes):
         full_path = self.get_full_path(object_id)
+        if os.path.exists(full_path):
+            logger.warning(f"will replace object: {object_id}")
+            
         self.write_sync(full_path, contents)
 
     def get(self, object_id: ObjectID) -> bytes:
         full_path = self.get_full_path(object_id)
+        if not os.path.exists(full_path):
+            return None
+        
         with open(full_path, "rb") as f:
             return f.read()
 
     def delete(self, object_id: ObjectID):
         full_path = self.get_full_path(object_id)
-        os.remove(full_path)
+        if os.path.exists(full_path):
+            os.remove(full_path)
 
     def exists(self, object_id: ObjectID) -> bool:
         full_path = self.get_full_path(object_id)

src/knowledge/object/object.py

@@ -1,6 +1,8 @@
 # define a object type enum
+from __future__ import annotations
 from abc import ABC, abstractmethod
 from enum import Enum
+
 from .object_id import ObjectID, ObjectType
 import hashlib
 import json
@@ -61,5 +63,5 @@ class KnowledgeObject(ABC):
         return pickle.dumps(self)
 
     @staticmethod
-    def decode(data: bytes):
+    def decode(data: bytes) -> "ImageObject":
         return pickle.loads(data)

src/knowledge/object_storage.py

@@ -1,33 +0,0 @@
-# import RDB LargeBinary
-from sqlalchemy import Column, String, LargeBinary, create_engine, sessionmaker, pickle
-from .object import KnowledgeObject
-
-# implement object storage with RDB
-# define object storage table
-class ObjectStorageTable(Base):
-    __tablename__ = 'object_storage'
-    id = Column(String, primary_key=True)
-    parent = Column(String, nullable=True)
-    object = Column(LargeBinary, nullable=False)
-
-    def __init__(self, id, parent, object): # pylint: disable=redefined-builtin
-        self.id = id
-        self.parent = parent
-        self.object = object
-
-# define object storage class
-class ObjectStorage:
-    async def __init__(self, db_url):
-        self.engine = create_engine(db_url)
-        self.session = sessionmaker(bind=self.engine)() # pylint: disable=not-callable
-
-    async def get(self, id) -> [KnowledgeObject, KnowledgeObject]:
-        obj = self.session.query(ObjectStorageTable).filter(ObjectStorageTable.id == id).first()
-        if obj is None:
-            return None
-        return pickle.loads(obj.object)
-    
-    # define insert method
-    async def insert(self, object, parent): # pylint: disable=redefined-builtin
-        obj = ObjectStorageTable(id, parent, pickle.dumps(object))
-        

src/knowledge/store.py

@@ -4,7 +4,7 @@ from .object import ObjectStore, ObjectRelationStore
 from .data import ChunkStore, ChunkTracker, ChunkListWriter, ChunkReader
 from .vector import ChromaVectorStore, VectorBase
 import logging
-import aios_kernel
+
 
 # KnowledgeStore class, which aggregates ChunkStore, ChunkTracker, and ObjectStore, and is a global singleton that makes it easy to use these three built-in store examples
 class KnowledgeStore:
@@ -13,6 +13,8 @@ class KnowledgeStore:
     def __new__(cls):
         if cls._instance is None:
             cls._instance = super().__new__(cls)
+            
+            import aios_kernel
             knowledge_dir = aios_kernel.storage.AIStorage().get_myai_dir() / "knowledge"
 
             if not os.path.exists(knowledge_dir):
@@ -60,5 +62,5 @@ class KnowledgeStore:
     
     def get_vector_store(self, model_name: str) -> VectorBase:
         if model_name not in self.vector_store:
-            self.vector_store[model_name] = ChromaVectorStore(model_name)
+            self.vector_store[model_name] = ChromaVectorStore(self.root, model_name)
         return self.vector_store[model_name]

src/requirements.txt

@@ -80,10 +80,61 @@ toml>=0.10.0
 protobuf
 grpcio
 grpcio-status
+h11==0.14.0
+httpcore==0.17.3
+httptools==0.6.0
+httpx==0.24.1
+huggingface-hub==0.16.4
+humanfriendly==10.0
+idna==3.4
+imageio==2.31.3
+imageio-ffmpeg==0.4.8
+importlib-resources==6.0.1
+mail-parser==3.15.0
+monotonic==1.6
+moviepy==1.0.0
+mpmath==1.3.0
+multidict==6.0.4
+numpy==1.25.2
+onnxruntime==1.15.1
+openai==0.28.0
+overrides==7.4.0
+packaging==23.1
+pandas==2.1.0
+Pillow==10.0.0
+posthog==3.0.2
+proglog==0.1.10
+prompt-toolkit==3.0.39
+proto-plus==1.22.3
+protobuf
+pulsar-client==3.3.0
+pyasn1==0.5.0
+pyasn1-modules==0.3.0
+pydantic==1.10.12
+PyPika==0.48.9
+pyreadline3==3.4.1
+python-dateutil==2.8.2
+python-dotenv==1.0.0
+python-telegram-bot==20.5
+pytz==2023.3.post1
+PyYAML==6.0.1
+requests==2.31.0
+rsa==4.9
+simplejson==3.19.1
+six==1.16.0
+sniffio==1.3.0
+soupsieve==2.5
+starlette==0.27.0
+sympy==1.12
+telegram==0.0.1
+tokenizers==0.14.0
+toml==0.10.0
 pysocks
 chardet
 pydub
 aiosqlite
 python-telegram-bot
 pydub
-stability_sdk
+stability_sdk
+sentence-transformers==2.2.2
+tiktoken

src/service/aios_shell/aios_shell.py

@@ -24,8 +24,7 @@ directory = os.path.dirname(__file__)
 sys.path.append(directory + '/../../')
 
 
-from aios_kernel import AIOS_Version,AgentMsgType,UserConfigItem,AIStorage,Workflow,AIAgent,AgentMsg,AgentMsgStatus,ComputeKernel,OpenAI_ComputeNode,AIBus,AIChatSession,AgentTunnel,TelegramTunnel,CalenderEnvironment,Environment,EmailTunnel,LocalLlama_ComputeNode,Local_Stability_ComputeNode,Stability_ComputeNode,PaintEnvironment
-from aios_kernel import ContactManager,Contact
+
 import proxy
 from aios_kernel import *
 
@@ -114,6 +113,10 @@ class AIOS_Shell:
 
             cm.add_family_member(self.username,owenr)
 
+        knowledge_env = KnowledgeEnvironment("knowledge")
+        Environment.set_env_by_id("knowledge",knowledge_env)
+
+
         cal_env = CalenderEnvironment("calender")
         await cal_env.start()
         Environment.set_env_by_id("calender",cal_env)
@@ -137,6 +140,7 @@ class AIOS_Shell:
             return False
         ComputeKernel.get_instance().add_compute_node(open_ai_node)
 
+
         if await AIStorage.get_instance().is_feature_enable("llama"):
             llama_ai_node = LocalLlama_ComputeNode()
             if await llama_ai_node.initial() is True:
@@ -146,6 +150,7 @@ class AIOS_Shell:
                 logger.error("llama node initial failed!")
                 await AIStorage.get_instance().set_feature_init_result("llama",False)
 
+
         if await AIStorage.get_instance().is_feature_enable("aigc"):
             try:
                 google_text_to_speech_node = GoogleTextToSpeechNode.get_instance()
@@ -160,13 +165,20 @@ class AIOS_Shell:
             #     logger.error("stability api node initial failed!")
             # ComputeKernel.get_instance().add_compute_node(stability_api_node)
 
-            local_sd_node = Local_Stability_ComputeNode.get_instance()
-            if await local_sd_node.initial() is True:
-                ComputeKernel.get_instance().add_compute_node(local_sd_node)
-            else:
-                logger.error("local stability node initial failed!")
-                await AIStorage.get_instance.set_feature_init_result("aigc",False)
-
+        
+        
+        local_st_text_compute_node = LocalSentenceTransformer_Text_ComputeNode()
+        if local_st_text_compute_node.initial() is not True:
+            logger.error("local sentence transformer text embedding node initial failed!")
+        else:
+            ComputeKernel.get_instance().add_compute_node(local_st_text_compute_node)
+            
+        local_st_image_compute_node = LocalSentenceTransformer_Image_ComputeNode()
+        if local_st_image_compute_node.initial() is not True:
+            logger.error("local sentence transformer image embedding node initial failed!")
+        else:
+            ComputeKernel.get_instance().add_compute_node(local_st_image_compute_node)
+       
 
         await ComputeKernel.get_instance().start()
 
@@ -308,8 +320,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 journal [$topn]\n"
-                              "/knowledge query $query\n")])
+                              "/knowledge journal [$topn]\n")])
         if len(args) < 1:
             return show_text
         sub_cmd = args[0]
@@ -346,13 +357,6 @@ class AIOS_Shell:
             topn = 10 if len(args) == 1 else int(args[1])
             journals = [str(journal) for journal in KnowledgePipline.get_instance().get_latest_journals(topn)]
             print_formatted_text("\r\n".join(journals))
-        if sub_cmd == "query":
-            if len(args) < 2:
-                return show_text
-            prompt = AgentPrompt()
-            prompt.messages.append({"role": "user", "content":" ".join(args[1:])})
-            result = await KnowledgeBase().query_prompt(prompt)
-            print_formatted_text(result.as_str())
 
     async def call_func(self,func_name, args):
         match func_name:
@@ -662,8 +666,7 @@ async def main():
                                '/connect $target',
                                '/contact $name',
                                '/knowledge add email | dir',
-                               '/knowledge journal [$topn]',
-                               '/knowledge query $query' 
+                               '/knowledge journal [$topn]', 
                                '/set_config $key',
                                '/enable $feature',
                                '/disable $feature',

test/test_chunk.py

@@ -59,7 +59,7 @@ class TestChunk(unittest.TestCase):
         with open(text_file, "r", encoding="utf-8") as file:
             text = file.read()
 
-        gen.create_chunk_list_from_text(text, 1024)
+        gen.create_chunk_list_from_text(text)
 
 
 if __name__ == "__main__":

test/test_embedding.py

@@ -0,0 +1,98 @@
+import sys
+import os
+import logging
+from sentence_transformers import SentenceTransformer, util
+
+
+dir_path = os.path.dirname(os.path.realpath(__file__))
+print(dir_path)
+
+sys.path.append("{}/../src/".format(dir_path))
+print(sys.path)
+
+root = logging.getLogger()
+root.setLevel(logging.DEBUG)
+handler = logging.StreamHandler(sys.stdout)
+handler.setLevel(logging.DEBUG)
+formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+handler.setFormatter(formatter)
+root.addHandler(handler)
+
+
+def test_st():
+    image_model = SentenceTransformer('clip-ViT-B-32-multilingual-v1')
+    model = SentenceTransformer("all-MiniLM-L6-v2")
+
+    # Our sentences we like to encode
+    sentences = [
+        "This framework generates embeddings for each input sentence",
+        "Sentences are passed as a list of string.",
+        "The quick brown fox jumps over the lazy dog.",
+    ]
+
+    # Sentences are encoded by calling model.encode()
+    sentence_embeddings = model.encode(sentences)
+
+    # Print the embeddings
+    for sentence, embedding in zip(sentences, sentence_embeddings):
+        print("Sentence:", sentence)
+        print("Embedding:", embedding)
+        print("")
+
+    # Single list of sentences
+
+    """
+    sentences = [
+        "The cat sits outside",
+        "A man is playing guitar",
+        "I love pasta",
+        "The new movie is awesome",
+        "The cat plays in the garden",
+        "A woman watches TV",
+        "The new movie is so great",
+        "Do you like pizza?",
+    ]
+    """
+    sentences = [
+        "猫坐在外面",
+        "狗坐在上面",
+        "狗坐在里面",
+        "一个男人在弹吉他",
+        "我爱意大利面",
+        "新电影太精彩了",
+        "猫在花园里玩耍",
+        "一个女人在看电视",
+        "新电影太棒了",
+        "你喜欢披萨吗？",
+    ]
+
+    # Compute embeddings
+    #embeddings = model.encode(sentences, convert_to_tensor=True)
+    embeddings = model.encode(sentences)
+    print("embeddings as follows: ")
+    print(embeddings)
+
+
+    # Compute cosine-similarities for each sentence with each other sentence
+    cosine_scores = util.cos_sim(embeddings, embeddings)
+
+    # Find the pairs with the highest cosine similarity scores
+    pairs = []
+    for i in range(len(cosine_scores) - 1):
+        for j in range(i + 1, len(cosine_scores)):
+            pairs.append({"index": [i, j], "score": cosine_scores[i][j]})
+
+    # Sort scores in decreasing order
+    pairs = sorted(pairs, key=lambda x: x["score"], reverse=True)
+
+    for pair in pairs[0:10]:
+        i, j = pair["index"]
+        print(
+            "{} \t\t {} \t\t Score: {:.4f}".format(
+                sentences[i], sentences[j], pair["score"]
+            )
+        )
+
+
+if __name__ == "__main__":
+    test_st()

test/test_object.py

@@ -24,13 +24,15 @@ from knowledge import (
     ObjectRelationStore,
     KnowledgeStore,
     EmailObject,
+    ImageObject,
 )
+from aios_kernel import LocalSentenceTransformer_Image_ComputeNode, ComputeTask
 import asyncio
 import unittest
 
 
-class TestVectorSTorage(unittest.TestCase):
-    def test_object(self):
+class TestVectorSTorage(unittest.IsolatedAsyncioTestCase):
+    async def test_object(self):
         data = HashValue.hash_data("1233".encode("utf-8"))
         print(data.to_base58())
         print(data.to_base36())
@@ -57,6 +59,41 @@ class TestVectorSTorage(unittest.TestCase):
         ret2 = obj.encode()
         self.assertEqual(ret, ret2)
         
+        images = email_object.get_rich_text().get_images()
+        image_keys = list(images.keys())
+        print("got image list: ", image_keys)
+        
+        image_id = images[image_keys[1]]
+        print(f"got image object: {image_keys[1]} {image_id.to_base58()}")
+        
+        node = LocalSentenceTransformer_Image_ComputeNode();
+        ret = node.initial()
+        self.assertEqual(ret, True)
+        
+        task = ComputeTask()
+        task.set_image_embedding_params(image_id)
+        ret = await node.execute_task(task)
+        print(ret)
+        '''
+        buf = KnowledgeStore().get_object_store().get_object(image_id)
+        image_obj= ImageObject.decode(buf)
+        file_size = image_obj.get_file_size()
+        print(f"got image object: {image_id.to_base58()}, size: {file_size}")
+        
+     
+        image_data = KnowledgeStore().get_chunk_reader().read_chunk_list_to_single_bytes(image_obj.get_chunk_list())
+        self.assertEqual(file_size, len(image_data))
+        
+        from PIL import Image
+        import io
+        image = Image.open(io.BytesIO(image_data))
+        image.show()
+        
+        from sentence_transformers import SentenceTransformer
+        #model = SentenceTransformer('clip-ViT-B-32-multilingual-v1')
+        model = SentenceTransformer('clip-ViT-B-32')
+        model.encode(image, convert_to_tensor=True)
+        '''
 
     def test_relation(self):
         obj1 = ObjectID.hash_data("12345".encode("utf-8"))