add knowledge pipeline manager
This commit is contained in:
tsukasa
@@ -2,4 +2,5 @@ from .object import *
|
||||
from .vector import *
|
||||
from .data import *
|
||||
from .store import KnowledgeStore
|
||||
from .core_object import *
|
||||
from .core_object import *
|
||||
from .pipeline import *
|
||||
@@ -0,0 +1 @@
|
||||
import local_dir
|
||||
@@ -0,0 +1,158 @@
|
||||
|
||||
class KnowledgeEmailSource:
|
||||
def __init__(self, config:dict):
|
||||
self.config = config
|
||||
self.config["type"] = "email"
|
||||
|
||||
def id(self):
|
||||
return self.config["address"]
|
||||
|
||||
@classmethod
|
||||
def user_config_items(cls):
|
||||
return [("address", "email address"),
|
||||
("password", "email password"),
|
||||
("imap_server", "imap server"),
|
||||
("imap_port", "imap port")
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def local_root(cls):
|
||||
user_data_dir = AIStorage.get_instance().get_myai_dir()
|
||||
return os.path.abspath(f"{user_data_dir}/knowledge/email")
|
||||
|
||||
async def run_once(self):
|
||||
# read config from toml file
|
||||
# and read from config config.local.toml if exists (config.local.toml is ignored by git)
|
||||
logging.debug(f"knowledge email source {self.id()} run once")
|
||||
filter = "ALL"
|
||||
self.client = self.email_client()
|
||||
await self.read_emails(imap_keyword=filter)
|
||||
|
||||
def email_client(self) -> imaplib.IMAP4_SSL:
|
||||
logging.info(f"read email config from {self.config.get('imap_server')}")
|
||||
client = imaplib.IMAP4_SSL(
|
||||
host=self.config.get('imap_server'),
|
||||
port=self.config.get('imap_port')
|
||||
)
|
||||
client.login(self.config.get('address'), self.config.get('password'))
|
||||
return client
|
||||
|
||||
async def read_emails(self, folder: str = 'INBOX', imap_keyword: str = "UNSEEN"):
|
||||
journal_client = KnowledgeJournalClient()
|
||||
latest_journal = journal_client.latest_journal(self.id())
|
||||
latest_uid = 0 if latest_journal is None else int(latest_journal.item_id)
|
||||
self.client.select(folder)
|
||||
_, data = self.client.uid('search', None, imap_keyword)
|
||||
|
||||
# get email uid list
|
||||
email_list = data[0].split()
|
||||
logging.info(f"got {len(email_list)} emails")
|
||||
journal_client = KnowledgeJournalClient()
|
||||
for uid in email_list:
|
||||
_uid = int.from_bytes(uid)
|
||||
if _uid > latest_uid:
|
||||
email_dir = self.check_email_saved(uid)
|
||||
if email_dir is not None:
|
||||
logging.info(f"email uid {uid} already saved")
|
||||
else:
|
||||
email_dir = self.read_and_save_email(uid)
|
||||
logging.info(f"email uid {uid} saved")
|
||||
email_object = EmailObjectBuilder({}, email_dir).build()
|
||||
await KnowledgeBase().insert_object(email_object)
|
||||
journal_client.insert(KnowledgeJournal("email", self.id(), str(int.from_bytes(uid)), str(email_object.calculate_id())))
|
||||
|
||||
|
||||
def read_and_save_email(self, uid: str) -> str:
|
||||
message_parts = "(BODY.PEEK[])"
|
||||
_, email_data = self.client.uid('fetch', uid, message_parts)
|
||||
mail = mailparser.parse_from_bytes(email_data[0][1])
|
||||
logging.info(f"got email subject [{mail.subject}]")
|
||||
self.save_email(mail)
|
||||
return self.get_local_dir_name(mail)
|
||||
|
||||
def get_local_dir_name(self, mail: mailparser.MailParser) -> str:
|
||||
dir = f"{self.local_root()}/{self.config.get('address')}"
|
||||
name = f"{mail.subject}__{mail.date}"
|
||||
name = hashlib.md5(name.encode('utf-8')).hexdigest()
|
||||
return f"{dir}/{name}"
|
||||
|
||||
def check_email_saved(self, uid: str) -> str:
|
||||
message_parts = "(BODY[HEADER])"
|
||||
_, email_data = self.client.uid('fetch', uid, message_parts)
|
||||
mail = mailparser.parse_from_bytes(email_data[0][1])
|
||||
logging.info(f"[{uid}]check email subject [{mail.subject}]")
|
||||
dir = self.get_local_dir_name(mail)
|
||||
logging.info(f"check email saved {dir}")
|
||||
file = f"{dir}/email.txt"
|
||||
if os.path.exists(file):
|
||||
return dir
|
||||
return None
|
||||
|
||||
# save email attachment(images)
|
||||
def save_email_attachment(self, mail: mailparser.MailParser, email_dir: str):
|
||||
for attachment in mail.attachments:
|
||||
if attachment['mail_content_type'] in ['image/png', 'image/jpeg', 'image/gif']:
|
||||
print('current mail have image attachment')
|
||||
img_dir = f"{email_dir}/image"
|
||||
if not os.path.exists(img_dir):
|
||||
os.makedirs(img_dir)
|
||||
filename = attachment['filename']
|
||||
filefullname = f"{img_dir}/{filename}"
|
||||
image_data = attachment['payload']
|
||||
try:
|
||||
image_data = base64.b64decode(image_data)
|
||||
except base64.binascii.Error:
|
||||
image_data = image_data.encode()
|
||||
with open(filefullname, 'wb') as f:
|
||||
f.write(image_data)
|
||||
logging.info(f"save email image {filename} success")
|
||||
|
||||
# save email body images(html content)
|
||||
def save_body_images(self, html_content: str, email_dir: str):
|
||||
# get all image urls
|
||||
soup = BeautifulSoup(html_content, 'html.parser')
|
||||
img_tags = soup.find_all('img')
|
||||
img_urls = [img['src'] for img in img_tags if 'src' in img.attrs]
|
||||
logging.info(f'Found {len(img_urls)} images in email body')
|
||||
|
||||
name_count = 0
|
||||
|
||||
if not os.path.exists(email_dir):
|
||||
os.makedirs(email_dir)
|
||||
|
||||
for img_url in img_urls:
|
||||
# keep the original image filename(last of url)
|
||||
ext = img_url.split('/')[-1].split('.')[-1]
|
||||
img_filename = os.path.join(email_dir, f"{name_count}.{ext}")
|
||||
name_count += 1
|
||||
# download image
|
||||
response = requests.get(img_url, stream=True)
|
||||
if response.status_code == 200:
|
||||
with open(img_filename, 'wb') as img_file:
|
||||
for chunk in response.iter_content(1024):
|
||||
img_file.write(chunk)
|
||||
logging.info(f'Downloaded {img_url} to {img_filename}')
|
||||
else:
|
||||
logging.info(f'Failed to download {img_url}')
|
||||
|
||||
# save email content to local dir
|
||||
def save_email(self, mail: mailparser.MailParser):
|
||||
dir = f"{self.local_root()}/{self.config.get('address')}"
|
||||
if not os.path.exists(dir):
|
||||
os.makedirs(dir)
|
||||
email_dir = self.get_local_dir_name(mail)
|
||||
logging.info(f"save email to {email_dir}")
|
||||
if not os.path.exists(email_dir):
|
||||
os.makedirs(email_dir)
|
||||
with open(f"{email_dir}/email.txt", "w", encoding='utf-8') as f:
|
||||
# soup = BeautifulSoup(mail.body, 'html.parser')
|
||||
f.write(mail.body)
|
||||
with open(f"{email_dir}/meta.json", "w", encoding='utf-8') as f:
|
||||
mail_dict = json.loads(mail.mail_json)
|
||||
if 'body' in mail_dict:
|
||||
del mail_dict['body']
|
||||
json.dump(mail_dict, f, ensure_ascii=False, indent=4)
|
||||
logging.info(f"save email meta info {f.name}")
|
||||
|
||||
self.save_email_attachment(mail, email_dir)
|
||||
self.save_body_images(mail.body, f"{email_dir}/body_image")
|
||||
@@ -0,0 +1,61 @@
|
||||
import os
|
||||
import aiofiles
|
||||
import chardet
|
||||
from knowledge import ImageObjectBuilder, DocumentObjectBuilder, KnowledgeBase
|
||||
|
||||
class KnowledgeDirSource:
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
config["path"] = os.path.abspath(config["path"])
|
||||
|
||||
# @classmethod
|
||||
# def user_config_items(cls):
|
||||
# return [("path", "local dir path")]
|
||||
|
||||
def path(self):
|
||||
return self.config["path"]
|
||||
|
||||
@staticmethod
|
||||
async def read_txt_file(file_path:str)->str:
|
||||
cur_encode = "utf-8"
|
||||
async with aiofiles.open(file_path,'rb') as f:
|
||||
cur_encode = chardet.detect(await f.read())['encoding']
|
||||
|
||||
async with aiofiles.open(file_path,'r',encoding=cur_encode) as f:
|
||||
return await f.read()
|
||||
|
||||
async def next(self):
|
||||
# logging.debug(f"knowledge dir source {self.id()} run once")
|
||||
# journal_client = KnowledgeJournalClient()
|
||||
# latest_journal = journal_client.latest_journal(self.id())
|
||||
# if latest_journal is not None:
|
||||
# if os.path.getmtime(self.path()) <= latest_journal.timestamp:
|
||||
# logging.debug(f"knowledge dir source {self.id()} ingnored for no update")
|
||||
# return
|
||||
while True:
|
||||
file_pathes = sorted(os.listdir(self.path()), key=lambda x: os.path.getctime(os.path.join(self.path(), x)))
|
||||
for rel_path in file_pathes:
|
||||
file_path = os.path.join(self.path(), rel_path)
|
||||
# timestamp = os.path.getctime(file_path)
|
||||
# if latest_journal is not None:
|
||||
# if timestamp <= latest_journal.timestamp:
|
||||
# continue
|
||||
ext = os.path.splitext(file_path)[1].lower()
|
||||
if ext in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']:
|
||||
# logging.info(f"knowledge dir source {self.id()} found image file {file_path}")
|
||||
image = ImageObjectBuilder({}, {}, file_path).build()
|
||||
await KnowledgeBase().insert_object(image)
|
||||
yield image.calculate_id()
|
||||
# journal_client.insert(KnowledgeJournal("dir", self.id(), rel_path, str(image.calculate_id()), timestamp))
|
||||
if ext in ['.txt']:
|
||||
# logging.info(f"knowledge dir source {self.id()} found text file {file_path}")
|
||||
text = await self.read_txt_file(file_path)
|
||||
document = DocumentObjectBuilder({}, {}, text).build()
|
||||
await KnowledgeBase().insert_object(document)
|
||||
yield document.calculate_id()
|
||||
# journal_client.insert(KnowledgeJournal("dir", self.id(), rel_path, str(document.calculate_id()), timestamp))
|
||||
yield 0
|
||||
|
||||
|
||||
def init(params: dict) -> KnowledgeDirSource:
|
||||
return KnowledgeDirSource(params)
|
||||
@@ -0,0 +1,93 @@
|
||||
# define a knowledge base class
|
||||
import json
|
||||
from aios_kernel import ComputeKernel, AIStorage
|
||||
from knowledge import *
|
||||
|
||||
|
||||
class EmbeddingParser:
|
||||
def __init__(self, params: dict):
|
||||
self.store = KnowledgeStore()
|
||||
self.compute_kernel = ComputeKernel.get_instance()
|
||||
self.knowledge_base = KnowledgeBase()
|
||||
self._default_text_model = "all-MiniLM-L6-v2"
|
||||
self._default_image_model = "clip-ViT-B-32"
|
||||
self.vector_store = ChromaVectorStore(AIStorage().get_myai_dir() / "knowledge", self._default_text_model)
|
||||
|
||||
def __get_vector_store(self, model_name: str) -> ChromaVectorStore:
|
||||
return ChromaVectorStore(AIStorage().get_myai_dir() / "knowledge", self._default_text_model)
|
||||
|
||||
async def __embedding_document(self, document: DocumentObject):
|
||||
for chunk_id in document.get_chunk_list():
|
||||
chunk = self.knowledge_base.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, self._default_text_model)
|
||||
if vector:
|
||||
await self.get_vector_store(self._default_text_model).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), self._default_text_model)
|
||||
vector = await self.compute_kernel.do_image_embedding(image.calculate_id(), self._default_image_model)
|
||||
if vector:
|
||||
await self.get_vector_store(self._default_image_model).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), self._default_text_model)
|
||||
await self.get_vector_store(self._default_text_model).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()), self._default_text_model)
|
||||
await self.get_vector_store(self._default_text_model).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
|
||||
|
||||
async def parse(self, object: ObjectID):
|
||||
obj = self.knowledge_base.load_object(object)
|
||||
await self.__do_embedding(obj)
|
||||
|
||||
def init(params: dict) -> EmbeddingParser:
|
||||
return EmbeddingParser(params)
|
||||
@@ -0,0 +1,165 @@
|
||||
|
||||
# class KnowledgePipelineTemplate
|
||||
import runpy
|
||||
import toml
|
||||
import datetime
|
||||
import sqlite3
|
||||
import os
|
||||
from . import ObjectID, KnowledgeStore
|
||||
import asyncio
|
||||
|
||||
class KnowledgePipelineJournal:
|
||||
def __init__(self, time: datetime.datetime, object_id: str, input: str, parser: str):
|
||||
self.time = time
|
||||
self.object_id = None if object_id is None else ObjectID.from_base58(object_id)
|
||||
self.input = input
|
||||
self.parser = parser
|
||||
|
||||
def is_finish(self) -> bool:
|
||||
self.object_id is None
|
||||
|
||||
|
||||
|
||||
# init sqlite3 client
|
||||
class KnowledgePipelineJournalClient:
|
||||
def __init__(self, pipeline_path: str = None):
|
||||
if not os.path.exists(pipeline_path):
|
||||
os.makedirs(pipeline_path)
|
||||
self.journal_path = os.path.join(pipeline_path, "journal.db")
|
||||
|
||||
conn = sqlite3.connect(self.journal_path)
|
||||
conn.execute(
|
||||
'''CREATE TABLE IF NOT EXISTS journal (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
time DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||||
object_id TEXT,
|
||||
input TEXT,
|
||||
parser TEXT)'''
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def insert(self, object_id: ObjectID, input: str, parser: str, timestamp: datetime.datetime = None):
|
||||
timestamp = datetime.datetime.now() if timestamp is None else timestamp
|
||||
conn = sqlite3.connect(self.journal_path)
|
||||
conn.execute(
|
||||
"INSERT INTO journal (time, object_id, input, parser) VALUES (?, ?, ?)",
|
||||
(timestamp, str(object_id), input, parser),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def latest_journals(self, topn) -> [KnowledgePipelineJournal]:
|
||||
conn = sqlite3.connect(self.journal_path)
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT * FROM journal ORDER BY id DESC LIMIT ?", (topn,))
|
||||
return [KnowledgePipelineJournal(time, object_id, input, parser) for (_, time, object_id, input, parser) in cursor.fetchall()]
|
||||
|
||||
class KnowledgePipelineEnvironment:
|
||||
def __init__(self, pipeline_path: str):
|
||||
self.knowledge_base = KnowledgeStore()
|
||||
self.pipeline_path = pipeline_path
|
||||
self.journal = KnowledgePipelineJournalClient(pipeline_path)
|
||||
|
||||
class KnowledgePipelineState(Enum):
|
||||
INIT = 0
|
||||
RUNNING = 1
|
||||
STOPPED = 2
|
||||
FINISHED = 3
|
||||
|
||||
class KnowledgePipeline:
|
||||
def __init__(self, name: str, env: KnowledgePipelineEnvironment, input_init, input_params, parser_init, parser_params):
|
||||
self.name = name
|
||||
self.state = KnowledgePipelineState.INIT
|
||||
self.input_init = input_init
|
||||
self.input_params = input_params
|
||||
self.parser_init = parser_init
|
||||
self.parser_params = parser_params
|
||||
self.env = env
|
||||
self.input = None
|
||||
self.parser = None
|
||||
|
||||
async def run(self):
|
||||
if self.state == KnowledgePipelineState.INIT:
|
||||
self.input = self.input_init(self.env, self.input_params)
|
||||
self.parser = self.parser_init(self.env, self.parser_params)
|
||||
self.state = KnowledgePipelineState.RUNNING
|
||||
if self.state == KnowledgePipelineState.RUNNING:
|
||||
for input in await self.input.next():
|
||||
if input is None:
|
||||
self.state = KnowledgePipelineState.FINISHED
|
||||
self.env.journal.insert(None, "finished", "finished")
|
||||
return
|
||||
(object_id, input_journal) = input
|
||||
if object_id is None:
|
||||
parser_journal = await self.parser.parse(object_id)
|
||||
self.env.journal.insert(object_id, input_journal, parser_journal)
|
||||
if self.state == KnowledgePipelineState.STOPPED:
|
||||
return
|
||||
if self.state == KnowledgePipelineState.FINISHED:
|
||||
return
|
||||
|
||||
class KnowledgePipelineManager:
|
||||
def __init__(self, root_dir: str):
|
||||
self.root_dir = root_dir
|
||||
self.input_modules = {}
|
||||
self.parser_modules = {}
|
||||
self.pipelines = {
|
||||
"names": {},
|
||||
"running": []
|
||||
}
|
||||
from .input import local_dir
|
||||
self.register_input("local_dir", local_dir.init)
|
||||
|
||||
def register_input(self, name: str, init_method):
|
||||
self.input_modules[name] = init_method
|
||||
|
||||
def register_parser(self, name: str, parser_method):
|
||||
self.parser_modules[name] = parser_method
|
||||
|
||||
def add_pipeline(self, config: dict, path: str):
|
||||
name = config["name"]
|
||||
if name in self.pipelines["names"]:
|
||||
return
|
||||
|
||||
input_module = self.input_modules[config["input"]["module"]]
|
||||
_, ext = os.path.splitext(input_module)
|
||||
if ext == ".py":
|
||||
input_module = os.path.abspath(path, input_module)
|
||||
input_init = runpy.run_path(input_module)["init"]
|
||||
else:
|
||||
input_init = self.input_modules.get(input_module)
|
||||
input_params = config["input"]["params"]
|
||||
|
||||
parser_module = self.parser_modules[config["parser"]["module"]]
|
||||
_, ext = os.path.splitext(parser_module)
|
||||
if ext == ".py":
|
||||
parser_module = os.path.abspath(path, parser_module)
|
||||
parser_init = runpy.run_path(parser_module)["init"]
|
||||
else:
|
||||
parser_init = self.parser_modules.get(parser_module)
|
||||
parser_params = config["parser"]["params"]
|
||||
|
||||
|
||||
data_path = self.root_dir / name
|
||||
env = KnowledgePipelineEnvironment(data_path)
|
||||
pipeline = KnowledgePipeline(name, env, input_init, input_params, parser_init, parser_params)
|
||||
self.pipelines["names"][name] = pipeline
|
||||
self.pipelines["running"].append(pipeline)
|
||||
|
||||
async def run(self):
|
||||
while True:
|
||||
for pipeline in self.pipelines["running"]:
|
||||
await pipeline.run()
|
||||
await asyncio.sleep(5)
|
||||
|
||||
def load_dir(self, root: str):
|
||||
config_path = os.path.join(root, "pipelines.toml")
|
||||
with open(config_path, "r") as f:
|
||||
config = toml.load(f)
|
||||
for path in config["pipelines"]:
|
||||
pipeline_path = os.path.join(root, path)
|
||||
with open(os.path.join(pipeline_path, "pipeline.toml")) as f:
|
||||
pipeline_config = toml.load(f)
|
||||
self.add_pipeline(pipeline_config, pipeline_path)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
|
||||
class KnowledgeEnvironment(Environment):
|
||||
def __init__(self, env_id: str) -> None:
|
||||
super().__init__(env_id)
|
||||
|
||||
query_param = {
|
||||
"tokens": "key words to query",
|
||||
"types": "prefered knowledge types, one or more of [text, image]",
|
||||
"index": "index of query result"
|
||||
}
|
||||
self.add_ai_function(SimpleAIFunction("query_knowledge",
|
||||
"vector query content from local knowledge base",
|
||||
self._query,
|
||||
query_param))
|
||||
async def query_objects(self, tokens: str, types: list[str], topk: int) -> [ObjectID]:
|
||||
texts = []
|
||||
if "text" in types:
|
||||
vector = await self.compute_kernel.do_text_embedding(tokens, self._default_text_model)
|
||||
texts = await self.store.get_vector_store(self._default_text_model).query(vector, topk)
|
||||
images = []
|
||||
if "image" in types:
|
||||
vector = await self.compute_kernel.do_text_embedding(tokens, self._default_image_model)
|
||||
images = await self.store.get_vector_store(self._default_image_model).query(vector, topk)
|
||||
return texts + images
|
||||
|
||||
|
||||
def tokens_from_objects(self, object_ids: [ObjectID]) -> list[str]:
|
||||
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 = ""
|
||||
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 = dict()
|
||||
desc["id"] = str(object_id)
|
||||
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"
|
||||
|
||||
return content
|
||||
|
||||
async def _query(self, tokens: str, types: list[str] = ["text"], index: str=0):
|
||||
index = int(index)
|
||||
object_ids = await KnowledgeBase().query_objects(tokens, types, 4)
|
||||
if len(object_ids) <= index:
|
||||
return "*** I have no more information for your reference.\n"
|
||||
else:
|
||||
content = "*** I have provided the following known information for your reference with json format:\n"
|
||||
return content + KnowledgeBase().tokens_from_objects(object_ids[index:index+1])
|
||||
+44
-9
@@ -1,8 +1,9 @@
|
||||
import os
|
||||
|
||||
from .object import ObjectStore, ObjectRelationStore
|
||||
from .object import ObjectStore, ObjectRelationStore, ObjectID, ObjectType, KnowledgeObject, DocumentObject, ImageObject, VideoObject, RichTextObject, EmailObject
|
||||
from core_object import DocumentObject, ImageObject, VideoObject, RichTextObject, EmailObject
|
||||
from .data import ChunkStore, ChunkTracker, ChunkListWriter, ChunkReader
|
||||
from .vector import ChromaVectorStore, VectorBase
|
||||
import json
|
||||
import logging
|
||||
|
||||
|
||||
@@ -17,7 +18,7 @@ class KnowledgeStore:
|
||||
cls._instance = super().__new__(cls)
|
||||
|
||||
import aios_kernel
|
||||
knowledge_dir = aios_kernel.storage.AIStorage().get_myai_dir() / "knowledge"
|
||||
knowledge_dir = aios_kernel.storage.AIStorage().get_myai_dir() / "knowledge" / "objects"
|
||||
|
||||
if not os.path.exists(knowledge_dir):
|
||||
os.makedirs(knowledge_dir)
|
||||
@@ -42,10 +43,8 @@ class KnowledgeStore:
|
||||
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
|
||||
|
||||
@@ -64,7 +63,43 @@ class KnowledgeStore:
|
||||
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(self.root, model_name)
|
||||
return self.vector_store[model_name]
|
||||
async def insert_object(self, object: KnowledgeObject):
|
||||
self.object_store.put_object(object.calculate_id(), object.encode())
|
||||
|
||||
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 parse_object_in_message(self, message: str) -> KnowledgeObject:
|
||||
# get message's first line
|
||||
logging.info(f"tg parse resp message: {message}")
|
||||
lines = message.split("\n")
|
||||
if len(lines) > 0:
|
||||
message = lines[0]
|
||||
try:
|
||||
desc = json.loads(message)
|
||||
if isinstance(desc, dict):
|
||||
object_id = desc["id"]
|
||||
else:
|
||||
object_id = desc[0]["id"]
|
||||
except Exception as e:
|
||||
return None
|
||||
|
||||
if object_id is not None:
|
||||
return self.load_object(ObjectID.from_base58(object_id))
|
||||
|
||||
|
||||
def bytes_from_object(self, object: KnowledgeObject) -> bytes:
|
||||
if object.get_object_type() == ObjectType.Image:
|
||||
image_object = object
|
||||
return self.get_chunk_reader().read_chunk_list_to_single_bytes(image_object.get_chunk_list())
|
||||
Reference in New Issue
Block a user