From 2761883ae81a2d8bb3048baef665c85b22b70039 Mon Sep 17 00:00:00 2001 From: liyaxing Date: Mon, 11 Sep 2023 02:05:24 +0800 Subject: [PATCH] Add some core objects impl --- src/knowledge/__init__.py | 3 +- src/knowledge/core_object/__init__.py | 4 + src/knowledge/core_object/document_object.py | 57 ++++++++++++ src/knowledge/core_object/email_object.py | 47 ++++++++++ src/knowledge/core_object/image_object.py | 89 +++++++++++++++++++ src/knowledge/core_object/rich_text_object.py | 80 +++++++++++++++++ src/knowledge/core_object/video_object.py | 84 +++++++++++++++++ src/knowledge/object/__init__.py | 2 +- src/knowledge/object/object.py | 31 +------ src/knowledge/object/object_id.py | 5 +- src/knowledge/store.py | 2 +- 11 files changed, 372 insertions(+), 32 deletions(-) create mode 100644 src/knowledge/core_object/__init__.py create mode 100644 src/knowledge/core_object/document_object.py create mode 100644 src/knowledge/core_object/email_object.py create mode 100644 src/knowledge/core_object/image_object.py create mode 100644 src/knowledge/core_object/rich_text_object.py create mode 100644 src/knowledge/core_object/video_object.py diff --git a/src/knowledge/__init__.py b/src/knowledge/__init__.py index e2cdeb2..680c632 100644 --- a/src/knowledge/__init__.py +++ b/src/knowledge/__init__.py @@ -1,4 +1,5 @@ from .object import * from .vector import * from .data import * -from .store import KnowledgeStore \ No newline at end of file +from .store import KnowledgeStore +from .core_object import * \ No newline at end of file diff --git a/src/knowledge/core_object/__init__.py b/src/knowledge/core_object/__init__.py new file mode 100644 index 0000000..065de74 --- /dev/null +++ b/src/knowledge/core_object/__init__.py @@ -0,0 +1,4 @@ +from .document_object import DocumentObject, DocumentObjectBuilder +from .image_object import ImageObject, ImageObjectBuilder +from .video_object import VideoObject, VideoObjectBuilder +from .email_object import EmailObject, EmailObjectBuilder \ No newline at end of file diff --git a/src/knowledge/core_object/document_object.py b/src/knowledge/core_object/document_object.py new file mode 100644 index 0000000..dd9bb8b --- /dev/null +++ b/src/knowledge/core_object/document_object.py @@ -0,0 +1,57 @@ +from ..object import KnowledgeObject +from ..data import ChunkList, ChunkListWriter +from ..object import ObjectType + + +# desc +# meta +# hash: "file-hash", +# tags: {} +# body +# chunk_list: [chunk_id, chunk_id, ...] + + +class DocumentObject(KnowledgeObject): + def __init__(self, meta: dict, tags: dict, chunk_list: ChunkList): + desc = dict() + body = dict() + desc["meta"] = meta + desc["tags"] = tags + desc["hash"] = chunk_list.hash.to_base58() + body["chunk_list"] = chunk_list.chunk_list + + super().__init__(ObjectType.Document, desc, body) + + def get_meta(self): + return self.desc["meta"] + + def get_tags(self): + return self.desc["tags"] + + def get_hash(self): + return self.desc["hash"] + + def get_chunk_list(self): + return self.body["chunk_list"] + + +class DocumentObjectBuilder: + def __init__(self, meta: dict, tags: dict, text: str): + self.meta = meta + self.tags = tags + self.text = text + + def set_meta(self, meta: dict): + self.meta = meta + return self + + def set_text(self, text: str): + self.text = text + return self + + def build(self) -> DocumentObject: + chunk_list = ChunkListWriter.create_chunk_list_from_text( + self.text, + 1024 * 4, + ) + return DocumentObject(self.meta, self.tags, chunk_list) diff --git a/src/knowledge/core_object/email_object.py b/src/knowledge/core_object/email_object.py new file mode 100644 index 0000000..fe5e0d4 --- /dev/null +++ b/src/knowledge/core_object/email_object.py @@ -0,0 +1,47 @@ +from .rich_text_object import RichTextObject, RichTextObjectBuilder +from ..object import ObjectID, ObjectType, KnowledgeObject + + +class EmailObject(KnowledgeObject): + def __init__(self, meta: dict, tags: dict, rich_text: RichTextObject): + desc = dict() + body = dict() + desc["meta"] = meta + desc["tags"] = tags + + # FIXME rich text content store in desc or body? which one is better? + body["content"] = rich_text + + super().__init__(ObjectType.Email, desc, body) + + def get_meta(self): + return self.desc["meta"] + + def get_tags(self): + return self.desc["tags"] + + def get_rich_text(self): + return self.body["rich_text"] + + +class EmailObjectBuilder: + def __init__(self, meta: dict, tags: dict, folder: str): + self.meta = meta + 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 + + def set_folder(self, folder: str): + self.folder = folder + return self + + def build(self) -> EmailObject: + content = RichTextObjectBuilder(self.folder).build() + return EmailObject(self.meta, self.tags, content) diff --git a/src/knowledge/core_object/image_object.py b/src/knowledge/core_object/image_object.py new file mode 100644 index 0000000..d6a0cde --- /dev/null +++ b/src/knowledge/core_object/image_object.py @@ -0,0 +1,89 @@ +from ..object import KnowledgeObject +from ..data import ChunkList, ChunkListWriter +from ..object import ObjectType + + +# desc +# meta +# tags +# hash: "file-hash", +# exif: {} +# body +# chunk_list: [chunk_id, chunk_id, ...] + + +class ImageObject(KnowledgeObject): + def __init__(self, meta: dict, tags: dict, exif: dict, chunk_list: ChunkList): + desc = dict() + body = dict() + desc["meta"] = meta + desc["exif"] = exif + desc + desc["hash"] = chunk_list.hash.to_base58() + body["chunk_list"] = chunk_list.chunk_list + + super().__init__(ObjectType.Image, desc, body) + + def get_meta(self): + return self.desc["meta"] + + def get_exif(self): + return self.desc["exif"] + + def get_tags(self): + return self.desc["tags"] + + def get_hash(self): + return self.desc["hash"] + + def get_chunk_list(self): + return self.body["chunk_list"] + + +from PIL import Image +from PIL.ExifTags import TAGS + + +def get_exif_data(image_path: str): + with Image.open(image_path) as image: + exif_data = image._getexif() + + if exif_data is not None: + return { + TAGS.get(key): exif_data[key] + for key in exif_data.keys() + if key in TAGS and isinstance(exif_data[key], (bytes, str)) + } + else: + return {} + + +class ImageObjectBuilder: + def __init__(self, meta: dict, tags: dict, image_file: str): + self.meta = meta + self.tags = tags + self.image_file = image_file + self.restore_file = False + + def set_meta(self, meta: dict): + self.meta = meta + return self + + def set_tags(self, tags: dict): + self.tags = tags + return self + + def set_image_file(self, image_file: str): + self.image_file = image_file + return self + + def set_restore_file(self, restore_file: bool): + self.restore_file = restore_file + return self + + def build(self) -> ImageObject: + chunk_list = ChunkListWriter.create_chunk_list_from_file( + self.image_file, 1024 * 1024 * 4, self.restore_file + ) + exif = get_exif_data(self.image_file) + return ImageObject(self.meta, self.tags, exif, chunk_list) diff --git a/src/knowledge/core_object/rich_text_object.py b/src/knowledge/core_object/rich_text_object.py new file mode 100644 index 0000000..ac3fcbe --- /dev/null +++ b/src/knowledge/core_object/rich_text_object.py @@ -0,0 +1,80 @@ +from knowledge.object.object_id import ObjectType +from ..object import KnowledgeObject +from ..data import ChunkList, ChunkListWriter +from ..object import ObjectType +from .video_object import VideoObjectBuilder, VideoObject +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 = {}}): + desc = dict() + desc.images = dict() + desc.videos = dict() + desc.documents = dict() + desc.rich_texts = dict() + + 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 + + def add_image(self, image_object: ImageObject): + self.desc.images[image_object.object_id()] = image_object + + def get_image_with_key(self, key) -> ImageObject: + return self.desc.images[key] + + def get_images(self) -> dict: + 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 + + def add_video(self, video_object: VideoObject): + self.desc.videos[video_object.object_id()] = video_object + + def get_video_with_key(self, key) -> VideoObject: + return self.desc.videos[key] + + def get_videos(self) -> dict: + 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 + + def add_document(self, document_object: DocumentObject): + self.desc.documents[document_object.object_id()] = document_object + + def get_document_with_key(self, key) -> DocumentObject: + return self.desc.documents[key] + + def get_documents(self) -> dict: + 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 + + def add_rich_text(self, 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] + + def get_rich_texts(self) -> dict: + return self.desc.rich_texts + + +class RichTextObjectBuilder: + def __init__(self, folder: str): + self.folder = folder + + def build(self) -> RichTextObject: + # TODO + return RichTextObject() \ No newline at end of file diff --git a/src/knowledge/core_object/video_object.py b/src/knowledge/core_object/video_object.py new file mode 100644 index 0000000..ab4c262 --- /dev/null +++ b/src/knowledge/core_object/video_object.py @@ -0,0 +1,84 @@ +from ..object import KnowledgeObject +from ..data import ChunkList, ChunkListWriter +from ..object import ObjectType + + +# desc +# meta +# tags +# hash: "file-hash", +# info: {} +# body +# chunk_list: [chunk_id, chunk_id, ...] + + +class VideoObject(KnowledgeObject): + def __init__(self, meta: dict, tags: dict, info: dict, chunk_list: ChunkList): + desc = dict() + body = dict() + desc["meta"] = meta + desc["tags"] = tags + desc["info"] = info + desc["hash"] = chunk_list.hash.to_base58() + body["chunk_list"] = chunk_list.chunk_list + + super().__init__(ObjectType.Video, desc, body) + + def get_meta(self): + return self.desc["meta"] + + def get_tags(self): + return self.desc["tags"] + + def get_info(self): + return self.desc["info"] + + def get_hash(self): + return self.desc["hash"] + + def get_chunk_list(self): + return self.body["chunk_list"] + + +from moviepy.editor import VideoFileClip + + +def get_video_info(video_path: str) -> dict: + clip = VideoFileClip(video_path) + return { + "duration": clip.duration, # Duration in seconds + "fps": clip.fps, # Frames per second + "nframes": clip.reader.nframes, # Total number of frames + "size": clip.size, # Size of the frames (width, height) + } + + +class VideoObjectBuilder: + def __init__(self, meta: dict, tags: dict, video_file: str): + self.meta = meta + self.tags = tags + self.video_file = video_file + self.restore_file = False + + def set_meta(self, meta: dict): + self.meta = meta + return self + + def set_tags(self, tags: dict): + self.tags = tags + return self + + def set_video_file(self, video_file: str): + self.video_file = video_file + return self + + def set_restore_file(self, restore_file: bool): + self.restore_file = restore_file + return self + + def build(self) -> VideoObject: + chunk_list = ChunkListWriter.create_chunk_list_from_file( + self.video_file, 1024 * 1024 * 4, self.restore_file + ) + info = get_video_info(self.video_file) + return VideoObject(self.meta, self.tags, info, chunk_list) diff --git a/src/knowledge/object/__init__.py b/src/knowledge/object/__init__.py index ef1db5c..7886d7a 100644 --- a/src/knowledge/object/__init__.py +++ b/src/knowledge/object/__init__.py @@ -1,4 +1,4 @@ -from .object import EmailObject, ImageObject, TextChunkObject, ObjectID +from .object import KnowledgeObject from .blob import FileBlobStorage from .hash import HashValue, hash_data from .object_store import ObjectStore diff --git a/src/knowledge/object/object.py b/src/knowledge/object/object.py index 0eff535..b3841ed 100644 --- a/src/knowledge/object/object.py +++ b/src/knowledge/object/object.py @@ -2,18 +2,18 @@ # define a object type enum from abc import ABC, abstractmethod from enum import Enum -from .object_id import ObjectID +from .object_id import ObjectID, ObjectType import hashlib import json import pickle class KnowledgeObject(ABC): - def __init__(self, object_type: int, desc: dict = {}, body: dict = {}): + def __init__(self, object_type: ObjectType, desc: dict = {}, body: dict = {}): self.desc = desc self.body = body self.object_type = object_type - def get_object_type(self): + def get_object_type(self) -> ObjectType: return self.object_type def object_id(self) -> ObjectID: @@ -50,30 +50,5 @@ class KnowledgeObject(ABC): @staticmethod def decode(data: bytes): return pickle.loads(data) - - - -# define a text chunk class -class TextChunkObject(KnowledgeObject): # pylint: disable=too-few-public-methods - def __init__(self, text: str): - super().__init__(ObjectType.TextChunk) - self.text = text - - -# define a image class -class ImageObject(KnowledgeObject): # pylint: disable=too-few-public-methods - def __init__(self, meta, path): - super().__init__(ObjectType.Image) - self.meta = meta - self.path = path - - -# define a email class -class EmailObject(KnowledgeObject): # pylint: disable=too-few-public-methods - def __init__(self, meta): - super().__init__(ObjectType.Email) - self.meta = meta - self.text = [ObjectID] - self.images = [ObjectID] diff --git a/src/knowledge/object/object_id.py b/src/knowledge/object/object_id.py index 27c0063..3c3ac24 100644 --- a/src/knowledge/object/object_id.py +++ b/src/knowledge/object/object_id.py @@ -10,7 +10,10 @@ class ObjectType(Enum): Chunk = 7 TextChunk = 100 Image = 101 - Email = 102 + Video = 102 + Document = 103 + RichText = 104 + Email = 105 # define a object ID class to identify a object diff --git a/src/knowledge/store.py b/src/knowledge/store.py index a4a0af6..ef697fa 100644 --- a/src/knowledge/store.py +++ b/src/knowledge/store.py @@ -4,7 +4,7 @@ from .data import ChunkStore, ChunkTracker import logging -# KnowledgeStore class, 聚合ChunkStore，ChunkTracker， ObjectStore，并且是一个全局单例，可以方便的使用这三个内置的store示例 +# KnowledgeStore class, which aggregates ChunkStore, ChunkTracker, and ObjectStore, and is a global singleton that makes it easy to use these three built-in store examples class KnowledgeStore: _instance = None