Add chromadb based vector store impl and test case

src/aios_kernel/knowledge/__init__.py

@@ -1,2 +0,0 @@
-from .object import EmailObject, ImageObject, TextChunkObject, ObjectID
-from .knowledge_base import KnowledgeBase

src/aios_kernel/knowledge/object.py

@@ -1,65 +0,0 @@
-
-# define a object type enum
-from abc import ABC
-
-
-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/aios_kernel/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/aios_kernel/knowledge/vector_base.py

@@ -1,14 +0,0 @@
-# import the ObjectID class
-from .object import ObjectID
-
-# define a vector base class
-class VectorBase:
-    def __init__(self, db_url, model_name) -> None:
-        self.db_url = db_url
-        self.model_name = model_name
-
-    async def insert(self, vector: [float], id: ObjectID):
-        pass
-
-    async def query(self, vector: [float], top_k: int) -> [ObjectID]:
-        pass

src/aios_kernel/knowledge_base.py

@@ -1,9 +1,8 @@
 # define a knowledge base class
-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
+from . import AgentPrompt, ComputeKernel
+from ..knowledge.object import KnowledgeObject, ObjectType, EmailObject, TextChunkObject, ImageObject
+from ..knowledge.object_storage import ObjectStorage
+from ..knowledge.vector.vector_base import VectorBase
 
 class KnowledgeBase:
     def __init__(self) -> None:
@@ -38,12 +37,11 @@ class KnowledgeBase:
                 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
+                        [object, email] = self.object_store.get(object_id)
+                        if object.object_type == ObjectType.Email:
+                            email: EmailObject = object
+                            prompt.append(AgentPrompt())
+                            prompt