Add chunk reader impl

src/knowledge/data/__init__.py

@@ -3,3 +3,4 @@ from .tracker import ChunkTracker
 from .chunk_store import ChunkStore
 from .writer import ChunkListWriter
 from .chunk_list import ChunkList
+from .reader import ChunkReader, Chunk

src/knowledge/data/chunk_store.py

@@ -22,5 +22,7 @@ class ChunkStore:
 
     def delete_chunk(self, chunk_id: ChunkID):
         self.blob.delete(chunk_id)
-        
+      
+    def get_chunk_file_path(self, chunk_id: ChunkID) -> str:
+        return self.blob.get_full_path(chunk_id, False)
     

src/knowledge/data/reader.py

@@ -0,0 +1,85 @@
+from .chunk import ChunkID, PositionType, PositionFileRange
+from .chunk_store import ChunkStore
+from .tracker import ChunkTracker
+from ..object import HashValue
+import logging
+from typing import List
+import hashlib
+
+class Chunk:
+    def __init__(self, file_path: str, range_start: int, size: int = -1):
+        self.file_path = file_path
+        self.range_start = range_start
+        self.size = size
+
+    def read(self):
+        with open(self.file_path, 'rb') as f:
+            f.seek(self.range_start)
+            return f.read(self.size)
+        
+        
+
+class ChunkReader:
+    def __init__(self, chunk_store: ChunkStore, chunk_tracker: ChunkTracker):
+        self.chunk_store = chunk_store
+        self.chunk_tracker = chunk_tracker
+
+    def get_chunk(self, chunk_id: ChunkID) -> Chunk:
+        positions = self.chunk_tracker.get_position(chunk_id)
+        if positions is None:
+            logging.warning(f"chunk not found: {chunk_id}")
+            return None 
+       
+        if len(positions) == 0:
+            logging.warning(f"chunk not found: {chunk_id}")
+            return None 
+
+        for pos in positions:
+            [position, position_type] = pos
+            logging.info(f"chunk position: {chunk_id}, {position}, {position_type}")
+            if position_type == PositionType.ChunkStore:
+                file_path = self.chunk_store.get_chunk_file_path(chunk_id)
+                return Chunk(file_path, 0, -1)
+            elif position_type == PositionType.File:
+                return Chunk(position, 0, -1)
+            elif position_type == PositionType.FileRange:
+                file_range = PositionFileRange.decode(position)
+                return Chunk(file_range.path, file_range.range_begin, file_range.range_end - file_range.range_begin)
+            else:
+                raise ValueError(f"invalid position type: {position_type}")
+
+        logging.error(f"chunk not found: {chunk_id}")
+        return None
+    
+    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]):
+        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()
+    
+    def read_text_chunk_list(self, chunk_ids: List[ChunkID]):
+        for chunk_id in chunk_ids:
+            chunk = self.get_chunk(chunk_id)
+            if chunk is None:
+                raise ValueError(f"text chunk not found: {chunk_id}")
+            
+            yield chunk.read().decode("utf-8")
+                
+    def calc_file_hash(self, file_path: str) -> HashValue:
+        hash_obj = hashlib.sha256()
+        with open(file_path, "rb") as file:
+            while True:
+                chunk = file.read(1024 * 1024)
+                if not chunk:
+                    break
+                hash_obj.update(chunk)
+        return HashValue(hash_obj.digest())
+
+    def calc_text_hash(self, text: str) -> HashValue:
+        hash_obj = hashlib.sha256()
+        hash_obj.update(text.encode("utf-8"))

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
 
 class ChunkTracker:
     def __init__(self, root_dir: str):
@@ -61,11 +61,11 @@ class ChunkTracker:
         )
         self.conn.commit()
 
-    def get_position(self, chunk_id: ChunkID) -> (str, PositionType):
+    def get_position(self, chunk_id: ChunkID) -> List[(str, PositionType)]:
         self.cursor.execute(
             """
             SELECT pos, pos_type FROM chunks WHERE id = ?
         """,
             (str(chunk_id),),
         )
-        return self.cursor.fetchone()
+        return self.cursor.fetchmany()

src/requirements.txt

@@ -2,4 +2,6 @@ chromadb==0.4
 openai==0.28
 toml==0.10
 Pillow==10.0
-moviepy==1.0
+moviepy==1.0
+base58==2.1
+base36==0.1