diff --git a/.gitignore b/.gitignore index f616e40..3bb2bad 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,5 @@ .vscode/ *.pyc -rootfs/data *.log rootfs/email/config.local.toml rootfs/data @@ -11,3 +10,4 @@ aios_shell_history.txt math_school_env.db workflows.db + diff --git a/aios_shell_history.txt b/aios_shell_history.txt new file mode 100644 index 0000000..34732fd --- /dev/null +++ b/aios_shell_history.txt @@ -0,0 +1,111 @@ + +# 2023-09-19 16:04:55.011656 ++tsukasa + +# 2023-09-19 16:05:27.714815 ++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 + +# 2023-09-20 14:48:06.181582 ++/connect + +# 2023-09-20 14:48:22.678945 ++/connect Jarvis email + +# 2023-09-20 14:54:35.101026 ++/connect Jarvis telegram + +# 2023-09-20 14:54:55.229801 ++token abcd + +# 2023-09-20 16:22:34.229084 ++/knowledge add email + +# 2023-09-20 16:22:46.797842 ++puotosssa@live.com + +# 2023-09-20 16:22:53.622015 ++p19870626 + +# 2023-09-20 16:22:59.186145 ++imap@live.com + +# 2023-09-20 16:23:02.334411 ++993 + +# 2023-09-20 16:40:29.996329 ++/knowledge add email + +# 2023-09-20 16:40:52.860700 ++puotosssa@live.com + +# 2023-09-20 16:40:55.324994 ++fdafda + +# 2023-09-20 16:40:58.167154 ++fdafdaf + +# 2023-09-20 16:40:59.757363 ++993 + +# 2023-09-20 16:42:55.288680 ++/knowledge add email + +# 2023-09-20 16:42:56.916408 ++fdfa + +# 2023-09-20 16:42:57.965232 ++fdfds + +# 2023-09-20 16:42:59.149516 ++fdafd + +# 2023-09-20 16:43:00.949104 ++993 + +# 2023-09-21 14:11:33.500467 ++/knowledge add c:/users/tsukasa/myai/photos + +# 2023-09-21 14:14:31.488794 ++/knowledge add dir + +# 2023-09-21 14:14:49.298635 ++c:/users/tsukasa/myai/photos + +# 2023-09-21 15:27:58.550658 ++/knowledge add dir + +# 2023-09-21 15:28:10.190467 ++c:/users/tsukasa/myai/photos + +# 2023-09-21 15:32:32.431305 ++/knowledge add dir + +# 2023-09-21 15:32:46.695566 ++c:/users/tsukasa/myai/photos + +# 2023-09-21 15:33:36.391465 ++/knowledge add dir + +# 2023-09-21 15:33:47.574853 ++c:/users/tsukasa/myai/photos + +# 2023-09-21 17:44:26.308692 ++/knowledge journals + +# 2023-09-21 17:45:42.925359 ++/knowledge journal + +# 2023-09-21 18:13:08.720054 ++/knowledge query what a see is + +# 2023-09-21 18:20:09.275556 ++/knowledge journal + +# 2023-09-21 18:20:10.746056 ++/knowledge query what a see is + +# 2023-09-21 18:29:48.391223 ++/knowledge journal + +# 2023-09-21 18:29:51.406532 ++/knowledge query what a see is diff --git a/src/aios_kernel/__init__.py b/src/aios_kernel/__init__.py index 7436dbf..7240b92 100644 --- a/src/aios_kernel/__init__.py +++ b/src/aios_kernel/__init__.py @@ -6,6 +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 KnowledgeEmailSource, KnowledgeDirSource, KnowledgePipline from .role import AIRole,AIRoleGroup from .workflow import Workflow from .bus import AIBus diff --git a/src/aios_kernel/compute_kernel.py b/src/aios_kernel/compute_kernel.py index 18ddeaf..97f9b9b 100644 --- a/src/aios_kernel/compute_kernel.py +++ b/src/aios_kernel/compute_kernel.py @@ -131,6 +131,7 @@ class ComputeKernel: 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) diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index b229864..11d6ce7 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -1,7 +1,9 @@ # define a knowledge base class import json import logging -from . import AgentPrompt, ComputeKernel +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 = [] @@ -239,9 +242,6 @@ class KnowledgeBase: prompt.messages.append({"role": "knowledge", "content": content}) return prompt - - - - + \ No newline at end of file diff --git a/src/aios_kernel/knowledge_pipeline.py b/src/aios_kernel/knowledge_pipeline.py new file mode 100644 index 0000000..c1f0a72 --- /dev/null +++ b/src/aios_kernel/knowledge_pipeline.py @@ -0,0 +1,368 @@ +""" +Capture your email locally, and parse out the pictures in the email body and the pictures, videos and other files in the attachment. Subsequently, it supports vectorized analysis of your personal data and serves as a knowledge base to enable large language model answers. Better results. + +An example of a local file is as follows: +├── data +│ └── alex0072@gmail.com +│ └── 5de3e52f3a6b90cabe6cbdd4ae3a5c5b +│ ├── email.txt +│ ├── meta.json +│ ├── image +│ │ ├── 0648B869@99C03070.DB94B354.jpg +│ └── body_image +│ ├── 11044884873.jpg +│ ├── 282985198265470.gif +│ └── dd-login-service-min.png + +""" +import asyncio +import datetime +import sqlite3 +import imaplib +import logging +import mailparser +import hashlib +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 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): + 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.client = self.email_client() + await self.read_emails() + + 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"): + 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") + email_list.reverse() + for uid in email_list: + if self.check_email_saved(uid): + logging.info(f"email uid {uid} already saved") + else: + self.read_and_save_email(uid) + 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]) + 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.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): + 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 False + return False + + # 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') + + 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) + img_filename = os.path.join(email_dir, img_url.split('/')[-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") as f: + 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") + + +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)) + + + + +# define singleton class knowledge pipline +class KnowledgePipline: + _instance = None + + @classmethod + def get_instance(cls): + if cls._instance is None: + cls._instance = 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): + 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 __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_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() diff --git a/src/knowledge/core_object/document_object.py b/src/knowledge/core_object/document_object.py index 7210f16..672b4a3 100644 --- a/src/knowledge/core_object/document_object.py +++ b/src/knowledge/core_object/document_object.py @@ -49,7 +49,7 @@ class DocumentObjectBuilder: self.text = text return self - def build(self, relation_store: ObjectRelationStore) -> DocumentObject: + def build(self) -> DocumentObject: chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_text( self.text, 1024 * 4, @@ -60,6 +60,6 @@ class DocumentObjectBuilder: # Add relation to store for chunk_id in chunk_list.chunk_list: - relation_store.add_relation(chunk_id, doc_id) + KnowledgeStore().get_relation_store().add_relation(chunk_id, doc_id) return doc diff --git a/src/knowledge/core_object/email_object.py b/src/knowledge/core_object/email_object.py index c86c3a2..37531ff 100644 --- a/src/knowledge/core_object/email_object.py +++ b/src/knowledge/core_object/email_object.py @@ -94,7 +94,7 @@ class EmailObjectBuilder: with open(content_file, "r", encoding="utf-8") as f: text = f.read() - document = DocumentObjectBuilder({}, {}, text).build(relation_store=relation) + document = DocumentObjectBuilder({}, {}, text).build() document_id = document.calculate_id() store.put_object(document_id, document.encode()) documents = {"email.txt": document_id} diff --git a/src/knowledge/core_object/image_object.py b/src/knowledge/core_object/image_object.py index 30b0e8a..2c02ad6 100644 --- a/src/knowledge/core_object/image_object.py +++ b/src/knowledge/core_object/image_object.py @@ -52,7 +52,7 @@ def get_exif_data(image_path: str): return { TAGS.get(key): exif_data[key] for key in exif_data.keys() - if key in TAGS and isinstance(exif_data[key], (bytes, str)) + if key in TAGS and isinstance(exif_data[key], str) } else: return {} diff --git a/src/knowledge/data/writer.py b/src/knowledge/data/writer.py index 3ac0ed1..ad6c520 100644 --- a/src/knowledge/data/writer.py +++ b/src/knowledge/data/writer.py @@ -46,7 +46,7 @@ class ChunkListWriter: ) file_hash = HashValue(hash_obj.digest()) - print(f"calc file hash: {file_path}, {file_hash}") + # print(f"calc file hash: {file_path}, {file_hash}") return ChunkList(chunk_list, file_hash) diff --git a/src/knowledge/object.py b/src/knowledge/object.py new file mode 100644 index 0000000..f7a127e --- /dev/null +++ b/src/knowledge/object.py @@ -0,0 +1,65 @@ + +# define a object type enum +from abc import ABC, abstractmethod +from enum import Enum + +class ObjectType(Enum): + TextChunk = 1 + Image = 2 + Email = 101 + + +# define a object ID class to identify a object +class ObjectID: # pylint: disable=too-few-public-methods + def __init__(self, object_type, digist): + self.object_type = object_type + self.digist = digist + + def __str__(self): + return f"{self.object_type.name}:{self.digist}" + + +# define a object class +class KnowledgeObject(ABC): # pylint: disable=too-few-public-methods + def __init__(self, object_type: ObjectType): + self.object_type = object_type + + @abstractmethod + def get_id(self) -> ObjectID: + pass + + # define a to binary method to convert object to binary + @abstractmethod + def to_binary(self) -> bytes: + pass + + # define a from binary method to convert binary to object + @abstractmethod + def from_binary(self, binary: bytes): + pass + + +# define a text chunk class +class TextChunkObject(KnowledgeObject): # pylint: disable=too-few-public-methods + def __init__(self, text: str): + super().__init__(ObjectType.TextChunk) + self.text = text + + +# define a image class +class ImageObject(KnowledgeObject): # pylint: disable=too-few-public-methods + def __init__(self, meta, path): + super().__init__(ObjectType.Image) + self.meta = meta + self.path = path + + +# define a email class +class EmailObject(KnowledgeObject): # pylint: disable=too-few-public-methods + def __init__(self, meta): + super().__init__(ObjectType.Email) + self.meta = meta + self.text = [ObjectID] + self.images = [ObjectID] + + diff --git a/src/knowledge/object_storage.py b/src/knowledge/object_storage.py new file mode 100644 index 0000000..1a6f85e --- /dev/null +++ b/src/knowledge/object_storage.py @@ -0,0 +1,33 @@ +# import RDB LargeBinary +from sqlalchemy import Column, String, LargeBinary, create_engine, sessionmaker, pickle +from .object import KnowledgeObject + +# implement object storage with RDB +# define object storage table +class ObjectStorageTable(Base): + __tablename__ = 'object_storage' + id = Column(String, primary_key=True) + parent = Column(String, nullable=True) + object = Column(LargeBinary, nullable=False) + + def __init__(self, id, parent, object): # pylint: disable=redefined-builtin + self.id = id + self.parent = parent + self.object = object + +# define object storage class +class ObjectStorage: + async def __init__(self, db_url): + self.engine = create_engine(db_url) + self.session = sessionmaker(bind=self.engine)() # pylint: disable=not-callable + + async def get(self, id) -> [KnowledgeObject, KnowledgeObject]: + obj = self.session.query(ObjectStorageTable).filter(ObjectStorageTable.id == id).first() + if obj is None: + return None + return pickle.loads(obj.object) + + # define insert method + async def insert(self, object, parent): # pylint: disable=redefined-builtin + obj = ObjectStorageTable(id, parent, pickle.dumps(object)) + \ No newline at end of file diff --git a/src/knowledge/store.py b/src/knowledge/store.py index 03e48d0..56ec64b 100644 --- a/src/knowledge/store.py +++ b/src/knowledge/store.py @@ -4,7 +4,7 @@ from .object import ObjectStore, ObjectRelationStore from .data import ChunkStore, ChunkTracker, ChunkListWriter, ChunkReader from .vector import ChromaVectorStore, VectorBase import logging - +import aios_kernel # 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: @@ -13,16 +13,12 @@ class KnowledgeStore: 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) + knowledge_dir = aios_kernel.storage.AIStorage().get_myai_dir() / "knowledge" - if not os.path.exists(directory): - os.makedirs(directory) + if not os.path.exists(knowledge_dir): + os.makedirs(knowledge_dir) - cls._instance.__singleton_init__(directory) + cls._instance.__singleton_init__(knowledge_dir) return cls._instance @@ -64,5 +60,5 @@ class KnowledgeStore: 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) + self.vector_store[model_name] = ChromaVectorStore(self.root, model_name) return self.vector_store[model_name] diff --git a/src/knowledge/vector/chroma_store.py b/src/knowledge/vector/chroma_store.py index 03006ae..260a8e8 100644 --- a/src/knowledge/vector/chroma_store.py +++ b/src/knowledge/vector/chroma_store.py @@ -6,16 +6,14 @@ import os class ChromaVectorStore(VectorBase): - def __init__(self, model_name: str) -> None: + def __init__(self, root_dir, 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" - ) + directory = os.path.join(root_dir, "vector") logging.info("will use vector store: {}".format(directory)) client = chromadb.PersistentClient( diff --git a/src/requirements.txt b/src/requirements.txt index 8427ade..0c88c8a 100644 --- a/src/requirements.txt +++ b/src/requirements.txt @@ -1,23 +1,83 @@ -chromadb==0.4 -moviepy==1.0 -base58==2.1 -base36==0.1 aiofiles==23.2.1 -aiohttp==3.7.0 +aiohttp==3.8.5 aioimaplib==1.0.1 +aiosignal==1.3.1 aiosmtplib==2.0.2 +anyio==4.0.0 +async-timeout==4.0.3 +attrs==23.1.0 +backoff==2.2.1 +base36==0.1.1 +base58==2.1.1 beautifulsoup4==4.12.2 -mail_parser==3.15.0 -prompt_toolkit==3.0.39 -pydantic==1.10.11 +cachetools==5.3.1 +certifi==2023.7.22 +charset-normalizer==3.2.0 +chroma-hnswlib==0.7.1 +chromadb==0.4.0 +click==8.1.7 +colorama==0.4.6 +coloredlogs==15.0.1 +decorator==4.4.2 +fastapi==0.99.1 +filelock==3.12.3 +flatbuffers==23.5.26 +frozenlist==1.4.0 +fsspec==2023.9.0 +google==3.0.0 +google-api-core==2.11.1 +google-auth==2.23.0 +google-cloud==0.34.0 +google-cloud-texttospeech==2.14.1 +googleapis-common-protos==1.60.0 +grpcio==1.58.0 +grpcio-status==1.58.0 +h11==0.14.0 +httpcore==0.17.3 +httptools==0.6.0 +httpx==0.24.1 +huggingface-hub==0.16.4 +humanfriendly==10.0 +idna==3.4 +imageio==2.31.3 +imageio-ffmpeg==0.4.8 +importlib-resources==6.0.1 +mail-parser==3.15.0 +monotonic==1.6 +moviepy==1.0.0 +mpmath==1.3.0 +multidict==6.0.4 +numpy==1.25.2 +onnxruntime==1.15.1 +openai==0.28.0 +overrides==7.4.0 +packaging==23.1 +pandas==2.1.0 +Pillow==10.0.0 +posthog==3.0.2 +proglog==0.1.10 +prompt-toolkit==3.0.39 +proto-plus==1.22.3 +protobuf==4.24.3 +pulsar-client==3.3.0 +pyasn1==0.5.0 +pyasn1-modules==0.3.0 +pydantic==1.10.12 +PyPika==0.48.9 +pyreadline3==3.4.1 +python-dateutil==2.8.2 +python-dotenv==1.0.0 python-telegram-bot==20.5 -Requests==2.31.0 -protobuf -stability_sdk -toml -base58 -google-cloud-texttospeech -openai -Pillow -aiosqlite -PySocks==1.7.1 +pytz==2023.3.post1 +PyYAML==6.0.1 +requests==2.31.0 +rsa==4.9 +simplejson==3.19.1 +six==1.16.0 +sniffio==1.3.0 +soupsieve==2.5 +starlette==0.27.0 +sympy==1.12 +telegram==0.0.1 +tokenizers==0.14.0 +toml==0.10.0 \ No newline at end of file diff --git a/src/service/aios_shell/aios_shell.py b/src/service/aios_shell/aios_shell.py index f5b9dec..3577d28 100644 --- a/src/service/aios_shell/aios_shell.py +++ b/src/service/aios_shell/aios_shell.py @@ -22,8 +22,11 @@ from prompt_toolkit.styles import Style directory = os.path.dirname(__file__) sys.path.append(directory + '/../../') -from aios_kernel import AIOS_Version,UserConfigItem,AIStorage,Workflow,AIAgent,AgentMsg,AgentMsgStatus,ComputeKernel,OpenAI_ComputeNode,AIBus,AIChatSession,AgentTunnel,TelegramTunnel,CalenderEnvironment,Environment,EmailTunnel,LocalLlama_ComputeNode + import proxy +from aios_kernel import * + + sys.path.append(directory + '/../../component/') from agent_manager import AgentManager @@ -115,7 +118,8 @@ class AIOS_Shell: await AgentTunnel.load_all_tunnels_from_config(tunnel_config) except Exception as e: logger.warning(f"load tunnels config from {tunnels_config_path} failed!") - + + KnowledgePipline.get_instance().initial() return True @@ -163,8 +167,7 @@ class AIOS_Shell: tunnel_config[key] = user_input - return tunnel_config - + return tunnel_config async def append_tunnel_config(self,tunnel_config): user_data_dir = AIStorage.get_instance().get_myai_dir() @@ -179,6 +182,55 @@ class AIOS_Shell: except Exception as e: logger.warning(f"load tunnels config from {tunnels_config_path} failed!") + async def handle_knowledge_commands(self, args): + show_text = FormattedText([("class:title", "sub command not support!\n" + "/knowledge add email | dir\n" + "/knowledge journal [$topn]\n" + "/knowledge query $query\n")]) + if len(args) < 1: + return show_text + sub_cmd = args[0] + if sub_cmd == "add": + if len(args) < 2: + return show_text + if args[1] == "email": + config = dict() + for key, item in KnowledgeEmailSource.user_config_items(): + user_input = await try_get_input(f"{key} : {item}") + if user_input is None: + return show_text + config[key] = user_input + error = KnowledgePipline.get_instance().add_email_source(KnowledgeEmailSource(config)) + if error is not None: + return FormattedText([("class:title", f"/knowledge add email failed {error}\n")]) + else: + KnowledgePipline.get_instance().save_config() + if args[1] == "dir": + config = dict() + for key, item in KnowledgeDirSource.user_config_items(): + user_input = await try_get_input(f"{key} : {item}") + if user_input is None: + return show_text + config[key] = user_input + error = KnowledgePipline.get_instance().add_dir_source(KnowledgeDirSource(config)) + if error is not None: + return FormattedText([("class:title", f"/knowledge add dir failed {error}\n")]) + else: + KnowledgePipline.get_instance().save_config() + else: + return show_text + if sub_cmd == "journal": + topn = 10 if len(args) == 1 else int(args[1]) + journals = [str(journal) for journal in KnowledgePipline.get_instance().get_latest_journals(topn)] + print_formatted_text("\r\n".join(journals)) + if sub_cmd == "query": + if len(args) < 2: + return show_text + prompt = AgentPrompt() + prompt.messages.append({"role": "user", "content":" ".join(args[1:])}) + result = await KnowledgeBase().query_prompt(prompt) + print_formatted_text(result.as_str()) + async def call_func(self,func_name, args): match func_name: case 'send': @@ -221,6 +273,8 @@ class AIOS_Shell: show_text = FormattedText([("class:title", f"connect to {tunnel_target} success!")]) return show_text + case 'knowledge': + return await self.handle_knowledge_commands(args) case 'open': if len(args) >= 1: target_id = args[0] @@ -400,6 +454,9 @@ async def main(): '/history $num $offset', '/login $username', '/connect $target', + '/knowledge add email | dir', + '/knowledge journal [$topn]', + '/knowledge query $query' '/set_config $key', '/list_config', '/show',