knowledge base basic flow

src/aios_kernel/knowledge/knowledge_base.py

@@ -1,12 +1,46 @@
 # define a knowledge base class
-from src.aios_kernel.agent import AgentPrompt
-from .object import KnowledgeObject
+from ..agent import AgentPrompt
+from ..compute_kernel import ComputeKernel
+from .object import KnowledgeObject, ObjectType, EmailObject, TextChunkObject, ImageObject
+from .object_storage import ObjectStorage
+from .vector_base import VectorBase
 
 class KnowledgeBase:
+    def __init__(self) -> None:
+        self.object_store = ObjectStorage()
+        self.vector_base = VectorBase()
+        self.compute_kernel = ComputeKernel()
+
     async def insert(self, object: KnowledgeObject):
+        if object.object_type == ObjectType.Email:
+            email: EmailObject = object
+            for text_id in email.text:
+                [text, _] = self.object_store.get(text_id)
+                text: TextChunkObject = text
+                vector = await self.compute_kernel.do_text_embedding(text.text)
+                self.vector_base.insert(vector, text_id)
+            
+            for image_id in email.images:
+                [image, _] = self.object_store.get(image_id)
+                image: ImageObject = image
+                vector = await self.compute_kernel.do_text_embedding(image.meta)
+                self.vector_base.insert(vector, image_id)
+            
+            vector = await self.compute_kernel.do_text_embedding(email.meta)
+            self.vector_base.insert(vector, email.get_id())
+        else:
             pass
 
     async def query(self, prompt: AgentPrompt) -> AgentPrompt:
-        pass
-
+        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)
+                for object_id in object_ids:
+                    if object_id.object_type == ObjectType.Email:
                     
+                    [object, email] = self.object_store.get(object_id)
+                    if object.object_type == ObjectType.Email:
+                        email: EmailObject = object
+                        prompt.append(AgentPrompt())
+                        prompt

src/aios_kernel/knowledge/object.py

@@ -21,7 +21,7 @@ class ObjectID: # pylint: disable=too-few-public-methods
 
 # define a object class
 class KnowledgeObject(ABC): # pylint: disable=too-few-public-methods
-    def __init__(self, object_type):
+    def __init__(self, object_type: ObjectType):
         self.object_type = object_type
 
     @abstractmethod
@@ -41,7 +41,7 @@ class KnowledgeObject(ABC): # pylint: disable=too-few-public-methods
 
 # define a text chunk class
 class TextChunkObject(KnowledgeObject): # pylint: disable=too-few-public-methods
-    def __init__(self, text):
+    def __init__(self, text: str):
         super().__init__(ObjectType.TextChunk)
         self.text = text
 
@@ -59,7 +59,7 @@ class EmailObject(KnowledgeObject): # pylint: disable=too-few-public-methods
     def __init__(self, meta):
         super().__init__(ObjectType.Email)
         self.meta = meta
-        self.text = []
-        self.images = []
+        self.text = [ObjectID]
+        self.images = [ObjectID]