Add object and chunk base impl and test cases
This commit is contained in:
liyaxing
@@ -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:
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
from .object import *
|
||||
from .vector import *
|
||||
from .data import *
|
||||
from .data import *
|
||||
from .store import KnowledgeStore
|
||||
@@ -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
|
||||
|
||||
@@ -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}"
|
||||
@@ -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
|
||||
@@ -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]
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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())
|
||||
@@ -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]
|
||||
|
||||
|
||||
@@ -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))
|
||||
@@ -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)
|
||||
@@ -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))
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user