From 805dc88de28d841ac1472803100e36a11ae2ae9c Mon Sep 17 00:00:00 2001 From: tsukasa Date: Mon, 18 Sep 2023 09:51:51 +0800 Subject: [PATCH] test prompt from knowledge object --- src/aios_kernel/compute_kernel.py | 2 +- src/aios_kernel/compute_task.py | 2 +- src/aios_kernel/knowledge_base.py | 151 ++++++++++--------- src/aios_kernel/open_ai_node.py | 2 +- src/knowledge/core_object/document_object.py | 1 + test/test_knowledge_base.py | 39 ++--- 6 files changed, 98 insertions(+), 99 deletions(-) diff --git a/src/aios_kernel/compute_kernel.py b/src/aios_kernel/compute_kernel.py index f626578..b1da61b 100644 --- a/src/aios_kernel/compute_kernel.py +++ b/src/aios_kernel/compute_kernel.py @@ -121,7 +121,7 @@ class ComputeKernel: def text_embedding(self,input:str,model_name:Optional[str] = None): task_req = ComputeTask() - task_req.set_text_embeding_params(input,model_name) + task_req.set_text_embedding_params(input,model_name) self.run(task_req) return task_req diff --git a/src/aios_kernel/compute_task.py b/src/aios_kernel/compute_task.py index 35a1afd..6916a2a 100644 --- a/src/aios_kernel/compute_task.py +++ b/src/aios_kernel/compute_task.py @@ -38,7 +38,7 @@ class ComputeTask: self.params["model_name"] = "gpt-4-0613" self.params["max_token_size"] = max_token_size - def set_text_embeding_params(self, input, model_name=None, callchain_id = None): + def set_text_embedding_params(self, input, model_name=None, callchain_id = None): self.task_type = "text_embedding" self.create_time = time.time() self.task_id = uuid.uuid4().hex diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index cabf005..dbb96e6 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -1,7 +1,7 @@ # define a knowledge base class import json from . import AgentPrompt, ComputeKernel -from ..knowledge import * +from knowledge import * class KnowledgeBase: @@ -62,7 +62,7 @@ class KnowledgeBase: 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.store.get_vector_store("default").insert(vector, email.calculate_id()) await self.__embedding_rich_text(email.get_rich_text()) @@ -80,89 +80,94 @@ class KnowledgeBase: else: pass - 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_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_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_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 + # 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, rich_text.encode()) + # self.store.get_object_store().put_object(rich_text_id, rich_text.encode()) - 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, email.encode()) + # 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, email.encode()) - 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 + # 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) + # self.__save_object(object) + await self.__do_embedding(object) + + async def query_prompt(self, prompt: AgentPrompt): + objects = await self.query_objects(prompt) + knowledge_prompt = self.prompt_from_objects(objects) + prompt.append(knowledge_prompt) async def query_objects(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.store.get_vector_store("default").query(vector, 10) + object_ids = await self.store.get_vector_store("default").query(vector, 10) results.append(object_ids) return results - async def __load_object(self, object_id: ObjectID) -> KnowledgeObject: + def __load_object(self, object_id: ObjectID) -> KnowledgeObject: if object_id.get_object_type() == ObjectType.Document: return DocumentObject.decode(self.store.get_object_store().get_object(object_id)) if object_id.get_object_type() == ObjectType.Image: @@ -177,10 +182,10 @@ class KnowledgeBase: pass - async def prompt_from_objects(self, object_ids: [ObjectID]) -> AgentPrompt: + def prompt_from_objects(self, object_ids: [ObjectID]) -> AgentPrompt: results = dict() for object_id in object_ids: - parents = self.store.get_relation_store().get_related_objects(object_id) + parents = self.store.get_relation_store().get_related_root_objects(object_id) # last parent is the root object root_object_id = parents[-1] if results[root_object_id] is None: @@ -192,9 +197,9 @@ class KnowledgeBase: result_desc = [] for result in results.values(): # first element in result is the root object - root_object = await self.__load_object(result[0]) - if root_object.get_object_type() == ObjectType.Email: - email = await self.__load_object(object_id) + root_object_id = result[0] + if root_object_id.get_object_type() == ObjectType.Email: + email = self.__load_object(root_object_id) desc = email.get_desc() desc["type"] = "email" desc["contents"] = [] @@ -208,12 +213,12 @@ class KnowledgeBase: if object_id.get_object_type() == ObjectType.Chunk: upper_list.append({"type": "text", "content": self.store.get_chunk_reader().get_chunk(object_id).read().decode("utf-8")}) if object_id.get_object_type() == ObjectType.Image: - image = await self.__load_object(object_id) + image = self.__load_object(object_id) desc = image.get_desc() desc["type"] = "image" upper_list.append(desc) if object_id.get_object_type() == ObjectType.Video: - video = await self.__load_object(object_id) + video = self.__load_object(object_id) desc = video.get_desc() desc["type"] = "video" upper_list.append(desc) diff --git a/src/aios_kernel/open_ai_node.py b/src/aios_kernel/open_ai_node.py index 27177f0..a37706e 100644 --- a/src/aios_kernel/open_ai_node.py +++ b/src/aios_kernel/open_ai_node.py @@ -82,7 +82,7 @@ class OpenAI_ComputeNode(ComputeNode): input = task.params["input"] logger.info(f"call openai {model_name} input: {input}") - resp = openai.Embeding.create(model=model_name, + resp = openai.Embedding.create(model=model_name, input=input) logger.info(f"openai response: {resp}") diff --git a/src/knowledge/core_object/document_object.py b/src/knowledge/core_object/document_object.py index 02056bc..ab02b8d 100644 --- a/src/knowledge/core_object/document_object.py +++ b/src/knowledge/core_object/document_object.py @@ -53,6 +53,7 @@ class DocumentObjectBuilder: chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_text( self.text, 1024 * 4, + "." ) doc = DocumentObject(self.meta, self.tags, chunk_list) doc_id = doc.calculate_id() diff --git a/test/test_knowledge_base.py b/test/test_knowledge_base.py index ebda00c..5c9d516 100644 --- a/test/test_knowledge_base.py +++ b/test/test_knowledge_base.py @@ -18,39 +18,32 @@ root.addHandler(handler) from knowledge import ObjectID, HashValue, EmailObjectBuilder -from aios_kernel import KnowledgeBase, AgentPrompt +from aios_kernel import KnowledgeBase, AgentPrompt, OpenAI_ComputeNode, ComputeKernel 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()) +async def test_embedding_email(test): + open_ai_node = OpenAI_ComputeNode() + open_ai_node.start() + ComputeKernel().add_compute_node(open_ai_node) + + email_folder = os.path.join(dir_path, "../rootfs/data/email/") + print("explore emails in folder ", email_folder) + for root, dirs, files in os.walk(email_folder): + for dir in dirs: + email_object = EmailObjectBuilder({}, os.path.join(root, dir)).build() + await KnowledgeBase().insert_object(email_object) - 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) + await KnowledgeBase().query_prompt(msg_prompt) + + class TestKnowledgeBase(unittest.TestCase): def test_embedding(self): - asyncio.run(test_embedding_email()) - - def test_query(self): - asyncio.run(test_query_email()) + asyncio.run(test_embedding_email(self)) if __name__ == "__main__":