From 58c0d74ca131c8c505aed52dc73102b4fa166370 Mon Sep 17 00:00:00 2001
From: liyaxing <liyaxing@buckyos.com>
Date: Mon, 11 Sep 2023 17:30:06 +0800
Subject: [PATCH] Add chunk reader impl

---
 src/knowledge/data/__init__.py    |  1 +
 src/knowledge/data/chunk_store.py |  4 +-
 src/knowledge/data/reader.py      | 85 +++++++++++++++++++++++++++++++
 src/knowledge/data/tracker.py     |  6 +--
 src/requirements.txt              |  4 +-
 5 files changed, 95 insertions(+), 5 deletions(-)
 create mode 100644 src/knowledge/data/reader.py

diff --git a/src/knowledge/data/__init__.py b/src/knowledge/data/__init__.py
index 2b391c9..3f470f0 100644
--- a/src/knowledge/data/__init__.py
+++ b/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
\ No newline at end of file
diff --git a/src/knowledge/data/chunk_store.py b/src/knowledge/data/chunk_store.py
index d06eb38..666b027 100644
--- a/src/knowledge/data/chunk_store.py
+++ b/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)
     
\ No newline at end of file
diff --git a/src/knowledge/data/reader.py b/src/knowledge/data/reader.py
new file mode 100644
index 0000000..c4cbb96
--- /dev/null
+++ b/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"))
\ No newline at end of file
diff --git a/src/knowledge/data/tracker.py b/src/knowledge/data/tracker.py
index 2905d11..57daa38 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
 
 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()
diff --git a/src/requirements.txt b/src/requirements.txt
index 811e8e1..1b90399 100644
--- a/src/requirements.txt
+++ b/src/requirements.txt
@@ -2,4 +2,6 @@ chromadb==0.4
 openai==0.28
 toml==0.10
 Pillow==10.0
-moviepy==1.0
\ No newline at end of file
+moviepy==1.0
+base58==2.1
+base36==0.1
\ No newline at end of file