Add object and chunk base impl and test cases

src/aios_kernel/knowledge_base.py

@@ -1,7 +1,7 @@
 # define a knowledge base class
 from . import AgentPrompt, ComputeKernel
 from ..knowledge.object import KnowledgeObject, ObjectType, EmailObject, TextChunkObject, ImageObject
-from ..knowledge.object_storage import ObjectStorage
+from ..knowledge.store import ObjectStorage
 from ..knowledge.vector.vector_base import VectorBase
 
 class KnowledgeBase:

src/knowledge/__init__.py

@@ -1,3 +1,4 @@
 from .object import *
 from .vector import *
 from .data import *
+from .store import KnowledgeStore

src/knowledge/data/__init__.py

@@ -1,4 +1,5 @@
 from .chunk import ChunkID, PositionType, PositionFileRange
 from .tracker import ChunkTracker
 from .chunk_store import ChunkStore
-from .gen import ChunkListGenerator
+from .writer import ChunkListWriter
+from .chunk_list import ChunkList

src/knowledge/data/chunk_list.py

@@ -0,0 +1,14 @@
+from ..object import HashValue
+from .chunk import ChunkID
+from typing import List
+
+class ChunkList:
+    def __init__(self, chunk_list: List[ChunkID], hash: HashValue):
+        self.chunk_list = chunk_list
+        self.hash = hash
+
+    def __str__(self):
+        return self.hash.to_base58()
+
+    def __repr__(self):
+        return f"chunk_list: {self.chunk_list}, hash: {self.hash}"

src/knowledge/data/writer.py

@@ -6,16 +6,16 @@ from .chunk_store import ChunkStore
 from .chunk import ChunkID, PositionFileRange, PositionType
 from ..object import HashValue
 from .tracker import ChunkTracker
+from .chunk_list import ChunkList
 
-
-class ChunkListGenerator:
+class ChunkListWriter:
     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]:
+    ) -> ChunkList:
         assert (
             chunk_size % (1024 * 1024) == 0
         ), "chunk size should be an integral multiple of 1MB"
@@ -47,7 +47,8 @@ class ChunkListGenerator:
 
         file_hash = HashValue(hash_obj.digest())
         print(f"calc file hash: {file_path}, {file_hash}")
