Merge pull request #55 from photosssa/MVP

Objected knowleadge base, a specialized implemention for emails
This commit is contained in:
Liu Zhicong
2023-09-18 00:37:48 -07:00
committed by GitHub
42 changed files with 3488 additions and 37 deletions
+3 -1
View File
@@ -1,10 +1,12 @@
.vscode/ .vscode/
*.pyc *.pyc
rootfs/data
*.log
rootfs/email/config.local.toml rootfs/email/config.local.toml
rootfs/data rootfs/data
venv venv
aios_shell.log aios_shell.log
history.txt history.txt
math_school_env.db math_school_env.db
workflows.db workflows.db
File diff suppressed because it is too large Load Diff
+20
View File
@@ -0,0 +1,20 @@
# Objected knowleadge base, a specialized implemention for emails
## Vectorized Knowledge
Large language models are trained on general corpora and without fine-tuning on user-specific data, they struggle to utilize user-related context effectively.
Users accumulate a vast amount of content that reflects their personality during their regular internet usage. This includes personal photos, tweets, Facebook posts, emails, etc. While it's possible to include all this content in the prompt during each interaction with the large language model, this approach is costly and can easily reach the token limit.
A common solution is to generate feature vectors from this content using word embedding techniques and store them in a vector database. During an interaction, the vector that is most relevant to the prompt is retrieved from the database, merged with the prompt, and then passed to the large language model.
We refer to this vectorized content as "knowledge".
## Objected knownleadge base
In a personal AI system, to build a user's own knowledge base, we first need to implement various spider programs to crawl and retrieve all user-related data. Modern web content is typically rich text, including text, images, videos, hyperlinks, etc. Organizing this rich text in a tree-like structure similar to HTML is necessary, hence the need to introduce an object structure to represent this content.
Different parts of this content cannot be vectorized using the same embedding model. For instance, text and images, as well as the content of an image and its EXIF information, need separate embeddings. This means that in the vector database, the same content may have multiple vector values, and a row can represent a whole content item or just a part of it.
We need a comprehensive object structure to represent the hierarchy and relationships of content, as well as to implement the indexing and storage of objects. In the Minimum Viable Product (MVP) version, we'll implement a specialized solution for email content. In future versions, we can generalize this to handle other types of content, such as Facebook posts, tweets, etc.
## Agent with knowleadge base
At the same time, we also need to explore the paradigm of using the knowledge base in Agents and workflows, so that the agent can better complete tasks in interaction with users through the context provided by the knowledge base.
+3 -2
View File
@@ -1,6 +1,6 @@
<mxfile host="65bd71144e" pages="3"> <mxfile host="65bd71144e" pages="3">
<diagram id="C5RBs43oDa-KdzZeNtuy" name="Page-1"> <diagram id="C5RBs43oDa-KdzZeNtuy" name="Page-1">
<mxGraphModel dx="2069" dy="1139" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="827" pageHeight="1169" math="0" shadow="0"> <mxGraphModel dx="500" dy="864" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="827" pageHeight="1169" math="0" shadow="0">
<root> <root>
<mxCell id="WIyWlLk6GJQsqaUBKTNV-0"/> <mxCell id="WIyWlLk6GJQsqaUBKTNV-0"/>
<mxCell id="WIyWlLk6GJQsqaUBKTNV-1" parent="WIyWlLk6GJQsqaUBKTNV-0"/> <mxCell id="WIyWlLk6GJQsqaUBKTNV-1" parent="WIyWlLk6GJQsqaUBKTNV-0"/>
@@ -123,7 +123,7 @@
</mxGraphModel> </mxGraphModel>
</diagram> </diagram>
<diagram id="kWxfmPxtNxOAf0a73TCG" name="Page-2"> <diagram id="kWxfmPxtNxOAf0a73TCG" name="Page-2">
<mxGraphModel dx="951" dy="944" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="850" pageHeight="1100" math="0" shadow="0"> <mxGraphModel dx="500" dy="864" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="850" pageHeight="1100" math="0" shadow="0">
<root> <root>
<mxCell id="0"/> <mxCell id="0"/>
<mxCell id="1" parent="0"/> <mxCell id="1" parent="0"/>
@@ -220,6 +220,7 @@
</mxGraphModel> </mxGraphModel>
</diagram> </diagram>
<diagram id="7NYJTgo0U9cdVshLy85U" name="Page-3"> <diagram id="7NYJTgo0U9cdVshLy85U" name="Page-3">
<mxGraphModel dx="1881" dy="676" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="850" pageHeight="1100" math="0" shadow="0"> <mxGraphModel dx="1881" dy="676" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="850" pageHeight="1100" math="0" shadow="0">
<root> <root>
<mxCell id="0"/> <mxCell id="0"/>
+1
View File
@@ -5,6 +5,7 @@ from .agent import AIAgent,AIAgentTemplete,AgentPrompt
from .compute_kernel import ComputeKernel,ComputeTask from .compute_kernel import ComputeKernel,ComputeTask
from .compute_node import ComputeNode,LocalComputeNode from .compute_node import ComputeNode,LocalComputeNode
from .open_ai_node import OpenAI_ComputeNode from .open_ai_node import OpenAI_ComputeNode
from .knowledge_base import KnowledgeBase
from .role import AIRole,AIRoleGroup from .role import AIRole,AIRoleGroup
from .workflow import Workflow from .workflow import Workflow
from .bus import AIBus from .bus import AIBus
+32
View File
@@ -117,3 +117,35 @@ class ComputeKernel:
return task_req.result return task_req.result
return "error!" return "error!"
def text_embedding(self,input:str,model_name:Optional[str] = None):
task_req = ComputeTask()
task_req.set_text_embedding_params(input,model_name)
self.run(task_req)
return task_req
async def do_text_embedding(self,input:str,model_name:Optional[str] = None) -> [float]:
task_req = self.text_embedding(input,model_name)
async def check_timer():
check_times = 0
while True:
if task_req.state == ComputeTaskState.DONE:
break
if task_req.state == ComputeTaskState.ERROR:
break
if check_times >= 20:
task_req.state = ComputeTaskState.ERROR
break
await asyncio.sleep(0.5)
check_times += 1
await asyncio.create_task(check_timer())
if task_req.state == ComputeTaskState.DONE:
return task_req.result.result
return "error!"
+11
View File
@@ -53,6 +53,17 @@ class ComputeTask:
if inner_functions is not None: if inner_functions is not None:
self.params["inner_functions"] = inner_functions self.params["inner_functions"] = inner_functions
def set_text_embedding_params(self, input, model_name=None, callchain_id = None):
self.task_type = "text_embedding"
self.create_time = time.time()
self.task_id = uuid.uuid4().hex
self.callchain_id = callchain_id
if model_name is not None:
self.params["model_name"] = model_name
else:
self.params["model_name"] = "text-embedding-ada-002"
self.params["input"] = input
def display(self) -> str: def display(self) -> str:
return f"ComputeTask: {self.task_id} {self.task_type} {self.state}" return f"ComputeTask: {self.task_id} {self.task_type} {self.state}"
+247
View File
@@ -0,0 +1,247 @@
# define a knowledge base class
import json
import logging
from . import AgentPrompt, ComputeKernel
from knowledge import *
class KnowledgeBase:
_instance = None
def __new__(cls):
if cls._instance is None:
cls._instance = super().__new__(cls)
cls._instance.__singleton_init__()
return cls._instance
def __singleton_init__(self) -> None:
self.store = KnowledgeStore()
self.compute_kernel = ComputeKernel()
async def __embedding_document(self, document: DocumentObject):
for chunk_id in document.get_chunk_list():
chunk = self.store.get_chunk_reader().get_chunk(chunk_id)
if chunk is None:
raise ValueError(f"text chunk not found: {chunk_id}")
text = chunk.read().decode("utf-8")
vector = await self.compute_kernel.do_text_embedding(text)
await self.store.get_vector_store("default").insert(vector, chunk_id)
async def __embedding_image(self, image: ImageObject):
desc = {}
if not not image.get_meta():
desc["meta"] = image.get_meta()
if not not image.get_exif():
desc["exif"] = image.get_exif()
if not not image.get_tags():
desc["tags"] = image.get_tags()
vector = await self.compute_kernel.do_text_embedding(json.dumps(desc))
await self.store.get_vector_store("default").insert(vector, image.calculate_id())
async def __embedding_video(self, vedio: VideoObject):
desc = {}
if not not vedio.get_meta():
desc["meta"] = vedio.get_meta()
if not not vedio.get_info():
desc["info"] = vedio.get_info()
if not not vedio.get_tags():
desc["tags"] = vedio.get_tags()
vector = await self.compute_kernel.do_text_embedding(json.dumps(desc))
await self.store.get_vector_store("default").insert(vector, vedio.calculate_id())
async def __embedding_rich_text(self, rich_text: RichTextObject):
for document_id in rich_text.get_documents().values():
document = DocumentObject.decode(self.store.get_object_store().get_object(document_id))
await self.__embedding_document(document)
for image_id in rich_text.get_images().values():
image = ImageObject.decode(self.store.get_object_store().get_object(image_id))
await self.__embedding_image(image)
for video_id in rich_text.get_videos().values():
video = VideoObject.decode(self.store.get_object_store().get_object(video_id))
await self.__embedding_video(video)
for rich_text_id in rich_text.get_rich_texts().values():
rich_text = RichTextObject.decode(self.store.get_object_store().get_object(rich_text_id))
await self.__embedding_rich_text(rich_text)
async def __embedding_email(self, email: EmailObject):
vector = await self.compute_kernel.do_text_embedding(json.dumps(email.get_desc()))
await self.store.get_vector_store("default").insert(vector, email.calculate_id())
await self.__embedding_rich_text(email.get_rich_text())
async def __do_embedding(self, object: KnowledgeObject):
if object.get_object_type() == ObjectType.Document:
await self.__embedding_document(object)
if object.get_object_type() == ObjectType.Image:
await self.__embedding_image(object)
if object.get_object_type() == ObjectType.Video:
await self.__embedding_video(object)
if object.get_object_type() == ObjectType.RichText:
await self.__embedding_rich_text(object)
if object.get_object_type() == ObjectType.Email:
await self.__embedding_email(object)
else:
pass
# def __save_document(self, document: DocumentObject):
# doc_id = document.calculate_id()
# self.store.get_object_store().put_object(doc_id, document.encode())
# for chunk_id in document.get_chunk_list():
# self.store.get_relation_store().add_relation(chunk_id, doc_id)
# def __save_image(self, image: ImageObject):
# image_id = image.calculate_id()
# self.store.get_object_store().put_object(image_id, image.encode())
# def __save_video(self, video: VideoObject):
# video_id = video.calculate_id()
# self.store.get_object_store().put_object(video_id, video.encode())
# def __save_rich_text(self, rich_text: RichTextObject):
# rich_text_id = rich_text.calculate_id()
# # rich_text_enc = dict()
# # rich_text_enc["desc"] = rich_text.desc
# # rich_text_enc["body"] = {"documents": {}, "images": {}, "videos": {}, "rich_texts": {}}
# for key, document in rich_text.get_documents().items():
# self.__save_document(document)
# doc_id = document.calculate_id()
# self.store.get_relation_store().add_relation(doc_id, rich_text_id)
# # rich_text_enc["body"]["documents"][key] = doc_id
# for key, image in rich_text.get_images().items():
# self.__save_image(image)
# image_id = image.calculate_id()
# self.store.get_relation_store().add_relation(image_id, rich_text_id)
# # rich_text_enc["body"]["images"][key] = image_id
# for key, video in rich_text.get_videos().items():
# self.__save_video(video)
# video_id = video.calculate_id()
# self.store.get_relation_store().add_relation(video_id, rich_text_id)
# # rich_text_enc["body"]["videos"][key] = video_id
# for key, rich_text in rich_text.get_rich_texts().items():
# self.__save_rich_text(rich_text)
# rich_text_id = rich_text.calculate_id()
# self.store.get_relation_store().add_relation(rich_text_id, rich_text_id)
# # rich_text_enc["body"]["rich_texts"][key] = rich_text_id
# self.store.get_object_store().put_object(rich_text_id, rich_text.encode())
# def __save_email(self, email: EmailObject):
# email_id = email.calculate_id()
# # email_enc = dict()
# # email_enc["desc"] = email.desc
# # email_enc["body"] = {"content": None}
# self.__save_rich_text(email.get_rich_text())
# rich_text_id = email.get_rich_text().calculate_id()
# self.store.get_relation_store().add_relation(rich_text_id, email_id)
# # email_enc["body"]["content"] = rich_text_id
# self.store.get_object_store().put_object(email_id, email.encode())
# def __save_object(self, object: KnowledgeObject):
# if object.get_object_type() == ObjectType.Document:
# self.__save_document(object)
# if object.get_object_type() == ObjectType.Image:
# self.__save_image(object)
# if object.get_object_type() == ObjectType.Video:
# self.__save_video(object)
# if object.get_object_type() == ObjectType.RichText:
# self.__save_rich_text(object)
# if object.get_object_type() == ObjectType.Email:
# self.__save_email(object)
# else:
# pass
async def insert_object(self, object: KnowledgeObject):
# self.__save_object(object)
await self.__do_embedding(object)
async def query_prompt(self, prompt: AgentPrompt):
logging.info(f"query_prompt: {prompt}")
objects = await self.query_objects(prompt)
knowledge_prompt = self.prompt_from_objects(objects)
logging.info(f"prompt_from_objects result: {knowledge_prompt.as_str()}")
prompt.append(knowledge_prompt)
async def query_objects(self, prompt: AgentPrompt) -> [ObjectID]:
results = []
for msg in prompt.messages:
if msg["role"] == "user":
vector = await self.compute_kernel.do_text_embedding(msg["content"])
object_ids = await self.store.get_vector_store("default").query(vector, 10)
results.extend(object_ids)
return results
def __load_object(self, object_id: ObjectID) -> KnowledgeObject:
if object_id.get_object_type() == ObjectType.Document:
return DocumentObject.decode(self.store.get_object_store().get_object(object_id))
if object_id.get_object_type() == ObjectType.Image:
return ImageObject.decode(self.store.get_object_store().get_object(object_id))
if object_id.get_object_type() == ObjectType.Video:
return VideoObject.decode(self.store.get_object_store().get_object(object_id))
if object_id.get_object_type() == ObjectType.RichText:
return RichTextObject.decode(self.store.get_object_store().get_object(object_id))
if object_id.get_object_type() == ObjectType.Email:
return EmailObject.decode(self.store.get_object_store().get_object(object_id))
else:
pass
def prompt_from_objects(self, object_ids: [ObjectID]) -> AgentPrompt:
results = dict()
for object_id in object_ids:
parents = self.store.get_relation_store().get_related_root_objects(object_id)
# last parent is the root object
root_object_id = parents[0] if parents else object_id
logging.info(f"object_id: {str(object_id)} root_object_id: {str(root_object_id)}")
if str(root_object_id) in results:
results[str(root_object_id)].append(object_id)
else:
results[str(root_object_id)] = [root_object_id, object_id]
content = "I found the following contents described with json format:\n"
result_desc = []
for result in results.values():
# first element in result is the root object
root_object_id = result[0]
if root_object_id.get_object_type() == ObjectType.Email:
email = self.__load_object(root_object_id)
desc = email.get_desc()
desc["type"] = "email"
desc["contents"] = []
result_desc.append(desc)
upper_list = desc["contents"]
result = result[1:]
else:
upper_list = result_desc
for object_id in result:
if object_id.get_object_type() == ObjectType.Chunk:
upper_list.append({"type": "text", "content": self.store.get_chunk_reader().get_chunk(object_id).read().decode("utf-8")})
if object_id.get_object_type() == ObjectType.Image:
image = self.__load_object(object_id)
desc = image.get_desc()
desc["type"] = "image"
upper_list.append(desc)
if object_id.get_object_type() == ObjectType.Video:
video = self.__load_object(object_id)
desc = video.get_desc()
desc["type"] = "video"
upper_list.append(desc)
else:
pass
content += json.dumps(result_desc)
content += ".\n"
prompt = AgentPrompt()
prompt.messages.append({"role": "knowledge", "content": content})
return prompt
+77 -30
View File
@@ -59,47 +59,86 @@ class OpenAI_ComputeNode(ComputeNode):
def _run_task(self, task: ComputeTask): def _run_task(self, task: ComputeTask):
task.state = ComputeTaskState.RUNNING task.state = ComputeTaskState.RUNNING
mode_name = task.params["model_name"] if task.task_type == "text_embedding":
# max_token_size = task.params["max_token_size"] model_name = task.params["model_name"]
prompts = task.params["prompts"] input = task.params["input"]
logger.info(f"call openai {model_name} input: {input}")
logger.info(f"call openai {mode_name} prompts: {prompts}") resp = openai.Embedding.create(model=model_name,
input=input)
if task.params.get("inner_functions") is None: # resp = {
resp = openai.ChatCompletion.create(model=mode_name, # "object": "list",
messages=prompts, # "data": [
max_tokens=task.params["max_token_size"], # {
temperature=0.7) # "object": "embedding",
else: # "index": 0,
resp = openai.ChatCompletion.create(model=mode_name, # "embedding": [
# -0.00930514745414257,
# 0.00765434792265296,
# -0.007167573552578688,
# -0.012373941019177437,
# -0.04884673282504082
# ]}]
# }
logger.info(f"openai response: {resp}")
result = ComputeTaskResult()
result.set_from_task(task)
result.worker_id = self.node_id
result.result = resp["data"][0]["embedding"]
return result
if task.task_type == "llm_completion":
mode_name = task.params["model_name"]
# max_token_size = task.params["max_token_size"]
prompts = task.params["prompts"]
mode_name = task.params["model_name"]
# max_token_size = task.params["max_token_size"]
prompts = task.params["prompts"]
logger.info(f"call openai {mode_name} prompts: {prompts}")
if task.params.get("inner_functions") is None:
resp = openai.ChatCompletion.create(model=mode_name,
messages=prompts, messages=prompts,
functions=task.params["inner_functions"],
max_tokens=task.params["max_token_size"], max_tokens=task.params["max_token_size"],
temperature=0.7) # TODO: add temperature to task params? temperature=0.7)
else:
resp = openai.ChatCompletion.create(model=mode_name,
messages=prompts,
functions=task.params["inner_functions"],
max_tokens=task.params["max_token_size"],
temperature=0.7) # TODO: add temperature to task params?
logger.info(f"openai response: {resp}") logger.info(f"openai response: {resp}")
result = ComputeTaskResult() result = ComputeTaskResult()
result.set_from_task(task) result.set_from_task(task)
status_code = resp["choices"][0]["finish_reason"] status_code = resp["choices"][0]["finish_reason"]
match status_code: match status_code:
case "function_call": case "function_call":
task.state = ComputeTaskState.DONE task.state = ComputeTaskState.DONE
case "stop": case "stop":
task.state = ComputeTaskState.DONE task.state = ComputeTaskState.DONE
case _: case _:
task.state = ComputeTaskState.ERROR task.state = ComputeTaskState.ERROR
task.error_str = f"The status code was {status_code}." task.error_str = f"The status code was {status_code}."
return None return None
result.worker_id = self.node_id result.worker_id = self.node_id
result.result_str = resp["choices"][0]["message"]["content"] result.result_str = resp["choices"][0]["message"]["content"]
result.result_message = resp["choices"][0]["message"] result.result_message = resp["choices"][0]["message"]
return result return result
def start(self): def start(self):
if self.is_start is True: if self.is_start is True:
return return
@@ -125,8 +164,16 @@ class OpenAI_ComputeNode(ComputeNode):
def get_capacity(self): def get_capacity(self):
pass pass
def is_support(self, task: ComputeTask) -> bool: def is_support(self, task: ComputeTask) -> bool:
return task.task_type == ComputeTaskType.LLM_COMPLETION and (not task.params["model_name"] or task.params["model_name"] == "gpt-4-0613") if task.task_type == ComputeTaskType.LLM_COMPLETION:
if (not task.params["model_name"] or task.params["model_name"] == "gpt-4-0613")
return True
if task.task_type == "text_embedding":
if task.params["model_name"] == "text-embedding-ada-002":
return True
return False
def is_local(self) -> bool: def is_local(self) -> bool:
return False return False
+5
View File
@@ -0,0 +1,5 @@
from .object import *
from .vector import *
from .data import *
from .store import KnowledgeStore
from .core_object import *
+5
View File
@@ -0,0 +1,5 @@
from .document_object import DocumentObject, DocumentObjectBuilder
from .image_object import ImageObject, ImageObjectBuilder
from .video_object import VideoObject, VideoObjectBuilder
from .rich_text_object import RichTextObject, RichTextObjectBuilder
from .email_object import EmailObject, EmailObjectBuilder
@@ -0,0 +1,65 @@
from ..object import KnowledgeObject, ObjectRelationStore
from ..data import ChunkList, ChunkListWriter
from ..object import ObjectType
from .. import KnowledgeStore
# desc
# meta
# hash: "file-hash",
# tags: {}
# body
# chunk_list: [chunk_id, chunk_id, ...]
class DocumentObject(KnowledgeObject):
def __init__(self, meta: dict, tags: dict, chunk_list: ChunkList):
desc = dict()
body = dict()
desc["meta"] = meta
desc["tags"] = tags
desc["hash"] = chunk_list.hash.to_base58()
body["chunk_list"] = chunk_list.chunk_list
super().__init__(ObjectType.Document, desc, body)
def get_meta(self):
return self.desc["meta"]
def get_tags(self):
return self.desc["tags"]
def get_hash(self):
return self.desc["hash"]
def get_chunk_list(self):
return self.body["chunk_list"]
class DocumentObjectBuilder:
def __init__(self, meta: dict, tags: dict, text: str):
self.meta = meta
self.tags = tags
self.text = text
def set_meta(self, meta: dict):
self.meta = meta
return self
def set_text(self, text: str):
self.text = text
return self
def build(self, relation_store: ObjectRelationStore) -> DocumentObject:
chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_text(
self.text,
1024 * 4,
".?!\n"
)
doc = DocumentObject(self.meta, self.tags, chunk_list)
doc_id = doc.calculate_id()
# Add relation to store
for chunk_id in chunk_list.chunk_list:
relation_store.add_relation(chunk_id, doc_id)
return doc
+160
View File
@@ -0,0 +1,160 @@
from .. import KnowledgeStore
from .rich_text_object import RichTextObject, RichTextObjectBuilder
from ..object import ObjectID, ObjectType, KnowledgeObject
from .document_object import DocumentObjectBuilder
from .image_object import ImageObjectBuilder
from .video_object import VideoObjectBuilder
import os
import json
import logging
class EmailObject(KnowledgeObject):
def __init__(self, meta: dict, tags: dict, rich_text: RichTextObject):
desc = dict()
body = dict()
desc["meta"] = meta
desc["tags"] = tags
# FIXME rich text content store in desc or body? which one is better?
body["content"] = rich_text
super().__init__(ObjectType.Email, desc, body)
def get_meta(self):
return self.desc["meta"]
def get_tags(self):
return self.desc["tags"]
def get_rich_text(self):
return self.body["content"]
"""
EmailObject folder structure:
.
├── email.txt
└── meta.json
├── image
│ ├── image1.jpg
│ ├── image2.jpg
│ └── ...
├── video
│ ├── video1.mp4
│ ├── video2.mv
│ └── ...
└── audio
├── audio1.m4a
├── audio2.flac
└── ...
EmailObjectBuilder will read the target folder and build the EmailObject
Store meta.json to meta in EmailObject
Store email.txt to DocumentObject and RichTextObject in EmailObject
Store very image file in image folder to ImageObject and RichTextObject in EmailObject, etc
"""
class EmailObjectBuilder:
def __init__(self, tags: dict, folder: str):
self.tags = tags
self.folder = folder
def set_tags(self, tags: dict):
self.tags = tags
return self
def set_folder(self, folder: str):
self.folder = folder
return self
def build(self) -> EmailObject:
# Just get the object store and relation store from global KnowledgeStore
store = KnowledgeStore().get_object_store()
relation = KnowledgeStore().get_relation_store()
# Read meta.json
meta = {}
meta_file = os.path.join(self.folder, "meta.json")
if os.path.exists(meta_file):
logging.info(f"Will read meta.json {meta_file}")
with open(meta_file, "r", encoding="utf-8") as f:
meta = json.load(f)
else:
logging.info(f"Meta file missing! {meta_file}")
# Read email.txt
documents = {}
content_file = os.path.join(self.folder, "email.txt")
if os.path.exists(content_file):
logging.info(f"Will read email.txt {content_file}")
try:
with open(content_file, "r", encoding="utf-8") as f:
text = f.read()
document = DocumentObjectBuilder({}, {}, text).build(relation_store=relation)
document_id = document.calculate_id()
store.put_object(document_id, document.encode())
documents = {"email.txt": document_id}
except Exception as e:
logging.error(f"Failed to read email.txt {content_file} {e}")
else:
logging.info(f"Content file missing! {content_file}")
# Process image files
images = {}
image_dir = os.path.join(self.folder, "image")
if os.path.exists(image_dir):
for image_file in os.listdir(image_dir):
image_path = os.path.join(image_dir, image_file)
logging.info(f"Will read image file {image_path}")
try:
image = ImageObjectBuilder({}, {}, image_path).build()
image_id = image.calculate_id()
store.put_object(image_id, image.encode())
images[image_file] = image_id
except Exception as e:
logging.error(f"Failed to read image file {image_path} {e}")
continue
# Process video files
videos = {}
video_dir = os.path.join(self.folder, "video")
if os.path.exists(video_dir):
for video_file in os.listdir(video_dir):
video_path = os.path.join(video_dir, video_file)
logging.info(f"Will read video file {video_path}")
try:
video = VideoObjectBuilder({}, {}, video_path).build()
video_id = video.calculate_id()
store.put_object(video_id, video.encode())
videos[video_file] = video_id
except Exception as e:
logging.error(f"Failed to read video file {video_path} {e}")
continue
# Create RichTextObject
rich_text = RichTextObject(images, videos, documents)
rich_text_id = rich_text.calculate_id()
# build relations with rich_text
for image_id in images.values():
relation.add_relation(image_id, rich_text_id)
for video_id in videos.values():
relation.add_relation(video_id, rich_text_id)
for document_id in documents.values():
relation.add_relation(document_id, rich_text_id)
# Create EmailObject
email_object = EmailObject(meta, {}, rich_text)
email_object_id = email_object.calculate_id()
store.put_object(email_object_id, email_object.encode())
# build relations with email_object
relation.add_relation(rich_text_id, email_object_id)
return email_object
+89
View File
@@ -0,0 +1,89 @@
from ..object import KnowledgeObject
from ..data import ChunkList, ChunkListWriter
from ..object import ObjectType
from .. import KnowledgeStore
# desc
# meta
# tags
# hash: "file-hash",
# exif: {}
# body
# chunk_list: [chunk_id, chunk_id, ...]
class ImageObject(KnowledgeObject):
def __init__(self, meta: dict, tags: dict, exif: dict, chunk_list: ChunkList):
desc = dict()
body = dict()
desc["meta"] = meta
desc["exif"] = exif
desc["tags"] = tags
desc["hash"] = chunk_list.hash.to_base58()
body["chunk_list"] = chunk_list.chunk_list
super().__init__(ObjectType.Image, desc, body)
def get_meta(self):
return self.desc["meta"]
def get_exif(self):
return self.desc["exif"]
def get_tags(self):
return self.desc["tags"]
def get_hash(self):
return self.desc["hash"]
def get_chunk_list(self):
return self.body["chunk_list"]
from PIL import Image
from PIL.ExifTags import TAGS
def get_exif_data(image_path: str):
with Image.open(image_path) as image:
exif_data = image._getexif()
if exif_data is not None:
return {
TAGS.get(key): exif_data[key]
for key in exif_data.keys()
if key in TAGS and isinstance(exif_data[key], (bytes, str))
}
else:
return {}
class ImageObjectBuilder:
def __init__(self, meta: dict, tags: dict, image_file: str):
self.meta = meta
self.tags = tags
self.image_file = image_file
self.restore_file = False
def set_meta(self, meta: dict):
self.meta = meta
return self
def set_tags(self, tags: dict):
self.tags = tags
return self
def set_image_file(self, image_file: str):
self.image_file = image_file
return self
def set_restore_file(self, restore_file: bool):
self.restore_file = restore_file
return self
def build(self) -> ImageObject:
chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_file(
self.image_file, 1024 * 1024 * 4, self.restore_file
)
exif = get_exif_data(self.image_file)
return ImageObject(self.meta, self.tags, exif, chunk_list)
@@ -0,0 +1,80 @@
from knowledge.object.object_id import ObjectType
from ..object import KnowledgeObject
from ..data import ChunkList, ChunkListWriter
from ..object import ObjectType
from .video_object import VideoObjectBuilder, VideoObject
from .image_object import ImageObjectBuilder, ImageObject
from .document_object import DocumentObjectBuilder, DocumentObject
class RichTextObject(KnowledgeObject):
def __init__(self, images: dict = {}, videos: dict = {}, documents: dict = {}, rich_texts: dict = {}):
desc = dict()
desc["images"] = images
desc["videos"] = videos
desc["documents"] = documents
desc["rich_texts"] = rich_texts
super().__init__(ObjectType.RichText, desc)
def add_image_with_key(self, key, image_object: ImageObject):
assert self.desc["images"][key] == None
self.desc["images"][key] = image_object
def add_image(self, image_object: ImageObject):
self.desc["images"][image_object.object_id()] = image_object
def get_image_with_key(self, key) -> ImageObject:
return self.desc["images"][key]
def get_images(self) -> dict:
return self.desc["images"]
def add_video_with_key(self, key, video_object: VideoObject):
assert self.desc["videos"][key] == None
self.desc["videos"][key] = video_object
def add_video(self, video_object: VideoObject):
self.desc["videos"][video_object.object_id()] = video_object
def get_video_with_key(self, key) -> VideoObject:
return self.desc["videos"][key]
def get_videos(self) -> dict:
return self.desc["videos"]
def add_document_with_key(self, key, document_object: DocumentObject):
assert self.desc["documents"][key] == None
self.desc["documents"][key] = document_object
def add_document(self, document_object: DocumentObject):
self.desc["documents"][document_object.object_id()] = document_object
def get_document_with_key(self, key) -> DocumentObject:
return self.desc["documents"][key]
def get_documents(self) -> dict:
return self.desc["documents"]
def add_rich_text_with_key(self, key, rich_text_object):
assert self.desc["rich_texts"][key] == None
self.desc["rich_texts"][key] = rich_text_object
def add_rich_text(self, rich_text_object):
self.desc["rich_texts"][rich_text_object.object_id()] = rich_text_object
def get_rich_text_with_key(self, key):
return self.desc["rich_texts"][key]
def get_rich_texts(self) -> dict:
return self.desc["rich_texts"]
class RichTextObjectBuilder:
def __init__(self, folder: str):
self.folder = folder
def build(self) -> RichTextObject:
# TODO
return RichTextObject()
+84
View File
@@ -0,0 +1,84 @@
from ..object import KnowledgeObject
from ..data import ChunkList, ChunkListWriter
from ..object import ObjectType
from .. import KnowledgeStore
# desc
# meta
# tags
# hash: "file-hash",
# info: {}
# body
# chunk_list: [chunk_id, chunk_id, ...]
class VideoObject(KnowledgeObject):
def __init__(self, meta: dict, tags: dict, info: dict, chunk_list: ChunkList):
desc = dict()
body = dict()
desc["meta"] = meta
desc["tags"] = tags
desc["info"] = info
desc["hash"] = chunk_list.hash.to_base58()
body["chunk_list"] = chunk_list.chunk_list
super().__init__(ObjectType.Video, desc, body)
def get_meta(self):
return self.desc["meta"]
def get_tags(self):
return self.desc["tags"]
def get_info(self):
return self.desc["info"]
def get_hash(self):
return self.desc["hash"]
def get_chunk_list(self):
return self.body["chunk_list"]
from moviepy.editor import VideoFileClip
def get_video_info(video_path: str) -> dict:
clip = VideoFileClip(video_path)
return {
"duration": clip.duration, # Duration in seconds
"fps": clip.fps, # Frames per second
"nframes": clip.reader.nframes, # Total number of frames
"size": clip.size, # Size of the frames (width, height)
}
class VideoObjectBuilder:
def __init__(self, meta: dict, tags: dict, video_file: str):
self.meta = meta
self.tags = tags
self.video_file = video_file
self.restore_file = False
def set_meta(self, meta: dict):
self.meta = meta
return self
def set_tags(self, tags: dict):
self.tags = tags
return self
def set_video_file(self, video_file: str):
self.video_file = video_file
return self
def set_restore_file(self, restore_file: bool):
self.restore_file = restore_file
return self
def build(self) -> VideoObject:
chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_file(
self.video_file, 1024 * 1024 * 4, self.restore_file
)
info = get_video_info(self.video_file)
return VideoObject(self.meta, self.tags, info, chunk_list)
+6
View File
@@ -0,0 +1,6 @@
from .chunk import ChunkID, PositionType, PositionFileRange
from .tracker import ChunkTracker
from .chunk_store import ChunkStore
from .writer import ChunkListWriter
from .chunk_list import ChunkList
from .reader import ChunkReader, Chunk
+43
View File
@@ -0,0 +1,43 @@
from enum import IntEnum
from ..object import ObjectID
ChunkID = ObjectID
class PositionType(IntEnum):
Unknown = 1
Device = 2
File = 3
FileRange = 4
ChunkStore = 5
class PositionFileRange:
def __init__(self, path: str, range_begin: int, range_end: int):
self.path = path
self.range_begin = range_begin
self.range_end = range_end
def encode(self):
return f"{self.range_begin}:{self.range_end}:{self.path}"
@staticmethod
def decode(value: str):
parts = value.split(":")
if len(parts) < 3:
raise ValueError("Invalid input string")
try:
range_begin = int(parts[0])
range_end = int(parts[1])
except ValueError as e:
raise ValueError("Invalid range_begin or range_end string") from e
path = ":".join(parts[2:])
return PositionFileRange(path, range_begin, range_end)
def __str__(self):
return self.encode()
@staticmethod
def from_string(value: str):
return PositionFileRange.decode(value)
+14
View File
@@ -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}"
+28
View File
@@ -0,0 +1,28 @@
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)
def get_chunk_file_path(self, chunk_id: ChunkID) -> str:
return self.blob.get_full_path(chunk_id, False)
+85
View File
@@ -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"))
+71
View File
@@ -0,0 +1,71 @@
import sqlite3
import time
import logging
import os
from .chunk import ChunkID, PositionType, PositionFileRange
from typing import List, Tuple
class ChunkTracker:
def __init__(self, root_dir: str):
if not os.path.exists(root_dir):
os.makedirs(root_dir)
file = os.path.join(root_dir, "chunk_tracker.db")
logging.info(f"will init chunk tracker, db={file}")
self.conn = sqlite3.connect(file)
self.cursor = self.conn.cursor()
self.cursor.execute(
"""
CREATE TABLE IF NOT EXISTS chunks (
id TEXT NOT NULL,
pos TEXT NOT NULL,
pos_type TINYINT NOT NULL,
insert_time UNSIGNED BIG INT NOT NULL,
update_time UNSIGNED BIG INT NOT NULL,
flags INTEGER DEFAULT 0,
PRIMARY KEY(id, pos, pos_type)
)
"""
)
self.conn.commit()
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())
self.cursor.execute(
"""
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,
),
)
self.conn.commit()
def remove_position(self, chunk_id: ChunkID):
logging.info(f"remove chunk position: {chunk_id}")
self.cursor.execute(
"""
DELETE FROM chunks WHERE id = ?
""",
(str(chunk_id),),
)
self.conn.commit()
def get_position(self, chunk_id: ChunkID) -> List[Tuple[str, PositionType]]:
self.cursor.execute(
"""
SELECT pos, pos_type FROM chunks WHERE id = ?
""",
(str(chunk_id),),
)
return self.cursor.fetchmany()
+93
View File
@@ -0,0 +1,93 @@
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
from .chunk_list import ChunkList
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
) -> ChunkList:
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, str(file_range), PositionType.FileRange
)
file_hash = HashValue(hash_obj.digest())
print(f"calc file hash: {file_path}, {file_hash}")
return ChunkList(chunk_list, file_hash)
def create_chunk_list_from_text(
self, text: str, chunk_max_words: int, separator_chars: str = ".,"
) -> ChunkList:
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 ChunkList(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
return sentences
+6
View File
@@ -0,0 +1,6 @@
from .object import KnowledgeObject
from .blob import FileBlobStorage
from .hash import HashValue, hash_data
from .relation import ObjectRelationStore
from .object_store import ObjectStore
from .object_id import ObjectID, ObjectType
+54
View File
@@ -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)
+42
View File
@@ -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())
+65
View File
@@ -0,0 +1,65 @@
# define a object type enum
from abc import ABC, abstractmethod
from enum import Enum
from .object_id import ObjectID, ObjectType
import hashlib
import json
import pickle
from typing import Any
class ObjectEnhancedJSONEncoder(json.JSONEncoder):
def default(self, o: Any) -> Any:
if isinstance(o, ObjectID):
return o.to_base58()
return super().default(o)
class KnowledgeObject(ABC):
def __init__(self, object_type: ObjectType, desc: dict = {}, body: dict = {}):
self.desc = desc
self.body = body
self.object_type = object_type
def get_object_type(self) -> ObjectType:
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},
cls=ObjectEnhancedJSONEncoder,
)
sha256 = hashlib.sha256()
sha256.update(data.encode())
hash_bytes = sha256.digest()
return ObjectID(bytes([self.object_type]) + hash_bytes[1:])
def encode(self) -> bytes:
return pickle.dumps(self)
@staticmethod
def decode(data: bytes):
return pickle.loads(data)
+58
View File
@@ -0,0 +1,58 @@
# define a object type enum
from abc import ABC, abstractmethod
from enum import IntEnum
from .hash import HashValue
import base58
import base36
class ObjectType(IntEnum):
Chunk = 7
Image = 101
Video = 102
Document = 103
RichText = 104
Email = 105
# define a object ID class to identify a object
class ObjectID: # pylint: disable=too-few-public-methods
def __init__(self, value: bytes):
assert len(value) == 32, "ObjectID must be 32 bytes long"
self.value = value
def __str__(self):
return self.to_base58()
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):
assert len(chunk_hash.value) == 32, "ObjectID must be 32 bytes long"
return ObjectID(bytes([ObjectType.Chunk]) + chunk_hash.value[1:])
def get_object_type(self) -> ObjectType:
return ObjectType(self.value[0])
@staticmethod
def hash_data(data: bytes):
return ObjectID.new_chunk_id(HashValue.hash_data(data))
def __eq__(self, other) -> bool:
return self.value == other.value
+24
View File
@@ -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)
+104
View File
@@ -0,0 +1,104 @@
# define a relation store class
from .object_id import ObjectID
import sqlite3
from typing import List, Tuple, Optional
import logging
import os
from enum import IntEnum
class ObjectRelationType(IntEnum):
Parent = 1
class ObjectRelationStore:
def __init__(self, root_dir: str):
if not os.path.exists(root_dir):
os.makedirs(root_dir)
file = os.path.join(root_dir, "relation.db")
logging.info(f"will init object relation store, db={file}")
self.conn = sqlite3.connect(file)
self.cursor = self.conn.cursor()
self.cursor.execute(
"""
CREATE TABLE IF NOT EXISTS relations (
object_id TEXT,
assoc_id TEXT,
relation_type TEXT,
PRIMARY KEY (object_id, assoc_id, relation_type)
)
"""
)
def add_relation(
self,
object_id: ObjectID,
assoc_id: ObjectID,
relation_type: ObjectRelationType = ObjectRelationType.Parent,
):
if relation_type == None:
relation_type = ObjectRelationType.Parent
self.cursor.execute(
"""
INSERT OR IGNORE INTO relations (object_id, assoc_id, relation_type)
VALUES (?, ?, ?)
""",
(str(object_id), str(assoc_id), relation_type.value),
)
self.conn.commit()
def get_related_objects(
self, object_id: ObjectID, relation_type: Optional[ObjectRelationType] = None
) -> List[ObjectID]:
if relation_type:
self.cursor.execute(
"""
SELECT assoc_id FROM relations WHERE object_id = ? AND relation_type = ?
""",
(str(object_id), relation_type.value),
)
else:
self.cursor.execute(
"""
SELECT assoc_id FROM relations WHERE object_id = ?
""",
(str(object_id),),
)
return [ObjectID.from_base58(row[0]) for row in self.cursor.fetchall()]
def get_related_root_objects(
self, object_id: ObjectID, relation_type: Optional[ObjectRelationType] = None
) -> List[ObjectID]:
root_objects = []
related_objects = self.get_related_objects(object_id, relation_type)
history = []
history.append(object_id)
while related_objects:
for obj in related_objects:
next_related_objects = self.get_related_objects(obj, relation_type)
if not next_related_objects:
if obj not in root_objects:
root_objects.append(obj)
else:
for related_object in next_related_objects:
if obj not in history:
related_objects.append(related_object)
else:
logging.warning(
f"loop detected: {obj} <-> {related_object}"
)
related_objects = next_related_objects
return root_objects
def delete_relation(self, object_id: ObjectID):
self.cursor.execute(
"""
DELETE FROM relations WHERE object_id = ?
""",
(str(object_id),),
)
self.conn.commit()
+68
View File
@@ -0,0 +1,68 @@
import os
from .object import ObjectStore, ObjectRelationStore
from .data import ChunkStore, ChunkTracker, ChunkListWriter, ChunkReader
from .vector import ChromaVectorStore, VectorBase
import logging
# KnowledgeStore class, which aggregates ChunkStore, ChunkTracker, and ObjectStore, and is a global singleton that makes it easy to use these three built-in store examples
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
relation_store_dir = os.path.join(root_dir, "relation")
self.relation_store = ObjectRelationStore(relation_store_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)
self.chunk_list_writer = ChunkListWriter(self.chunk_store, self.chunk_tracker)
self.chunk_reader = ChunkReader(self.chunk_store, self.chunk_tracker)
self.vector_store = {}
def get_relation_store(self) -> ObjectRelationStore:
return self.relation_store
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
def get_chunk_list_writer(self) -> ChunkListWriter:
return self.chunk_list_writer
def get_chunk_reader(self) -> ChunkReader:
return self.chunk_reader
def get_vector_store(self, model_name: str) -> VectorBase:
if model_name not in self.vector_store:
self.vector_store[model_name] = ChromaVectorStore(model_name)
return self.vector_store[model_name]
+2
View File
@@ -0,0 +1,2 @@
from .vector_base import VectorBase
from .chroma_store import ChromaVectorStore
+52
View File
@@ -0,0 +1,52 @@
from .vector_base import VectorBase
from ..object import ObjectID
import chromadb
import logging
import os
class ChromaVectorStore(VectorBase):
def __init__(self, model_name: str) -> None:
super().__init__(model_name)
logging.info(
"will init chroma vector store, model={}".format(model_name)
)
directory = os.path.join(
os.path.dirname(__file__), "../../../rootfs/data/vector"
)
logging.info("will use vector store: {}".format(directory))
client = chromadb.PersistentClient(
path=directory, settings=chromadb.Settings(anonymized_telemetry=False)
)
# client = chromadb.Client()
collection_name = "coll_{}".format(model_name)
logging.info("will init chroma colletion: %s", collection_name)
collection = client.get_or_create_collection(collection_name)
self.collection = collection
async def insert(self, vector: [float], id: ObjectID):
logging.info(f"will insert vector: {vector} id: {str(id)}")
self.collection.add(
embeddings=vector,
ids=str(id),
)
async def query(self, vector: [float], top_k: int) -> [ObjectID]:
ret = self.collection.query(
query_embeddings=vector,
n_results=top_k,
)
logging.info(f"query result {ret}")
if len(ret['ids']) == 0:
return []
return list(map(ObjectID.from_base58, ret["ids"][0]))
async def delete(self, id: ObjectID):
self.collection.delete(
ids=id,
)
+16
View File
@@ -0,0 +1,16 @@
# import the ObjectID class
from ..object import ObjectID
# define a vector base class
class VectorBase:
def __init__(self, model_name) -> None:
self.model_name = model_name
async def insert(self, vector: [float], id: ObjectID):
pass
async def query(self, vector: [float], top_k: int) -> [ObjectID]:
pass
async def delete(self, id: ObjectID):
pass
+10 -2
View File
@@ -1,3 +1,10 @@
chromadb==0.4
openai==0.28
toml==0.10
moviepy==1.0
base58==2.1
base36==0.1
aiofiles==23.2.1 aiofiles==23.2.1
aiohttp==3.7.0 aiohttp==3.7.0
aioimaplib==1.0.1 aioimaplib==1.0.1
@@ -5,11 +12,12 @@ aiosmtplib==2.0.2
beautifulsoup4==4.12.2 beautifulsoup4==4.12.2
mail_parser==3.15.0 mail_parser==3.15.0
openai==0.27.10 openai==0.27.10
Pillow==10.0.1 Pillow
prompt_toolkit==3.0.39 prompt_toolkit==3.0.39
protobuf protobuf
pydantic==1.10.11 pydantic==1.10.11
python-telegram-bot==20.5 python-telegram-bot==20.5
Requests==2.31.0 Requests==2.31.0
stability_sdk==0.8.4 stability_sdk
toml==0.10.2 toml==0.10.2
+24
View File
@@ -0,0 +1,24 @@
from aios_kernel.knowledge import KnowledgeBase, EmailObject
# define a email converter class
class EmailConverter:
# define init method
def __init__(self, local_dir, knowledge_base: KnowledgeBase) -> None:
pass
async def run(self):
# convert the email to knowledge object
for email_dir in self._next():
# convert the email to knowledge object
knowledge_object = self._convert(email_dir)
# insert the knowledge object to knowledge base
await self.knowledge_base.insert(knowledge_object)
def _next(self) -> str:
pass
def _convert(self, email_dir) -> EmailObject:
pass
+12
View File
@@ -0,0 +1,12 @@
import asyncio
from .spider import EmailSpider, EmailConverter
if __name__ == "__main__":
spider = EmailSpider("smtp.163.com","user","pwd","./email")
asyncio.run(spider.run())
converter = EmailConverter("./email",KnowledgeBase())
asyncio.run(converter.run())
+17
View File
@@ -0,0 +1,17 @@
# define a email spider class
class EmailSpider:
def __init__(self, address, account, pwd, local_dir) -> None:
pass
async def run(self):
# spide the email from the email server
for email_link in self._next():
# save the email to local directory
self._save(email_link)
def _next(self):
pass
def _save(self, email_link) -> str:
pass
+66
View File
@@ -0,0 +1,66 @@
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)
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,
ChunkListWriter,
)
import asyncio
import unittest
class TestChunk(unittest.TestCase):
def test_chunk_tracker(self):
tracker = KnowledgeStore().get_chunk_tracker()
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 = 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:
text = file.read()
gen.create_chunk_list_from_text(text, 1024)
if __name__ == "__main__":
unittest.main()
+50
View File
@@ -0,0 +1,50 @@
import sys
import os
import logging
dir_path = os.path.dirname(os.path.realpath(__file__))
print(dir_path)
sys.path.append("{}/../src/".format(dir_path))
print(sys.path)
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 ObjectID, HashValue, EmailObjectBuilder
from aios_kernel import KnowledgeBase, AgentPrompt, OpenAI_ComputeNode, ComputeKernel
import asyncio
import unittest
async def test_embedding_email(test):
open_ai_node = OpenAI_ComputeNode()
open_ai_node.start()
ComputeKernel().add_compute_node(open_ai_node)
email_folder = os.path.join(dir_path, "../rootfs/data/email/")
print("explore emails in folder ", email_folder)
for root, dirs, files in os.walk(email_folder):
for dir in dirs:
email_object = EmailObjectBuilder({}, os.path.join(root, dir)).build()
await KnowledgeBase().insert_object(email_object)
msg_prompt = AgentPrompt()
msg_prompt.messages = [{"role":"user","content":"abcdef"}]
await KnowledgeBase().query_prompt(msg_prompt)
class TestKnowledgeBase(unittest.TestCase):
def test_embedding(self):
asyncio.run(test_embedding_email(self))
if __name__ == "__main__":
unittest.main()
+87
View File
@@ -0,0 +1,87 @@
import sys
import os
import logging
dir_path = os.path.dirname(os.path.realpath(__file__))
print(dir_path)
sys.path.append("{}/../src/".format(dir_path))
print(sys.path)
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 (
ObjectID,
HashValue,
EmailObjectBuilder,
ObjectRelationStore,
KnowledgeStore,
EmailObject,
)
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())
email_folder = "F:\\system\\Downloads\\8081ffdb80925f5bff9c6ab9c4756c7d"
email_object = EmailObjectBuilder({}, email_folder).build()
id = email_object.calculate_id()
print(f"got email object: {id.to_base58()}")
# test encode & decode
ret = email_object.encode()
obj = EmailObject.decode(ret)
id2 = obj.calculate_id()
print(f"got email object: {id2.to_base58()}")
self.assertEqual(id.to_base58(), id2.to_base58())
ret2 = obj.encode()
self.assertEqual(ret, ret2)
def test_relation(self):
obj1 = ObjectID.hash_data("12345".encode("utf-8"))
obj2 = ObjectID.hash_data("67890".encode("utf-8"))
obj3 = ObjectID.hash_data("abcde".encode("utf-8"))
obj4 = ObjectID.hash_data("fghij".encode("utf-8"))
print(obj1.to_base58(), obj2.to_base58(), obj3.to_base58())
relation_store = KnowledgeStore().get_relation_store()
relation_store.add_relation(obj1, obj2)
relation_store.add_relation(obj1, obj2)
relation_store.add_relation(obj2, obj3)
relation_store.add_relation(obj1, obj3)
relation_store.add_relation(obj1, obj4)
objs = relation_store.get_related_objects(obj2)
self.assertEqual(len(objs), 1)
self.assertEqual(objs[0], obj3)
objs = relation_store.get_related_root_objects(obj1)
self.assertEqual(len(objs), 2)
self.assertEqual(obj3 in objs, True)
self.assertEqual(obj4 in objs, True)
# self.assertCountEqual(objs, [obj3, obj4])
if __name__ == "__main__":
unittest.main()
+29
View File
@@ -0,0 +1,29 @@
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 ChromaVectorStore
import asyncio
import unittest
async def test_vector():
storage = ChromaVectorStore("test")
await storage.insert([1, 2, 3], "test")
ids = await storage.query([1, 2, 3], 10)
print(ids)
class TestVectorStorage(unittest.TestCase):
def test_run(self):
asyncio.run(test_vector())
if __name__ == "__main__":
unittest.main()