shell knowledge commands

This commit is contained in:
tsukasa
2023-09-21 18:32:17 +08:00
parent edbd1c41da
commit a0a45b8998
11 changed files with 409 additions and 74 deletions
+1 -1
View File
@@ -6,7 +6,7 @@ from .compute_kernel import ComputeKernel,ComputeTask
from .compute_node import ComputeNode,LocalComputeNode
from .open_ai_node import OpenAI_ComputeNode
from .knowledge_base import KnowledgeBase
from .knowledge_pipeline import EmailSpider
from .knowledge_pipeline import KnowledgeEmailSource, KnowledgeDirSource, KnowledgePipline
from .role import AIRole,AIRoleGroup
from .workflow import Workflow
from .bus import AIBus
+6 -3
View File
@@ -1,7 +1,9 @@
# define a knowledge base class
import json
import logging
from . import AgentPrompt, ComputeKernel, AIStorage
from .agent import AgentPrompt
from .compute_kernel import ComputeKernel
from .storage import AIStorage
from knowledge import *
@@ -17,7 +19,7 @@ class KnowledgeBase:
def __singleton_init__(self) -> None:
self.store = KnowledgeStore()
self.compute_kernel = ComputeKernel()
self.compute_kernel = ComputeKernel.get_instance()
async def __embedding_document(self, document: DocumentObject):
for chunk_id in document.get_chunk_list():
@@ -155,7 +157,7 @@ class KnowledgeBase:
# pass
async def insert_object(self, object: KnowledgeObject):
# self.__save_object(object)
self.store.get_object_store().put_object(object.calculate_id(), object.encode())
await self.__do_embedding(object)
async def query_prompt(self, prompt: AgentPrompt):
@@ -164,6 +166,7 @@ class KnowledgeBase:
knowledge_prompt = self.prompt_from_objects(objects)
logging.info(f"prompt_from_objects result: {knowledge_prompt.as_str()}")
prompt.append(knowledge_prompt)
return prompt
async def query_objects(self, prompt: AgentPrompt) -> [ObjectID]:
results = []
+223 -47
View File
@@ -15,10 +15,10 @@ An example of a local file is as follows:
│ └── dd-login-service-min.png
"""
import asyncio
import datetime
import sqlite3
import imaplib
import os
import toml
import logging
import mailparser
import hashlib
@@ -26,63 +26,139 @@ import json
import base64
from bs4 import BeautifulSoup
import requests
import os
import toml
from .storage import AIStorage, UserConfigItem
from .knowledge_base import KnowledgeBase, ImageObjectBuilder, ObjectID, ObjectType, DocumentObjectBuilder
class EmailSpider:
class KnowledgeJournal:
def __init__(self, source_type: str, source_id: str, item_id: str, object_id: str, timestamp=None):
# define a timestamp variable
self.timestamp = datetime.datetime.now() if timestamp is None else timestamp
self.object_id = object_id
self.source_type = source_type
self.source_id = source_id
self.item_id = item_id
def __str__(self) -> str:
if self.source_type == "dir":
object_id = ObjectID.from_base58(self.object_id)
object_type = None
if object_id.get_object_type() == ObjectType.Image:
object_type = "image"
else:
pass
return f"Add {object_type} from {os.path.join(self.source_id, self.item_id)}"
if self.source_type == "email":
pass
# init sqlite3 client
class KnowledgeJournalClient:
def __init__(self):
# logger config
self.logger = logging.getLogger('email spider')
self.logger.setLevel(logging.DEBUG)
ch = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s [%(name)s] [%(levelname)s] %(message)s')
ch.setFormatter(formatter)
self.logger.addHandler(ch)
knowledge_dir = os.path.join(AIStorage.get_instance().get_myai_dir(), "knowledge")
if not os.path.exists(knowledge_dir):
os.makedirs(knowledge_dir)
self.journal_path = os.path.join(knowledge_dir, "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,
source_type TEXT,
source_id TEXT,
item_id TEXT,
object_id TEXT)'''
)
conn.commit()
def insert(self, journal: KnowledgeJournal):
conn = sqlite3.connect(self.journal_path)
conn.execute(
"INSERT INTO journal (time, source_type, source_id, item_id, object_id) VALUES (?, ?, ?, ?, ?)",
(journal.timestamp, journal.source_type, journal.source_id, journal.item_id, journal.object_id),
)
conn.commit()
def latest_journal(self, source_id: str) -> KnowledgeJournal:
conn = sqlite3.connect(self.journal_path)
cursor = conn.cursor()
cursor.execute("SELECT * FROM journal WHERE source_id = ? ORDER BY id DESC LIMIT 1", (source_id,))
result = cursor.fetchone()
if result is None:
return None
else:
(_, timestamp, source_type, sorce_id, item_id, object_id) = result
return KnowledgeJournal(source_type, sorce_id, item_id, object_id, timestamp)
def latest_journals(self, topn) -> [KnowledgeJournal]:
conn = sqlite3.connect(self.journal_path)
cursor = conn.cursor()
cursor.execute("SELECT * FROM journal ORDER BY id DESC LIMIT ?", (topn,))
return [KnowledgeJournal(source_type, sorce_id, item_id, object_id, timestamp) for (_, timestamp, source_type, sorce_id, item_id, object_id) in cursor.fetchall()]
class KnowledgeEmailSource:
def __init__(self, config:dict):
self.config = config
self.config["type"] = "email"
def id(self):
"::".join([self.config["imap_server"], 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}/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)
self.config = toml.load('./rootfs/email/config.toml')
if os.path.exists('./rootfs/email/config.local.toml'):
self.config = toml.load('./rootfs/email/config.local.toml')
self.client = self.email_client()
await self.read_emails()
def email_client(self) -> imaplib.IMAP4_SSL:
self.logger.info(f"read email config from {self.config.get('EMAIL_IMAP_SERVER')}")
logging.info(f"read email config from {self.config.get('imap_server')}")
client = imaplib.IMAP4_SSL(
host=self.config.get('EMAIL_IMAP_SERVER'),
port=self.config.get('EMAIL_IMAP_PORT')
host=self.config.get('imap_server'),
port=self.config.get('imap_port')
)
client.login(self.config.get('EMAIL_ADDRESS'), self.config.get('EMAIL_PASSWORD'))
client.login(self.config.get('address'), self.config.get('password'))
return client
def list_box(self):
_, mailbox_list = self.client.list()
for mailbox in mailbox_list:
print(mailbox.decode())
def read_emails(self, folder: str = 'INBOX', imap_keyword: str = "UNSEEN"):
async def read_emails(self, folder: str = 'INBOX', imap_keyword: str = "UNSEEN"):
self.client.select(folder)
_, data = self.client.uid('search', None, imap_keyword)
# get email uid list
email_list = data[0].split()
self.logger.info(f"got {len(email_list)} emails")
logging.info(f"got {len(email_list)} emails")
email_list.reverse()
for uid in email_list:
if self.check_email_saved(uid):
self.logger.info(f"email uid {uid} already saved")
logging.info(f"email uid {uid} already saved")
else:
self.read_and_save_email(uid)
self.logger.info(f"email uid {uid} saved")
logging.info(f"email uid {uid} saved")
def read_and_save_email(self, uid: str):
message_parts = "(BODY.PEEK[])"
_, email_data = self.client.uid('fetch', uid, message_parts)
mail = mailparser.parse_from_bytes(email_data[0][1])
self.logger.info(f"got email subject [{mail.subject}]")
logging.info(f"got email subject [{mail.subject}]")
self.save_email(mail)
def get_local_dir_name(self, mail: mailparser.MailParser) -> str:
dir = f"{self.config.get('LOCAL_DIR')}/{self.config.get('EMAIL_ADDRESS')}"
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}"
@@ -91,9 +167,9 @@ class EmailSpider:
message_parts = "(BODY[HEADER])"
_, email_data = self.client.uid('fetch', uid, message_parts)
mail = mailparser.parse_from_bytes(email_data[0][1])
self.logger.info(f"[{uid}]check email subject [{mail.subject}]")
logging.info(f"[{uid}]check email subject [{mail.subject}]")
dir = self.get_local_dir_name(mail)
self.logger.info(f"check email saved {dir}")
logging.info(f"check email saved {dir}")
file = f"{dir}/email.txt"
if os.path.exists(file):
return False
@@ -116,7 +192,7 @@ class EmailSpider:
image_data = image_data.encode()
with open(filefullname, 'wb') as f:
f.write(image_data)
self.logger.info(f"save email image {filename} success")
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):
@@ -124,7 +200,7 @@ class EmailSpider:
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]
self.logger.info(f'Found {len(img_urls)} images in email body')
logging.info(f'Found {len(img_urls)} images in email body')
if not os.path.exists(email_dir):
os.makedirs(email_dir)
@@ -138,17 +214,17 @@ class EmailSpider:
with open(img_filename, 'wb') as img_file:
for chunk in response.iter_content(1024):
img_file.write(chunk)
self.logger.info(f'Downloaded {img_url} to {img_filename}')
logging.info(f'Downloaded {img_url} to {img_filename}')
else:
self.logger.info(f'Failed to download {img_url}')
logging.info(f'Failed to download {img_url}')
# save email content to local dir
def save_email(self, mail: mailparser.MailParser):
dir = f"{self.config.get('LOCAL_DIR')}/{self.config.get('EMAIL_ADDRESS')}"
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)
self.logger.info(f"save email to {email_dir}")
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") as f:
@@ -158,14 +234,58 @@ class EmailSpider:
if 'body' in mail_dict:
del mail_dict['body']
json.dump(mail_dict, f, ensure_ascii=False, indent=4)
self.logger.info(f"save email meta info {f.name}")
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")
class KnowledgeDirSource:
def __init__(self, config):
self.config = config
self.config["type"] = "dir"
@classmethod
def user_config_items(cls):
return [("path", "local dir path")]
def id(self):
return self.config["path"]
def path(self):
return self.config["path"]
async def run_once(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
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)
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}")
with open(file_path, "r", encoding="utf-8") as f:
text = f.read()
document = DocumentObjectBuilder({}, {}, text).build()
await KnowledgeBase().insert_object(document)
journal_client.insert(KnowledgeJournal("dir", self.id(), rel_path, str(document.calculate_id()), timestamp))
from . import AIStorage, KnowledgeBase
# define singleton class knowledge pipline
class KnowledgePipline:
@@ -178,15 +298,71 @@ class KnowledgePipline:
cls._instance.__singleton_init__()
return cls._instance
def initial(self):
config_path = self.__config_path()
logging.info(f"initial knowledge pipline from {config_path}")
if os.path.exists(config_path):
config = toml.load(self.__config_path())
for source_config in config["sources"]:
if source_config['type'] == 'email':
self.add_email_source(KnowledgeEmailSource(source_config))
if source_config['type'] == 'dir':
self.add_dir_source(KnowledgeDirSource(source_config))
def __singleton_init__(self) -> None:
def __singleton_init__(self):
self.knowledge_base = KnowledgeBase()
self.email_sources = dict()
self.dir_sources = dict()
self.source_queue = list()
self.run_lock = asyncio.Lock()
asyncio.create_task(self.run_loop())
def save_config(self):
config = dict()
config["sources"] = [source.config for source in self.source_queue]
with open(self.__config_path(), "w") as f:
toml.dump(config, f)
@classmethod
def declare_user_config(cls):
user_config = AIStorage.get_instance().get_user_config()
user_config.add_user_config("email_spiders","email addresses to build knowledge base",True,None,"list")
user_config.add_user_config("personal_dirs", "personal directories to build knowledge base", True, None, "list")
def __config_path(cls) -> str:
user_data_dir = AIStorage.get_instance().get_myai_dir()
return os.path.abspath(f"{user_data_dir}/etc/knowledge.cfg.toml")
def add_email_address(self, ) -> None:
pass
def add_email_source(self, source: KnowledgeEmailSource):
if self.email_sources.get(source.id()) is not None:
return "already exists"
self.email_sources[source.id()] = source
self.source_queue.append(source)
return None
def add_dir_source(self, source: KnowledgeDirSource):
if self.dir_sources.get(source.id()) is not None:
logging.info(f"knowledge add source {source.id()} failed for already exists")
return "already exists"
logging.info(f"knowledge added source {source.id()}")
self.dir_sources[source.id()] = source
self.source_queue.append(source)
return None
def get_latest_journals(self, topn) -> [KnowledgeJournal]:
return KnowledgeJournalClient().latest_journals(topn)
async def run_loop(self):
while True:
await self.run_once()
await asyncio.sleep(5)
async def run_once(self):
logging.info(f"knowledge pipeline started")
# sources = list()
# async with self.run_lock:
# for source in self.source_queue:
# sources.append(source)
# for source in sources:
# await source.run_once()
for source in self.source_queue:
await source.run_once()