-        return chunk_list, file_hash
+        
+        return ChunkList(chunk_list, file_hash)
 
     def create_chunk_list_from_text(
         self, text: str, chunk_max_words: int, separator_chars: str = ".,"
@@ -57,18 +58,16 @@ class ChunkListGenerator:
         hash_obj = hashlib.sha256()
 
         for text in text_list:
-            chunk_bytes  = text.encode("utf-8")
+            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_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
+        return ChunkList(chunk_list, hash)
 
     @staticmethod
     def _split_text_list(
@@ -91,4 +90,3 @@ class ChunkListGenerator:
         if chunk:
             chunk_list.append(" ".join(chunk))
         return chunk_list
-

src/knowledge/object.py

@@ -1,65 +0,0 @@
-
-# define a object type enum
-from abc import ABC, abstractmethod
-from enum import Enum
-
-class ObjectType(Enum):
-    TextChunk = 1
-    Image = 2
-    Email = 101
-
-
-# define a object ID class to identify a object
-class ObjectID: # pylint: disable=too-few-public-methods
-    def __init__(self, object_type, digist):
-        self.object_type = object_type
-        self.digist = digist
-
-    def __str__(self):
-        return f"{self.object_type.name}:{self.digist}"
-    
-
-# define a object class
-class KnowledgeObject(ABC): # pylint: disable=too-few-public-methods
-    def __init__(self, object_type: ObjectType):
-        self.object_type = object_type
-
-    @abstractmethod
-    def get_id(self) -> ObjectID:
-        pass
-
-    # define a to binary method to convert object to binary
-    @abstractmethod
-    def to_binary(self) -> bytes:
-        pass
-
-    # define a from binary method to convert binary to object
-    @abstractmethod
-    def from_binary(self, binary: bytes):
-        pass
-
-
-# define a text chunk class
-class TextChunkObject(KnowledgeObject): # pylint: disable=too-few-public-methods
-    def __init__(self, text: str):
-        super().__init__(ObjectType.TextChunk)
-        self.text = text
-
-
-# define a image class
-class ImageObject(KnowledgeObject): # pylint: disable=too-few-public-methods
-    def __init__(self, meta, path):
-        super().__init__(ObjectType.Image)
-        self.meta = meta
-        self.path = path
-
-
-# define a email class
-class EmailObject(KnowledgeObject): # pylint: disable=too-few-public-methods
-    def __init__(self, meta):
-        super().__init__(ObjectType.Email)
-        self.meta = meta
-        self.text = [ObjectID]
-        self.images = [ObjectID]
-
-    

src/knowledge/object/__init__.py

@@ -0,0 +1,5 @@
+from .object import EmailObject, ImageObject, TextChunkObject, ObjectID
+from .blob import FileBlobStorage
+from .hash import HashValue, hash_data
+from .object_store import ObjectStore
+from .object_id import ObjectID, ObjectType

src/knowledge/object/blob.py

@@ -0,0 +1,54 @@
+import os
+import shutil
+from .object import ObjectID
+
+
+class FileBlobStorage:
+    def __init__(self, root):
+        self.root = root
+
+    def get_full_path(self, object_id: ObjectID, auto_create: bool = True):
+        if os.name == "nt":  # Windows
+            hash_str = object_id.to_base36()
+            len = 3
+        else:
+            hash_str = str(object_id)
+            len = 2
+
+        tmp, first = hash_str[:-len], hash_str[-len:]
+        second = tmp[-len:]
+
+        if os.name == "nt":  # Windows
+            if second in ["con", "aux", "nul", "prn"]:
+                second = tmp[-(len + 1) :]
+            if first in ["con", "aux", "nul", "prn"]:
+                first = f"{first}_"
+
+        path = os.path.join(self.root, first, second)
+        if auto_create and not os.path.exists(path):
+            os.makedirs(path)
+
+        path = os.path.join(path, hash_str)
+
+        return path
+
+    def write_sync(self, path: str, contents: bytes):
+        with open(path, "wb") as f:
+            f.write(contents)
+
+    def put(self, object_id: ObjectID, contents: bytes):
+        full_path = self.get_full_path(object_id)
+        self.write_sync(full_path, contents)
+
+    def get(self, object_id: ObjectID) -> bytes:
+        full_path = self.get_full_path(object_id)
+        with open(full_path, "rb") as f:
+            return f.read()
+
+    def delete(self, object_id: ObjectID):
+        full_path = self.get_full_path(object_id)
+        os.remove(full_path)
+
+    def exists(self, object_id: ObjectID) -> bool:
+        full_path = self.get_full_path(object_id)
+        return os.path.exists(full_path)

src/knowledge/object/hash.py

@@ -0,0 +1,42 @@
+import hashlib
+import base58
+import base36
+
+class HashValue:
+    def __init__(self, value: bytes):
+        assert len(value) == 32, "HashValue must be 32 bytes long"
+        self.value = value
+
+    def __str__(self) -> str:
+        return self.to_base58()
+        
+    @staticmethod
+    def hash_data(data):
+        return hash_data(data)
+
+    def to_base58(self):
+        return base58.b58encode(self.value).decode()
+
+    @staticmethod
+    def from_base58(s):
+        return HashValue(base58.b58decode(s))
+
+    def to_base36(self):
+        # Convert the bytes to int before encoding
+        num = int.from_bytes(self.value, 'big')
+        return base36.dumps(num)
+
+    @staticmethod
+    def from_base36(s):
+        # Decode to int and then convert to bytes
+        num = base36.loads(s)
+        return HashValue(num.to_bytes((num.bit_length() + 7) // 8, 'big'))
+    
+    
+HASH_VALUE_LEN = 32
+
+
+def hash_data(data: bytes):
+    sha256 = hashlib.sha256()
+    sha256.update(data)
+    return HashValue(sha256.digest())

src/knowledge/object/object.py

@@ -0,0 +1,79 @@
+
+# define a object type enum
+from abc import ABC, abstractmethod
+from enum import Enum
+from .object_id import ObjectID
+import hashlib
+import json
+import pickle
+
+class KnowledgeObject(ABC):
+    def __init__(self, object_type: int, desc: dict = {}, body: dict = {}):
+        self.desc = desc
+        self.body = body
+        self.object_type = object_type
+
+    def get_object_type(self):
+        return self.object_type
+
+    def object_id(self) -> ObjectID:
+        return self.calculate_id()
+    
+    def set_desc_with_key_value(self, key, value):
+        self.desc[key] = value
+
+    def get_desc_with_key(self, key):
+        return self.desc.get(key)
+
+    def get_desc(self) -> dict:
+        return self.desc
+    
+    def set_body_with_key_value(self, key, value):
+        self.body[key] = value
+
+    def get_body_with_key(self, key):
+        return self.body.get(key)
+
+    def get_body(self) -> dict:
+        return self.body
+    
+    def calculate_id(self):
+        # Convert the object_type and desc to string and compute the SHA256 hash
+        data = json.dumps({"object_type": self.object_type, "desc": self.desc})
+        sha256 = hashlib.sha256()
+        sha256.update(data.encode())
+        return ObjectID(sha256.digest())
+
+    def encode(self) -> bytes:
+        return pickle.dumps(self)
+
+    @staticmethod
+    def decode(data: bytes):
+        return pickle.loads(data)
+    
+
+        
+# define a text chunk class
+class TextChunkObject(KnowledgeObject): # pylint: disable=too-few-public-methods
+    def __init__(self, text: str):
+        super().__init__(ObjectType.TextChunk)
+        self.text = text
+
+
+# define a image class
+class ImageObject(KnowledgeObject): # pylint: disable=too-few-public-methods
+    def __init__(self, meta, path):
+        super().__init__(ObjectType.Image)
+        self.meta = meta
+        self.path = path
+
+
+# define a email class
+class EmailObject(KnowledgeObject): # pylint: disable=too-few-public-methods
+    def __init__(self, meta):
+        super().__init__(ObjectType.Email)
+        self.meta = meta
+        self.text = [ObjectID]
+        self.images = [ObjectID]
+
+    

src/knowledge/object/object_id.py

@@ -0,0 +1,53 @@
+# define a object type enum
+from abc import ABC, abstractmethod
+from enum import Enum
+from .hash import HashValue
+import base58
+import base36
+
+
+class ObjectType(Enum):
+    Chunk = 7
+    TextChunk = 100
+    Image = 101
+    Email = 102
+
+
+# define a object ID class to identify a object
+class ObjectID:  # pylint: disable=too-few-public-methods
+    def __init__(self, object_type: ObjectType, value: bytes):
+        assert len(value) == 32, "ObjectID must be 32 bytes long"
+        self.object_type = object_type
+        self.value = value
+
+    def __str__(self):
+        return self.to_base58()
+
+    def get_object_type(self):
+        return self.object_type
+    
+    def to_base58(self):
+        return base58.b58encode(self.value).decode()
+
+    @staticmethod
+    def from_base58(s):
+        return ObjectID(base58.b58decode(s))
+
+    def to_base36(self):
+        # Convert the bytes to int before encoding
+        num = int.from_bytes(self.value, "big")
+        return base36.dumps(num)
+
+    @staticmethod
+    def from_base36(s):
+        # Decode to int and then convert to bytes
+        num = base36.loads(s)
+        return ObjectID(num.to_bytes((num.bit_length() + 7) // 8, "big"))
+
+    @staticmethod
+    def new_chunk_id(chunk_hash: HashValue):
+        return ObjectID(ObjectType.Chunk, chunk_hash.value)
+    
+    @staticmethod
+    def hash_data(data: bytes):
+        return ObjectID.new_chunk_id(HashValue.hash_data(data))

src/knowledge/object/object_store.py

@@ -0,0 +1,24 @@
+import os
+import logging
+from .blob import FileBlobStorage
+from .object_id import ObjectID
+
+
+class ObjectStore:
+    def __init__(self, root_dir: str):
+        logging.info(f"will init object blob store, root_dir={root_dir}")
+
+        blob_dir = os.path.join(root_dir, "blob")
+        if not os.path.exists(blob_dir):
+            logging.info(f"will create blob dir: {blob_dir}")
+            os.makedirs(blob_dir)
+        self.blob = FileBlobStorage(blob_dir)
+
+    def put_object(self, object_id: ObjectID, contents: bytes):
+        self.blob.put(object_id, contents)
+
+    def get_object(self, object_id: ObjectID) -> bytes:
+        return self.blob.get(object_id)
+
+    def delete_object(self, object_id: ObjectID):
+        self.blob.delete(object_id)

src/knowledge/object_storage.py

@@ -1,33 +0,0 @@
-# import RDB LargeBinary
-from sqlalchemy import Column, String, LargeBinary, create_engine, sessionmaker, pickle
-from .object import KnowledgeObject
-
-# implement object storage with RDB
-# define object storage table
-class ObjectStorageTable(Base):
-    __tablename__ = 'object_storage'
-    id = Column(String, primary_key=True)
-    parent = Column(String, nullable=True)
-    object = Column(LargeBinary, nullable=False)
-
-    def __init__(self, id, parent, object): # pylint: disable=redefined-builtin
-        self.id = id
-        self.parent = parent
-        self.object = object
-
-# define object storage class
-class ObjectStorage:
-    async def __init__(self, db_url):
-        self.engine = create_engine(db_url)
-        self.session = sessionmaker(bind=self.engine)() # pylint: disable=not-callable
-
-    async def get(self, id) -> [KnowledgeObject, KnowledgeObject]:
-        obj = self.session.query(ObjectStorageTable).filter(ObjectStorageTable.id == id).first()
-        if obj is None:
-            return None
-        return pickle.loads(obj.object)
-    
-    # define insert method
-    async def insert(self, object, parent): # pylint: disable=redefined-builtin
-        obj = ObjectStorageTable(id, parent, pickle.dumps(object))
-        

src/knowledge/store.py

@@ -0,0 +1,46 @@
+import os
+from .object import ObjectStore
+from .data import ChunkStore, ChunkTracker
+import logging
+
+
+# KnowledgeStore class, 聚合ChunkStore，ChunkTracker， ObjectStore，并且是一个全局单例，可以方便的使用这三个内置的store示例
+class KnowledgeStore:
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+            directory = os.path.join(
+                os.path.dirname(__file__), "../../rootfs/data/"
+            )
+            directory = os.path.normpath(directory)
+            print(directory)
+
+            if not os.path.exists(directory):
+                os.makedirs(directory)
+
+            cls._instance.__singleton_init__(directory)
+
+        return cls._instance
+
+    def __singleton_init__(self, root_dir: str):
+        logging.info(f"will init knowledge store, root_dir={root_dir}")
+
+        self.root = root_dir
+
+        object_store_dir = os.path.join(root_dir, "object")
+        self.object_store = ObjectStore(object_store_dir)
+
+        chunk_store_dir = os.path.join(root_dir, "chunk")
+        self.chunk_store = ChunkStore(chunk_store_dir)
+        self.chunk_tracker = ChunkTracker(chunk_store_dir)
+
+    def get_object_store(self) -> ObjectStore:
+        return self.object_store
+
+    def get_chunk_store(self) -> ChunkStore:
+        return self.chunk_store
+
+    def get_chunk_tracker(self) -> ChunkTracker:
+        return self.chunk_tracker

test/test_chunk.py

@@ -15,12 +15,19 @@ 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')
+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
+from knowledge import (
+    ChunkTracker,
+    ChunkID,
+    HashValue,
+    PositionType,
+    KnowledgeStore,
+    ChunkListWriter,
+)
 import asyncio
 import unittest
 
@@ -29,7 +36,7 @@ class TestChunk(unittest.TestCase):
     def test_chunk_tracker(self):
         tracker = KnowledgeStore().get_chunk_tracker()
 
-        hash = HashValue.hash_data("1234567890".encode("utf-8"));
+        hash = HashValue.hash_data("1234567890".encode("utf-8"))
         cid = ChunkID.new_chunk_id(hash)
         print(cid)
 
@@ -42,15 +49,18 @@ class TestChunk(unittest.TestCase):
         self.assertEqual(ret, None)
 
     def test_chunk(self):
-        gen = ChunkListGenerator(KnowledgeStore().get_chunk_store(), KnowledgeStore().get_chunk_tracker())
+        gen = ChunkListWriter(
+            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:
+        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()

test/test_object.py

@@ -0,0 +1,31 @@
+import sys
+import os
+
+dir_path = os.path.dirname(os.path.realpath(__file__))
+print(dir_path)
+
+sys.path.append("{}/../src/".format(dir_path))
+print(sys.path)
+
+from knowledge import ObjectID, HashValue
+
+
+import asyncio
+import unittest
+
+
+class TestVectorSTorage(unittest.TestCase):
+    def test_object(self):
+        data = HashValue.hash_data("1233".encode("utf-8"));
+        print(data.to_base58())
+        print(data.to_base36())
+        
+        data2 = HashValue.from_base58(data.to_base58())
+        self.assertEqual(data.to_base36(), data2.to_base36())
+        
+        data2 = HashValue.from_base36(data.to_base36())
+        self.assertEqual(data.to_base58(), data2.to_base58())
+
+
+if __name__ == "__main__":
+    unittest.main()