embedding email object

src/aios_kernel/__init__.py

@@ -4,7 +4,8 @@ from .chatsession import AIChatSession
 from .agent import AIAgent,AIAgentTemplete,AgentPrompt
 from .compute_kernel import ComputeKernel,ComputeTask
 from .compute_node import ComputeNode,LocalComputeNode
-from .open_ai_node import OpenAI_ComputeNode 
+from .open_ai_node import OpenAI_ComputeNode
+from .knowledge_base import KnowledgeBase 
 from .role import AIRole,AIRoleGroup
 from .workflow import Workflow
 from .bus import AIBus

src/aios_kernel/knowledge_base.py

@@ -1,47 +1,96 @@
 # define a knowledge base class
+import json
 from . import AgentPrompt, ComputeKernel
-from ..knowledge.object import KnowledgeObject, ObjectType, EmailObject, TextChunkObject, ImageObject
+from ..knowledge import *
-from ..knowledge.store import ObjectStorage
+
-from ..knowledge.vector.vector_base import VectorBase
 
 class KnowledgeBase:
-    def __init__(self) -> None:
+    _instance = None
-        self.object_store = ObjectStorage()
+
-        self.vector_base = VectorBase()
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            cls._instance.__singleton_init__()
+
+        return cls._instance
+
+    def __singleton_init__(self) -> None:
+        self.store = KnowledgeStore()
         self.compute_kernel = ComputeKernel()
 
-    async def insert(self, object: KnowledgeObject):
+    async def __embedding_document(self, document: DocumentObject):
-        if object.object_type == ObjectType.Email:
+        for chunk_id in document.get_chunk_list():
-            email: EmailObject = object
+            chunk = self.store.get_chunk_reader().get_chunk(chunk_id)
-            for text_id in email.text:
+            if chunk is None:
-                [text, _] = self.object_store.get(text_id)
+                raise ValueError(f"text chunk not found: {chunk_id}")
-                text: TextChunkObject = text
+        
-                vector = await self.compute_kernel.do_text_embedding(text.text)
+            text = chunk.read().decode("utf-8")
-                self.vector_base.insert(vector, text_id)
+            vector = await self.compute_kernel.do_text_embedding(text)
-            
+            self.store.get_vector_store("default").insert(vector, chunk_id)
-            for image_id in email.images:
+
-                [image, _] = self.object_store.get(image_id)
+    async def __embedding_image(self, image: ImageObject):
-                image: ImageObject = image
+        desc = {}
-                vector = await self.compute_kernel.do_text_embedding(image.meta)
+        if not image.get_meta():
-                self.vector_base.insert(vector, image_id)
+            desc["meta"] = image.get_meta()
-            
+        if not image.get_exif():
-            vector = await self.compute_kernel.do_text_embedding(email.meta)
+            desc["exif"] = image.get_exif()
-            self.vector_base.insert(vector, email.get_id())
+        if not image.get_tags():
+            desc["tags"] = image.get_tags()
+        vector = await self.compute_kernel.do_text_embedding(json.dumps(desc))
+        self.store.get_vector_store("default").insert(vector, image.calculate_id())
+
+    async def __embedding_vedio(self, vedio: VideoObject):
+        desc = {}
+        if not vedio.get_meta():
+            desc["meta"] = vedio.get_meta()
+        if not vedio.get_info():
+            desc["info"] = vedio.get_info()
+        if not vedio.get_tags():
+            desc["tags"] = vedio.get_tags()
+        vector = await self.compute_kernel.do_text_embedding(json.dumps(desc))
+        self.store.get_vector_store("default").insert(vector, vedio.calculate_id())
+
+    async def __embedding_rich_text(self, rich_text: RichTextObject):
+        for document in rich_text.get_documents().values():
+            await self.__embedding_document(document)
+        for image in rich_text.get_images().values():
+            await self.__embedding_image(image)
+        for vedio in rich_text.get_videos().values():
+            await self.__embedding_vedio(vedio)
+        for rich_text in rich_text.get_rich_texts().values():
+            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()))
+        self.store.get_vector_store("default").insert(vector, email.calculate_id())
+        await self.__embedding_rich_text(email.get_rich_text())
+
+    async def do_embedding(self, object: KnowledgeObject):
+        if object.get_object_type() == ObjectType.Document:
+            await self.__embedding_document(object)
+        if object.get_object_type() == ObjectType.Image:
+            await self.__embedding_image(object)
+        if object.get_object_type() == ObjectType.Video:
+            await self.__embedding_vedio(object)
+        if object.get_object_type() == ObjectType.RichText:
+            await self.__embedding_rich_text(object)
+        if object.get_object_type() == ObjectType.Email:
+            await self.__embedding_email(object)
         else:
             pass
 
