From ddc103af38c79c562a94e895babcb6a357df3b48 Mon Sep 17 00:00:00 2001 From: liyaxing Date: Mon, 11 Sep 2023 20:32:52 +0800 Subject: [PATCH] Add EmailObjectBuilder impl and fix some bugs --- src/knowledge/core_object/document_object.py | 4 +- src/knowledge/core_object/email_object.py | 102 ++++++++++++++++-- src/knowledge/core_object/image_object.py | 4 +- src/knowledge/core_object/rich_text_object.py | 50 ++++----- src/knowledge/core_object/video_object.py | 4 +- src/knowledge/data/tracker.py | 4 +- src/knowledge/data/writer.py | 2 +- src/knowledge/store.py | 11 +- test/test_object.py | 14 ++- 9 files changed, 150 insertions(+), 45 deletions(-) diff --git a/src/knowledge/core_object/document_object.py b/src/knowledge/core_object/document_object.py index dd9bb8b..6d6270e 100644 --- a/src/knowledge/core_object/document_object.py +++ b/src/knowledge/core_object/document_object.py @@ -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, ) diff --git a/src/knowledge/core_object/email_object.py b/src/knowledge/core_object/email_object.py index fe5e0d4..8b8efb5 100644 --- a/src/knowledge/core_object/email_object.py +++ b/src/knowledge/core_object/email_object.py @@ -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) diff --git a/src/knowledge/core_object/image_object.py b/src/knowledge/core_object/image_object.py index d6a0cde..83ab434 100644 --- a/src/knowledge/core_object/image_object.py +++ b/src/knowledge/core_object/image_object.py @@ -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) diff --git a/src/knowledge/core_object/rich_text_object.py b/src/knowledge/core_object/rich_text_object.py index ac3fcbe..7791689 100644 --- a/src/knowledge/core_object/rich_text_object.py +++ b/src/knowledge/core_object/rich_text_object.py @@ -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: diff --git a/src/knowledge/core_object/video_object.py b/src/knowledge/core_object/video_object.py index ab4c262..56e9829 100644 --- a/src/knowledge/core_object/video_object.py +++ b/src/knowledge/core_object/video_object.py @@ -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) diff --git a/src/knowledge/data/tracker.py b/src/knowledge/data/tracker.py index 57daa38..66024c9 100644 --- a/src/knowledge/data/tracker.py +++ b/src/knowledge/data/tracker.py @@ -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 = ? diff --git a/src/knowledge/data/writer.py b/src/knowledge/data/writer.py index 7f79447..7d38e9d 100644 --- a/src/knowledge/data/writer.py +++ b/src/knowledge/data/writer.py @@ -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()) diff --git a/src/knowledge/store.py b/src/knowledge/store.py index ef697fa..dc6b17a 100644 --- a/src/knowledge/store.py +++ b/src/knowledge/store.py @@ -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 diff --git a/test/test_object.py b/test/test_object.py index 6a603fd..1e1bc56 100644 --- a/test/test_object.py +++ b/test/test_object.py @@ -1,5 +1,6 @@ import sys import os +import logging dir_path = os.path.dirname(os.path.realpath(__file__)) print(dir_path) @@ -7,9 +8,16 @@ print(dir_path) sys.path.append("{}/../src/".format(dir_path)) print(sys.path) -from knowledge import ObjectID, HashValue +root = logging.getLogger() +root.setLevel(logging.DEBUG) +handler = logging.StreamHandler(sys.stdout) +handler.setLevel(logging.DEBUG) +formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") +handler.setFormatter(formatter) +root.addHandler(handler) +from knowledge import ObjectID, HashValue, EmailObjectBuilder import asyncio import unittest @@ -26,6 +34,8 @@ class TestVectorSTorage(unittest.TestCase): data2 = HashValue.from_base36(data.to_base36()) self.assertEqual(data.to_base58(), data2.to_base58()) - + email_folder = "F:\\system\\Downloads\\8081ffdb80925f5bff9c6ab9c4756c7d" + email_object = EmailObjectBuilder({}, email_folder).build() + if __name__ == "__main__": unittest.main()