diff --git a/src/aios_kernel/knowledge/knowledge_base.py b/src/aios_kernel/knowledge_base.py similarity index 74% rename from src/aios_kernel/knowledge/knowledge_base.py rename to src/aios_kernel/knowledge_base.py index bb06401..a882ba7 100644 --- a/src/aios_kernel/knowledge/knowledge_base.py +++ b/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 diff --git a/src/aios_kernel/knowledge/__init__.py b/src/knowledge/__init__.py similarity index 63% rename from src/aios_kernel/knowledge/__init__.py rename to src/knowledge/__init__.py index 13bb147..c879518 100644 --- a/src/aios_kernel/knowledge/__init__.py +++ b/src/knowledge/__init__.py @@ -1,2 +1,2 @@ from .object import EmailObject, ImageObject, TextChunkObject, ObjectID -from .knowledge_base import KnowledgeBase +from .vector import * \ No newline at end of file diff --git a/src/aios_kernel/knowledge/object.py b/src/knowledge/object.py similarity index 96% rename from src/aios_kernel/knowledge/object.py rename to src/knowledge/object.py index 815d4e2..f7a127e 100644 --- a/src/aios_kernel/knowledge/object.py +++ b/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 diff --git a/src/aios_kernel/knowledge/object_storage.py b/src/knowledge/object_storage.py similarity index 100% rename from src/aios_kernel/knowledge/object_storage.py rename to src/knowledge/object_storage.py diff --git a/src/knowledge/vector/__init__.py b/src/knowledge/vector/__init__.py new file mode 100644 index 0000000..b3f0d78 --- /dev/null +++ b/src/knowledge/vector/__init__.py @@ -0,0 +1,2 @@ +from .vector_base import VectorBase +from .chroma_store import ChromaVectorStore diff --git a/src/knowledge/vector/chroma_store.py b/src/knowledge/vector/chroma_store.py new file mode 100644 index 0000000..626cd4f --- /dev/null +++ b/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, + ) diff --git a/src/aios_kernel/knowledge/vector_base.py b/src/knowledge/vector/vector_base.py similarity index 76% rename from src/aios_kernel/knowledge/vector_base.py rename to src/knowledge/vector/vector_base.py index b2cf442..354c21e 100644 --- a/src/aios_kernel/knowledge/vector_base.py +++ b/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 \ No newline at end of file diff --git a/test/test_vector_storage.py b/test/test_vector_storage.py new file mode 100644 index 0000000..7357d65 --- /dev/null +++ b/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()