embedding email object
This commit is contained in:
@@ -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
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
from .document_object import DocumentObject, DocumentObjectBuilder
|
||||
from .image_object import ImageObject, ImageObjectBuilder
|
||||
from .video_object import VideoObject, VideoObjectBuilder
|
||||
from .rich_text_object import RichTextObject, RichTextObjectBuilder
|
||||
from .email_object import EmailObject, EmailObjectBuilder
|
||||
@@ -18,7 +18,7 @@ class ImageObject(KnowledgeObject):
|
||||
body = dict()
|
||||
desc["meta"] = meta
|
||||
desc["exif"] = exif
|
||||
desc
|
||||
desc["tags"] = tags
|
||||
desc["hash"] = chunk_list.hash.to_base58()
|
||||
body["chunk_list"] = chunk_list.chunk_list
|
||||
|
||||
|
||||
@@ -8,7 +8,6 @@ import base36
|
||||
|
||||
class ObjectType(Enum):
|
||||
Chunk = 7
|
||||
TextChunk = 100
|
||||
Image = 101
|
||||
Video = 102
|
||||
Document = 103
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import os
|
||||
from .object import ObjectStore
|
||||
from .data import ChunkStore, ChunkTracker, ChunkListWriter, ChunkReader
|
||||
from .vector import ChromaVectorStore, VectorBase
|
||||
import logging
|
||||
|
||||
|
||||
@@ -37,6 +38,7 @@ class KnowledgeStore:
|
||||
self.chunk_tracker = ChunkTracker(chunk_store_dir)
|
||||
self.chunk_list_writer = ChunkListWriter(self.chunk_store, self.chunk_tracker)
|
||||
self.chunk_reader = ChunkReader(self.chunk_store, self.chunk_tracker)
|
||||
self.vector_store = {}
|
||||
|
||||
|
||||
def get_object_store(self) -> ObjectStore:
|
||||
@@ -53,3 +55,8 @@ class KnowledgeStore:
|
||||
|
||||
def get_chunk_reader(self) -> ChunkReader:
|
||||
return self.chunk_reader
|
||||
|
||||
def get_vector_store(self, model_name: str) -> VectorBase:
|
||||
if model_name not in self.vector_store:
|
||||
self.vector_store[model_name] = ChromaVectorStore(model_name)
|
||||
return self.vector_store[model_name]
|
||||
|
||||
@@ -6,11 +6,11 @@ import os
|
||||
|
||||
|
||||
class ChromaVectorStore(VectorBase):
|
||||
def __init__(self, db_url, model_name: str) -> None:
|
||||
super().__init__(db_url, model_name)
|
||||
def __init__(self, model_name: str) -> None:
|
||||
super().__init__(model_name)
|
||||
|
||||
logging.info(
|
||||
"will init chroma vector store, db={}, model={}".format(db_url, model_name)
|
||||
"will init chroma vector store, model={}".format(model_name)
|
||||
)
|
||||
|
||||
directory = os.path.join(
|
||||
|
||||
@@ -3,8 +3,7 @@ from ..object import ObjectID
|
||||
|
||||
# define a vector base class
|
||||
class VectorBase:
|
||||
def __init__(self, db_url, model_name) -> None:
|
||||
self.db_url = db_url
|
||||
def __init__(self, model_name) -> None:
|
||||
self.model_name = model_name
|
||||
|
||||
async def insert(self, vector: [float], id: ObjectID):
|
||||
|
||||
Reference in New Issue
Block a user