Add chunk store and chunk list helper

src/knowledge/data/__init__.py

@@ -1,2 +1,4 @@
 from .chunk import ChunkID, PositionType, PositionFileRange
-from .tracker import ChunkTracker
+from .tracker import ChunkTracker
+from .chunk_store import ChunkStore
+from .gen import ChunkListGenerator

src/knowledge/data/chunk.py

@@ -8,7 +8,7 @@ class PositionType(Enum):
     Device = 2
     File = 3
     FileRange = 4
-    ChunkManager = 5
+    ChunkStore = 5
 
 
 class PositionFileRange:

src/knowledge/data/chunk_store.py

@@ -0,0 +1,26 @@
+import os
+import logging
+from ..object import FileBlobStorage
+from .chunk import ChunkID
+
+
+class ChunkStore:
+    def __init__(self, root_dir: str):
+        logging.info(f"will init chunk store, root_dir={root_dir}")
+
+        if not os.path.exists(root_dir):
+            os.makedirs(root_dir)
+            
+        self.root = root_dir
+        self.blob = FileBlobStorage(root_dir)
+
+    def put_chunk(self, chunk_id: ChunkID, contents: bytes):
+        self.blob.put(chunk_id, contents)
+
+    def get_chunk(self, chunk_id: ChunkID) -> bytes:
+        return self.blob.get(chunk_id)
+
+    def delete_chunk(self, chunk_id: ChunkID):
+        self.blob.delete(chunk_id)
+        
+    

src/knowledge/data/gen.py

@@ -0,0 +1,94 @@
+import os
+import hashlib
+import re
+from typing import Tuple, List
+from .chunk_store import ChunkStore
+from .chunk import ChunkID, PositionFileRange, PositionType
+from ..object import HashValue
+from .tracker import ChunkTracker
+
+
+class ChunkListGenerator:
+    def __init__(self, chunk_store: ChunkStore, chunk_tracker: ChunkTracker):
+        self.chunk_store = chunk_store
+        self.chunk_tracker = chunk_tracker
+
+    def create_chunk_list_from_file(
+        self, file_path: str, chunk_size: int, restore: bool
+    ) -> Tuple[List[ChunkID], HashValue]:
+        assert (
+            chunk_size % (1024 * 1024) == 0
+        ), "chunk size should be an integral multiple of 1MB"
+        chunk_list = []
+        hash_obj = hashlib.sha256()
+
+        with open(file_path, "rb") as file:
+            while True:
+                chunk = file.read(chunk_size)
+                if not chunk:
+                    break
+                chunk_id = ChunkID.hash_data(chunk)
+                chunk_list.append(chunk_id)
+
+                hash_obj.update(chunk)
+
+                if restore:
+                    self.chunk_tracker.add_position(
+                        chunk_id, file_path, PositionType.ChunkStore
+                    )
+                    self.chunk_store.put_chunk(chunk_id, chunk)
+                else:
+                    file_range = PositionFileRange(
+                        file_path, file.tell() - chunk_size, chunk_size
+                    )
+                    self.chunk_tracker.add_position(
+                        chunk_id, file_range, PositionType.FileRange
+                    )
+
+        file_hash = HashValue(hash_obj.digest())
+        print(f"calc file hash: {file_path}, {file_hash}")
+        return chunk_list, file_hash
+
+    def create_chunk_list_from_text(
+        self, text: str, chunk_max_words: int, separator_chars: str = ".,"
+    ) -> Tuple[List[ChunkID], HashValue]:
+        text_list = self._split_text_list(text, chunk_max_words, separator_chars)
+        chunk_list = []
+        hash_obj = hashlib.sha256()
+        
+        for text in text_list:
+            chunk_bytes  = text.encode("utf-8")
+            hash_obj.update(chunk_bytes)
+            
+            chunk_id = ChunkID.hash_data(chunk_bytes)
+            chunk_list.append(chunk_id)
+            self.chunk_tracker.add_position(
+                chunk_id, "", PositionType.ChunkStore
+            )
+            self.chunk_store.put_chunk(chunk_id, chunk_bytes)
+        
+        hash = HashValue(hash_obj.digest())
+        return chunk_list, hash
+
+    @staticmethod
+    def _split_text_list(
+        text: str, chunk_max_words: int, separator_chars: str = ".,"
+    ) -> List[str]:
+        sentences = re.split(f"[{separator_chars}]", text)
+        chunk_list = []
+        chunk = []
+        word_count = 0
+        for sentence in sentences:
+            words = sentence.split()
+            for word in words:
+                if word_count < chunk_max_words:
+                    chunk.append(word)
+                    word_count += 1
+                else:
+                    chunk_list.append(" ".join(chunk))
+                    chunk = [word]
+                    word_count = 1
+        if chunk:
+            chunk_list.append(" ".join(chunk))
+        return chunk_list
+

