From 6e5beb4cdffb42c85eeb720ab0035b7c102ac58f Mon Sep 17 00:00:00 2001 From: tsukasa Date: Wed, 13 Sep 2023 17:55:27 +0800 Subject: [PATCH] save relationship of objects --- src/aios_kernel/knowledge_base.py | 107 +++++++++++++++++++++- src/knowledge/core_object/email_object.py | 2 +- src/knowledge/object/__init__.py | 1 + src/knowledge/object/relation.py | 16 ++++ src/knowledge/store.py | 10 +- test/test_knowledge_base.py | 2 +- test/test_vector_storage.py | 4 +- 7 files changed, 131 insertions(+), 11 deletions(-) create mode 100644 src/knowledge/object/relation.py diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index 52b2943..ce689e6 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -1,5 +1,6 @@ # define a knowledge base class import json +import pickle from . import AgentPrompt, ComputeKernel from ..knowledge import * @@ -39,7 +40,7 @@ class KnowledgeBase: 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): + async def __embedding_video(self, vedio: VideoObject): desc = {} if not vedio.get_meta(): desc["meta"] = vedio.get_meta() @@ -56,7 +57,7 @@ class KnowledgeBase: 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) + await self.__embedding_video(vedio) for rich_text in rich_text.get_rich_texts().values(): await self.__embedding_rich_text(rich_text) @@ -65,13 +66,14 @@ class KnowledgeBase: 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): + + 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) + await self.__embedding_video(object) if object.get_object_type() == ObjectType.RichText: await self.__embedding_rich_text(object) if object.get_object_type() == ObjectType.Email: @@ -79,7 +81,80 @@ class KnowledgeBase: else: pass - async def query(self, prompt: AgentPrompt) -> [ObjectID]: + def __save_document(self, document: DocumentObject): + doc_id = document.calculate_id() + self.store.get_object_store().put_object(doc_id, document.encode()) + for chunk_id in document.get_chunk_list(): + self.store.get_relation_store().add_relation(chunk_id, doc_id) + + def __save_image(self, image: ImageObject): + image_id = image.calculate_id() + self.store.get_object_store().put_object(image_id, image.encode()) + + def __save_video(self, video: VideoObject): + video_id = video.calculate_id() + self.store.get_object_store().put_object(video_id, video.encode()) + + def __save_rich_text(self, rich_text: RichTextObject): + rich_text_id = rich_text.calculate_id() + rich_text_enc = dict() + rich_text_enc["desc"] = rich_text.desc + rich_text_enc["body"] = {"documents": {}, "images": {}, "videos": {}, "rich_texts": {}} + for key, document in rich_text.get_documents().items(): + self.__save_document(document) + doc_id = document.calculate_id() + self.store.get_relation_store().add_relation(doc_id, rich_text_id) + rich_text_enc["body"]["documents"][key] = doc_id + for key, image in rich_text.get_images().items(): + self.__save_image(image) + image_id = image.calculate_id() + self.store.get_relation_store().add_relation(image_id, rich_text_id) + rich_text_enc["body"]["images"][key] = image_id + for key, video in rich_text.get_videos().items(): + self.__save_video(video) + video_id = video.calculate_id() + self.store.get_relation_store().add_relation(video_id, rich_text_id) + rich_text_enc["body"]["videos"][key] = video_id + for key, rich_text in rich_text.get_rich_texts().items(): + self.__save_rich_text(rich_text) + rich_text_id = rich_text.calculate_id() + self.store.get_relation_store().add_relation(rich_text_id, rich_text_id) + rich_text_enc["body"]["rich_texts"][key] = rich_text_id + + + self.store.get_object_store().put_object(rich_text_id, pickle.dumps(rich_text_enc)) + + def __save_email(self, email: EmailObject): + email_id = email.calculate_id() + email_enc = dict() + email_enc["desc"] = email.desc + email_enc["body"] = {"content": None} + self.__save_rich_text(email.get_rich_text()) + rich_text_id = email.get_rich_text().calculate_id() + self.store.get_relation_store().add_relation(rich_text_id, email_id) + email_enc["body"]["content"] = rich_text_id + self.store.get_object_store().put_object(email_id, pickle.dumps(email_enc)) + + + def __save_object(self, object: KnowledgeObject): + if object.get_object_type() == ObjectType.Document: + self.__save_document(object) + if object.get_object_type() == ObjectType.Image: + self.__save_image(object) + if object.get_object_type() == ObjectType.Video: + self.__save_video(object) + if object.get_object_type() == ObjectType.RichText: + self.__save_rich_text(object) + if object.get_object_type() == ObjectType.Email: + self.__save_email(object) + else: + pass + + async def insert_object(self, object: KnowledgeObject): + self.__save_object(object) + self.__do_embedding(object) + + async def query_objects(self, prompt: AgentPrompt) -> [ObjectID]: results = [] for msg in prompt.messages: if msg.role == "user": @@ -89,6 +164,28 @@ class KnowledgeBase: return results + async def __prompt_from_objects(self, object_ids: [ObjectID]) -> AgentPrompt: + prompt = AgentPrompt() + for object_id in object_ids: + object = self.store.get_object_reader().get_object(object_id) + if object is None: + raise ValueError(f"object not found: {object_id}") + if object.get_object_type() == ObjectType.Document: + document = object + prompt.messages.append({"role": "agent", "content": document.get_body()}) + if object.get_object_type() == ObjectType.Image: + image = object + prompt.messages.append({"role": "agent", "content": json.dumps(image.get_desc())}) + if object.get_object_type() == ObjectType.Video: + video = object + prompt.messages.append({"role": "agent", "content": json.dumps(video.get_desc())}) + if object.get_object_type() == ObjectType.RichText: + rich_text = object + prompt.messages.append({"role": "agent", "content": json.dumps(rich_text.get_desc())}) + if object.get_object_type() == ObjectType.Email: + email = object + prompt.messages.append({"role": "agent", "content": json.dumps(email.get_desc())}) + return prompt diff --git a/src/knowledge/core_object/email_object.py b/src/knowledge/core_object/email_object.py index 8b8efb5..f563052 100644 --- a/src/knowledge/core_object/email_object.py +++ b/src/knowledge/core_object/email_object.py @@ -28,7 +28,7 @@ class EmailObject(KnowledgeObject): return self.desc["tags"] def get_rich_text(self): - return self.body["rich_text"] + return self.body["content"] """ diff --git a/src/knowledge/object/__init__.py b/src/knowledge/object/__init__.py index 7886d7a..291715f 100644 --- a/src/knowledge/object/__init__.py +++ b/src/knowledge/object/__init__.py @@ -1,5 +1,6 @@ from .object import KnowledgeObject from .blob import FileBlobStorage from .hash import HashValue, hash_data +from .relation import ObjectRelationStore from .object_store import ObjectStore from .object_id import ObjectID, ObjectType diff --git a/src/knowledge/object/relation.py b/src/knowledge/object/relation.py new file mode 100644 index 0000000..1b56e48 --- /dev/null +++ b/src/knowledge/object/relation.py @@ -0,0 +1,16 @@ +# define a relation store class +from .object_id import ObjectID + + +class ObjectRelationStore: + def __init__(self, root_dir: str): + pass + + def add_relation(self, object: ObjectID, parent: ObjectID): + pass + + def get_related_objects(self, object: ObjectID) -> [ObjectID]: + pass + + def delete_relation(self, object_id: ObjectID) -> [ObjectID]: + pass \ No newline at end of file diff --git a/src/knowledge/store.py b/src/knowledge/store.py index c0cc6b5..03e48d0 100644 --- a/src/knowledge/store.py +++ b/src/knowledge/store.py @@ -1,5 +1,6 @@ import os -from .object import ObjectStore + +from .object import ObjectStore, ObjectRelationStore from .data import ChunkStore, ChunkTracker, ChunkListWriter, ChunkReader from .vector import ChromaVectorStore, VectorBase import logging @@ -30,6 +31,9 @@ class KnowledgeStore: self.root = root_dir + relation_store_dir = os.path.join(root_dir, "relation") + self.relation_store = ObjectRelationStore(relation_store_dir) + object_store_dir = os.path.join(root_dir, "object") self.object_store = ObjectStore(object_store_dir) @@ -39,7 +43,9 @@ class KnowledgeStore: 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_relation_store(self) -> ObjectRelationStore: + return self.relation_store def get_object_store(self) -> ObjectStore: return self.object_store diff --git a/test/test_knowledge_base.py b/test/test_knowledge_base.py index efed1bd..ebda00c 100644 --- a/test/test_knowledge_base.py +++ b/test/test_knowledge_base.py @@ -45,7 +45,7 @@ async def test_query_email(): KnowledgeBase().query(msg_prompt) -class TestVectorSTorage(unittest.TestCase): +class TestKnowledgeBase(unittest.TestCase): def test_embedding(self): asyncio.run(test_embedding_email()) diff --git a/test/test_vector_storage.py b/test/test_vector_storage.py index 6023ab1..ff4dc72 100644 --- a/test/test_vector_storage.py +++ b/test/test_vector_storage.py @@ -14,13 +14,13 @@ import asyncio import unittest -async def test_embedding_email(): +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): +class TestVectorStorage(unittest.TestCase): def test_run(self): asyncio.run(test_vector())