knowledge pipeline manager init

rootfs/knowledge_pipelines/local_embedding/pipeline.toml

@@ -0,0 +1,6 @@
+name = "LocalEmbedding"
+input.module = "local_dir"
+input.params.path = "${myai_dir}/data"
+parser.module = "embedding"
+parser.params.path = "${myai_dir}/knowledge/indices/embedding"
+

rootfs/knowledge_pipelines/local_photos_embedding/pipeline.toml

@@ -1,6 +0,0 @@
-name = "LocalPhotoEmbedding"
-input.module = "local_dir"
-input.params.path = "~/myai/photos"
-parser.module = "embedding"
-parser.params.path = "~/myai/knowledge/indices/photo_embedding"
-

rootfs/knowledge_pipelines/pipelines.toml

@@ -1,2 +1,3 @@
-[[pipelines]]
+pipelines = [
-"local_photos_embedding"
+    "local_embedding"
+]

src/component/knowledge_manager/__init__.py

@@ -0,0 +1 @@
+from .pipeline import KnowledgePipelineManager

src/component/knowledge_manager/input/local_dir.py

@@ -0,0 +1,65 @@
+import os
+import aiofiles
+import chardet
+import logging
+import string
+from knowledge import ImageObjectBuilder, DocumentObjectBuilder, KnowledgePipelineEnvironment, KnowledgePipelineJournal
+from aios_kernel.storage import AIStorage
+
+class KnowledgeDirSource:
+    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  
+
+    # @classmethod
+    # def user_config_items(cls):
+    #     return [("path", "local dir path")]
+
+    def path(self):
+        return self.config["path"]
+    
+    @staticmethod
+    async def read_txt_file(file_path:str)->str:
+        cur_encode = "utf-8"
+        async with aiofiles.open(file_path,'rb') as f:
+            cur_encode = chardet.detect(await f.read())['encoding']
+        
+        async with aiofiles.open(file_path,'r',encoding=cur_encode) as f:
+            return await f.read()
+        
+    async def next(self):
+        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
+            from_time = os.path.getctime(latest_journal.get_input())
+            if os.path.getmtime(self.path()) <= from_time:
+                yield (None, None)
+        while True:
+            file_pathes = sorted(os.listdir(self.path()), key=lambda x: os.path.getctime(os.path.join(self.path(), x)))
+            for rel_path in file_pathes:
+                file_path = os.path.join(self.path(), rel_path)
+                timestamp = os.path.getctime(file_path)
+                if timestamp <= from_time:
+                    continue
+                ext = os.path.splitext(file_path)[1].lower()
+                if ext in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']:
+                    logging.info(f"knowledge dir source found image file {file_path}")
+                    image = ImageObjectBuilder({}, {}, file_path).build(self.env.get_knowledge_store())
+                    await self.env.get_knowledge_store().insert_object(image)
+                    yield (image.calculate_id(), file_path)
+                if ext in ['.txt']:
+                    logging.info(f"knowledge dir source found text file {file_path}")
+                    text = await self.read_txt_file(file_path)
+                    document = DocumentObjectBuilder({}, {}, text).build(self.env.get_knowledge_store())
+                    await self.env.get_knowledge_store().insert_object(document)
+                    yield (document.calculate_id(), file_path)
+            yield (None, None)
+            
+
+def init(env: KnowledgePipelineEnvironment, params: dict) -> KnowledgeDirSource:
+    return KnowledgeDirSource(env, params)

src/component/knowledge_manager/parser/embedding.py

@@ -1,31 +1,39 @@
 # define a knowledge base class
 import json
+import string
 from aios_kernel import ComputeKernel, AIStorage
 from knowledge import *
 
 
 class EmbeddingParser:
-    def __init__(self, params: dict):
+    def __init__(self, env: KnowledgePipelineEnvironment, config: dict):
-        self.store = KnowledgeStore()
-        self.compute_kernel = ComputeKernel.get_instance()
-        self.knowledge_base = KnowledgeBase()
         self._default_text_model = "all-MiniLM-L6-v2"
         self._default_image_model = "clip-ViT-B-32"
