Merge pull request #60 from photosssa/MVP

Add knowledge commands to aios shell

.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
 
+

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 

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

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)

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 = []
@@ -242,6 +245,3 @@ class KnowledgeBase:
 
 
     
-                
-
-    

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()

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

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}

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 {}

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)
 

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]
+
+    

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))
+        

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]

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(

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

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
@@ -116,6 +119,7 @@ class AIOS_Shell:
         except Exception as e:
             logger.warning(f"load tunnels config from {tunnels_config_path} failed!")
         
+        KnowledgePipline.get_instance().initial()
         return True 
         
 
@@ -165,7 +169,6 @@ class AIOS_Shell:
 
         return tunnel_config        
 
-
     async def append_tunnel_config(self,tunnel_config):
         user_data_dir = AIStorage.get_instance().get_myai_dir()
         tunnels_config_path = os.path.abspath(f"{user_data_dir}/etc/tunnels.cfg.toml")
@@ -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',