
import os
import chardet
import importlib
from pathlib import Path
from WebUI.text_splitter import zh_title_enhance as func_zh_title_enhance
from WebUI.Server.document_loaders import RapidOCRPDFLoader, RapidOCRLoader
import langchain.document_loaders
from langchain.docstore.document import Document
from langchain.text_splitter import TextSplitter
from WebUI.configs.basicconfig import (GetKbConfig, GetKbRootPath, GetTextSplitterDict)
from WebUI.Server.utils import run_in_thread_pool, get_model_worker_config
from WebUI.Server.document_loaders import *
from typing import List, Union,Dict, Tuple, Generator

TEXT_SPLITTER_NAME = "ChineseRecursiveTextSplitter"
CHUNK_SIZE = 500
OVERLAP_SIZE = 100
ZH_TITLE_ENHANCE = False
VECTOR_SEARCH_TOP_K = 5
SCORE_THRESHOLD = 1.5

LOADER_DICT = {"UnstructuredHTMLLoader": ['.html'],
               "MHTMLLoader": ['.mhtml'],
               "UnstructuredMarkdownLoader": ['.md'],
               "JSONLoader": [".json"],
               "JSONLinesLoader": [".jsonl"],
               "CSVLoader": [".csv"],
               # "FilteredCSVLoader": [".csv"], # 需要自己指定，目前还没有支持
               "RapidOCRPDFLoader": [".pdf"],
               "RapidOCRLoader": ['.png', '.jpg', '.jpeg', '.bmp'],
               "UnstructuredEmailLoader": ['.eml', '.msg'],
               "UnstructuredEPubLoader": ['.epub'],
               "UnstructuredExcelLoader": ['.xlsx', '.xls', '.xlsd'],
               "NotebookLoader": ['.ipynb'],
               "UnstructuredODTLoader": ['.odt'],
               "PythonLoader": ['.py'],
               "UnstructuredRSTLoader": ['.rst'],
               "UnstructuredRTFLoader": ['.rtf'],
               "SRTLoader": ['.srt'],
               "TomlLoader": ['.toml'],
               "UnstructuredTSVLoader": ['.tsv'],
               "UnstructuredWordDocumentLoader": ['.docx', '.doc'],
               "UnstructuredXMLLoader": ['.xml'],
               "UnstructuredPowerPointLoader": ['.ppt', '.pptx'],
               "EverNoteLoader": ['.enex'],
               "UnstructuredFileLoader": ['.txt'],
               }
SUPPORTED_EXTS = [ext for sublist in LOADER_DICT.values() for ext in sublist]

def validate_kb_name(knowledge_base_id: str) -> bool:
    if "../" in knowledge_base_id:
        return False
    return True


def get_kb_path(knowledge_base_name: str):
    kb_config = GetKbConfig()
    kb_root_path = GetKbRootPath(kb_config)
    return os.path.join(kb_root_path, knowledge_base_name)


def get_doc_path(knowledge_base_name: str):
    return os.path.join(get_kb_path(knowledge_base_name), "content")


def get_vs_path(knowledge_base_name: str, vector_name: str):
    return os.path.join(get_kb_path(knowledge_base_name), "vector_store", vector_name)


def get_file_path(knowledge_base_name: str, doc_name: str):
    return os.path.join(get_doc_path(knowledge_base_name), doc_name)

def list_files_from_folder(kb_name: str):
    doc_path = get_doc_path(kb_name)
    result = []

    def is_skiped_path(path: str):
        tail = os.path.basename(path).lower()
        for x in ["temp", "tmp", ".", "~$"]:
            if tail.startswith(x):
                return True
        return False

    def process_entry(entry):
        if is_skiped_path(entry.path):
            return

        if entry.is_symlink():
            target_path = os.path.realpath(entry.path)
            with os.scandir(target_path) as target_it:
                for target_entry in target_it:
                    process_entry(target_entry)
        elif entry.is_file():
            file_path = (Path(os.path.relpath(entry.path, doc_path)).as_posix())
            result.append(file_path)
        elif entry.is_dir():
            with os.scandir(entry.path) as it:
                for sub_entry in it:
                    process_entry(sub_entry)

    with os.scandir(doc_path) as it:
        for entry in it:
            process_entry(entry)
    return result