-    async def query(self, prompt: AgentPrompt) -> AgentPrompt:
+    async def query(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 = self.vector_base.query(vector, 10)
+                object_ids = self.store.get_vector_store("default").query(vector, 10)
-                for object_id in object_ids:
+                results.append(object_ids)
-                    if object_id.object_type == ObjectType.Email:
+        return results
-                        [object, email] = self.object_store.get(object_id)
+
-                        if object.object_type == ObjectType.Email:
+
-                            email: EmailObject = object
+
-                            prompt.append(AgentPrompt())
+                        
-                            prompt
                     
                 
 

src/knowledge/core_object/__init__.py

@@ -1,4 +1,5 @@
 from .document_object import DocumentObject, DocumentObjectBuilder
 from .image_object import ImageObject, ImageObjectBuilder
 from .video_object import VideoObject, VideoObjectBuilder
+from .rich_text_object import RichTextObject, RichTextObjectBuilder
 from .email_object import EmailObject, EmailObjectBuilder

src/knowledge/core_object/image_object.py

@@ -18,7 +18,7 @@ class ImageObject(KnowledgeObject):
         body = dict()
         desc["meta"] = meta
         desc["exif"] = exif
-        desc
+        desc["tags"] = tags
         desc["hash"] = chunk_list.hash.to_base58()
         body["chunk_list"] = chunk_list.chunk_list
 

src/knowledge/object/object_id.py

@@ -8,7 +8,6 @@ import base36
 
 class ObjectType(Enum):
     Chunk = 7
-    TextChunk = 100
     Image = 101
     Video = 102
     Document = 103

src/knowledge/store.py

@@ -1,6 +1,7 @@
 import os
 from .object import ObjectStore
 from .data import ChunkStore, ChunkTracker, ChunkListWriter, ChunkReader
+from .vector import ChromaVectorStore, VectorBase
 import logging
 
 
@@ -37,6 +38,7 @@ class KnowledgeStore:
         self.chunk_tracker = ChunkTracker(chunk_store_dir)
         self.chunk_list_writer = ChunkListWriter(self.chunk_store, self.chunk_tracker)
         self.chunk_reader = ChunkReader(self.chunk_store, self.chunk_tracker)
+        self.vector_store = {}
         
 
     def get_object_store(self) -> ObjectStore:
@@ -53,3 +55,8 @@ class KnowledgeStore:
     
     def get_chunk_reader(self) -> ChunkReader:
         return self.chunk_reader
+    
+    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)
+        return self.vector_store[model_name]

src/knowledge/vector/chroma_store.py

@@ -6,11 +6,11 @@ import os
 
 
 class ChromaVectorStore(VectorBase):
-    def __init__(self, db_url, model_name: str) -> None:
+    def __init__(self, model_name: str) -> None:
-        super().__init__(db_url, model_name)
+        super().__init__(model_name)
 
         logging.info(
-            "will init chroma vector store, db={}, model={}".format(db_url, model_name)
+            "will init chroma vector store, model={}".format(model_name)
         )
 
         directory = os.path.join(

src/knowledge/vector/vector_base.py

@@ -3,8 +3,7 @@ from ..object import ObjectID
 
 # define a vector base class
 class VectorBase:
-    def __init__(self, db_url, model_name) -> None:
+    def __init__(self, model_name) -> None:
-        self.db_url = db_url
         self.model_name = model_name
 
     async def insert(self, vector: [float], id: ObjectID):

test/test_knowledge_base.py

@@ -0,0 +1,57 @@
+import sys
+import os
+import logging
+
+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)
+
+
+from knowledge import ObjectID, HashValue, EmailObjectBuilder
+from aios_kernel import KnowledgeBase, AgentPrompt
+import asyncio
+import unittest
+
+async def test_embedding_email():
+    data = HashValue.hash_data("1233".encode("utf-8"));
+    print(data.to_base58())
+    print(data.to_base36())
+    
+    data2 = HashValue.from_base58(data.to_base58())
+    self.assertEqual(data.to_base36(), data2.to_base36())
+    
+    data2 = HashValue.from_base36(data.to_base36())
+    self.assertEqual(data.to_base58(), data2.to_base58())
+
+    email_folder = "F:\\system\\Downloads\\8081ffdb80925f5bff9c6ab9c4756c7d"
+    email_object = EmailObjectBuilder({}, email_folder).build()
+
+    await KnowledgeBase().do_embedding(email_object)
+
+
+async def test_query_email():
+    msg_prompt = AgentPrompt()
+    msg_prompt.messages = [{"role":"user","content":"abcdef"}]
+       
+    KnowledgeBase().query(msg_prompt)
+
+class TestVectorSTorage(unittest.TestCase):
+    def test_embedding(self):
+        asyncio.run(test_embedding_email())
+
+    def test_query(self):
+        asyncio.run(test_query_email())
+        
+        
+if __name__ == "__main__":
+    unittest.main()

test/test_vector_storage.py

@@ -14,8 +14,8 @@ import asyncio
 import unittest
 
 
-async def test_vector():
+async def test_embedding_email():
-    storage = ChromaVectorStore("", "test")
+    storage = ChromaVectorStore("test")
     await storage.insert([1, 2, 3], "test")
     ids = await storage.query([1, 2, 3], 10)
     print(ids)