diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index a882ba7..4f4ba4f 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/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: diff --git a/src/knowledge/__init__.py b/src/knowledge/__init__.py index 6edb83e..e2cdeb2 100644 --- a/src/knowledge/__init__.py +++ b/src/knowledge/__init__.py @@ -1,3 +1,4 @@ from .object import * from .vector import * -from .data import * \ No newline at end of file +from .data import * +from .store import KnowledgeStore \ No newline at end of file diff --git a/src/knowledge/data/__init__.py b/src/knowledge/data/__init__.py index f8e69aa..2b391c9 100644 --- a/src/knowledge/data/__init__.py +++ b/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 \ No newline at end of file +from .writer import ChunkListWriter +from .chunk_list import ChunkList diff --git a/src/knowledge/data/chunk_list.py b/src/knowledge/data/chunk_list.py new file mode 100644 index 0000000..969b03f --- /dev/null +++ b/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}" \ No newline at end of file diff --git a/src/knowledge/data/gen.py b/src/knowledge/data/writer.py similarity index 89% rename from src/knowledge/data/gen.py rename to src/knowledge/data/writer.py index a6d9425..7f79447 100644 --- a/src/knowledge/data/gen.py +++ b/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 = ".," @@ -55,20 +56,18 @@ class ChunkListGenerator: 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") + 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( @@ -90,5 +89,4 @@ class ChunkListGenerator: word_count = 1 if chunk: chunk_list.append(" ".join(chunk)) - return chunk_list - + return chunk_list \ No newline at end of file diff --git a/src/knowledge/object.py b/src/knowledge/object.py deleted file mode 100644 index f7a127e..0000000 --- a/src/knowledge/object.py +++ /dev/null @@ -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] - - diff --git a/src/knowledge/object/__init__.py b/src/knowledge/object/__init__.py new file mode 100644 index 0000000..ef1db5c --- /dev/null +++ b/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 diff --git a/src/knowledge/object/blob.py b/src/knowledge/object/blob.py new file mode 100644 index 0000000..b005e62 --- /dev/null +++ b/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) diff --git a/src/knowledge/object/hash.py b/src/knowledge/object/hash.py new file mode 100644 index 0000000..5063814 --- /dev/null +++ b/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()) diff --git a/src/knowledge/object/object.py b/src/knowledge/object/object.py new file mode 100644 index 0000000..0eff535 --- /dev/null +++ b/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] + + diff --git a/src/knowledge/object/object_id.py b/src/knowledge/object/object_id.py new file mode 100644 index 0000000..27c0063 --- /dev/null +++ b/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)) \ No newline at end of file diff --git a/src/knowledge/object/object_store.py b/src/knowledge/object/object_store.py new file mode 100644 index 0000000..d903b8f --- /dev/null +++ b/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) diff --git a/src/knowledge/object_storage.py b/src/knowledge/object_storage.py deleted file mode 100644 index 1a6f85e..0000000 --- a/src/knowledge/object_storage.py +++ /dev/null @@ -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)) - \ No newline at end of file diff --git a/src/knowledge/store.py b/src/knowledge/store.py new file mode 100644 index 0000000..a4a0af6 --- /dev/null +++ b/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 diff --git a/test/test_chunk.py b/test/test_chunk.py index b2f2e7e..f1f496b 100644 --- a/test/test_chunk.py +++ b/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 @@ -28,29 +35,32 @@ import unittest 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) - + tracker.add_position(cid, "/tmp/1", PositionType.File) ret = tracker.get_position(cid) print(ret[0]) - + 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 = 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() diff --git a/test/test_object.py b/test/test_object.py new file mode 100644 index 0000000..6a603fd --- /dev/null +++ b/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()