diff --git a/rootfs/knowledge_pipelines/local_embedding/pipeline.toml b/rootfs/knowledge_pipelines/local_embedding/pipeline.toml new file mode 100644 index 0000000..ee8c5b0 --- /dev/null +++ b/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" + diff --git a/rootfs/knowledge_pipelines/local_photos_embedding/pipeline.toml b/rootfs/knowledge_pipelines/local_photos_embedding/pipeline.toml deleted file mode 100644 index e26e7a8..0000000 --- a/rootfs/knowledge_pipelines/local_photos_embedding/pipeline.toml +++ /dev/null @@ -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" - diff --git a/rootfs/knowledge_pipelines/pipelines.toml b/rootfs/knowledge_pipelines/pipelines.toml index 99d5b59..e6bf006 100644 --- a/rootfs/knowledge_pipelines/pipelines.toml +++ b/rootfs/knowledge_pipelines/pipelines.toml @@ -1,2 +1,3 @@ -[[pipelines]] -"local_photos_embedding" \ No newline at end of file +pipelines = [ + "local_embedding" +] \ No newline at end of file diff --git a/src/component/knowledge_manager/__init__.py b/src/component/knowledge_manager/__init__.py new file mode 100644 index 0000000..8f6fd1f --- /dev/null +++ b/src/component/knowledge_manager/__init__.py @@ -0,0 +1 @@ +from .pipeline import KnowledgePipelineManager \ No newline at end of file diff --git a/src/knowledge/input/email.py b/src/component/knowledge_manager/input/email.py similarity index 100% rename from src/knowledge/input/email.py rename to src/component/knowledge_manager/input/email.py diff --git a/src/component/knowledge_manager/input/local_dir.py b/src/component/knowledge_manager/input/local_dir.py new file mode 100644 index 0000000..76c1137 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/src/knowledge/parser/embedding.py b/src/component/knowledge_manager/parser/embedding.py similarity index 55% rename from src/knowledge/parser/embedding.py rename to src/component/knowledge_manager/parser/embedding.py index e512eea..cb8b5f8 100644 --- a/src/knowledge/parser/embedding.py +++ b/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): - self.store = KnowledgeStore() - self.compute_kernel = ComputeKernel.get_instance() - self.knowledge_base = KnowledgeBase() + def __init__(self, env: KnowledgePipelineEnvironment, config: dict): 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) - await self.get_vector_store(self._default_text_model).insert(vector, vedio.calculate_id()) + 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()) 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) - await self.get_vector_store(self._default_text_model).insert(vector, email.calculate_id()) + 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.__embedding_rich_text(email.get_rich_text()) @@ -85,9 +93,10 @@ class EmbeddingParser: else: pass - async def parse(self, object: ObjectID): - obj = self.knowledge_base.load_object(object) + async def parse(self, object: ObjectID) -> str: + obj = self.env.get_knowledge_store().load_object(object) await self.__do_embedding(obj) + return "insert into vector store" -def init(params: dict) -> EmbeddingParser: - return EmbeddingParser(params) \ No newline at end of file +def init(env: KnowledgePipelineEnvironment, params: dict) -> EmbeddingParser: + return EmbeddingParser(env, params) \ No newline at end of file diff --git a/src/component/knowledge_manager/pipeline.py b/src/component/knowledge_manager/pipeline.py new file mode 100644 index 0000000..653ea9c --- /dev/null +++ b/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) diff --git a/src/knowledge/query/embedding.py b/src/component/knowledge_manager/query/embedding.py similarity index 100% rename from src/knowledge/query/embedding.py rename to src/component/knowledge_manager/query/embedding.py diff --git a/src/knowledge/input/__init__.py b/src/knowledge/input/__init__.py deleted file mode 100644 index 0e9e2a5..0000000 --- a/src/knowledge/input/__init__.py +++ /dev/null @@ -1 +0,0 @@ -import local_dir \ No newline at end of file diff --git a/src/knowledge/input/local_dir.py b/src/knowledge/input/local_dir.py deleted file mode 100644 index 2bbafe1..0000000 --- a/src/knowledge/input/local_dir.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/src/knowledge/parser/__init__.py b/src/knowledge/parser/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/knowledge/pipeline.py b/src/knowledge/pipeline.py index e051925..e637fb7 100644 --- a/src/knowledge/pipeline.py +++ b/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) - \ No newline at end of file diff --git a/src/knowledge/store.py b/src/knowledge/store.py index 803c3b6..a2ebd77 100644 --- a/src/knowledge/store.py +++ b/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 diff --git a/src/service/aios_shell/aios_shell.py b/src/service/aios_shell/aios_shell.py index 191e0d1..0dc47ea 100644 --- a/src/service/aios_shell/aios_shell.py +++ b/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.load_dir(AIStorage().get_instance().get_system_app_dir() / "knowledge_pipelines") - pipelines.load_dir(AIStorage().get_instance().get_myai_dir() / "knowledge_pipelines") + pipelines = KnowledgePipelineManager(os.path.join(AIStorage().get_instance().get_myai_dir(), "knowledge/pipelines")) + pipelines.load_dir(os.path.join(AIStorage().get_instance().get_system_app_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() diff --git a/src/service/email_spider/converter.py b/src/service/email_spider/converter.py deleted file mode 100644 index 193b220..0000000 --- a/src/service/email_spider/converter.py +++ /dev/null @@ -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 - - \ No newline at end of file diff --git a/src/service/email_spider/main.py b/src/service/email_spider/main.py deleted file mode 100644 index cb75abd..0000000 --- a/src/service/email_spider/main.py +++ /dev/null @@ -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()) - - diff --git a/src/service/email_spider/spider.py b/src/service/email_spider/spider.py deleted file mode 100644 index 0b5a212..0000000 --- a/src/service/email_spider/spider.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/src/service/spider/email_spider.py b/src/service/spider/email_spider.py deleted file mode 100644 index cd44097..0000000 --- a/src/service/spider/email_spider.py +++ /dev/null @@ -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) \ No newline at end of file