From 8b862c0075bf9e1ba2ff9cfe4bf1ff72b6694970 Mon Sep 17 00:00:00 2001 From: liyaxing Date: Tue, 26 Sep 2023 02:30:59 +0800 Subject: [PATCH] Improve chunk relate code and add read test case --- src/knowledge/core_object/email_object.py | 6 +++--- src/knowledge/core_object/image_object.py | 21 ++++++++++++++------- src/knowledge/data/reader.py | 18 ++++++++++++++---- src/knowledge/data/writer.py | 5 ++++- src/knowledge/object/object.py | 4 +++- test/test_object.py | 17 +++++++++++++++++ 6 files changed, 55 insertions(+), 16 deletions(-) diff --git a/src/knowledge/core_object/email_object.py b/src/knowledge/core_object/email_object.py index 37531ff..108c9cc 100644 --- a/src/knowledge/core_object/email_object.py +++ b/src/knowledge/core_object/email_object.py @@ -21,13 +21,13 @@ class EmailObject(KnowledgeObject): super().__init__(ObjectType.Email, desc, body) - def get_meta(self): + def get_meta(self) -> dict: return self.desc["meta"] - def get_tags(self): + def get_tags(self) -> dict: return self.desc["tags"] - def get_rich_text(self): + def get_rich_text(self) -> RichTextObject: return self.body["content"] diff --git a/src/knowledge/core_object/image_object.py b/src/knowledge/core_object/image_object.py index 2c02ad6..340bcc2 100644 --- a/src/knowledge/core_object/image_object.py +++ b/src/knowledge/core_object/image_object.py @@ -2,6 +2,7 @@ from ..object import KnowledgeObject from ..data import ChunkList, ChunkListWriter from ..object import ObjectType from .. import KnowledgeStore +import os # desc # meta @@ -13,30 +14,34 @@ from .. import KnowledgeStore class ImageObject(KnowledgeObject): - def __init__(self, meta: dict, tags: dict, exif: dict, chunk_list: ChunkList): + def __init__(self, meta: dict, tags: dict, exif: dict, file_size: int, chunk_list: ChunkList): desc = dict() body = dict() desc["meta"] = meta desc["exif"] = exif desc["tags"] = tags desc["hash"] = chunk_list.hash.to_base58() + desc["file_size"] = file_size body["chunk_list"] = chunk_list.chunk_list super().__init__(ObjectType.Image, desc, body) - def get_meta(self): + def get_meta(self) -> dict: return self.desc["meta"] - def get_exif(self): + def get_exif(self) -> dict: return self.desc["exif"] - def get_tags(self): + def get_tags(self) -> dict: return self.desc["tags"] - def get_hash(self): + def get_hash(self) -> str: return self.desc["hash"] - def get_chunk_list(self): + def get_file_size(self) -> int: + return self.desc["file_size"] + + def get_chunk_list(self) -> ChunkList: return self.body["chunk_list"] @@ -82,8 +87,10 @@ class ImageObjectBuilder: return self def build(self) -> ImageObject: + + file_size = os.path.getsize(self.image_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) - return ImageObject(self.meta, self.tags, exif, chunk_list) + return ImageObject(self.meta, self.tags, exif, file_size, chunk_list) diff --git a/src/knowledge/data/reader.py b/src/knowledge/data/reader.py index c4cbb96..ede2648 100644 --- a/src/knowledge/data/reader.py +++ b/src/knowledge/data/reader.py @@ -12,7 +12,7 @@ class Chunk: self.range_start = range_start self.size = size - def read(self): + def read(self) -> bytes: with open(self.file_path, 'rb') as f: f.seek(self.range_start) return f.read(self.size) @@ -26,6 +26,8 @@ class ChunkReader: def get_chunk(self, chunk_id: ChunkID) -> Chunk: positions = self.chunk_tracker.get_position(chunk_id) + logging.info(f"chunk positions: {chunk_id}, {positions}") + if positions is None: logging.warning(f"chunk not found: {chunk_id}") return None @@ -54,15 +56,23 @@ class ChunkReader: def get_chunk_list(self, chunk_list: List[ChunkID]) -> List[Chunk]: return [self.get_chunk(chunk_id) for chunk_id in chunk_list] - def read_chunk_list(self, chunk_ids: List[ChunkID]): + def read_chunk_list(self, chunk_ids: List[ChunkID]) -> bytes: for chunk_id in chunk_ids: chunk = self.get_chunk(chunk_id) if chunk is None: raise ValueError(f"chunk not found: {chunk_id}") - yield from chunk.read() + yield chunk.read() - def read_text_chunk_list(self, chunk_ids: List[ChunkID]): + def read_chunk_list_to_single_bytes(self, chunk_ids: List[ChunkID]) -> bytes: + chunks = [] + for chunk in self.read_chunk_list(chunk_ids): + chunks.append(chunk) + + image_data = b''.join(chunks) + return image_data + + def read_text_chunk_list(self, chunk_ids: List[ChunkID]) -> str: for chunk_id in chunk_ids: chunk = self.get_chunk(chunk_id) if chunk is None: diff --git a/src/knowledge/data/writer.py b/src/knowledge/data/writer.py index ad6c520..93b9ecc 100644 --- a/src/knowledge/data/writer.py +++ b/src/knowledge/data/writer.py @@ -27,6 +27,8 @@ class ChunkListWriter: chunk = file.read(chunk_size) if not chunk: break + + chunk_len = len(chunk) chunk_id = ChunkID.hash_data(chunk) chunk_list.append(chunk_id) @@ -38,8 +40,9 @@ class ChunkListWriter: ) self.chunk_store.put_chunk(chunk_id, chunk) else: + pos = file.tell() file_range = PositionFileRange( - file_path, file.tell() - chunk_size, chunk_size + file_path, pos - chunk_len, pos ) self.chunk_tracker.add_position( chunk_id, str(file_range), PositionType.FileRange diff --git a/src/knowledge/object/object.py b/src/knowledge/object/object.py index a63e30a..ac6f2af 100644 --- a/src/knowledge/object/object.py +++ b/src/knowledge/object/object.py @@ -1,6 +1,8 @@ # define a object type enum +from __future__ import annotations from abc import ABC, abstractmethod from enum import Enum + from .object_id import ObjectID, ObjectType import hashlib import json @@ -61,5 +63,5 @@ class KnowledgeObject(ABC): return pickle.dumps(self) @staticmethod - def decode(data: bytes): + def decode(data: bytes) -> "ImageObject": return pickle.loads(data) diff --git a/test/test_object.py b/test/test_object.py index 655e2a2..53dc6b0 100644 --- a/test/test_object.py +++ b/test/test_object.py @@ -24,6 +24,7 @@ from knowledge import ( ObjectRelationStore, KnowledgeStore, EmailObject, + ImageObject, ) import asyncio import unittest @@ -57,6 +58,22 @@ class TestVectorSTorage(unittest.TestCase): ret2 = obj.encode() self.assertEqual(ret, ret2) + images = email_object.get_rich_text().get_images() + image_keys = list(images.keys()) + print("got image list: ", image_keys) + + image_id = images[image_keys[1]] + print(f"got image object: {image_keys[1]} {image_id.to_base58()}") + + buf = KnowledgeStore().get_object_store().get_object(image_id) + image_obj= ImageObject.decode(buf) + file_size = image_obj.get_file_size() + print(f"got image object: {image_id.to_base58()}, size: {file_size}") + + + image_data = KnowledgeStore().get_chunk_reader().read_chunk_list_to_single_bytes(image_obj.get_chunk_list()) + self.assertEqual(file_size, len(image_data)) + def test_relation(self): obj1 = ObjectID.hash_data("12345".encode("utf-8"))