Improve chunk relate code and add read test case

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"]
 
 

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)

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:

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

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)