def get_LoaderClass(file_extension):
    for LoaderClass, extensions in LOADER_DICT.items():
        if file_extension in extensions:
            return LoaderClass
        
def get_loader(loader_name: str, file_path: str, loader_kwargs: Dict = None):
    loader_kwargs = loader_kwargs or {}
    try:
        if loader_name == "RapidOCRPDFLoader":
            DocumentLoader = RapidOCRPDFLoader
        elif loader_name == "RapidOCRLoader":
            DocumentLoader = RapidOCRLoader
        else:
            document_loaders_module = importlib.import_module('langchain.document_loaders')
            DocumentLoader = getattr(document_loaders_module, loader_name)
    except Exception as e:
        msg = f"for file {file_path} search loader {loader_name} failed: {e}"
        print(f'{e.__class__.__name__}: {msg}')
        document_loaders_module = importlib.import_module('langchain.document_loaders')
        DocumentLoader = getattr(document_loaders_module, "UnstructuredFileLoader")

    if loader_name == "UnstructuredFileLoader":
        loader_kwargs.setdefault("autodetect_encoding", True)
    elif loader_name == "CSVLoader":
        if not loader_kwargs.get("encoding"):
            with open(file_path, 'rb') as struct_file:
                encode_detect = chardet.detect(struct_file.read())
            if encode_detect is None:
                encode_detect = {"encoding": "utf-8"}
            loader_kwargs["encoding"] = encode_detect["encoding"]

    elif loader_name == "JSONLoader":
        loader_kwargs.setdefault("jq_schema", ".")
        loader_kwargs.setdefault("text_content", False)
    elif loader_name == "JSONLinesLoader":
        loader_kwargs.setdefault("jq_schema", ".")
        loader_kwargs.setdefault("text_content", False)

    loader = DocumentLoader(file_path, **loader_kwargs)
    return loader

def make_text_splitter(
        splitter_name: str = TEXT_SPLITTER_NAME,
        chunk_size: int = CHUNK_SIZE,
        chunk_overlap: int = OVERLAP_SIZE,
        llm_model: str = "",
):
    splitter_name = splitter_name or "SpacyTextSplitter"
    try:
        text_splitter_dict = GetTextSplitterDict()
        if splitter_name == "MarkdownHeaderTextSplitter":
            headers_to_split_on = text_splitter_dict[splitter_name]['headers_to_split_on']
            text_splitter = langchain.text_splitter.MarkdownHeaderTextSplitter(
                headers_to_split_on=headers_to_split_on)
        elif splitter_name == "ChineseRecursiveTextSplitter":
            text_splitter_module = importlib.import_module('text_splitter')
            TextSplitter = getattr(text_splitter_module, splitter_name)
            text_splitter = TextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        else:
            try:
                text_splitter_module = importlib.import_module('text_splitter')
                TextSplitter = getattr(text_splitter_module, splitter_name)
            except:
                text_splitter_module = importlib.import_module('langchain.text_splitter')
                TextSplitter = getattr(text_splitter_module, splitter_name)

            if text_splitter_dict[splitter_name]["source"] == "tiktoken":
                try:
                    text_splitter = TextSplitter.from_tiktoken_encoder(
                        encoding_name=text_splitter_dict[splitter_name]["tokenizer_name_or_path"],
                        pipeline="zh_core_web_sm",
                        chunk_size=chunk_size,
                        chunk_overlap=chunk_overlap
                    )
                except:
                    text_splitter = TextSplitter.from_tiktoken_encoder(
                        encoding_name=text_splitter_dict[splitter_name]["tokenizer_name_or_path"],
                        chunk_size=chunk_size,
                        chunk_overlap=chunk_overlap
                    )
            elif text_splitter_dict[splitter_name]["source"] == "huggingface":
                if text_splitter_dict[splitter_name]["tokenizer_name_or_path"] == "":
                    config = get_model_worker_config(llm_model)
                    text_splitter_dict[splitter_name]["tokenizer_name_or_path"] = \
                        config.get("model_path")

                if text_splitter_dict[splitter_name]["tokenizer_name_or_path"] == "gpt2":
                    from transformers import GPT2TokenizerFast
                    from langchain.text_splitter import CharacterTextSplitter
                    tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
                else:
                    from transformers import AutoTokenizer
                    tokenizer = AutoTokenizer.from_pretrained(
                        text_splitter_dict[splitter_name]["tokenizer_name_or_path"],
                        trust_remote_code=True)
                text_splitter = TextSplitter.from_huggingface_tokenizer(
                    tokenizer=tokenizer,
                    chunk_size=chunk_size,
                    chunk_overlap=chunk_overlap
                )
            else:
                try:
                    text_splitter = TextSplitter(
                        pipeline="zh_core_web_sm",
                        chunk_size=chunk_size,
                        chunk_overlap=chunk_overlap
                    )
                except:
                    text_splitter = TextSplitter(
                        chunk_size=chunk_size,
                        chunk_overlap=chunk_overlap
                    )
    except Exception as e:
        print(e)
        text_splitter_module = importlib.import_module('langchain.text_splitter')
        TextSplitter = getattr(text_splitter_module, "RecursiveCharacterTextSplitter")
        text_splitter = TextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        splitter_name = "RecursiveCharacterTextSplitter"
    return text_splitter, splitter_name