-        self.vector_store = ChromaVectorStore(AIStorage().get_myai_dir() / "knowledge", self._default_text_model)
+
+        path = string.Template(config["path"]).substitute(myai_dir=AIStorage.get_instance().get_myai_dir())
+        if not os.path.exists(path):
+            os.makedirs(path)
+        config["path"] = path
+        
+        self.env = env
+        self.config = config
+
+    def get_path(self) -> str:
+        return self.config["path"]
 
     def __get_vector_store(self, model_name: str) -> ChromaVectorStore:
-        return ChromaVectorStore(AIStorage().get_myai_dir() / "knowledge", self._default_text_model)
+        return ChromaVectorStore(self.get_path(), model_name)
 
     async def __embedding_document(self, document: DocumentObject):
         for chunk_id in document.get_chunk_list():
-            chunk = self.knowledge_base.store.get_chunk_reader().get_chunk(chunk_id)
+            chunk = self.env.get_knowledge_store().get_chunk_reader().get_chunk(chunk_id)
             if chunk is None:
                 raise ValueError(f"text chunk not found: {chunk_id}")
         
             text = chunk.read().decode("utf-8")
-            vector = await self.compute_kernel.do_text_embedding(text, self._default_text_model)
+            vector = await ComputeKernel.get_instance().do_text_embedding(text, self._default_text_model)
             if vector:
-                await self.get_vector_store(self._default_text_model).insert(vector, chunk_id)
+                await self.__get_vector_store(self._default_text_model).insert(vector, chunk_id)
 
     async def __embedding_image(self, image: ImageObject):
         # desc = {}
@@ -36,9 +44,9 @@ class EmbeddingParser:
         # if not not image.get_tags():
         #     desc["tags"] = image.get_tags()
         # vector = await self.compute_kernel.do_text_embedding(json.dumps(desc), self._default_text_model)
-        vector = await self.compute_kernel.do_image_embedding(image.calculate_id(), self._default_image_model)
+        vector = await ComputeKernel.get_instance().do_image_embedding(image.calculate_id(), self._default_image_model)
         if vector:
-            await self.get_vector_store(self._default_image_model).insert(vector, image.calculate_id())
+            await self.__get_vector_store(self._default_image_model).insert(vector, image.calculate_id())
 
     async def __embedding_video(self, vedio: VideoObject):
         desc = {}
@@ -48,26 +56,26 @@ class EmbeddingParser:
             desc["info"] = vedio.get_info()
         if not not vedio.get_tags():
             desc["tags"] = vedio.get_tags()
-        vector = await self.compute_kernel.do_text_embedding(json.dumps(desc), self._default_text_model)
+        vector = await ComputeKernel.get_instance().do_text_embedding(json.dumps(desc), self._default_text_model)
-        await self.get_vector_store(self._default_text_model).insert(vector, vedio.calculate_id())
+        await self.__get_vector_store(self._default_text_model).insert(vector, vedio.calculate_id())
 
     async def __embedding_rich_text(self, rich_text: RichTextObject):
         for document_id in rich_text.get_documents().values():
-            document = DocumentObject.decode(self.store.get_object_store().get_object(document_id))
+            document = DocumentObject.decode(self.env.get_knowledge_store().get_object_store().get_object(document_id))
             await self.__embedding_document(document)
         for image_id in rich_text.get_images().values():
-            image = ImageObject.decode(self.store.get_object_store().get_object(image_id))
+            image = ImageObject.decode(self.env.get_knowledge_store().get_object_store().get_object(image_id))
             await self.__embedding_image(image)
         for video_id in rich_text.get_videos().values():
-            video = VideoObject.decode(self.store.get_object_store().get_object(video_id))
+            video = VideoObject.decode(self.env.get_knowledge_store().get_object_store().get_object(video_id))
             await self.__embedding_video(video)
         for rich_text_id in rich_text.get_rich_texts().values():
-            rich_text = RichTextObject.decode(self.store.get_object_store().get_object(rich_text_id))
+            rich_text = RichTextObject.decode(self.env.get_knowledge_store().get_object_store().get_object(rich_text_id))
             await self.__embedding_rich_text(rich_text)
 
     async def __embedding_email(self, email: EmailObject):
-        vector = await self.compute_kernel.do_text_embedding(json.dumps(email.get_desc()), self._default_text_model)
+        vector = await ComputeKernel.get_instance().do_text_embedding(json.dumps(email.get_desc()), self._default_text_model)
-        await self.get_vector_store(self._default_text_model).insert(vector, email.calculate_id())
+        await self.__get_vector_store(self._default_text_model).insert(vector, email.calculate_id())
         await self.__embedding_rich_text(email.get_rich_text())
 
 
