diff --git a/src/aios_kernel/__init__.py b/src/aios_kernel/__init__.py index bd1434b..9817b76 100644 --- a/src/aios_kernel/__init__.py +++ b/src/aios_kernel/__init__.py @@ -4,7 +4,8 @@ from .chatsession import AIChatSession from .agent import AIAgent,AIAgentTemplete,AgentPrompt from .compute_kernel import ComputeKernel,ComputeTask from .compute_node import ComputeNode,LocalComputeNode -from .open_ai_node import OpenAI_ComputeNode +from .open_ai_node import OpenAI_ComputeNode +from .knowledge_base import KnowledgeBase from .role import AIRole,AIRoleGroup from .workflow import Workflow from .bus import AIBus \ No newline at end of file diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index 4f4ba4f..52b2943 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -1,47 +1,96 @@ # define a knowledge base class +import json from . import AgentPrompt, ComputeKernel -from ..knowledge.object import KnowledgeObject, ObjectType, EmailObject, TextChunkObject, ImageObject -from ..knowledge.store import ObjectStorage -from ..knowledge.vector.vector_base import VectorBase +from ..knowledge import * + class KnowledgeBase: - def __init__(self) -> None: - self.object_store = ObjectStorage() - self.vector_base = VectorBase() + _instance = None + + def __new__(cls): + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance.__singleton_init__() + + return cls._instance + + def __singleton_init__(self) -> None: + self.store = KnowledgeStore() self.compute_kernel = ComputeKernel() - async def insert(self, object: KnowledgeObject): - if object.object_type == ObjectType.Email: - email: EmailObject = object - for text_id in email.text: - [text, _] = self.object_store.get(text_id) - text: TextChunkObject = text - vector = await self.compute_kernel.do_text_embedding(text.text) - self.vector_base.insert(vector, text_id) - - for image_id in email.images: - [image, _] = self.object_store.get(image_id) - image: ImageObject = image - vector = await self.compute_kernel.do_text_embedding(image.meta) - self.vector_base.insert(vector, image_id) - - vector = await self.compute_kernel.do_text_embedding(email.meta) - self.vector_base.insert(vector, email.get_id()) + async def __embedding_document(self, document: DocumentObject): + for chunk_id in document.get_chunk_list(): + chunk = self.store.get_chunk_reader().get_chunk(chunk_id) + if chunk is None: + raise ValueError(f"text chunk not found: {chunk_id}") + + text = chunk.read().decode("utf-8") + vector = await self.compute_kernel.do_text_embedding(text) + self.store.get_vector_store("default").insert(vector, chunk_id) + + async def __embedding_image(self, image: ImageObject): + desc = {} + if not image.get_meta(): + desc["meta"] = image.get_meta() + if not image.get_exif(): + desc["exif"] = image.get_exif() + if not image.get_tags(): + desc["tags"] = image.get_tags() + vector = await self.compute_kernel.do_text_embedding(json.dumps(desc)) + self.store.get_vector_store("default").insert(vector, image.calculate_id()) + + async def __embedding_vedio(self, vedio: VideoObject): + desc = {} + if not vedio.get_meta(): + desc["meta"] = vedio.get_meta() + if not vedio.get_info(): + desc["info"] = vedio.get_info() + if not vedio.get_tags(): + desc["tags"] = vedio.get_tags() + vector = await self.compute_kernel.do_text_embedding(json.dumps(desc)) + self.store.get_vector_store("default").insert(vector, vedio.calculate_id()) + + async def __embedding_rich_text(self, rich_text: RichTextObject): + for document in rich_text.get_documents().values(): + await self.__embedding_document(document) + for image in rich_text.get_images().values(): + await self.__embedding_image(image) + for vedio in rich_text.get_videos().values(): + await self.__embedding_vedio(vedio) + for rich_text in rich_text.get_rich_texts().values(): + await self.__embedding_rich_text(rich_text) + + async def __embedding_email(self, email: EmailObject): + vector = await self.compute_kernel.do_text_embedding(json.dumps(email.get_desc())) + self.store.get_vector_store("default").insert(vector, email.calculate_id()) + await self.__embedding_rich_text(email.get_rich_text()) + + async def do_embedding(self, object: KnowledgeObject): + if object.get_object_type() == ObjectType.Document: + await self.__embedding_document(object) + if object.get_object_type() == ObjectType.Image: + await self.__embedding_image(object) + if object.get_object_type() == ObjectType.Video: + await self.__embedding_vedio(object) + if object.get_object_type() == ObjectType.RichText: + await self.__embedding_rich_text(object) + if object.get_object_type() == ObjectType.Email: + await self.__embedding_email(object) else: pass - async def query(self, prompt: AgentPrompt) -> AgentPrompt: + async def query(self, prompt: AgentPrompt) -> [ObjectID]: + results = [] for msg in prompt.messages: if msg.role == "user": vector = await self.compute_kernel.do_text_embedding(msg.content) - 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_ids = self.store.get_vector_store("default").query(vector, 10) + results.append(object_ids) + return results + + + + diff --git a/src/knowledge/core_object/__init__.py b/src/knowledge/core_object/__init__.py index 065de74..dee13e2 100644 --- a/src/knowledge/core_object/__init__.py +++ b/src/knowledge/core_object/__init__.py @@ -1,4 +1,5 @@ from .document_object import DocumentObject, DocumentObjectBuilder from .image_object import ImageObject, ImageObjectBuilder from .video_object import VideoObject, VideoObjectBuilder +from .rich_text_object import RichTextObject, RichTextObjectBuilder from .email_object import EmailObject, EmailObjectBuilder \ No newline at end of file diff --git a/src/knowledge/core_object/image_object.py b/src/knowledge/core_object/image_object.py index 83ab434..30b0e8a 100644 --- a/src/knowledge/core_object/image_object.py +++ b/src/knowledge/core_object/image_object.py @@ -18,7 +18,7 @@ class ImageObject(KnowledgeObject): body = dict() desc["meta"] = meta desc["exif"] = exif - desc + desc["tags"] = tags desc["hash"] = chunk_list.hash.to_base58() body["chunk_list"] = chunk_list.chunk_list diff --git a/src/knowledge/object/object_id.py b/src/knowledge/object/object_id.py index 3c3ac24..6ce77ae 100644 --- a/src/knowledge/object/object_id.py +++ b/src/knowledge/object/object_id.py @@ -8,7 +8,6 @@ import base36 class ObjectType(Enum): Chunk = 7 - TextChunk = 100 Image = 101 Video = 102 Document = 103 diff --git a/src/knowledge/store.py b/src/knowledge/store.py index dc6b17a..c0cc6b5 100644 --- a/src/knowledge/store.py +++ b/src/knowledge/store.py @@ -1,6 +1,7 @@ import os from .object import ObjectStore from .data import ChunkStore, ChunkTracker, ChunkListWriter, ChunkReader +from .vector import ChromaVectorStore, VectorBase import logging @@ -37,6 +38,7 @@ class KnowledgeStore: self.chunk_tracker = ChunkTracker(chunk_store_dir) self.chunk_list_writer = ChunkListWriter(self.chunk_store, self.chunk_tracker) self.chunk_reader = ChunkReader(self.chunk_store, self.chunk_tracker) + self.vector_store = {} def get_object_store(self) -> ObjectStore: @@ -53,3 +55,8 @@ class KnowledgeStore: def get_chunk_reader(self) -> ChunkReader: return self.chunk_reader + + def get_vector_store(self, model_name: str) -> VectorBase: + if model_name not in self.vector_store: + self.vector_store[model_name] = ChromaVectorStore(model_name) + return self.vector_store[model_name] diff --git a/src/knowledge/vector/chroma_store.py b/src/knowledge/vector/chroma_store.py index 626cd4f..0056375 100644 --- a/src/knowledge/vector/chroma_store.py +++ b/src/knowledge/vector/chroma_store.py @@ -6,11 +6,11 @@ import os class ChromaVectorStore(VectorBase): - def __init__(self, db_url, model_name: str) -> None: - super().__init__(db_url, model_name) + def __init__(self, model_name: str) -> None: + super().__init__(model_name) logging.info( - "will init chroma vector store, db={}, model={}".format(db_url, model_name) + "will init chroma vector store, model={}".format(model_name) ) directory = os.path.join( diff --git a/src/knowledge/vector/vector_base.py b/src/knowledge/vector/vector_base.py index 354c21e..83276ed 100644 --- a/src/knowledge/vector/vector_base.py +++ b/src/knowledge/vector/vector_base.py @@ -3,8 +3,7 @@ from ..object import ObjectID # define a vector base class class VectorBase: - def __init__(self, db_url, model_name) -> None: - self.db_url = db_url + def __init__(self, model_name) -> None: self.model_name = model_name async def insert(self, vector: [float], id: ObjectID): diff --git a/test/test_knowledge_base.py b/test/test_knowledge_base.py new file mode 100644 index 0000000..efed1bd --- /dev/null +++ b/test/test_knowledge_base.py @@ -0,0 +1,57 @@ +import sys +import os +import logging + +dir_path = os.path.dirname(os.path.realpath(__file__)) +print(dir_path) + +sys.path.append("{}/../src/".format(dir_path)) +print(sys.path) + +root = logging.getLogger() +root.setLevel(logging.DEBUG) +handler = logging.StreamHandler(sys.stdout) +handler.setLevel(logging.DEBUG) +formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") +handler.setFormatter(formatter) +root.addHandler(handler) + + +from knowledge import ObjectID, HashValue, EmailObjectBuilder +from aios_kernel import KnowledgeBase, AgentPrompt +import asyncio +import unittest + +async def test_embedding_email(): + data = HashValue.hash_data("1233".encode("utf-8")); + print(data.to_base58()) + print(data.to_base36()) + + data2 = HashValue.from_base58(data.to_base58()) + self.assertEqual(data.to_base36(), data2.to_base36()) + + data2 = HashValue.from_base36(data.to_base36()) + self.assertEqual(data.to_base58(), data2.to_base58()) + + email_folder = "F:\\system\\Downloads\\8081ffdb80925f5bff9c6ab9c4756c7d" + email_object = EmailObjectBuilder({}, email_folder).build() + + await KnowledgeBase().do_embedding(email_object) + + +async def test_query_email(): + msg_prompt = AgentPrompt() + msg_prompt.messages = [{"role":"user","content":"abcdef"}] + + KnowledgeBase().query(msg_prompt) + +class TestVectorSTorage(unittest.TestCase): + def test_embedding(self): + asyncio.run(test_embedding_email()) + + def test_query(self): + asyncio.run(test_query_email()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_vector_storage.py b/test/test_vector_storage.py index 7357d65..6023ab1 100644 --- a/test/test_vector_storage.py +++ b/test/test_vector_storage.py @@ -14,8 +14,8 @@ import asyncio import unittest -async def test_vector(): - storage = ChromaVectorStore("", "test") +async def test_embedding_email(): + storage = ChromaVectorStore("test") await storage.insert([1, 2, 3], "test") ids = await storage.query([1, 2, 3], 10) print(ids)