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 ..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
                     
                 
 

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:
@@ -11,4 +11,7 @@ class VectorBase:
         pass
 
     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()