@@ -85,9 +93,10 @@ class EmbeddingParser:
         else:
             pass
     
-    async def parse(self, object: ObjectID):
+    async def parse(self, object: ObjectID) -> str:
-        obj = self.knowledge_base.load_object(object)
+        obj = self.env.get_knowledge_store().load_object(object)
         await self.__do_embedding(obj)
+        return "insert into vector store"
 
-def init(params: dict) -> EmbeddingParser:
+def init(env: KnowledgePipelineEnvironment, params: dict) -> EmbeddingParser:
-    return EmbeddingParser(params)
+    return EmbeddingParser(env, params)

src/component/knowledge_manager/pipeline.py

@@ -0,0 +1,73 @@
+import os
+import runpy
+import toml
+import asyncio
+from knowledge import KnowledgePipelineEnvironment, KnowledgePipeline
+
+class KnowledgePipelineManager:
+    def __init__(self, root_dir: str):
+        self.root_dir = root_dir
+        self.input_modules = {}
+        self.parser_modules = {}
+        self.pipelines = {
+            "names": {},
+            "running": []
+        }
+        from .input import local_dir
+        self.register_input("local_dir", local_dir.init)
+        from .parser import embedding
+        self.register_parser("embedding", embedding.init)
+
+    def register_input(self, name: str, init_method):
+        self.input_modules[name] = init_method
+        
+    def register_parser(self, name: str, parser_method):
+        self.parser_modules[name] = parser_method
+
+    def add_pipeline(self, config: dict, path: str):
+        name = config["name"]
+        if name in self.pipelines["names"]:
+            return
+        
+        input_module = config["input"]["module"]
+        _, ext = os.path.splitext(input_module)
+        if ext == ".py":
+            input_module = os.path.abspath(path, input_module)
+            input_init = runpy.run_path(input_module)["init"]
+        else:
+            input_init = self.input_modules.get(input_module)
+        input_params = config["input"]["params"]
+
+        parser_module = config["parser"]["module"]
+        _, ext = os.path.splitext(parser_module)
+        if ext == ".py":
+            parser_module = os.path.abspath(path, parser_module)
+            parser_init = runpy.run_path(parser_module)["init"]
+        else:
+            parser_init = self.parser_modules.get(parser_module)
+        parser_params = config["parser"]["params"]
+
+
+        data_path = os.path.join(self.root_dir, name)
+        env = KnowledgePipelineEnvironment(data_path)
+        pipeline = KnowledgePipeline(name, env, input_init, input_params, parser_init, parser_params)
+        self.pipelines["names"][name] = pipeline
+        self.pipelines["running"].append(pipeline)
+
+    async def run(self):
+        while True:
+            for pipeline in self.pipelines["running"]:
+                await pipeline.run()
+            await asyncio.sleep(5)
+
+    def load_dir(self, root: str):
+        config_path = os.path.join(root, "pipelines.toml")
+        if not os.path.exists(config_path):
+            return 
+        with open(config_path, "r") as f:
+            config = toml.load(f)
+        for path in config["pipelines"]:
+            pipeline_path = os.path.join(root, path)
+            with open(os.path.join(pipeline_path, "pipeline.toml")) as f:
+                pipeline_config = toml.load(f)
+                self.add_pipeline(pipeline_config, pipeline_path)

src/knowledge/input/__init__.py

@@ -1 +0,0 @@
-import local_dir

src/knowledge/input/local_dir.py

