save relationship of objects

src/aios_kernel/knowledge_base.py

@@ -1,5 +1,6 @@
 # define a knowledge base class
 import json
+import pickle
 from . import AgentPrompt, ComputeKernel
 from ..knowledge import *
 
@@ -39,7 +40,7 @@ class KnowledgeBase:
         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):
+    async def __embedding_video(self, vedio: VideoObject):
         desc = {}
         if not vedio.get_meta():
             desc["meta"] = vedio.get_meta()
@@ -56,7 +57,7 @@ class KnowledgeBase:
         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)
+            await self.__embedding_video(vedio)
         for rich_text in rich_text.get_rich_texts().values():
             await self.__embedding_rich_text(rich_text)
 
@@ -65,13 +66,14 @@ class KnowledgeBase:
         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):
+
+    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)
+            await self.__embedding_video(object)
         if object.get_object_type() == ObjectType.RichText:
             await self.__embedding_rich_text(object)
         if object.get_object_type() == ObjectType.Email:
@@ -79,7 +81,80 @@ class KnowledgeBase:
         else:
             pass
 
-    async def query(self, prompt: AgentPrompt) -> [ObjectID]:
+    def __save_document(self, document: DocumentObject):
+        doc_id = document.calculate_id()
+        self.store.get_object_store().put_object(doc_id, document.encode())
+        for chunk_id in document.get_chunk_list():
+            self.store.get_relation_store().add_relation(chunk_id, doc_id)
+
+    def __save_image(self, image: ImageObject):
+        image_id = image.calculate_id()
+        self.store.get_object_store().put_object(image_id, image.encode())
+
+    def __save_video(self, video: VideoObject):
+        video_id = video.calculate_id()
+        self.store.get_object_store().put_object(video_id, video.encode())
+
+    def __save_rich_text(self, rich_text: RichTextObject):
+        rich_text_id = rich_text.calculate_id()
+        rich_text_enc = dict()
+        rich_text_enc["desc"] = rich_text.desc
+        rich_text_enc["body"] = {"documents": {}, "images": {}, "videos": {}, "rich_texts": {}}
+        for key, document in rich_text.get_documents().items():
+            self.__save_document(document)
+            doc_id = document.calculate_id()
+            self.store.get_relation_store().add_relation(doc_id, rich_text_id)
+            rich_text_enc["body"]["documents"][key] = doc_id
+        for key, image in rich_text.get_images().items():
+            self.__save_image(image)
+            image_id = image.calculate_id()
+            self.store.get_relation_store().add_relation(image_id, rich_text_id)
+            rich_text_enc["body"]["images"][key] = image_id
+        for key, video in rich_text.get_videos().items():
+            self.__save_video(video)
+            video_id = video.calculate_id()
+            self.store.get_relation_store().add_relation(video_id, rich_text_id)
+            rich_text_enc["body"]["videos"][key] = video_id
+        for key, rich_text in rich_text.get_rich_texts().items():
+            self.__save_rich_text(rich_text)
+            rich_text_id = rich_text.calculate_id()
+            self.store.get_relation_store().add_relation(rich_text_id, rich_text_id)
+            rich_text_enc["body"]["rich_texts"][key] = rich_text_id
+
+
+        self.store.get_object_store().put_object(rich_text_id, pickle.dumps(rich_text_enc))
+
+    def __save_email(self, email: EmailObject):
+        email_id = email.calculate_id()
+        email_enc = dict()
+        email_enc["desc"] = email.desc
+        email_enc["body"] = {"content": None}
+        self.__save_rich_text(email.get_rich_text())
+        rich_text_id = email.get_rich_text().calculate_id()
+        self.store.get_relation_store().add_relation(rich_text_id, email_id)
+        email_enc["body"]["content"] = rich_text_id
+        self.store.get_object_store().put_object(email_id, pickle.dumps(email_enc))
+
+    
+    def __save_object(self, object: KnowledgeObject):
+        if object.get_object_type() == ObjectType.Document:
+            self.__save_document(object)
+        if object.get_object_type() == ObjectType.Image:
+            self.__save_image(object)
+        if object.get_object_type() == ObjectType.Video:
+            self.__save_video(object)
+        if object.get_object_type() == ObjectType.RichText:
+            self.__save_rich_text(object)
+        if object.get_object_type() == ObjectType.Email:
+            self.__save_email(object)
+        else:
+            pass
+
+    async def insert_object(self, object: KnowledgeObject):
+        self.__save_object(object)
+        self.__do_embedding(object)
+
+    async def query_objects(self, prompt: AgentPrompt) -> [ObjectID]:
         results = []
         for msg in prompt.messages:
             if msg.role == "user":
