Add EmailObjectBuilder impl and fix some bugs
This commit is contained in:
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -3,7 +3,7 @@ import time
|
||||
import logging
|
||||
import os
|
||||
from .chunk import ChunkID, PositionType, PositionFileRange
|
||||
from typing import List
|
||||
from typing import List, Tuple
|
||||
|
||||
class ChunkTracker:
|
||||
def __init__(self, root_dir: str):
|
||||
@@ -61,7 +61,7 @@ class ChunkTracker:
|
||||
)
|
||||
self.conn.commit()
|
||||
|
||||
def get_position(self, chunk_id: ChunkID) -> List[(str, PositionType)]:
|
||||
def get_position(self, chunk_id: ChunkID) -> List[Tuple[str, PositionType]]:
|
||||
self.cursor.execute(
|
||||
"""
|
||||
SELECT pos, pos_type FROM chunks WHERE id = ?
|
||||
|
||||
@@ -42,7 +42,7 @@ class ChunkListWriter:
|
||||
file_path, file.tell() - chunk_size, chunk_size
|
||||
)
|
||||
self.chunk_tracker.add_position(
|
||||
chunk_id, file_range, PositionType.FileRange
|
||||
chunk_id, str(file_range), PositionType.FileRange
|
||||
)
|
||||
|
||||
file_hash = HashValue(hash_obj.digest())
|
||||
|
||||
+10
-1
@@ -1,6 +1,6 @@
|
||||
import os
|
||||
from .object import ObjectStore
|
||||
from .data import ChunkStore, ChunkTracker
|
||||
from .data import ChunkStore, ChunkTracker, ChunkListWriter, ChunkReader
|
||||
import logging
|
||||
|
||||
|
||||
@@ -35,6 +35,9 @@ class KnowledgeStore:
|
||||
chunk_store_dir = os.path.join(root_dir, "chunk")
|
||||
self.chunk_store = ChunkStore(chunk_store_dir)
|
||||
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)
|
||||
|
||||
|
||||
def get_object_store(self) -> ObjectStore:
|
||||
return self.object_store
|
||||
@@ -44,3 +47,9 @@ class KnowledgeStore:
|
||||
|
||||
def get_chunk_tracker(self) -> ChunkTracker:
|
||||
return self.chunk_tracker
|
||||
|
||||
def get_chunk_list_writer(self) -> ChunkListWriter:
|
||||
return self.chunk_list_writer
|
||||
|
||||
def get_chunk_reader(self) -> ChunkReader:
|
||||
return self.chunk_reader
|
||||
|
||||
Reference in New Issue
Block a user