@@ -1,61 +0,0 @@
-import os
-import aiofiles
-import chardet
-from knowledge import ImageObjectBuilder, DocumentObjectBuilder, KnowledgeBase
-
-class KnowledgeDirSource:
-    def __init__(self, config):
-        self.config = config
-        config["path"] = os.path.abspath(config["path"])
-
-    # @classmethod
-    # def user_config_items(cls):
-    #     return [("path", "local dir path")]
-
-    def path(self):
-        return self.config["path"]
-    
-    @staticmethod
-    async def read_txt_file(file_path:str)->str:
-        cur_encode = "utf-8"
-        async with aiofiles.open(file_path,'rb') as f:
-            cur_encode = chardet.detect(await f.read())['encoding']
-        
-        async with aiofiles.open(file_path,'r',encoding=cur_encode) as f:
-            return await f.read()
-        
-    async def next(self):
-        # logging.debug(f"knowledge dir source {self.id()} run once")
-        # journal_client = KnowledgeJournalClient()
-        # latest_journal = journal_client.latest_journal(self.id())
-        # if latest_journal is not None:
-        #     if os.path.getmtime(self.path()) <= latest_journal.timestamp:
-        #         logging.debug(f"knowledge dir source {self.id()} ingnored for no update")
-        #         return
-        while True:
-            file_pathes = sorted(os.listdir(self.path()), key=lambda x: os.path.getctime(os.path.join(self.path(), x)))
-            for rel_path in file_pathes:
-                file_path = os.path.join(self.path(), rel_path)
-                # timestamp = os.path.getctime(file_path)
-                # if latest_journal is not None:
-                #     if timestamp <= latest_journal.timestamp:
-                #         continue
-                ext = os.path.splitext(file_path)[1].lower()
-                if ext in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']:
-                    # logging.info(f"knowledge dir source {self.id()} found image file {file_path}")
-                    image = ImageObjectBuilder({}, {}, file_path).build()
-                    await KnowledgeBase().insert_object(image)
-                    yield image.calculate_id()
-                    # journal_client.insert(KnowledgeJournal("dir", self.id(), rel_path, str(image.calculate_id()), timestamp))
-                if ext in ['.txt']:
-                    # logging.info(f"knowledge dir source {self.id()} found text file {file_path}")
-                    text = await self.read_txt_file(file_path)
-                    document = DocumentObjectBuilder({}, {}, text).build()
-                    await KnowledgeBase().insert_object(document)
-                    yield document.calculate_id()
-                    # journal_client.insert(KnowledgeJournal("dir", self.id(), rel_path, str(document.calculate_id()), timestamp))
-            yield 0
-            
-
-def init(params: dict) -> KnowledgeDirSource:
-    return KnowledgeDirSource(params)

src/knowledge/pipeline.py

@@ -1,12 +1,7 @@
-
-# class KnowledgePipelineTemplate
-import runpy
-import toml
 import datetime
 import sqlite3
 import os
 from . import ObjectID, KnowledgeStore
-import asyncio
 from enum import Enum
 
 class KnowledgePipelineJournal:
@@ -17,9 +12,13 @@ class KnowledgePipelineJournal:
         self.parser = parser
     
     def is_finish(self) -> bool:
-        self.object_id is None
+        return self.object_id is None
 
+    def get_input(self) -> str:
+        return self.input
     
+    def get_parser(self) -> str:
+        return self.parser
 
 # init sqlite3 client
 class KnowledgePipelineJournalClient:
@@ -43,7 +42,7 @@ class KnowledgePipelineJournalClient:
         timestamp = datetime.datetime.now() if timestamp is None else timestamp
         conn = sqlite3.connect(self.journal_path)
         conn.execute(
-            "INSERT INTO journal (time, object_id, input, parser) VALUES (?, ?, ?)",
+            "INSERT INTO journal (time, object_id, input, parser) VALUES (?, ?, ?, ?)",
             (timestamp, str(object_id), input, parser),
         )
         conn.commit()
@@ -56,10 +55,18 @@ class KnowledgePipelineJournalClient:
 
 class KnowledgePipelineEnvironment:
     def __init__(self, pipeline_path: str):
-        self.knowledge_base = KnowledgeStore()
+        self.knowledge_store = KnowledgeStore()
+        if not os.path.exists(pipeline_path):
+            os.makedirs(pipeline_path)
         self.pipeline_path = pipeline_path
         self.journal = KnowledgePipelineJournalClient(pipeline_path)
 
+    def get_journal(self) -> KnowledgePipelineJournalClient:
+        return self.journal
+    
+    def get_knowledge_store(self) -> KnowledgeStore:
+        return self.knowledge_store
+
 class KnowledgePipelineState(Enum):
     INIT = 0
     RUNNING = 1
@@ -84,13 +91,13 @@ class KnowledgePipeline:
             self.parser = self.parser_init(self.env, self.parser_params)
             self.state = KnowledgePipelineState.RUNNING
         if self.state == KnowledgePipelineState.RUNNING:
-            for input in await self.input.next():
+            async for input in self.input.next():
                 if input is None:
                     self.state = KnowledgePipelineState.FINISHED
                     self.env.journal.insert(None, "finished", "finished")
                     return 
                 (object_id, input_journal) = input
