diff --git a/src/aios_kernel/knowledge/knowledge_base.py b/src/aios_kernel/knowledge/knowledge_base.py index a576280..5150ede 100644 --- a/src/aios_kernel/knowledge/knowledge_base.py +++ b/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): - pass + 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 - - \ No newline at end of file + 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 diff --git a/src/aios_kernel/knowledge/object.py b/src/aios_kernel/knowledge/object.py index 320cd8f..815d4e2 100644 --- a/src/aios_kernel/knowledge/object.py +++ b/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]