Add chromadb based vector store impl and test case

src/knowledge/__init__.py

@@ -0,0 +1,2 @@
+from .object import EmailObject, ImageObject, TextChunkObject, ObjectID
+from .vector import *

src/knowledge/object.py

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

@@ -0,0 +1,33 @@
+# 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/knowledge/vector/__init__.py

@@ -0,0 +1,2 @@
+from .vector_base import VectorBase
+from .chroma_store import ChromaVectorStore

src/knowledge/vector/chroma_store.py

@@ -0,0 +1,49 @@
+from .vector_base import VectorBase
+from ..object import ObjectID
+import chromadb
+import logging
+import os
+
+
+class ChromaVectorStore(VectorBase):
+    def __init__(self, db_url, model_name: str) -> None:
+        super().__init__(db_url, model_name)
+
+        logging.info(
+            "will init chroma vector store, db={}, model={}".format(db_url, model_name)
+        )
+
+        directory = os.path.join(
+            os.path.dirname(__file__), "../../../rootfs/data/vector"
+        )
+        logging.info("will use vector store: {}".format(directory))
+
+        client = chromadb.PersistentClient(
+            path=directory, settings=chromadb.Settings(anonymized_telemetry=False)
+        )
+        # client = chromadb.Client()
+
+        collection_name = "coll_{}".format(model_name)
+        logging.info("will init chroma colletion: %s", collection_name)
+
+        collection = client.get_or_create_collection(collection_name)
+        self.collection = collection
+
+    async def insert(self, vector: [float], id: ObjectID):
+        self.collection.add(
+            embeddings=vector,
+            ids=id,
+        )
+
+    async def query(self, vector: [float], top_k: int) -> [ObjectID]:
+        ret = self.collection.query(
+            query_embeddings=vector,
+            n_results=top_k,
+        )
+
+        return ret["ids"]
+
+    async def delete(self, id: ObjectID):
+        self.collection.delete(
+            ids=id,
+        )

src/knowledge/vector/vector_base.py

@@ -0,0 +1,17 @@
+# 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
+    
+    async def delete(self, id: ObjectID):
+        pass