Merge pull request #55 from photosssa/MVP
Objected knowleadge base, a specialized implemention for emails
This commit is contained in:
+3
-1
@@ -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
@@ -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.
|
||||||
|
|
||||||
@@ -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"/>
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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!"
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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}"
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
@@ -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
|
||||||
|
|||||||
@@ -0,0 +1,5 @@
|
|||||||
|
from .object import *
|
||||||
|
from .vector import *
|
||||||
|
from .data import *
|
||||||
|
from .store import KnowledgeStore
|
||||||
|
from .core_object import *
|
||||||
@@ -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
|
||||||
@@ -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
|
||||||
@@ -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()
|
||||||
@@ -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)
|
||||||
@@ -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
|
||||||
@@ -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)
|
||||||
@@ -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}"
|
||||||
@@ -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)
|
||||||
|
|
||||||
@@ -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"))
|
||||||
@@ -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()
|
||||||
@@ -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
|
||||||
@@ -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
|
||||||
@@ -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,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)
|
||||||
@@ -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
|
||||||
@@ -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)
|
||||||
@@ -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()
|
||||||
@@ -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]
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
from .vector_base import VectorBase
|
||||||
|
from .chroma_store import ChromaVectorStore
|
||||||
@@ -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,
|
||||||
|
)
|
||||||
@@ -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
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|
||||||
@@ -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())
|
||||||
|
|
||||||
|
|
||||||
@@ -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
|
||||||
@@ -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()
|
||||||
@@ -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()
|
||||||
@@ -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()
|
||||||
@@ -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()
|
||||||
Reference in New Issue
Block a user