embedding email object

This commit is contained in:
tsukasa
2023-09-12 15:28:59 +08:00
parent ddc103af38
commit 8ac6065332
10 changed files with 155 additions and 42 deletions
+2 -1
View File
@@ -4,7 +4,8 @@ from .chatsession import AIChatSession
from .agent import AIAgent,AIAgentTemplete,AgentPrompt
from .compute_kernel import ComputeKernel,ComputeTask
from .compute_node import ComputeNode,LocalComputeNode
from .open_ai_node import OpenAI_ComputeNode
from .open_ai_node import OpenAI_ComputeNode
from .knowledge_base import KnowledgeBase
from .role import AIRole,AIRoleGroup
from .workflow import Workflow
from .bus import AIBus
+81 -32
View File
@@ -1,47 +1,96 @@
# define a knowledge base class
import json
from . import AgentPrompt, ComputeKernel
from ..knowledge.object import KnowledgeObject, ObjectType, EmailObject, TextChunkObject, ImageObject
from ..knowledge.store import ObjectStorage
from ..knowledge.vector.vector_base import VectorBase
from ..knowledge import *
class KnowledgeBase:
def __init__(self) -> None:
self.object_store = ObjectStorage()
self.vector_base = VectorBase()
_instance = None
def __new__(cls):
if cls._instance is None:
cls._instance = super().__new__(cls)
cls._instance.__singleton_init__()
return cls._instance
def __singleton_init__(self) -> None:
self.store = KnowledgeStore()
self.compute_kernel = ComputeKernel()
async def insert(self, object: KnowledgeObject):
if object.object_type == ObjectType.Email:
email: EmailObject = object
for text_id in email.text:
[text, _] = self.object_store.get(text_id)
text: TextChunkObject = text
vector = await self.compute_kernel.do_text_embedding(text.text)
self.vector_base.insert(vector, text_id)
for image_id in email.images:
[image, _] = self.object_store.get(image_id)
image: ImageObject = image
vector = await self.compute_kernel.do_text_embedding(image.meta)
self.vector_base.insert(vector, image_id)
vector = await self.compute_kernel.do_text_embedding(email.meta)
self.vector_base.insert(vector, email.get_id())
async def __embedding_document(self, document: DocumentObject):
for chunk_id in document.get_chunk_list():
chunk = self.store.get_chunk_reader().get_chunk(chunk_id)
if chunk is None:
raise ValueError(f"text chunk not found: {chunk_id}")
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)
async def __embedding_image(self, image: ImageObject):
desc = {}
if not image.get_meta():
desc["meta"] = image.get_meta()
if not image.get_exif():
desc["exif"] = image.get_exif()
if 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())
async def __embedding_vedio(self, vedio: VideoObject):
desc = {}
if not vedio.get_meta():
desc["meta"] = vedio.get_meta()
if not vedio.get_info():
desc["info"] = vedio.get_info()
if 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())
async def __embedding_rich_text(self, rich_text: RichTextObject):
for document in rich_text.get_documents().values():
await self.__embedding_document(document)
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)
for rich_text in rich_text.get_rich_texts().values():
await self.__embedding_rich_text(rich_text)
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.__embedding_rich_text(email.get_rich_text())
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)
if object.get_object_type() == ObjectType.RichText:
await self.__embedding_rich_text(object)
if object.get_object_type() == ObjectType.Email:
await self.__embedding_email(object)
else:
pass
async def query(self, prompt: AgentPrompt) -> AgentPrompt:
async def query(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.vector_base.query(vector, 10)
for object_id in object_ids:
if object_id.object_type == ObjectType.Email:
[object, email] = self.object_store.get(object_id)
if object.object_type == ObjectType.Email:
email: EmailObject = object
prompt.append(AgentPrompt())
prompt
object_ids = self.store.get_vector_store("default").query(vector, 10)
results.append(object_ids)
return results