def list_kbs_from_folder():
    kb_config = GetKbConfig()
    kb_root_path = GetKbRootPath(kb_config)
    kb_list = []
    try:
        dirs = os.listdir(kb_root_path)
        for f in dirs:
            if os.path.isdir(os.path.join(kb_root_path, f)):
                kb_list.append(f)
    except Exception as e:
        pass
    return kb_list

class KnowledgeFile:
    def __init__(
            self,
            filename: str,
            knowledge_base_name: str,
            loader_kwargs: Dict = {},
    ):
        self.kb_name = knowledge_base_name
        self.filename = str(Path(filename).as_posix())
        self.ext = os.path.splitext(filename)[-1].lower()
        if self.ext not in SUPPORTED_EXTS:
            raise ValueError(f"Not support file format: {self.filename}")
        self.loader_kwargs = loader_kwargs
        self.filepath = get_file_path(knowledge_base_name, filename)
        self.docs = None
        self.splited_docs = None
        self.document_loader_name = get_LoaderClass(self.ext)
        self.text_splitter_name = TEXT_SPLITTER_NAME

    def file2docs(self, refresh: bool = False):
        if self.docs is None or refresh:
            print(f"{self.document_loader_name} used for {self.filepath}")
            loader = get_loader(loader_name=self.document_loader_name,
                                file_path=self.filepath,
                                loader_kwargs=self.loader_kwargs)
            self.docs = loader.load()
        return self.docs

    def docs2texts(
            self,
            docs: List[Document] = None,
            zh_title_enhance: bool = ZH_TITLE_ENHANCE,
            refresh: bool = False,
            chunk_size: int = CHUNK_SIZE,
            chunk_overlap: int = OVERLAP_SIZE,
            text_splitter: TextSplitter = None,
    ):
        docs = docs or self.file2docs(refresh=refresh)
        if not docs:
            return []
        if self.ext not in [".csv"]:
            if text_splitter is None:
                text_splitter, new_text_splitter_name = make_text_splitter(splitter_name=self.text_splitter_name, chunk_size=chunk_size,
                                                   chunk_overlap=chunk_overlap)
                if new_text_splitter_name != self.text_splitter_name:
                    self.text_splitter_name = new_text_splitter_name
            if self.text_splitter_name == "MarkdownHeaderTextSplitter":
                docs = text_splitter.split_text(docs[0].page_content)
            else:
                docs = text_splitter.split_documents(docs)

        if not docs:
            return []

        print(f"Document split samples: {docs[0]}")
        if zh_title_enhance:
            docs = func_zh_title_enhance(docs)
        self.splited_docs = docs
        return self.splited_docs

    def file2text(
            self,
            zh_title_enhance: bool = ZH_TITLE_ENHANCE,
            refresh: bool = False,
            chunk_size: int = CHUNK_SIZE,
            chunk_overlap: int = OVERLAP_SIZE,
            text_splitter: TextSplitter = None,
    ):
        if self.splited_docs is None or refresh:
            docs = self.file2docs()
            self.splited_docs = self.docs2texts(docs=docs,
                                                zh_title_enhance=zh_title_enhance,
                                                refresh=refresh,
                                                chunk_size=chunk_size,
                                                chunk_overlap=chunk_overlap,
                                                text_splitter=text_splitter)
        return self.splited_docs

    def file_exist(self):
        return os.path.isfile(self.filepath)

    def get_mtime(self):
        return os.path.getmtime(self.filepath)

    def get_size(self):
        return os.path.getsize(self.filepath)


