test prompt from knowledge object

This commit is contained in:
tsukasa
2023-09-18 11:28:21 +08:00
parent 805dc88de2
commit 441583b7f8
7 changed files with 71 additions and 38 deletions
+24 -16
View File
@@ -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
+16
View File
@@ -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()
+1 -1
View File
@@ -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()
+17 -16
View File
@@ -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
# 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
+2 -1
View File
@@ -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)
+5 -1
View File
@@ -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):
+6 -3
View File
@@ -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(