prompt from knowledge object

This commit is contained in:
tsukasa
2023-09-14 16:42:14 +08:00
parent 9cbf29e0b0
commit a0cb5480df
+57 -21
View File
@@ -165,30 +165,66 @@ class KnowledgeBase:
async def __load_object(self, object_id: ObjectID) -> KnowledgeObject: async def __load_object(self, object_id: ObjectID) -> KnowledgeObject:
if object_id.get_object_type() == ObjectType.Document: if object_id.get_object_type() == ObjectType.Document:
return DocumentObject.decode(self.store.get_object_store().get_object(object_id)) return DocumentObject.decode(self.store.get_object_store().get_object(object_id))
object = self.store.get_object_store().get_object(object_id) if object_id.get_object_type() == ObjectType.Image:
return ImageObject.decode(self.store.get_object_store().get_object(object_id))
if object_id.get_object_type() == ObjectType.Video:
return VideoObject.decode(self.store.get_object_store().get_object(object_id))
if object_id.get_object_type() == ObjectType.RichText:
return RichTextObject.decode(self.store.get_object_store().get_object(object_id))
if object_id.get_object_type() == ObjectType.Email:
return EmailObject.decode(self.store.get_object_store().get_object(object_id))
else:
pass
async def __prompt_from_objects(self, object_ids: [ObjectID]) -> AgentPrompt: async def prompt_from_objects(self, object_ids: [ObjectID]) -> AgentPrompt:
prompt = AgentPrompt() results = dict()
for object_id in object_ids: for object_id in object_ids:
object = self.store.get_object_reader().get_object(object_id) parents = self.store.get_relation_store().get_related_objects(object_id)
if object is None: # last parent is the root object
raise ValueError(f"object not found: {object_id}") root_object_id = parents[-1]
if object.get_object_type() == ObjectType.Document: if results[root_object_id] is None:
document = object results[root_object_id] = [root_object_id, object_id]
prompt.messages.append({"role": "agent", "content": document.get_body()}) else:
if object.get_object_type() == ObjectType.Image: results[root_object_id].append(object_id)
image = object
prompt.messages.append({"role": "agent", "content": json.dumps(image.get_desc())}) content = "I found the following contents described with json format:\n"
if object.get_object_type() == ObjectType.Video: result_desc = []
video = object for result in results.values():
prompt.messages.append({"role": "agent", "content": json.dumps(video.get_desc())}) # first element in result is the root object
if object.get_object_type() == ObjectType.RichText: root_object = await self.__load_object(result[0])
rich_text = object if root_object.get_object_type() == ObjectType.Email:
prompt.messages.append({"role": "agent", "content": json.dumps(rich_text.get_desc())}) email = await self.__load_object(object_id)
if object.get_object_type() == ObjectType.Email: desc = email.get_desc()
email = object desc["type"] = "email"
prompt.messages.append({"role": "agent", "content": json.dumps(email.get_desc())}) desc["contents"] = []
result_desc.append(desc)
upper_list = desc["contents"]
result = result[1:]
else:
upper_list = result_desc
for object_id in result:
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)
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)
desc = video.get_desc()
desc["type"] = "video"
upper_list.append(desc)
else:
pass
content += json.dumps(result_desc)
content += ".\n"
prompt = AgentPrompt()
prompt.add_message("knowledge", content)
return prompt return prompt