Add EmailObjectBuilder impl and fix some bugs

This commit is contained in:
liyaxing
2023-09-11 20:32:52 +08:00
parent 58c0d74ca1
commit ddc103af38
9 changed files with 150 additions and 45 deletions
+2 -2
View File
@@ -1,7 +1,7 @@
from ..object import KnowledgeObject
from ..data import ChunkList, ChunkListWriter
from ..object import ObjectType
from .. import KnowledgeStore
# desc
# meta
@@ -50,7 +50,7 @@ class DocumentObjectBuilder:
return self
def build(self) -> DocumentObject:
chunk_list = ChunkListWriter.create_chunk_list_from_text(
chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_text(
self.text,
1024 * 4,
)
+94 -8
View File
@@ -1,5 +1,12 @@
from .rich_text_object import RichTextObject, RichTextObjectBuilder
from ..object import ObjectID, ObjectType, KnowledgeObject
from .document_object import DocumentObjectBuilder
from .image_object import ImageObjectBuilder
from .video_object import VideoObjectBuilder
from .rich_text_object import RichTextObjectBuilder
import os
import json
import logging
class EmailObject(KnowledgeObject):
@@ -24,16 +31,35 @@ class EmailObject(KnowledgeObject):
return self.body["rich_text"]
"""
EmailObject folder structure:
.
├── email.txt
└── meta.json
├── image
│ ├── image1.jpg
│ ├── image2.jpg
│ └── ...
├── video
│ ├── video1.mp4
│ ├── video2.mv
│ └── ...
└── audio
├── audio1.m4a
├── audio2.flac
└── ...
EmailObjectBuilder will read the target folder and build the EmailObject
Store meta.json to meta in EmailObject
Store email.txt to DocumentObject and RichTextObject in EmailObject
Store very image file in image folder to ImageObject and RichTextObject in EmailObject, etc
"""
class EmailObjectBuilder:
def __init__(self, meta: dict, tags: dict, folder: str):
self.meta = meta
def __init__(self, tags: dict, folder: str):
self.tags = tags
self.folder = folder
def set_meta(self, meta: dict):
self.meta = meta
return self
def set_tags(self, tags: dict):
self.tags = tags
return self
@@ -43,5 +69,65 @@ class EmailObjectBuilder:
return self
def build(self) -> EmailObject:
content = RichTextObjectBuilder(self.folder).build()
return EmailObject(self.meta, self.tags, content)
# Read meta.json
meta = {}
meta_file = os.path.join(self.folder, "meta.json")
if os.path.exists(meta_file):
logging.info(f"Will read meta.json {meta_file}")
with open(meta_file, "r", encoding="utf-8") as f:
meta = json.load(f)
else:
logging.info(f"Meta file missing! {meta_file}")
# Read email.txt
documents = {}
content_file = os.path.join(self.folder, "email.txt")
if os.path.exists(content_file):
logging.info(f"Will read email.txt {content_file}")
try:
with open(content_file, "r", encoding="utf-8") as f:
text = f.read()
document = DocumentObjectBuilder({}, {}, text).build()
documents = {"email.txt": document}
except Exception as e:
logging.error(f"Failed to read email.txt {content_file} {e}")
else:
logging.info(f"Content file missing! {content_file}")
# Process image files
images = {}
image_dir = os.path.join(self.folder, "image")
if os.path.exists(image_dir):
for image_file in os.listdir(image_dir):
image_path = os.path.join(image_dir, image_file)
logging.info(f"Will read image file {image_path}")
try:
image = ImageObjectBuilder({}, {}, image_path).build()
images[image_file] = image
except Exception as e:
logging.error(f"Failed to read image file {image_path} {e}")
continue
# Process video files
videos = {}
video_dir = os.path.join(self.folder, "video")
if os.path.exists(video_dir):
for video_file in os.listdir(video_dir):
video_path = os.path.join(video_dir, video_file)
logging.info(f"Will read video file {video_path}")
try:
video = VideoObjectBuilder({}, {}, video_path).build()
videos[video_file] = video
except Exception as e:
logging.error(f"Failed to read video file {video_path} {e}")
continue
# Create RichTextObject
rich_text = RichTextObject(images, videos, documents)
# Create EmailObject
return EmailObject(meta, {}, rich_text)
+2 -2
View File
@@ -1,7 +1,7 @@
from ..object import KnowledgeObject
from ..data import ChunkList, ChunkListWriter
from ..object import ObjectType
from .. import KnowledgeStore
# desc
# meta
@@ -82,7 +82,7 @@ class ImageObjectBuilder:
return self
def build(self) -> ImageObject:
chunk_list = ChunkListWriter.create_chunk_list_from_file(
chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_file(
self.image_file, 1024 * 1024 * 4, self.restore_file
)
exif = get_exif_data(self.image_file)
+25 -25
View File
@@ -7,68 +7,68 @@ from .image_object import ImageObjectBuilder, ImageObject
from .document_object import DocumentObjectBuilder, DocumentObject
class RichTextObject(KnowledgeObject):
def __init__(self, images: dict = {}, videos: dict = {}, documents: dict = {}, rich_texts: dict = {}}):
def __init__(self, images: dict = {}, videos: dict = {}, documents: dict = {}, rich_texts: dict = {}):
desc = dict()
desc.images = dict()
desc.videos = dict()
desc.documents = dict()
desc.rich_texts = dict()
desc["images"] = images
desc["videos"] = videos
desc["documents"] = documents
desc["rich_texts"] = rich_texts
super().__init__(ObjectType.RichText, desc)
def add_image_with_key(self, key, image_object: ImageObject):
assert self.desc.images[key] == None
self.desc.images[key] = image_object
assert self.desc["images"][key] == None
self.desc["images"][key] = image_object
def add_image(self, image_object: ImageObject):
self.desc.images[image_object.object_id()] = image_object
self.desc["images"][image_object.object_id()] = image_object
def get_image_with_key(self, key) -> ImageObject:
return self.desc.images[key]
return self.desc["images"][key]
def get_images(self) -> dict:
return self.desc.images
return self.desc["images"]
def add_video_with_key(self, key, video_object: VideoObject):
assert self.desc.videos[key] == None
self.desc.videos[key] = video_object
assert self.desc["videos"][key] == None
self.desc["videos"][key] = video_object
def add_video(self, video_object: VideoObject):
self.desc.videos[video_object.object_id()] = video_object
self.desc["videos"][video_object.object_id()] = video_object
def get_video_with_key(self, key) -> VideoObject:
return self.desc.videos[key]
return self.desc["videos"][key]
def get_videos(self) -> dict:
return self.desc.videos
return self.desc["videos"]
def add_document_with_key(self, key, document_object: DocumentObject):
assert self.desc.documents[key] == None
self.desc.documents[key] = document_object
assert self.desc["documents"][key] == None
self.desc["documents"][key] = document_object
def add_document(self, document_object: DocumentObject):
self.desc.documents[document_object.object_id()] = document_object
self.desc["documents"][document_object.object_id()] = document_object
def get_document_with_key(self, key) -> DocumentObject:
return self.desc.documents[key]
return self.desc["documents"][key]
def get_documents(self) -> dict:
return self.desc.documents
return self.desc["documents"]
def add_rich_text_with_key(self, key, rich_text_object):
assert self.desc.rich_texts[key] == None
self.desc.rich_texts[key] = rich_text_object
assert self.desc["rich_texts"][key] == None
self.desc["rich_texts"][key] = rich_text_object
def add_rich_text(self, rich_text_object):
self.desc.rich_texts[rich_text_object.object_id()] = rich_text_object
self.desc["rich_texts"][rich_text_object.object_id()] = rich_text_object
def get_rich_text_with_key(self, key):
return self.desc.rich_texts[key]
return self.desc["rich_texts"][key]
def get_rich_texts(self) -> dict:
return self.desc.rich_texts
return self.desc["rich_texts"]
class RichTextObjectBuilder:
+2 -2
View File
@@ -1,7 +1,7 @@
from ..object import KnowledgeObject
from ..data import ChunkList, ChunkListWriter
from ..object import ObjectType
from .. import KnowledgeStore
# desc
# meta
@@ -77,7 +77,7 @@ class VideoObjectBuilder:
return self
def build(self) -> VideoObject:
chunk_list = ChunkListWriter.create_chunk_list_from_file(
chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_file(
self.video_file, 1024 * 1024 * 4, self.restore_file
)
info = get_video_info(self.video_file)