-                if object_id is None:
+                if object_id is not None:
                     parser_journal = await self.parser.parse(object_id)
                     self.env.journal.insert(object_id, input_journal, parser_journal)
         if self.state == KnowledgePipelineState.STOPPED:
@@ -98,69 +105,5 @@ class KnowledgePipeline:
         if self.state == KnowledgePipelineState.FINISHED:
             return 
 
-class KnowledgePipelineManager:
-    def __init__(self, root_dir: str):
-        self.root_dir = root_dir
-        self.input_modules = {}
-        self.parser_modules = {}
-        self.pipelines = {
-            "names": {},
-            "running": []
-        }
-        from .input import local_dir
-        self.register_input("local_dir", local_dir.init)
-
-    def register_input(self, name: str, init_method):
-        self.input_modules[name] = init_method
-        
-    def register_parser(self, name: str, parser_method):
-        self.parser_modules[name] = parser_method
-
-    def add_pipeline(self, config: dict, path: str):
-        name = config["name"]
-        if name in self.pipelines["names"]:
-            return
-        
-        input_module = self.input_modules[config["input"]["module"]]
-        _, ext = os.path.splitext(input_module)
-        if ext == ".py":
-            input_module = os.path.abspath(path, input_module)
-            input_init = runpy.run_path(input_module)["init"]
-        else:
-            input_init = self.input_modules.get(input_module)
-        input_params = config["input"]["params"]
-
-        parser_module = self.parser_modules[config["parser"]["module"]]
-        _, ext = os.path.splitext(parser_module)
-        if ext == ".py":
-            parser_module = os.path.abspath(path, parser_module)
-            parser_init = runpy.run_path(parser_module)["init"]
-        else:
-            parser_init = self.parser_modules.get(parser_module)
-        parser_params = config["parser"]["params"]
-
-
-        data_path = self.root_dir / name
-        env = KnowledgePipelineEnvironment(data_path)
-        pipeline = KnowledgePipeline(name, env, input_init, input_params, parser_init, parser_params)
-        self.pipelines["names"][name] = pipeline
-        self.pipelines["running"].append(pipeline)
-
-    async def run(self):
-        while True:
-            for pipeline in self.pipelines["running"]:
-                await pipeline.run()
-            await asyncio.sleep(5)
-
-    def load_dir(self, root: str):
-        config_path = os.path.join(root, "pipelines.toml")
-        with open(config_path, "r") as f:
-            config = toml.load(f)
-        for path in config["pipelines"]:
-            pipeline_path = os.path.join(root, path)
-            with open(os.path.join(pipeline_path, "pipeline.toml")) as f:
-                pipeline_config = toml.load(f)
-                self.add_pipeline(pipeline_config, pipeline_path)
-
 
     

src/knowledge/store.py

@@ -68,15 +68,15 @@ class KnowledgeStore:
     
     def load_object(self, object_id: ObjectID) -> KnowledgeObject:
         if object_id.get_object_type() == ObjectType.Document:
-            return DocumentObject.decode(self.store.get_object_store().get_object(object_id))
+            return DocumentObject.decode(self.object_store.get_object(object_id))
         if object_id.get_object_type() == ObjectType.Image:
-            return ImageObject.decode(self.store.get_object_store().get_object(object_id))
+            return ImageObject.decode(self.object_store.get_object(object_id))
         if object_id.get_object_type() == ObjectType.Video:
-            return VideoObject.decode(self.store.get_object_store().get_object(object_id))
+            return VideoObject.decode(self.object_store.get_object(object_id))
         if object_id.get_object_type() == ObjectType.RichText:
-            return RichTextObject.decode(self.store.get_object_store().get_object(object_id))
+            return RichTextObject.decode(self.object_store.get_object(object_id))
         if object_id.get_object_type() == ObjectType.Email:
-            return EmailObject.decode(self.store.get_object_store().get_object(object_id))
+            return EmailObject.decode(self.object_store.get_object(object_id))
         else:
             pass
     

src/service/aios_shell/aios_shell.py

@@ -32,6 +32,7 @@ from aios_kernel import *
 sys.path.append(directory + '/../../component/')
 from agent_manager import AgentManager
 from workflow_manager import WorkflowManager
+from knowledge_manager import KnowledgePipelineManager
 
 
 logger = logging.getLogger(__name__)