src/knowledge/data/tracker.py

@@ -2,31 +2,17 @@ import sqlite3
 import time
 import logging
 import os
-from .chunk import  ChunkID, PositionType, PositionFileRange
+from .chunk import ChunkID, PositionType, PositionFileRange
 
 
 class ChunkTracker:
-    _instance = None
+    def __init__(self, root_dir: str):
-    
+        if not os.path.exists(root_dir):
-    def __new__(cls):
+            os.makedirs(root_dir)
-        if cls._instance is None:
+        file = os.path.join(root_dir, "chunk_tracker.db")
-            cls._instance = super().__new__(cls)
+        logging.info(f"will init chunk tracker, db={file}")
-            directory = os.path.join(
-                os.path.dirname(__file__), "../../../rootfs/data/chunk/"
-            )
-            
-            if not os.path.exists(directory):
-                os.mkdir(directory)
-            file = os.path.join(directory, "chunk_tracker.db")
-            logging.info(f"will init chunk tracker, db={file}")
-            cls._instance.__singleton_init__(file)
-            
-        return cls._instance
 
-    def __singleton_init__(self, db_path: str):
+        self.conn = sqlite3.connect(file)
-        logging.info(f"will init chunk tracker, db={db_path}")
-
-        self.conn = sqlite3.connect(db_path)
         self.cursor = self.conn.cursor()
         self.cursor.execute(
             """
@@ -43,7 +29,9 @@ class ChunkTracker:
         )
         self.conn.commit()
 
-    def add_position(self, chunk_id: ChunkID, position: str, position_type: PositionType):
+    def add_position(
+        self, chunk_id: ChunkID, position: str, position_type: PositionType
+    ):
         logging.debug(f"add chunk position: {chunk_id}, {position}, {position_type}")
 
         insert_time = update_time = int(time.time())
@@ -52,7 +40,13 @@ class ChunkTracker:
             INSERT OR REPLACE INTO chunks (id, pos, pos_type, insert_time, update_time)
             VALUES (?, ?, ?, ?, ?)
         """,
-            (str(chunk_id), position, position_type.value, insert_time, update_time,),
+            (
+                str(chunk_id),
+                position,
+                position_type.value,
+                insert_time,
+                update_time,
+            ),
         )
         self.conn.commit()
 

test/test_chunk.py

@@ -7,16 +7,27 @@ print(dir_path)
 sys.path.append("{}/../src/".format(dir_path))
 print(sys.path)
 
-from knowledge import ChunkTracker, ChunkID, HashValue, PositionType
+
+import logging
+import sys
+
+root = logging.getLogger()
+root.setLevel(logging.DEBUG)
+handler = logging.StreamHandler(sys.stdout)
+handler.setLevel(logging.DEBUG)
+formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+handler.setFormatter(formatter)
+root.addHandler(handler)
 
 
+from knowledge import ChunkTracker, ChunkID, HashValue, PositionType, KnowledgeStore, ChunkListGenerator
 import asyncio
 import unittest
 
 
 class TestChunk(unittest.TestCase):
     def test_chunk_tracker(self):
-        tracker = ChunkTracker()
+        tracker = KnowledgeStore().get_chunk_tracker()
         
         hash = HashValue.hash_data("1234567890".encode("utf-8"));
         cid = ChunkID.new_chunk_id(hash)
@@ -29,8 +40,17 @@ class TestChunk(unittest.TestCase):
         tracker.remove_position(cid)
         ret = tracker.get_position(cid)
         self.assertEqual(ret, None)
-    
+        
+    def test_chunk(self):
+        gen = ChunkListGenerator(KnowledgeStore().get_chunk_store(), KnowledgeStore().get_chunk_tracker())
+        gen.create_chunk_list_from_file("H:/test", 1024 * 1024, True)
+        
+        # Read the file
+        text_file = "H:/test.txt"
+        with open(text_file, 'r', encoding = "utf-8") as file:
+            text = file.read()
 
+        gen.create_chunk_list_from_text(text, 1024)
 
 if __name__ == "__main__":
     unittest.main()