define bas environment
This commit is contained in:
tsukasa
@@ -0,0 +1,671 @@
|
||||
# import os
|
||||
# import aiofiles
|
||||
# import chardet
|
||||
# import logging
|
||||
# import string
|
||||
# from knowledge import ImageObjectBuilder, DocumentObjectBuilder, KnowledgePipelineEnvironment, KnowledgePipelineJournal
|
||||
# from aios_kernel.storage import AIStorage
|
||||
|
||||
|
||||
import os
|
||||
import aiofiles
|
||||
import chardet
|
||||
import logging
|
||||
import string
|
||||
import sqlite3
|
||||
import json
|
||||
import threading
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from typing import Optional, List
|
||||
from knowledge import ImageObjectBuilder, DocumentObjectBuilder, KnowledgePipelineEnvironment, KnowledgePipelineJournal
|
||||
from aios_kernel import AIStorage, SimpleEnvironment
|
||||
|
||||
|
||||
class ScanLocalDocument:
|
||||
def __init__(self, env: KnowledgePipelineEnvironment, config):
|
||||
self.env = env
|
||||
path = string.Template(config["path"]).substitute(myai_dir=AIStorage.get_instance().get_myai_dir())
|
||||
config["path"] = path
|
||||
self.config = config
|
||||
|
||||
def path(self):
|
||||
return self.config["path"]
|
||||
|
||||
async def next(self):
|
||||
while True:
|
||||
journals = self.env.journal.latest_journals(1)
|
||||
from_time = 0
|
||||
if len(journals) == 1:
|
||||
latest_journal = journals[0]
|
||||
if latest_journal.is_finish():
|
||||
yield None
|
||||
continue
|
||||
from_time = os.path.getctime(latest_journal.get_input())
|
||||
if os.path.getmtime(self.path()) <= from_time:
|
||||
yield (None, None)
|
||||
continue
|
||||
|
||||
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 timestamp <= from_time:
|
||||
continue
|
||||
ext = os.path.splitext(file_path)[1].lower()
|
||||
if ext in ['.pdf', '.md', '.txt']:
|
||||
logging.info(f"knowledge dir source found document file {file_path}")
|
||||
yield (file_path, file_path)
|
||||
yield (None, None)
|
||||
|
||||
class MetaDatabase:
|
||||
def __init__(self,db_path:str):
|
||||
self.db_path = db_path
|
||||
self._get_conn()
|
||||
|
||||
def _get_conn(self):
|
||||
""" get db connection """
|
||||
local = threading.local()
|
||||
if not hasattr(local, 'conn'):
|
||||
local.conn = self._create_connection(self.db_path)
|
||||
return local.conn
|
||||
|
||||
|
||||
def _create_connection(self, db_file):
|
||||
""" create a database connection to a SQLite database """
|
||||
conn = None
|
||||
try:
|
||||
conn = sqlite3.connect(db_file)
|
||||
except Exception as e:
|
||||
logger.error("Error occurred while connecting to database: %s", e)
|
||||
return None
|
||||
|
||||
if conn:
|
||||
self._create_tables(conn)
|
||||
|
||||
return conn
|
||||
|
||||
def _create_tables(self,conn):
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS documents (
|
||||
doc_path TEXT PRIMARY KEY,
|
||||
length INTEGER,
|
||||
last_modify TEXT,
|
||||
doc_hash TEXT,
|
||||
create_time TEXT
|
||||
)
|
||||
''')
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS knowledge (
|
||||
doc_hash TEXT PRIMARY KEY,
|
||||
title TEXT,
|
||||
summary TEXT,
|
||||
content TEXT,
|
||||
catalogs TEXT,
|
||||
tags TEXT,
|
||||
llm_title TEXT,
|
||||
llm_summary TEXT,
|
||||
create_time TEXT
|
||||
)
|
||||
''')
|
||||
|
||||
cursor.execute('''
|
||||
CREATE INDEX IF NOT EXISTS idx_documents_doc_hash
|
||||
ON documents (doc_hash)
|
||||
''')
|
||||
|
||||
cursor.execute('''
|
||||
CREATE INDEX IF NOT EXISTS idx_knowledge_tags
|
||||
ON knowledge (tags)
|
||||
''')
|
||||
|
||||
conn.commit()
|
||||
|
||||
def add_doc(self, doc_path: str, length: int, last_modify: str, doc_hash: Optional[str] = None):
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
create_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
cursor.execute('''
|
||||
INSERT INTO documents (doc_path, length, last_modify, doc_hash,create_time)
|
||||
VALUES (?, ?, ?, ?,?)
|
||||
''', (doc_path, length, last_modify, doc_hash,create_time))
|
||||
conn.commit()
|
||||
|
||||
def is_doc_exist(self, doc_path: str) -> bool:
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
SELECT doc_path
|
||||
FROM documents
|
||||
WHERE doc_path = ?
|
||||
''', (doc_path,))
|
||||
return len(cursor.fetchall()) > 0
|
||||
|
||||
def set_doc_hash(self, doc_path: str, doc_hash: str):
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
UPDATE documents
|
||||
SET doc_hash = ?
|
||||
WHERE doc_path = ?
|
||||
''', (doc_hash, doc_path))
|
||||
conn.commit()
|
||||
|
||||
def get_docs_without_hash(self,limit:int=1024) -> List[str]:
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
SELECT doc_path
|
||||
FROM documents
|
||||
WHERE doc_hash IS NULL OR doc_hash = ''
|
||||
ORDER BY create_time DESC
|
||||
LIMIT ?
|
||||
''',(limit,))
|
||||
return [row[0] for row in cursor.fetchall()]
|
||||
|
||||
#metadata["summary"]
|
||||
#metadata["catelogs"]
|
||||
#metadata["tags"]
|
||||
def add_knowledge(self, doc_hash: str, title: str, metadata: dict,content:str = None,):
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
|
||||
create_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
summary = metadata.get("summary", "")
|
||||
catalogs = metadata.get("catalogs","")
|
||||
tags = ','.join(metadata.get("tags", []))
|
||||
|
||||
cursor.execute('''
|
||||
INSERT INTO knowledge (doc_hash, title , summary , catalogs , tags,create_time)
|
||||
VALUES (?, ?, ?, ?, ?,?)
|
||||
''', (doc_hash, title, summary, catalogs, tags,create_time))
|
||||
conn.commit()
|
||||
|
||||
#llm_result["summary"]
|
||||
#llm_result["tags"]
|
||||
#llm_result["catelog"]
|
||||
def set_knowledge_llm_result(self, doc_hash: str, llm_result: dict):
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
|
||||
title = llm_result.get("title", "")
|
||||
summary = llm_result.get("summary", "")
|
||||
catalogs = json.dumps(llm_result.get("catalogs", {}))
|
||||
tags = ','.join(llm_result.get("tags", []))
|
||||
|
||||
cursor.execute('''
|
||||
UPDATE knowledge
|
||||
SET llm_title = ?,llm_summary = ?, catalogs = ?, tags = ?
|
||||
WHERE doc_hash = ?
|
||||
''', (title,summary, catalogs, tags, doc_hash))
|
||||
conn.commit()
|
||||
|
||||
def get_hash_by_doc_path(self, doc_path: str) -> Optional[str]:
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
SELECT doc_hash
|
||||
FROM documents
|
||||
WHERE doc_path = ?
|
||||
''', (doc_path,))
|
||||
row = cursor.fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
return row[0]
|
||||
|
||||
def get_knowledge(self, doc_hash: str) -> Optional[dict]:
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
SELECT title, summary, catalogs, tags, llm_title, llm_summary
|
||||
FROM knowledge
|
||||
WHERE doc_hash = ?
|
||||
''', (doc_hash,))
|
||||
row = cursor.fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
|
||||
# get doc path
|
||||
cursor.execute('''
|
||||
SELECT doc_path
|
||||
FROM documents
|
||||
WHERE doc_hash = ?
|
||||
''', (doc_hash,))
|
||||
row2 = cursor.fetchone()
|
||||
if row2 is None:
|
||||
return None
|
||||
doc_path = row2[0]
|
||||
|
||||
|
||||
return {
|
||||
"full_path": doc_path,
|
||||
"title": row[0],
|
||||
"summary": row[1],
|
||||
"catalogs": row[2],
|
||||
"tags": row[3],
|
||||
"llm_title" : row[4],
|
||||
"llm_summary" : row[5],
|
||||
}
|
||||
|
||||
def get_knowledge_without_llm_title(self,limit:int=16) -> List[str]:
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
SELECT doc_hash
|
||||
FROM knowledge
|
||||
WHERE llm_title IS NULL OR llm_title = ''
|
||||
ORDER BY create_time DESC
|
||||
LIMIT ?
|
||||
''',(limit,))
|
||||
return [row[0] for row in cursor.fetchall()]
|
||||
|
||||
def query_docs_by_tag(self, tag: str) -> List[str]:
|
||||
conn = self._get_conn()
|
||||
cursor = conn.cursor()
|
||||
tag_json = json.dumps(tag) # 将标签转换为 JSON 字符串
|
||||
cursor.execute('''
|
||||
SELECT documents.doc_path
|
||||
FROM documents
|
||||
JOIN knowledge ON documents.doc_hash = knowledge.doc_hash
|
||||
WHERE json_extract(knowledge.tags, '$') LIKE ?
|
||||
''', (tag))
|
||||
return [row[0] for row in cursor.fetchall()]
|
||||
|
||||
|
||||
class DocumentKnowledgeBase(SimpleEnvironment):
|
||||
async def get_knowledege_catalog(self,path:str=None,only_dir =True,max_depth:int=5)->str:
|
||||
if path:
|
||||
full_path = f"{self.root_path}/knowledge/{path}"
|
||||
else:
|
||||
full_path = f"{self.root_path}/knowledge"
|
||||
|
||||
catlogs,file_count = await self.get_directory_structure(full_path,max_depth,only_dir)
|
||||
return catlogs
|
||||
|
||||
async def get_directory_structure(self,root_dir, max_depth:int=4, only_dir=True, indent=1):
|
||||
file_count = 0
|
||||
structure_str = ''
|
||||
if os.path.isdir(root_dir):
|
||||
sub_files = []
|
||||
with os.scandir(root_dir) as it:
|
||||
for entry in it:
|
||||
if entry.is_dir():
|
||||
sub_structure, sub_count = await self.get_directory_structure(entry.path, max_depth, only_dir, indent + 1)
|
||||
if sub_structure:
|
||||
structure_str += sub_structure
|
||||
file_count += sub_count
|
||||
else:
|
||||
file_count += 1
|
||||
sub_files.append(entry.name)
|
||||
|
||||
if only_dir is False:
|
||||
for file_name in sub_files:
|
||||
structure_str = structure_str + ' ' * (indent+1) + file_name + '\n'
|
||||
|
||||
dir_name = os.path.basename(root_dir)
|
||||
dir_info = f"{dir_name} <count: {file_count}>"
|
||||
|
||||
|
||||
structure_str = ' ' * indent + dir_info + '\n' + structure_str
|
||||
|
||||
if indent - 1 >= max_depth:
|
||||
return None, file_count
|
||||
else:
|
||||
return structure_str, file_count
|
||||
|
||||
# inner_function
|
||||
async def get_knowledge(self,path:str) -> str:
|
||||
full_path = f"{self.root_path}/knowledge/{path}"
|
||||
if os.islink(full_path):
|
||||
org_path = os.readlink(full_path)
|
||||
hash = self.kb_db.get_hash_by_doc_path(org_path)
|
||||
if hash:
|
||||
return self.kb_db.get_knowledge(org_path)
|
||||
|
||||
return "not found"
|
||||
|
||||
|
||||
class ParseLocalDocument:
|
||||
def _parse_pdf_bookmarks(self,bookmarks, parent:list):
|
||||
|
||||
for item in bookmarks:
|
||||
if isinstance(item,list):
|
||||
self._parse_pdf_bookmarks(item,parent)
|
||||
else:
|
||||
if item.title:
|
||||
new_item = {}
|
||||
new_item["page"] = item.page.idnum
|
||||
new_item["title"] = item.title
|
||||
my_childs = []
|
||||
if item.childs:
|
||||
if len(item.childs) > 0:
|
||||
self._parse_pdf_bookmarks(item.childs, my_childs)
|
||||
new_item["childs"] = my_childs
|
||||
parent.append(new_item)
|
||||
else:
|
||||
logger.warning("parse pdf bookmarks failed: item.title is None!")
|
||||
|
||||
return
|
||||
|
||||
def _parse_pdf(self,doc_path:str):
|
||||
metadata = {}
|
||||
with open(doc_path, 'rb') as file:
|
||||
reader = PyPDF2.PdfReader(file)
|
||||
try:
|
||||
doc_info = reader.metadata
|
||||
if doc_info:
|
||||
if doc_info.title:
|
||||
metadata["title"] = doc_info.title
|
||||
if doc_info.author:
|
||||
metadata["authors"] = doc_info.author
|
||||
except Exception as e:
|
||||
logger.warn("parse pdf metadata failed:%s",e)
|
||||
|
||||
dir_path = os.path.dirname(doc_path)
|
||||
base_name = os.path.basename(doc_path)
|
||||
text_content_path = f"{dir_path}/.{base_name}.txt"
|
||||
full_text = ""
|
||||
|
||||
for page in reader.pages:
|
||||
text = page.extract_text()
|
||||
full_text += text
|
||||
with open(text_content_path, 'w', encoding='utf-8') as f:
|
||||
f.write(full_text)
|
||||
|
||||
try:
|
||||
bookmarks = reader.outline
|
||||
if bookmarks:
|
||||
catalogs = []
|
||||
self._parse_pdf_bookmarks(bookmarks,catalogs)
|
||||
metadata["catalogs"] = json.dumps(catalogs)
|
||||
except Exception as e:
|
||||
logger.warn("parse pdf bookmarks failed:%s",e)
|
||||
|
||||
return metadata
|
||||
|
||||
def _parse_txt(self,doc_path:str):
|
||||
return {}
|
||||
|
||||
def _parse_md(self,doc_path:str):
|
||||
metadata = {}
|
||||
cur_encode = "utf-8"
|
||||
with open(doc_path,'rb') as f:
|
||||
cur_encode = chardet.detect(f.read(1024))['encoding']
|
||||
|
||||
with open(doc_path, mode='r', encoding=cur_encode) as f:
|
||||
content = f.read()
|
||||
match = re.search(r'^# (.*)', content, re.MULTILINE)
|
||||
if match:
|
||||
metadata['title'] = match.group(1).strip()
|
||||
md = Markdown(extensions=['toc'])
|
||||
html_str = md.convert(content)
|
||||
toc = md.toc
|
||||
if toc:
|
||||
metadata['catalogs'] = toc
|
||||
|
||||
return metadata
|
||||
|
||||
def _parse_document(self,doc_path:str):
|
||||
hash_result = None
|
||||
title = os.path.basename(doc_path)
|
||||
meta_data = {}
|
||||
|
||||
with open(doc_path, "rb") as f:
|
||||
hash_md5 = hashlib.md5()
|
||||
for chunk in iter(lambda: f.read(1024*1024), b""):
|
||||
hash_md5.update(chunk)
|
||||
hash_result = hash_md5.hexdigest()
|
||||
try:
|
||||
if doc_path.endswith(".md"):
|
||||
meta_data = self._parse_md(doc_path)
|
||||
elif doc_path.endswith(".pdf"):
|
||||
meta_data = self._parse_pdf(doc_path)
|
||||
except Exception as e:
|
||||
logger.error("parse document %s failed:%s",doc_path,e)
|
||||
traceback.print_exc()
|
||||
|
||||
if meta_data.get("title"):
|
||||
title = meta_data["title"]
|
||||
logger.info("parse document %s!",doc_path)
|
||||
return hash_result,title,meta_data
|
||||
|
||||
async def parse(self, file_path: str) -> str:
|
||||
|
||||
|
||||
|
||||
# async def get_knowledege_catalog(self,path:str=None,only_dir =True,max_depth:int=5)->str:
|
||||
# if path:
|
||||
# full_path = f"{self.root_path}/knowledge/{path}"
|
||||
# else:
|
||||
# full_path = f"{self.root_path}/knowledge"
|
||||
|
||||
# catlogs,file_count = await self.get_directory_structure(full_path,max_depth,only_dir)
|
||||
# return catlogs
|
||||
|
||||
# async def get_directory_structure(self,root_dir, max_depth:int=4, only_dir=True, indent=1):
|
||||
# file_count = 0
|
||||
# structure_str = ''
|
||||
# if os.path.isdir(root_dir):
|
||||
# sub_files = []
|
||||
# with os.scandir(root_dir) as it:
|
||||
# for entry in it:
|
||||
# if entry.is_dir():
|
||||
# sub_structure, sub_count = await self.get_directory_structure(entry.path, max_depth, only_dir, indent + 1)
|
||||
# if sub_structure:
|
||||
# structure_str += sub_structure
|
||||
# file_count += sub_count
|
||||
# else:
|
||||
# file_count += 1
|
||||
# sub_files.append(entry.name)
|
||||
|
||||
# if only_dir is False:
|
||||
# for file_name in sub_files:
|
||||
# structure_str = structure_str + ' ' * (indent+1) + file_name + '\n'
|
||||
|
||||
# dir_name = os.path.basename(root_dir)
|
||||
# dir_info = f"{dir_name} <count: {file_count}>"
|
||||
|
||||
|
||||
# structure_str = ' ' * indent + dir_info + '\n' + structure_str
|
||||
|
||||
# if indent - 1 >= max_depth:
|
||||
# return None, file_count
|
||||
# else:
|
||||
# return structure_str, file_count
|
||||
|
||||
# # inner_function
|
||||
# async def get_knowledge(self,path:str) -> str:
|
||||
# full_path = f"{self.root_path}/knowledge/{path}"
|
||||
# if os.islink(full_path):
|
||||
# org_path = os.readlink(full_path)
|
||||
# hash = self.kb_db.get_hash_by_doc_path(org_path)
|
||||
# if hash:
|
||||
# return self.kb_db.get_knowledge(org_path)
|
||||
|
||||
# return "not found"
|
||||
|
||||
async def load_knowledge_content(self,path:str,pos:int=0,length:int=None) -> str:
|
||||
if path.endswith("pdf"):
|
||||
logger.info("load_knowledge_content:pdf")
|
||||
dir_path = os.path.dirname(path)
|
||||
base_name = os.path.basename(path)
|
||||
text_content_path = f"{dir_path}/.{base_name}.txt"
|
||||
if os.path.exists(text_content_path) is False:
|
||||
return None
|
||||
async with aiofiles.open(path, mode='r', encoding=cur_encode) as f:
|
||||
await f.seek(pos)
|
||||
content = await f.read(length)
|
||||
return content
|
||||
else:
|
||||
async with aiofiles.open(path,'rb') as f:
|
||||
cur_encode = chardet.detect(await f.read())['encoding']
|
||||
|
||||
async with aiofiles.open(path, mode='r', encoding=cur_encode) as f:
|
||||
await f.seek(pos)
|
||||
content = await f.read(length)
|
||||
return content
|
||||
|
||||
return "load content failed."
|
||||
|
||||
def _add_document_dir(self,path:str):
|
||||
self.doc_dirs[path] = 0
|
||||
|
||||
|
||||
def _parse_pdf_bookmarks(self,bookmarks, parent:list):
|
||||
|
||||
for item in bookmarks:
|
||||
if isinstance(item,list):
|
||||
self._parse_pdf_bookmarks(item,parent)
|
||||
else:
|
||||
if item.title:
|
||||
new_item = {}
|
||||
new_item["page"] = item.page.idnum
|
||||
new_item["title"] = item.title
|
||||
my_childs = []
|
||||
if item.childs:
|
||||
if len(item.childs) > 0:
|
||||
self._parse_pdf_bookmarks(item.childs, my_childs)
|
||||
new_item["childs"] = my_childs
|
||||
parent.append(new_item)
|
||||
else:
|
||||
logger.warning("parse pdf bookmarks failed: item.title is None!")
|
||||
|
||||
return
|
||||
|
||||
def _parse_pdf(self,doc_path:str):
|
||||
metadata = {}
|
||||
with open(doc_path, 'rb') as file:
|
||||
reader = PyPDF2.PdfReader(file)
|
||||
try:
|
||||
doc_info = reader.metadata
|
||||
if doc_info:
|
||||
if doc_info.title:
|
||||
metadata["title"] = doc_info.title
|
||||
if doc_info.author:
|
||||
metadata["authors"] = doc_info.author
|
||||
except Exception as e:
|
||||
logger.warn("parse pdf metadata failed:%s",e)
|
||||
|
||||
dir_path = os.path.dirname(doc_path)
|
||||
base_name = os.path.basename(doc_path)
|
||||
text_content_path = f"{dir_path}/.{base_name}.txt"
|
||||
full_text = ""
|
||||
|
||||
for page in reader.pages:
|
||||
text = page.extract_text()
|
||||
full_text += text
|
||||
with open(text_content_path, 'w', encoding='utf-8') as f:
|
||||
f.write(full_text)
|
||||
|
||||
try:
|
||||
bookmarks = reader.outline
|
||||
if bookmarks:
|
||||
catalogs = []
|
||||
self._parse_pdf_bookmarks(bookmarks,catalogs)
|
||||
metadata["catalogs"] = json.dumps(catalogs)
|
||||
except Exception as e:
|
||||
logger.warn("parse pdf bookmarks failed:%s",e)
|
||||
|
||||
return metadata
|
||||
|
||||
def _parse_txt(self,doc_path:str):
|
||||
return {}
|
||||
|
||||
def _parse_md(self,doc_path:str):
|
||||
metadata = {}
|
||||
cur_encode = "utf-8"
|
||||
with open(doc_path,'rb') as f:
|
||||
cur_encode = chardet.detect(f.read(1024))['encoding']
|
||||
|
||||
with open(doc_path, mode='r', encoding=cur_encode) as f:
|
||||
content = f.read()
|
||||
match = re.search(r'^# (.*)', content, re.MULTILINE)
|
||||
if match:
|
||||
metadata['title'] = match.group(1).strip()
|
||||
md = Markdown(extensions=['toc'])
|
||||
html_str = md.convert(content)
|
||||
toc = md.toc
|
||||
if toc:
|
||||
metadata['catalogs'] = toc
|
||||
|
||||
return metadata
|
||||
|
||||
def _parse_document(self,doc_path:str):
|
||||
hash_result = None
|
||||
title = os.path.basename(doc_path)
|
||||
meta_data = {}
|
||||
|
||||
with open(doc_path, "rb") as f:
|
||||
hash_md5 = hashlib.md5()
|
||||
for chunk in iter(lambda: f.read(1024*1024), b""):
|
||||
hash_md5.update(chunk)
|
||||
hash_result = hash_md5.hexdigest()
|
||||
try:
|
||||
if doc_path.endswith(".md"):
|
||||
meta_data = self._parse_md(doc_path)
|
||||
elif doc_path.endswith(".pdf"):
|
||||
meta_data = self._parse_pdf(doc_path)
|
||||
except Exception as e:
|
||||
logger.error("parse document %s failed:%s",doc_path,e)
|
||||
traceback.print_exc()
|
||||
|
||||
if meta_data.get("title"):
|
||||
title = meta_data["title"]
|
||||
logger.info("parse document %s!",doc_path)
|
||||
return hash_result,title,meta_data
|
||||
|
||||
|
||||
def _support_file(self,file_name:str) -> bool:
|
||||
if file_name.startswith("."):
|
||||
return False
|
||||
|
||||
if file_name.endswith(".pdf"):
|
||||
return True
|
||||
if file_name.endswith(".md"):
|
||||
return True
|
||||
if file_name.endswith(".txt"):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _scan_dir(self):
|
||||
while True:
|
||||
time.sleep(10)
|
||||
for directory in self.doc_dirs.keys():
|
||||
now = time.time()
|
||||
if now - self.doc_dirs[directory] > 60*15:
|
||||
self.doc_dirs[directory] = time.time()
|
||||
else:
|
||||
continue
|
||||
|
||||
for root, dirs, files in os.walk(directory):
|
||||
for file in files:
|
||||
if self._support_file(file):
|
||||
full_path = os.path.join(root, file)
|
||||
full_path = os.path.normpath(full_path)
|
||||
if self.kb_db.is_doc_exist(full_path):
|
||||
continue
|
||||
|
||||
file_stat = os.stat(full_path)
|
||||
if file_stat.st_size < 1:
|
||||
continue
|
||||
|
||||
if file_stat.st_size < 1024*1024*8:
|
||||
#parse and insert
|
||||
hash,title,meta_data = self._parse_document(full_path)
|
||||
self.kb_db.add_doc(full_path,file_stat.st_size,file_stat.st_mtime,hash)
|
||||
self.kb_db.add_knowledge(hash,title,meta_data)
|
||||
|
||||
else:
|
||||
self.kb_db.add_doc(full_path,file_stat.st_size,file_stat.st_mtime)
|
||||
|
||||
def _scan_document(self):
|
||||
while True:
|
||||
time.sleep(10)
|
||||
parse_queue = self.kb_db.get_docs_without_hash()
|
||||
for doc_path in parse_queue:
|
||||
hash,title,meta_data = self._parse_document(doc_path)
|
||||
self.kb_db.set_doc_hash(doc_path,hash)
|
||||
self.kb_db.add_knowledge(hash,title,meta_data)
|
||||
|
||||
|
||||
@@ -0,0 +1,214 @@
|
||||
# 尝试自我学习，会主动获取、读取资料并进行整理
|
||||
# LLM的本质能力是处理海量知识，应该让LLM能基于知识把自己的工作处理的更好
|
||||
async def do_self_learn(self) -> None:
|
||||
# 不同的workspace是否应该有不同的学习方法？
|
||||
workspace = self.get_workspace_by_msg(None)
|
||||
hash_list = workspace.kb_db.get_knowledge_without_llm_title()
|
||||
for hash in hash_list:
|
||||
if self.agent_energy <= 0:
|
||||
break
|
||||
|
||||
knowledge = workspace.kb_db.get_knowledge(hash)
|
||||
if knowledge is None:
|
||||
continue
|
||||
|
||||
full_path = knowledge.get("full_path")
|
||||
if full_path is None:
|
||||
continue
|
||||
|
||||
if os.path.exists(full_path) is False:
|
||||
logger.warning(f"do_self_learn: knowledge {full_path} is not exists!")
|
||||
continue
|
||||
|
||||
#TODO 可以用v-db 对不同目录的名字进行选择后，先进行一次快速的插入。有时间再慢慢用LLM整理
|
||||
result_obj = await self._llm_read_article(knowledge,full_path)
|
||||
|
||||
#根据结果更新knowledge
|
||||
if result_obj is not None:
|
||||
workspace.kb_db.set_knowledge_llm_result(hash,result_obj)
|
||||
# 在知识库中创建软链接
|
||||
path_list = result_obj.get("path")
|
||||
new_title = result_obj.get("title")
|
||||
if path_list:
|
||||
for new_path in path_list:
|
||||
full_new_path = f"/knowledge{new_path}/{new_title}"
|
||||
await workspace.symlink(full_path,full_new_path)
|
||||
logger.info(f"create soft link {full_path} -> {full_new_path}")
|
||||
|
||||
|
||||
self.agent_energy -= 1
|
||||
|
||||
# match item.type():
|
||||
# case "book":
|
||||
# self.llm_read_book(kb,item)
|
||||
# learn_power -= 1
|
||||
# case "article":
|
||||
#
|
||||
# self.llm_read_article(kb,item)
|
||||
# learn_power -= 1
|
||||
# case "video":
|
||||
# self.llm_watch_video(kb,item)
|
||||
# learn_power -= 1
|
||||
# case "audio":
|
||||
# self.llm_listen_audio(kb,item)
|
||||
# learn_power -= 1
|
||||
# case "code_project":
|
||||
# self.llm_read_code_project(kb,item)
|
||||
# learn_power -= 1
|
||||
# case "image":
|
||||
# self.llm_view_image(kb,item)
|
||||
# learn_power -= 1
|
||||
# case "other":
|
||||
# self.llm_read_other(kb,item)
|
||||
# learn_power -= 1
|
||||
# case _:
|
||||
# self.llm_learn_any(kb,item)
|
||||
# pass
|
||||
|
||||
|
||||
async def do_blance_knowledge_base(selft):
|
||||
# 整理自己的知识库(让分类更平衡，更由于自己以后的工作)，并尝试更新学习目标
|
||||
current_path = "/"
|
||||
current_list = kb.get_list(current_path)
|
||||
self_assessment_with_goal = self.get_self_assessment_with_goal()
|
||||
learn_goal = {}
|
||||
|
||||
|
||||
llm_blance_knowledge_base(current_path,current_list,self_assessment_with_goal,learn_goal,learn_power)
|
||||
|
||||
# 主动学习
|
||||
# 方法目前只有使用搜索引擎一种？
|
||||
for goal in learn_goal.items():
|
||||
self.llm_learn_with_search_engine(kb,goal,learn_power)
|
||||
if learn_power <= 0:
|
||||
break
|
||||
|
||||
|
||||
def parser_learn_llm_result(self,llm_result:LLMResult):
|
||||
pass
|
||||
|
||||
async def gen_known_info_for_knowledge_prompt(self,knowledge_item:dict,temp_meta = None,need_catalogs = False) -> AgentPrompt:
|
||||
workspace =self.get_workspace_by_msg(None)
|
||||
kb_tree = await workspace.get_knowledege_catalog()
|
||||
|
||||
|
||||
known_obj = {}
|
||||
title = knowledge_item.get("title")
|
||||
if title:
|
||||
known_obj["title"] = title
|
||||
summary = knowledge_item.get("summary")
|
||||
if summary:
|
||||
known_obj["summary"] = summary
|
||||
tags = knowledge_item.get("tags")
|
||||
if tags:
|
||||
known_obj["tags"] = tags
|
||||
if need_catalogs:
|
||||
catalogs = knowledge_item.get("catalogs")
|
||||
if catalogs:
|
||||
known_obj["catalogs"] = catalogs
|
||||
|
||||
if temp_meta:
|
||||
for key in temp_meta.keys():
|
||||
known_obj[key] = temp_meta[key]
|
||||
|
||||
org_path = knowledge_item.get("full_path")
|
||||
known_obj["orginal_path"] = org_path
|
||||
know_info_str = f"# Known information:\n## Current directory structure:\n{kb_tree}\n## Knowlege Metadata:\n{json.dumps(known_obj)}\n"
|
||||
return AgentPrompt(know_info_str)
|
||||
|
||||
async def _llm_read_article(self,knowledge_item:dict,full_path:str) -> ComputeTaskResult:
|
||||
# Objectives:
|
||||
# Obtain better titles, abstracts, table of contents (if necessary), tags
|
||||
# Determine the appropriate place to put it (in line with the organization's goals)
|
||||
# Known information:
|
||||
# The reason why the target service's learn_prompt is being sorted
|
||||
# Summary of the organization's work (if any)
|
||||
# The current structure of the knowledge base (note the size control) gen_kb_tree_prompt (when empty, LLM should generate an appropriate initial directory structure)
|
||||
# Original path, current title, abstract, table of contents
|
||||
|
||||
# Sorting long files (general tricks)
|
||||
# Indicate that the input is part of the content, let LLM generate intermediate results for the task
|
||||
# Enter the content in sequence, when the last content block is input, LLM gets the result
|
||||
|
||||
|
||||
#full_content = item.get_article_full_content()
|
||||
workspace = self.get_workspace_by_msg(None)
|
||||
full_content_len = self.token_len(full_content)
|
||||
|
||||
if full_content_len < self.get_llm_learn_token_limit():
|
||||
|
||||
# 短文章不用总结catelog
|
||||
#path_list,summary = llm_get_summary(summary,full_content)
|
||||
#prompt = self.get_agent_role_prompt()
|
||||
prompt = AgentPrompt()
|
||||
prompt.append(self.get_learn_prompt())
|
||||
known_info_prompt = await self.gen_known_info_for_knowledge_prompt(knowledge_item)
|
||||
prompt.append(known_info_prompt)
|
||||
content_prompt = AgentPrompt(full_content)
|
||||
prompt.append(content_prompt)
|
||||
env_functions = None
|
||||
#env_functions,function_len = workspace.get_knowledge_base_ai_functions()
|
||||
task_result:ComputeTaskResult = await self.do_llm_complection(prompt,is_json_resp=True)
|
||||
if task_result.result_code != ComputeTaskResultCode.OK:
|
||||
result_obj = {}
|
||||
result_obj["error_str"] = task_result.error_str
|
||||
return result_obj
|
||||
|
||||
result_obj = json.loads(task_result.result_str)
|
||||
return result_obj
|
||||
|
||||
else:
|
||||
logger.warning(f"llm_read_article: article {full_path} use LLM loop learn!")
|
||||
pos = 0
|
||||
read_len = int(self.get_llm_learn_token_limit() * 1.2)
|
||||
|
||||
temp_meta_data = {}
|
||||
is_final = False
|
||||
while pos < str_len:
|
||||
_content = full_content[pos:pos+read_len]
|
||||
part_cotent_len = len(_content)
|
||||
if part_cotent_len < read_len:
|
||||
# last chunk
|
||||
is_final = True
|
||||
part_content = f"<<Final Part:start at {pos}>>\n{_content}"
|
||||
else:
|
||||
part_content = f"<<Part:start at {pos}>>\n{_content}"
|
||||
|
||||
pos = pos + read_len
|
||||
prompt = AgentPrompt()
|
||||
prompt.append(self.get_learn_prompt())
|
||||
known_info_prompt = await self.gen_known_info_for_knowledge_prompt(knowledge_item,temp_meta_data)
|
||||
prompt.append(known_info_prompt)
|
||||
content_prompt = AgentPrompt(part_content)
|
||||
prompt.append(content_prompt)
|
||||
#env_functions,function_len = workspace.get_knowledge_base_ai_functions()
|
||||
task_result:ComputeTaskResult = await self.do_llm_complection(prompt,is_json_resp=True)
|
||||
if task_result.result_code != ComputeTaskResultCode.OK:
|
||||
result_obj = {}
|
||||
result_obj["error_str"] = task_result.error_str
|
||||
return result_obj
|
||||
|
||||
result_obj = json.loads(task_result.result_str)
|
||||
temp_meta_data = result_obj
|
||||
if is_final:
|
||||
return result_obj
|
||||
|
||||
return None
|
||||
|
||||
|
||||
async def do_self_think(self):
|
||||
session_id_list = AIChatSession.list_session(self.agent_id,self.chat_db)
|
||||
for session_id in session_id_list:
|
||||
if self.agent_energy <= 0:
|
||||
break
|
||||
used_energy = await self.think_chatsession(session_id)
|
||||
self.agent_energy -= used_energy
|
||||
|
||||
todo_logs = await self.get_todo_logs()
|
||||
for todo_log in todo_logs:
|
||||
if self.agent_energy <= 0:
|
||||
break
|
||||
used_energy = await self.think_todo_log(todo_log)
|
||||
self.agent_energy -= used_energy
|
||||
|
||||
return
|
||||
@@ -44,14 +44,19 @@ class KnowledgePipelineManager:
|
||||
input_init = self.input_modules.get(input_module)
|
||||
input_params = config["input"].get("params")
|
||||
|
||||
parser_module = config["parser"]["module"]
|
||||
_, ext = os.path.splitext(parser_module)
|
||||
if ext == ".py":
|
||||
parser_module = os.path.join(path, parser_module)
|
||||
parser_init = runpy.run_path(parser_module)["init"]
|
||||
parser_config = config.get("parser")
|
||||
if parser_config is None:
|
||||
parser_init = None
|
||||
parser_params = None
|
||||
else:
|
||||
parser_init = self.parser_modules.get(parser_module)
|
||||
parser_params = config["parser"].get("params")
|
||||
parser_module = parser_config["module"]
|
||||
_, ext = os.path.splitext(parser_module)
|
||||
if ext == ".py":
|
||||
parser_module = os.path.join(path, parser_module)
|
||||
parser_init = runpy.run_path(parser_module)["init"]
|
||||
else:
|
||||
parser_init = self.parser_modules.get(parser_module)
|
||||
parser_params = parser_config.get("params")
|
||||
|
||||
|
||||
data_path = os.path.join(self.root_dir, name)
|
||||
|
||||
@@ -22,7 +22,7 @@ class LocalEmail:
|
||||
if latest_journal.is_finish():
|
||||
yield None
|
||||
continue
|
||||
parsed = str(latest_journal.get_object_id())
|
||||
parsed = latest_journal.get_input()
|
||||
|
||||
mail_id = self.mail_storage.next_mail_id(parsed)
|
||||
if mail_id is None:
|
||||
|
||||
@@ -7,17 +7,20 @@ import datetime
|
||||
from bs4 import BeautifulSoup
|
||||
import sqlite3
|
||||
import html2text
|
||||
from urllib.parse import urlparse
|
||||
from aios import *
|
||||
|
||||
|
||||
|
||||
class Mail:
|
||||
def __init__(self, **kwargs) -> None:
|
||||
self.from_addr = kwargs.get("From")
|
||||
self.to_addr = kwargs.get("To")
|
||||
self.subject = kwargs.get("Subject")
|
||||
self.date = kwargs.get("Date")
|
||||
self.bcc = kwargs.get("BCC")
|
||||
self.cc = kwargs.get("CC")
|
||||
self.reply_to = None
|
||||
self.from_addr = kwargs.get("from")
|
||||
self.to_addr = kwargs.get("to")
|
||||
self.subject = kwargs.get("subject")
|
||||
self.date = kwargs.get("date")
|
||||
self.bcc = kwargs.get("bcc")
|
||||
self.cc = kwargs.get("cc")
|
||||
self.reply_to = kwargs.get("reply_to")
|
||||
self.id: str = None
|
||||
self.content: str = None
|
||||
|
||||
@@ -192,20 +195,36 @@ class MailStorage:
|
||||
self.conn.commit()
|
||||
await asyncio.sleep(10)
|
||||
|
||||
def download(self, uid, mail: mailparser.MailParser):
|
||||
def download(self, uid, parser: mailparser.MailParser,
|
||||
save_image=True,
|
||||
from_field="From",
|
||||
to_field="To",
|
||||
subject_field="Subject",
|
||||
date_field="Date",
|
||||
reply_to_field="In-Reply-To",
|
||||
cc_field="CC",
|
||||
bcc_field="BCC"):
|
||||
mail_dir = self.mail_dir(uid)
|
||||
os.makedirs(dir)
|
||||
if not os.path.exists(mail_dir):
|
||||
os.makedirs(mail_dir)
|
||||
|
||||
meta = json.loads(mail.mail_json)
|
||||
mail = Mail(**meta)
|
||||
reply_to = meta.get("In-Reply-To")
|
||||
src_meta = json.loads(parser.mail_json)
|
||||
meta = {}
|
||||
meta["from"] = src_meta.get(from_field)
|
||||
meta["to"] = src_meta.get(to_field)
|
||||
meta["subject"] = src_meta.get(subject_field)
|
||||
meta["date"] = src_meta.get(date_field)
|
||||
meta["bcc"] = src_meta.get(bcc_field)
|
||||
meta["cc"] = src_meta.get(cc_field)
|
||||
reply_to = src_meta.get(reply_to_field)
|
||||
if reply_to:
|
||||
mail.reply_to = self.uid_to_object_id(reply_to)
|
||||
meta["reply_to"] = self.uid_to_object_id(reply_to)
|
||||
mail = Mail(**meta)
|
||||
|
||||
h = html2text.HTML2Text()
|
||||
h.ignore_links = True
|
||||
h.ignore_images = True
|
||||
mail_content = h.handle(mail.body)
|
||||
mail_content = h.handle(parser.body)
|
||||
mail.content = mail_content
|
||||
|
||||
mail.calculate_id()
|
||||
@@ -216,41 +235,52 @@ class MailStorage:
|
||||
with open(f"{mail_dir}/mail.txt", "w", encoding='utf-8') as f:
|
||||
f.write(mail_content)
|
||||
|
||||
for attachment in mail.attachments:
|
||||
if attachment['mail_content_type'] in ['image/png', 'image/jpeg', 'image/gif']:
|
||||
filename = attachment['filename']
|
||||
filefullname = f"{mail_dir}/{filename}"
|
||||
image_data = attachment['payload']
|
||||
if save_image:
|
||||
for attachment in parser.attachments:
|
||||
if attachment['mail_content_type'] in ['image/png', 'image/jpg', 'image/jpeg', 'image/gif', 'image/svg']:
|
||||
filename = attachment['filename']
|
||||
filefullname = f"{mail_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")
|
||||
|
||||
# get all image urls
|
||||
soup = BeautifulSoup(parser.body, '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
|
||||
|
||||
for img_url in img_urls:
|
||||
# keep the original image filename(last of url)
|
||||
url_result = urlparse(img_url)
|
||||
if url_result.scheme not in ['http', 'https']:
|
||||
continue
|
||||
ext = url_result.path.split('/')[-1].split('.')[-1]
|
||||
if ext in ['png', 'jpg', 'jpeg', 'gif', 'svg']:
|
||||
img_filename = os.path.join(mail_dir, f"{name_count}.{ext}")
|
||||
else :
|
||||
img_filename = os.path.join(mail_dir, f"{name_count}")
|
||||
name_count += 1
|
||||
# download image
|
||||
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")
|
||||
|
||||
# get all image urls
|
||||
soup = BeautifulSoup(mail.body, '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
|
||||
|
||||
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(mail_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}')
|
||||
response = requests.get(img_url, stream=True)
|
||||
except requests.exceptions.RequestException as e:
|
||||
logging.error(f'Failed to download {img_url}: {e}')
|
||||
continue
|
||||
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.error(f'Failed to download {img_url}')
|
||||
|
||||
cursor = self.conn.cursor()
|
||||
cursor.execute(
|
||||
@@ -260,5 +290,8 @@ class MailStorage:
|
||||
""",
|
||||
(uid, mail.id, mail.date, mail.from_addr),
|
||||
)
|
||||
self.conn.commit()
|
||||
|
||||
return mail.id
|
||||
|
||||
|
||||
@@ -1,9 +1,13 @@
|
||||
import os
|
||||
import logging
|
||||
import json
|
||||
import string
|
||||
import imaplib
|
||||
import mailparser
|
||||
from aios import *
|
||||
|
||||
from knowledge import *
|
||||
from aios_kernel.storage import AIStorage
|
||||
from .mail import Mail, MailStorage
|
||||
|
||||
|
||||
class EmailSpider:
|
||||
@@ -16,14 +20,22 @@ class EmailSpider:
|
||||
port=self.config.get('imap_port')
|
||||
)
|
||||
self.client.login(self.config.get('address'), self.config.get('password'))
|
||||
self.mail_local_root = os.path.join(self.env.pipeline_path, self.config.get("address"))
|
||||
os.makedirs(self.mail_local_root)
|
||||
self.client.select("INBOX")
|
||||
local_path = string.Template(config["path"]).substitute(myai_dir=AIStorage.get_instance().get_myai_dir())
|
||||
local_path = os.path.join(local_path, self.config.get('address'))
|
||||
self.mail_storage = MailStorage(local_path)
|
||||
|
||||
|
||||
async def next(self):
|
||||
while True:
|
||||
_, data = self.client.uid('search', None, "ALL")
|
||||
try:
|
||||
_, data = self.client.uid('search', None, "ALL")
|
||||
except Exception as e:
|
||||
self.env.get_logger().error(f"email spider error: {e}")
|
||||
yield (None, None)
|
||||
continue
|
||||
uid_list = data[0].split()
|
||||
if uid_list.len() == 0:
|
||||
if len(uid_list) == 0:
|
||||
yield (None, None)
|
||||
continue
|
||||
|
||||
@@ -43,9 +55,16 @@ class EmailSpider:
|
||||
_uid = int.from_bytes(uid)
|
||||
if _uid > from_uid:
|
||||
message_parts = "(BODY.PEEK[])"
|
||||
_, email_data = self.client.uid('fetch', uid, message_parts)
|
||||
mail = mailparser.parse_from_bytes(email_data[0][1])
|
||||
self.save_email(_uid, mail)
|
||||
try:
|
||||
_, email_data = self.client.uid('fetch', uid, message_parts)
|
||||
mail = mailparser.parse_from_bytes(email_data[0][1])
|
||||
id = self.mail_storage.download(_uid, mail)
|
||||
except Exception as e:
|
||||
self.env.get_logger().error(f"email spider error: {e}")
|
||||
yield (None, None)
|
||||
break
|
||||
yield (ObjectID.from_base58(id), str(_uid))
|
||||
|
||||
|
||||
yield (None, None)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user