@@ -185,9 +186,9 @@ class AIOS_Shell:
         AIBus().get_default_bus().register_message_handler(self.username,self._user_process_msg)
         
         
-        pipelines = KnowledgePipelineManager(AIStorage().get_instance().get_myai_dir() / "knowledge" / "pipelines")
+        pipelines = KnowledgePipelineManager(os.path.join(AIStorage().get_instance().get_myai_dir(), "knowledge/pipelines"))
-        pipelines.load_dir(AIStorage().get_instance().get_system_app_dir() / "knowledge_pipelines")
+        pipelines.load_dir(os.path.join(AIStorage().get_instance().get_system_app_dir(), "knowledge_pipelines"))
-        pipelines.load_dir(AIStorage().get_instance().get_myai_dir() / "knowledge_pipelines")
+        pipelines.load_dir(os.path.join(AIStorage().get_instance().get_myai_dir(), "knowledge_pipelines"))
         asyncio.create_task(pipelines.run())
 
         TelegramTunnel.register_to_loader()

src/service/email_spider/converter.py

@@ -1,24 +0,0 @@
-from aios_kernel.knowledge import KnowledgeBase, EmailObject
-
-# define a email converter class
-
-class EmailConverter:
-    # define init method
-    def __init__(self, local_dir, knowledge_base: KnowledgeBase) -> None:
-        pass
-
-    async def run(self):
-        # convert the email to knowledge object
-        for email_dir in self._next():
-            # convert the email to knowledge object
-            knowledge_object = self._convert(email_dir)
-            # insert the knowledge object to knowledge base
-            await self.knowledge_base.insert(knowledge_object)
-
-    def _next(self) -> str:
-        pass
-
-    def _convert(self, email_dir) -> EmailObject:
-        pass
-
-    

src/service/email_spider/main.py

@@ -1,12 +0,0 @@
-import asyncio
-from .spider import EmailSpider, EmailConverter
-
-
-if __name__ == "__main__":
-    spider = EmailSpider("smtp.163.com","user","pwd","./email")
-    asyncio.run(spider.run())
-
-    converter = EmailConverter("./email",KnowledgeBase())
-    asyncio.run(converter.run())
-
-    

src/service/email_spider/spider.py

@@ -1,17 +0,0 @@
-# define a email spider class
-
-class EmailSpider:
-    def __init__(self, address, account, pwd, local_dir) -> None:
-        pass
-
-    async def run(self): 
-        # spide the email from the email server
-        for email_link in self._next():
-            # save the email to local directory
-            self._save(email_link)
-
-    def _next(self):
-        pass
-
-    def _save(self, email_link) -> str:
-        pass

src/service/spider/email_spider.py

@@ -1,171 +0,0 @@
-"""
-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 imaplib
-import os
-import toml
-import logging
-import mailparser
-import hashlib
-import json
-import base64
-from bs4 import BeautifulSoup
-import requests
-
-class EmailSpider:
-    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)
-
-        # 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()
-
-    def email_client(self) -> imaplib.IMAP4_SSL:
-        self.logger.info(f"read email config from {self.config.get('EMAIL_IMAP_SERVER')}")
-        client = imaplib.IMAP4_SSL(
-            host=self.config.get('EMAIL_IMAP_SERVER'), 
-            port=self.config.get('EMAIL_IMAP_PORT')
-        )
-        client.login(self.config.get('EMAIL_ADDRESS'), self.config.get('EMAIL_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"):
-        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")
-        email_list.reverse()
-        for uid in email_list:
-            if self.check_email_saved(uid):
-                self.logger.info(f"email uid {uid} already saved")
-            else:
-                self.read_and_save_email(uid)
-                self.logger.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}]")
-        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')}"
-        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])
-        self.logger.info(f"[{uid}]check email subject [{mail.subject}]")
-        dir = self.get_local_dir_name(mail)
-        self.logger.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)
-                    self.logger.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]
-        self.logger.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)
-                self.logger.info(f'Downloaded {img_url} to {img_filename}')
-            else:
-                self.logger.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')}"
-        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}")
-        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)
-            self.logger.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")
-
-
-if __name__ == "__main__":
-    spider = EmailSpider()
-    folder = 'INBOX'
-    imap_keyword = "ALL"
-    spider.read_emails(folder, imap_keyword)