diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index dbb96e6..b229864 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 logging from . import AgentPrompt, ComputeKernel from knowledge import * @@ -26,7 +27,7 @@ class KnowledgeBase: 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) + await self.store.get_vector_store("default").insert(vector, chunk_id) async def __embedding_image(self, image: ImageObject): desc = {} @@ -37,7 +38,7 @@ class KnowledgeBase: if not 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()) + await self.store.get_vector_store("default").insert(vector, image.calculate_id()) async def __embedding_video(self, vedio: VideoObject): desc = {} @@ -48,16 +49,20 @@ class KnowledgeBase: if not 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()) + await 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(): + for document_id in rich_text.get_documents().values(): + document = DocumentObject.decode(self.store.get_object_store().get_object(document_id)) await self.__embedding_document(document) - for image in rich_text.get_images().values(): + for image_id in rich_text.get_images().values(): + image = ImageObject.decode(self.store.get_object_store().get_object(image_id)) await self.__embedding_image(image) - for vedio in rich_text.get_videos().values(): - await self.__embedding_video(vedio) - for rich_text in rich_text.get_rich_texts().values(): + for video_id in rich_text.get_videos().values(): + video = VideoObject.decode(self.store.get_object_store().get_object(video_id)) + await self.__embedding_video(video) + for rich_text_id in rich_text.get_rich_texts().values(): + rich_text = RichTextObject.decode(self.store.get_object_store().get_object(rich_text_id)) await self.__embedding_rich_text(rich_text) async def __embedding_email(self, email: EmailObject): @@ -154,17 +159,19 @@ class KnowledgeBase: await self.__do_embedding(object) async def query_prompt(self, prompt: AgentPrompt): + logging.info(f"query_prompt: {prompt}") objects = await self.query_objects(prompt) knowledge_prompt = self.prompt_from_objects(objects) + logging.info(f"prompt_from_objects result: {knowledge_prompt.as_str()}") 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) + if msg["role"] == "user": + vector = await self.compute_kernel.do_text_embedding(msg["content"]) object_ids = await self.store.get_vector_store("default").query(vector, 10) - results.append(object_ids) + results.extend(object_ids) return results def __load_object(self, object_id: ObjectID) -> KnowledgeObject: @@ -187,11 +194,12 @@ class KnowledgeBase: for object_id in object_ids: 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: - results[root_object_id] = [root_object_id, object_id] + root_object_id = parents[0] if parents else object_id + logging.info(f"object_id: {str(object_id)} root_object_id: {str(root_object_id)}") + if str(root_object_id) in results: + results[str(root_object_id)].append(object_id) else: - results[root_object_id].append(object_id) + results[str(root_object_id)] = [root_object_id, object_id] content = "I found the following contents described with json format:\n" result_desc = [] @@ -228,7 +236,7 @@ class KnowledgeBase: content += ".\n" prompt = AgentPrompt() - prompt.add_message("knowledge", content) + prompt.messages.append({"role": "knowledge", "content": content}) return prompt diff --git a/src/aios_kernel/open_ai_node.py b/src/aios_kernel/open_ai_node.py index a37706e..1c39c96 100644 --- a/src/aios_kernel/open_ai_node.py +++ b/src/aios_kernel/open_ai_node.py @@ -84,6 +84,22 @@ class OpenAI_ComputeNode(ComputeNode): resp = openai.Embedding.create(model=model_name, input=input) + + # resp = { + # "object": "list", + # "data": [ + # { + # "object": "embedding", + # "index": 0, + # "embedding": [ + # -0.00930514745414257, + # 0.00765434792265296, + # -0.007167573552578688, + # -0.012373941019177437, + # -0.04884673282504082 + # ]}] + # } + logger.info(f"openai response: {resp}") result = ComputeTaskResult() diff --git a/src/knowledge/core_object/document_object.py b/src/knowledge/core_object/document_object.py index ab02b8d..7210f16 100644 --- a/src/knowledge/core_object/document_object.py +++ b/src/knowledge/core_object/document_object.py @@ -53,7 +53,7 @@ class DocumentObjectBuilder: chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_text( self.text, 1024 * 4, - "." + ".?!\n" ) doc = DocumentObject(self.meta, self.tags, chunk_list) doc_id = doc.calculate_id() diff --git a/src/knowledge/data/writer.py b/src/knowledge/data/writer.py index 599c51b..3ac0ed1 100644 --- a/src/knowledge/data/writer.py +++ b/src/knowledge/data/writer.py @@ -74,19 +74,20 @@ class ChunkListWriter: text: str, chunk_max_words: int, separator_chars: str = ".," ) -> List[str]: sentences = re.split(f"[{separator_chars}]", text) - chunk_list = [] - chunk = [] - word_count = 0 - for sentence in sentences: - words = sentence.split() - for word in words: - if word_count < chunk_max_words: - chunk.append(word) - word_count += 1 - else: - chunk_list.append(" ".join(chunk)) - chunk = [word] - word_count = 1 - if chunk: - chunk_list.append(" ".join(chunk)) - return chunk_list \ No newline at end of file + # chunk_list = [] + # chunk = [] + # word_count = 0 + # for sentence in sentences: + # words = sentence.split() + # for word in words: + # if word_count < chunk_max_words: + # chunk.append(word) + # word_count += 1 + # else: + # chunk_list.append(" ".join(chunk)) + # chunk = [word] + # word_count = 1 + # if chunk: + # chunk_list.append(" ".join(chunk)) + # return chunk_list + return sentences \ No newline at end of file diff --git a/src/knowledge/object/object.py b/src/knowledge/object/object.py index 490555b..a63e30a 100644 --- a/src/knowledge/object/object.py +++ b/src/knowledge/object/object.py @@ -54,7 +54,8 @@ class KnowledgeObject(ABC): ) sha256 = hashlib.sha256() sha256.update(data.encode()) - return ObjectID(sha256.digest()) + hash_bytes = sha256.digest() + return ObjectID(bytes([self.object_type]) + hash_bytes[1:]) def encode(self) -> bytes: return pickle.dumps(self) diff --git a/src/knowledge/object/object_id.py b/src/knowledge/object/object_id.py index 42f313a..46955f4 100644 --- a/src/knowledge/object/object_id.py +++ b/src/knowledge/object/object_id.py @@ -44,7 +44,11 @@ class ObjectID: # pylint: disable=too-few-public-methods @staticmethod def new_chunk_id(chunk_hash: HashValue): - return ObjectID(chunk_hash.value) + assert len(chunk_hash.value) == 32, "ObjectID must be 32 bytes long" + return ObjectID(bytes([ObjectType.Chunk]) + chunk_hash.value[1:]) + + def get_object_type(self) -> ObjectType: + return ObjectType(self.value[0]) @staticmethod def hash_data(data: bytes): diff --git a/src/knowledge/vector/chroma_store.py b/src/knowledge/vector/chroma_store.py index 0056375..03006ae 100644 --- a/src/knowledge/vector/chroma_store.py +++ b/src/knowledge/vector/chroma_store.py @@ -30,9 +30,10 @@ class ChromaVectorStore(VectorBase): self.collection = collection async def insert(self, vector: [float], id: ObjectID): + logging.info(f"will insert vector: {vector} id: {str(id)}") self.collection.add( embeddings=vector, - ids=id, + ids=str(id), ) async def query(self, vector: [float], top_k: int) -> [ObjectID]: @@ -40,8 +41,10 @@ class ChromaVectorStore(VectorBase): query_embeddings=vector, n_results=top_k, ) - - return ret["ids"] + logging.info(f"query result {ret}") + if len(ret['ids']) == 0: + return [] + return list(map(ObjectID.from_base58, ret["ids"][0])) async def delete(self, id: ObjectID): self.collection.delete(