def files2docs_in_thread(
        files: List[Union[KnowledgeFile, Tuple[str, str], Dict]],
        chunk_size: int = CHUNK_SIZE,
        chunk_overlap: int = OVERLAP_SIZE,
        zh_title_enhance: bool = ZH_TITLE_ENHANCE,
    ) -> Generator:
    def file2docs(*, file: KnowledgeFile, **kwargs) -> Tuple[bool, Tuple[str, str, List[Document]]]:
        try:
            return True, (file.kb_name, file.filename, file.file2text(**kwargs))
        except Exception as e:
            msg = f"from {file.kb_name}/{file.filename} load failed: {e}"
            return False, (file.kb_name, file.filename, msg)

    kwargs_list = []
    for i, file in enumerate(files):
        kwargs = {}
        try:
            if isinstance(file, tuple) and len(file) >= 2:
                filename = file[0]
                kb_name = file[1]
                file = KnowledgeFile(filename=filename, knowledge_base_name=kb_name)
            elif isinstance(file, dict):
                filename = file.pop("filename")
                kb_name = file.pop("kb_name")
                kwargs.update(file)
                file = KnowledgeFile(filename=filename, knowledge_base_name=kb_name)
            kwargs["file"] = file
            kwargs["chunk_size"] = chunk_size
            kwargs["chunk_overlap"] = chunk_overlap
            kwargs["zh_title_enhance"] = zh_title_enhance
            kwargs_list.append(kwargs)
        except Exception as e:
            yield False, (kb_name, filename, str(e))

    for result in run_in_thread_pool(func=file2docs, params=kwargs_list):
        yield result
