Add chromadb based vector store impl and test case

src/aios_kernel/knowledge_base.py

@@ -1,9 +1,8 @@
 # define a knowledge base class
-from ..agent import AgentPrompt
+from . import AgentPrompt, ComputeKernel
-from ..compute_kernel import ComputeKernel
+from ..knowledge.object import KnowledgeObject, ObjectType, EmailObject, TextChunkObject, ImageObject
-from .object import KnowledgeObject, ObjectType, EmailObject, TextChunkObject, ImageObject
+from ..knowledge.object_storage import ObjectStorage
-from .object_storage import ObjectStorage
+from ..knowledge.vector.vector_base import VectorBase
-from .vector_base import VectorBase
 
 class KnowledgeBase:
     def __init__(self) -> None:
@@ -38,7 +37,6 @@ 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

src/knowledge/__init__.py

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

src/knowledge/object.py

@@ -1,7 +1,7 @@
 
 # define a object type enum
-from abc import ABC
+from abc import ABC, abstractmethod
-
+from enum import Enum
 
 class ObjectType(Enum):
     TextChunk = 1

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

@@ -1,5 +1,5 @@
 # import the ObjectID class
-from .object import ObjectID
+from ..object import ObjectID
 
 # define a vector base class
 class VectorBase:
@@ -12,3 +12,6 @@ class VectorBase:
 
     async def query(self, vector: [float], top_k: int) -> [ObjectID]:
         pass
+    
+    async def delete(self, id: ObjectID):
+        pass

test/test_vector_storage.py

@@ -0,0 +1,29 @@
+import sys
+import os
+
+dir_path = os.path.dirname(os.path.realpath(__file__))
+print(dir_path)
+
+sys.path.append("{}/../src/".format(dir_path))
+print(sys.path)
+
+from knowledge import ChromaVectorStore
+
+
+import asyncio
+import unittest
+
+
+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):
+    def test_run(self):
+        asyncio.run(test_vector())
+
+
+if __name__ == "__main__":
+    unittest.main()