import asyncio import os import re from configs import ( KB_ROOT_PATH, CHUNK_SIZE, OVERLAP_SIZE, ZH_TITLE_ENHANCE, logger, log_verbose, text_splitter_dict, LLM_MODELS, TEXT_SPLITTER_NAME, TEXT_SPLITTER_MAP ) import importlib from server.chat.policy_fun_iast import get_llm_model_response_async from server.knowledge_base import kb_service as tr from server.knowledge_base.TexkRank import TextRank from text_splitter import zh_title_enhance as func_zh_title_enhance import langchain.document_loaders from langchain.docstore.document import Document from langchain.text_splitter import TextSplitter from pathlib import Path from server.utils import run_in_thread_pool, get_model_worker_config import json from typing import List, Union,Dict, Tuple, Generator import chardet import time def get_split_time(f): def inner(*arg,**kwarg): s_time = time.time() res = f(*arg,**kwarg) e_time = time.time() print('切片耗时：{}秒'.format(e_time - s_time)) return res return inner 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): 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_kbs_from_folder(): return [f for f in os.listdir(KB_ROOT_PATH) if os.path.isdir(os.path.join(KB_ROOT_PATH, f))] 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()) # 路径统一为 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 LOADER_DICT = {"GCYHTMLLoader": ['.html', '.htm'], "GCYWordLoader2": ['.docx', '.doc'], # "GCYWordLoader": ['.docx'], # .doc 解析目前有点问题，暂时关掉 "MHTMLLoader": ['.mhtml'], "TextLoader": ['.md', '.txt'], "JSONLoader": [".json"], "JSONLinesLoader": [".jsonl"], "RapidOCRCSVLoader": [".csv"], # "CSVLoader": [".csv"], # "FilteredCSVLoader": [".csv"], 如果使用自定义分割csv "PyMuPDFLoader": [".pdf"], #"RapidOCRDocLoader": ['.docx', '.doc'], "RapidOCRPPTLoader": ['.ppt', '.pptx', ], "RapidOCRLoader": ['.png', '.jpg', '.jpeg', '.bmp'], "UnstructuredFileLoader": ['.eml', '.msg', '.rst', '.rtf', '.xml', '.epub', '.odt','.tsv'], "UnstructuredEmailLoader": ['.eml', '.msg'], "UnstructuredEPubLoader": ['.epub'], "ExcelLoader": ['.xlsx', '.xls', '.xlsd'], "NotebookLoader": ['.ipynb'], "UnstructuredODTLoader": ['.odt'], "PythonLoader": ['.py'], "UnstructuredRSTLoader": ['.rst'], "UnstructuredRTFLoader": ['.rtf'], "SRTLoader": ['.srt'], "TomlLoader": ['.toml'], "UnstructuredTSVLoader": ['.tsv'], "UnstructuredXMLLoader": ['.xml'], "UnstructuredPowerPointLoader": ['.ppt', '.pptx'], "EverNoteLoader": ['.enex'], } SUPPORTED_EXTS = [ext for sublist in LOADER_DICT.values() for ext in sublist] # patch json.dumps to disable ensure_ascii def _new_json_dumps(obj, **kwargs): kwargs["ensure_ascii"] = False return _origin_json_dumps(obj, **kwargs) if json.dumps is not _new_json_dumps: _origin_json_dumps = json.dumps json.dumps = _new_json_dumps class JSONLinesLoader(langchain.document_loaders.JSONLoader): ''' 行式 Json 加载器，要求文件扩展名为 .jsonl ''' def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._json_lines = True langchain.document_loaders.JSONLinesLoader = JSONLinesLoader def get_LoaderClass(file_extension): for LoaderClass, extensions in LOADER_DICT.items(): if file_extension in extensions: return LoaderClass def get_SplitterClass(file_extension): """ 根据文件类型获取文本分块器类型 """ # print('get Splitter Class', file_extension) for SplitterClass, extensions in TEXT_SPLITTER_MAP.items(): if file_extension in extensions: return SplitterClass print(f'未找到文件类型"{file_extension}"对应的切分器，使用默认值') return TEXT_SPLITTER_NAME def get_loader(loader_name: str, file_path: str, loader_kwargs: Dict = None): ''' 根据loader_name和文件路径或内容返回文档加载器。 ''' loader_kwargs = loader_kwargs or {} try: # print(loader_name) if loader_name in ["RapidOCRLoader", "FilteredCSVLoader", "GCYWordLoader","GCYWordLoader2", "GCYHTMLLoader", "RapidOCRPPTLoader", "RapidOCRCSVLoader","ExcelLoader"]: document_loaders_module = importlib.import_module('document_loaders') else: document_loaders_module = importlib.import_module('langchain.document_loaders') DocumentLoader = getattr(document_loaders_module, loader_name) except Exception as e: msg = f"为文件{file_path}查找加载器{loader_name}时出错：{e}" logger.error(f'{e.__class__.__name__}: {msg}', exc_info=e if log_verbose else None) document_loaders_module = importlib.import_module('langchain.document_loaders') DocumentLoader = getattr(document_loaders_module, "UnstructuredFileLoader") if loader_name in ["UnstructuredFileLoader", "TextLoader"]: loader_kwargs.setdefault("autodetect_encoding", True) elif loader_name == "CSVLoader": if not loader_kwargs.get("encoding"): # 如果未指定 encoding，自动识别文件编码类型，避免langchain loader 加载文件报编码错误 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 = LLM_MODELS[0], ): """ 根据参数获取特定的分词器 """ # print('spliter name', splitter_name) splitter_name = splitter_name or "SpacyTextSplitter" try: # if splitter_name == "GCYMarkdownTextSplitter": # MarkdownHeaderTextSplitter特殊判定 if splitter_name == "MarkdownTextSplitter": # MarkdownHeaderTextSplitter特殊判定 text_splitter_module = importlib.import_module('text_splitter') TextSplitter = getattr(text_splitter_module, splitter_name) headers_to_split_on = text_splitter_dict[splitter_name]['headers_to_split_on'] text_splitter = TextSplitter( headers_to_split_on=headers_to_split_on, strip_headers=False, # 不要将标题从分块文本中去掉 promote_headers=True ) else: try: ## 优先使用用户自定义的text_splitter text_splitter_module = importlib.import_module('text_splitter') TextSplitter = getattr(text_splitter_module, splitter_name) except: ## 否则使用langchain的text_splitter text_splitter_module = importlib.import_module('langchain.text_splitter') TextSplitter = getattr(text_splitter_module, splitter_name) if text_splitter_dict[splitter_name]["source"] == "tiktoken": ## 从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": ## 从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 ) elif text_splitter_dict[splitter_name]["source"] == "no_tokenizer": # IAST 0429: 目前不需要使用分词器 text_splitter = TextSplitter( 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) # If you use SpacyTextSplitter you can use GPU to do split likes Issue #1287 # text_splitter._tokenizer.max_length = 37016792 # text_splitter._tokenizer.prefer_gpu() return text_splitter 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"暂未支持的文件格式 {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 = get_SplitterClass(self.ext) def get_full_text(self) -> Dict[str, str]: """ 获取文件的全文内容，并返回文件名和全文内容的结构。 """ try: docs = self.file2docs() full_text = "".join([doc.page_content for doc in docs]) result = json.dumps( { "filename": self.filename, "full_text": full_text }, ensure_ascii=False) return result except Exception as e: logger.error(f"获取文件全文内容时出错：{e}", exc_info=e if log_verbose else None) return { "filename": self.filename, "full_text": "加载文件失败或文件内容为空" } async def get_llm_result(self) -> Dict[str, str]: """ 根据文件的全文内容，异步调用模型生成文章摘要、关键词和章节速览。 """ try: # full_text_data = self.get_full_text() # full_text = full_text_data.get("full_text", "") loop = asyncio.get_event_loop() full_text_data = await loop.run_in_executor(None, self.get_full_text) # full_text = full_text_data.get("full_text", "") try: # 将 JSON 字符串解析为字典 full_text_dict = json.loads(full_text_data) full_text = full_text_dict.get("full_text", "") except json.JSONDecodeError: print("解析 JSON 数据时出错") full_text = "" if len(full_text) > 40000: # 判断英文占比 # english_chars = re.findall(r'[a-zA-Z]', full_text) # english_ratio = len(english_chars) / len(full_text) if len(full_text) > 0 else 0 # if english_ratio > 0.9 and len(full_text) > 50000: full_text = full_text[:40000] # logger.info(f'=============文章长度{len(full_text)}') # full_text_80 = TextRank(full_text, 80) # logger.info(f'=============按80句压缩后文章长度{len(full_text_80)}') # if len(full_text_80) > 55000: # full_text_10 = TextRank(full_text_80, num_sentences=10) # logger.info(f'=============按10句压缩后文章长度{len(full_text_10)}') # full_text = full_text_10 # else: # full_text = full_text_80 else: pass # 异步调用模型 from asyncio import gather llm_time = time.time() abstract_task = get_llm_model_response_async( strategy_name="gen_abstract", llm_model_name=LLM_MODELS[1], template_prompt_name="gen_abstract", prompt_param_dict={"context": full_text}, temperature=0.7, max_tokens=4096 ) keywords_task = get_llm_model_response_async( strategy_name="gen_keywords", llm_model_name=LLM_MODELS[1], template_prompt_name="gen_keywords", prompt_param_dict={"context": full_text}, temperature=0.7, max_tokens=512 ) paragraph_task = get_llm_model_response_async( strategy_name="gen_paragraph", llm_model_name=LLM_MODELS[0], template_prompt_name="gen_paragraph", prompt_param_dict={"context": full_text}, temperature=0.7, max_tokens=8192 ) # 并行执行任务 article_abstract, article_keywords, article_paragraph = await gather( abstract_task, keywords_task, paragraph_task ) logger.info(f'生成导读用时：{time.time() - llm_time}') return { "filename": self.filename, "full_text": full_text, "article_abstract": article_abstract, "article_keywords": article_keywords, "article_paragraph": article_paragraph } except Exception as e: logger.error(f"生成LLM结果时出错：{e}", exc_info=e if log_verbose else None) return { "filename": self.filename, "article_abstract": "生成摘要失败", "article_keywords": "生成关键词失败", "article_paragraph": "生成章节速览失败" } def file2docs(self, refresh: bool = False): if self.docs is None or refresh: try: logger.info(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() except Exception as e: if self.document_loader_name == 'GCYWordLoader': loader = get_loader(loader_name='GCYWordLoader2', file_path=self.filepath, loader_kwargs=self.loader_kwargs) else: logger.error(f"加载文件 {self.filepath} 时出错：{e}", exc_info=e if log_verbose else None) self.docs = loader.load() return self.docs @get_split_time 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) # debug 0429 # print('docs2texts',docs ) if not docs: return [] if text_splitter is None: self.text_splitter_name = get_SplitterClass(self.ext) text_splitter = make_text_splitter(splitter_name=self.text_splitter_name, chunk_size=chunk_size, chunk_overlap=chunk_overlap) # if self.text_splitter_name == "GCYMarkdownTextSplitter": if self.text_splitter_name == "MarkdownTextSplitter": doc_source = (docs[0].metadata)["source"] docs = text_splitter.split_markdown_text(docs[0].page_content, doc_source) else: docs = text_splitter.split_documents(docs) if not docs: return [] # 检查切分好的文档是否有'h1'标题字段，如果没有，就加上。为之后入库其它有h1的文件做准备 if 'h1' not in docs[0].metadata: for doc in docs: doc.metadata['h1'] = '' print(f"文档切分示例：{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: ''' 利用多线程批量将磁盘文件转化成langchain Document. 如果传入参数是Tuple，形式为(filename, kb_name) 生成器返回值为 status, (kb_name, file_name, docs | error) ''' 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"从文件 {file.kb_name}/{file.filename} 加载文档时出错：{e}" logger.error(f'{e.__class__.__name__}: {msg}', exc_info=e if log_verbose else None) 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 if __name__ == "__main__": from pprint import pprint kb_file = KnowledgeFile( filename="/home/congyin/Code/Project_Langchain_0814/Langchain-Chatchat/knowledge_base/csv1/content/gm.csv", knowledge_base_name="samples") # kb_file.text_splitter_name = "RecursiveCharacterTextSplitter" docs = kb_file.file2docs() # pprint(docs[-1])