langchain.text_splitter.TextSplitter.from_huggingface_tokenizer langchain.text_splitter.TextSplitter.from_tiktoken_encoder langchain.text_splitter.TextSplitter
[((2330, 2343), 'WebUI.configs.basicconfig.GetKbConfig', 'GetKbConfig', ([], {}), '()\n', (2341, 2343), False, 'from WebUI.configs.basicconfig import GetKbConfig, GetKbRootPath, GetTextSplitterDict\n'), ((2363, 2387), 'WebUI.configs.basicconfig.GetKbRootPath', 'GetKbRootPath', (['kb_config'], {}), '(kb_config)\n', (2376, 2387), False, 'from WebUI.configs.basicconfig import GetKbConfig, GetKbRootPath, GetTextSplitterDict\n'), ((2399, 2446), 'os.path.join', 'os.path.join', (['kb_root_path', 'knowledge_base_name'], {}), '(kb_root_path, knowledge_base_name)\n', (2411, 2446), False, 'import os\n'), ((9822, 9835), 'WebUI.configs.basicconfig.GetKbConfig', 'GetKbConfig', ([], {}), '()\n', (9833, 9835), False, 'from WebUI.configs.basicconfig import GetKbConfig, GetKbRootPath, GetTextSplitterDict\n'), ((9855, 9879), 'WebUI.configs.basicconfig.GetKbRootPath', 'GetKbRootPath', (['kb_config'], {}), '(kb_config)\n', (9868, 9879), False, 'from WebUI.configs.basicconfig import GetKbConfig, GetKbRootPath, GetTextSplitterDict\n'), ((15110, 15164), 'WebUI.Server.utils.run_in_thread_pool', 'run_in_thread_pool', ([], {'func': 'file2docs', 'params': 'kwargs_list'}), '(func=file2docs, params=kwargs_list)\n', (15128, 15164), False, 'from WebUI.Server.utils import run_in_thread_pool, get_model_worker_config\n'), ((3792, 3812), 'os.scandir', 'os.scandir', (['doc_path'], {}), '(doc_path)\n', (3802, 3812), False, 'import os\n'), ((6015, 6036), 'WebUI.configs.basicconfig.GetTextSplitterDict', 'GetTextSplitterDict', ([], {}), '()\n', (6034, 6036), False, 'from WebUI.configs.basicconfig import GetKbConfig, GetKbRootPath, GetTextSplitterDict\n'), ((9921, 9945), 'os.listdir', 'os.listdir', (['kb_root_path'], {}), '(kb_root_path)\n', (9931, 9945), False, 'import os\n'), ((13436, 13465), 'os.path.isfile', 'os.path.isfile', (['self.filepath'], {}), '(self.filepath)\n', (13450, 13465), False, 'import os\n'), ((13507, 13538), 'os.path.getmtime', 'os.path.getmtime', (['self.filepath'], {}), '(self.filepath)\n', (13523, 13538), False, 'import os\n'), ((13579, 13609), 'os.path.getsize', 'os.path.getsize', (['self.filepath'], {}), '(self.filepath)\n', (13594, 13609), False, 'import os\n'), ((3297, 3325), 'os.path.realpath', 'os.path.realpath', (['entry.path'], {}), '(entry.path)\n', (3313, 3325), False, 'import os\n'), ((4750, 4803), 'importlib.import_module', 'importlib.import_module', (['"""langchain.document_loaders"""'], {}), "('langchain.document_loaders')\n", (4773, 4803), False, 'import importlib\n'), ((9453, 9503), 'importlib.import_module', 'importlib.import_module', (['"""langchain.text_splitter"""'], {}), "('langchain.text_splitter')\n", (9476, 9503), False, 'import importlib\n'), ((9615, 9679), 'langchain.text_splitter.TextSplitter', 'TextSplitter', ([], {'chunk_size': 'chunk_size', 'chunk_overlap': 'chunk_overlap'}), '(chunk_size=chunk_size, chunk_overlap=chunk_overlap)\n', (9627, 9679), False, 'from langchain.text_splitter import TextSplitter\n'), ((12478, 12505), 'WebUI.text_splitter.zh_title_enhance', 'func_zh_title_enhance', (['docs'], {}), '(docs)\n', (12499, 12505), True, 'from WebUI.text_splitter import zh_title_enhance as func_zh_title_enhance\n'), ((2990, 3012), 'os.path.basename', 'os.path.basename', (['path'], {}), '(path)\n', (3006, 3012), False, 'import os\n'), ((3343, 3366), 'os.scandir', 'os.scandir', (['target_path'], {}), '(target_path)\n', (3353, 3366), False, 'import os\n'), ((4434, 4487), 'importlib.import_module', 'importlib.import_module', (['"""langchain.document_loaders"""'], {}), "('langchain.document_loaders')\n", (4457, 4487), False, 'import importlib\n'), ((6420, 6460), 'importlib.import_module', 'importlib.import_module', (['"""text_splitter"""'], {}), "('text_splitter')\n", (6443, 6460), False, 'import importlib\n'), ((6561, 6625), 'langchain.text_splitter.TextSplitter', 'TextSplitter', ([], {'chunk_size': 'chunk_size', 'chunk_overlap': 'chunk_overlap'}), '(chunk_size=chunk_size, chunk_overlap=chunk_overlap)\n', (6573, 6625), False, 'from langchain.text_splitter import TextSplitter\n'), ((9998, 10027), 'os.path.join', 'os.path.join', (['kb_root_path', 'f'], {}), '(kb_root_path, f)\n', (10010, 10027), False, 'import os\n'), ((6696, 6736), 'importlib.import_module', 'importlib.import_module', (['"""text_splitter"""'], {}), "('text_splitter')\n", (6719, 6736), False, 'import importlib\n'), ((10362, 10376), 'pathlib.Path', 'Path', (['filename'], {}), '(filename)\n', (10366, 10376), False, 'from pathlib import Path\n'), ((10408, 10434), 'os.path.splitext', 'os.path.splitext', (['filename'], {}), '(filename)\n', (10424, 10434), False, 'import os\n'), ((3670, 3692), 'os.scandir', 'os.scandir', (['entry.path'], {}), '(entry.path)\n', (3680, 3692), False, 'import os\n'), ((6872, 6922), 'importlib.import_module', 'importlib.import_module', (['"""langchain.text_splitter"""'], {}), "('langchain.text_splitter')\n", (6895, 6922), False, 'import importlib\n'), ((7131, 7328), 'langchain.text_splitter.TextSplitter.from_tiktoken_encoder', 'TextSplitter.from_tiktoken_encoder', ([], {'encoding_name': "text_splitter_dict[splitter_name]['tokenizer_name_or_path']", 'pipeline': '"""zh_core_web_sm"""', 'chunk_size': 'chunk_size', 'chunk_overlap': 'chunk_overlap'}), "(encoding_name=text_splitter_dict[\n    splitter_name]['tokenizer_name_or_path'], pipeline='zh_core_web_sm',\n    chunk_size=chunk_size, chunk_overlap=chunk_overlap)\n", (7165, 7328), False, 'from langchain.text_splitter import TextSplitter\n'), ((8731, 8848), 'langchain.text_splitter.TextSplitter.from_huggingface_tokenizer', 'TextSplitter.from_huggingface_tokenizer', ([], {'tokenizer': 'tokenizer', 'chunk_size': 'chunk_size', 'chunk_overlap': 'chunk_overlap'}), '(tokenizer=tokenizer, chunk_size=\n    chunk_size, chunk_overlap=chunk_overlap)\n', (8770, 8848), False, 'from langchain.text_splitter import TextSplitter\n'), ((3536, 3573), 'os.path.relpath', 'os.path.relpath', (['entry.path', 'doc_path'], {}), '(entry.path, doc_path)\n', (3551, 3573), False, 'import os\n'), ((7498, 7668), 'langchain.text_splitter.TextSplitter.from_tiktoken_encoder', 'TextSplitter.from_tiktoken_encoder', ([], {'encoding_name': "text_splitter_dict[splitter_name]['tokenizer_name_or_path']", 'chunk_size': 'chunk_size', 'chunk_overlap': 'chunk_overlap'}), "(encoding_name=text_splitter_dict[\n    splitter_name]['tokenizer_name_or_path'], chunk_size=chunk_size,\n    chunk_overlap=chunk_overlap)\n", (7532, 7668), False, 'from langchain.text_splitter import TextSplitter\n'), ((7948, 7982), 'WebUI.Server.utils.get_model_worker_config', 'get_model_worker_config', (['llm_model'], {}), '(llm_model)\n', (7971, 7982), False, 'from WebUI.Server.utils import run_in_thread_pool, get_model_worker_config\n'), ((8380, 8421), 'transformers.GPT2TokenizerFast.from_pretrained', 'GPT2TokenizerFast.from_pretrained', (['"""gpt2"""'], {}), "('gpt2')\n", (8413, 8421), False, 'from transformers import GPT2TokenizerFast\n'), ((8535, 8654), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (["text_splitter_dict[splitter_name]['tokenizer_name_or_path']"], {'trust_remote_code': '(True)'}), "(text_splitter_dict[splitter_name][\n    'tokenizer_name_or_path'], trust_remote_code=True)\n", (8564, 8654), False, 'from transformers import AutoTokenizer\n'), ((8997, 9092), 'langchain.text_splitter.TextSplitter', 'TextSplitter', ([], {'pipeline': '"""zh_core_web_sm"""', 'chunk_size': 'chunk_size', 'chunk_overlap': 'chunk_overlap'}), "(pipeline='zh_core_web_sm', chunk_size=chunk_size,\n    chunk_overlap=chunk_overlap)\n", (9009, 9092), False, 'from langchain.text_splitter import TextSplitter\n'), ((9243, 9307), 'langchain.text_splitter.TextSplitter', 'TextSplitter', ([], {'chunk_size': 'chunk_size', 'chunk_overlap': 'chunk_overlap'}), '(chunk_size=chunk_size, chunk_overlap=chunk_overlap)\n', (9255, 9307), False, 'from langchain.text_splitter import TextSplitter\n')]