@@ -89,6 +164,28 @@ class KnowledgeBase:
         return results
 
 
+    async def __prompt_from_objects(self, object_ids: [ObjectID]) -> AgentPrompt:
+        prompt = AgentPrompt()
+        for object_id in object_ids:
+            object = self.store.get_object_reader().get_object(object_id)
+            if object is None:
+                raise ValueError(f"object not found: {object_id}")
+            if object.get_object_type() == ObjectType.Document:
+                document = object
+                prompt.messages.append({"role": "agent", "content": document.get_body()})
+            if object.get_object_type() == ObjectType.Image:
+                image = object
+                prompt.messages.append({"role": "agent", "content": json.dumps(image.get_desc())})
+            if object.get_object_type() == ObjectType.Video:
+                video = object
+                prompt.messages.append({"role": "agent", "content": json.dumps(video.get_desc())})
+            if object.get_object_type() == ObjectType.RichText:
+                rich_text = object
+                prompt.messages.append({"role": "agent", "content": json.dumps(rich_text.get_desc())})
+            if object.get_object_type() == ObjectType.Email:
+                email = object
+                prompt.messages.append({"role": "agent", "content": json.dumps(email.get_desc())})
+        return prompt
 
                         
                     

src/knowledge/core_object/email_object.py

@@ -28,7 +28,7 @@ class EmailObject(KnowledgeObject):
         return self.desc["tags"]
 
     def get_rich_text(self):
-        return self.body["rich_text"]
+        return self.body["content"]
 
 
 """

src/knowledge/object/__init__.py

@@ -1,5 +1,6 @@
 from .object import KnowledgeObject
 from .blob import FileBlobStorage
 from .hash import HashValue, hash_data
+from .relation import ObjectRelationStore
 from .object_store import ObjectStore
 from .object_id import ObjectID, ObjectType

src/knowledge/object/relation.py

@@ -0,0 +1,16 @@
+# define a relation store class
+from .object_id import ObjectID
+
+
+class ObjectRelationStore:
+    def __init__(self, root_dir: str):
+        pass
+
+    def add_relation(self, object: ObjectID, parent: ObjectID):
+        pass
+
+    def get_related_objects(self, object: ObjectID) -> [ObjectID]:
+        pass
+
+    def delete_relation(self, object_id: ObjectID) -> [ObjectID]:
+        pass

src/knowledge/store.py

@@ -1,5 +1,6 @@
 import os
-from .object import ObjectStore
+
+from .object import ObjectStore, ObjectRelationStore
 from .data import ChunkStore, ChunkTracker, ChunkListWriter, ChunkReader
 from .vector import ChromaVectorStore, VectorBase
 import logging
@@ -30,6 +31,9 @@ class KnowledgeStore:
 
         self.root = root_dir
 
+        relation_store_dir = os.path.join(root_dir, "relation")
+        self.relation_store = ObjectRelationStore(relation_store_dir)
+
         object_store_dir = os.path.join(root_dir, "object")
         self.object_store = ObjectStore(object_store_dir)
 
@@ -39,7 +43,9 @@ class KnowledgeStore:
         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_relation_store(self) -> ObjectRelationStore:
+        return self.relation_store
 
     def get_object_store(self) -> ObjectStore:
         return self.object_store

test/test_knowledge_base.py

@@ -45,7 +45,7 @@ async def test_query_email():
        
     KnowledgeBase().query(msg_prompt)
 
-class TestVectorSTorage(unittest.TestCase):
+class TestKnowledgeBase(unittest.TestCase):
     def test_embedding(self):
         asyncio.run(test_embedding_email())
 

test/test_vector_storage.py

@@ -14,13 +14,13 @@ import asyncio
 import unittest
 
 
-async def test_embedding_email():
+async def test_vector():
     storage = ChromaVectorStore("test")
     await storage.insert([1, 2, 3], "test")
     ids = await storage.query([1, 2, 3], 10)
     print(ids)
 
-class TestVectorSTorage(unittest.TestCase):
+class TestVectorStorage(unittest.TestCase):
     def test_run(self):
         asyncio.run(test_vector())