import os import sys import re import json sys.path.append(os.path.dirname(os.path.dirname(__file__))) import tempfile from dotenv import load_dotenv, find_dotenv from embedding.call_embedding import get_embedding from langchain.document_loaders import UnstructuredFileLoader from langchain.document_loaders import UnstructuredMarkdownLoader from langchain.text_splitter import RecursiveCharacterTextSplitter, TextSplitter from langchain.document_loaders import PyMuPDFLoader from langchain.document_loaders import UnstructuredWordDocumentLoader from langchain.vectorstores import Chroma from langchain.schema import Document # 禁用 Pebblo 安全模块 # os.environ["PEBBLO_DISABLED"] = "1" # 新增环境变量 # 设置模型缓存目录 CACHE_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "models") # os.environ['TRANSFORMERS_CACHE'] = CACHE_DIR # 修改环境变量设置（保留一个即可） os.environ['HF_HOME'] = CACHE_DIR # 首先实现基本配置 # 原代码 DEFAULT_DB_PATH = "./knowledge_db/sanguo_characters" DEFAULT_PERSIST_PATH = "./vector_db/chroma_sanguo" class CharacterTextSplitter(TextSplitter): """专门用于处理角色JSON数据的文本分割器""" def split_text(self, text: str) -> list[str]: # 使用更健壮的正则表达式匹配每个角色的JSON数据 pattern = r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}' matches = re.finditer(pattern, text) # 将每个匹配的JSON字符串转换为文本块 chunks = [] for match in matches: try: # 解析JSON数据 char_data = json.loads(match.group()) # 检查必要字段 if 'name' not in char_data: print(f"警告：发现缺少name字段的JSON数据: {match.group()[:100]}...") continue # 处理技能数据，将stamina_cost转换为endurance_cost if 'skills' in char_data: for skill in char_data['skills']: if 'stamina_cost' in skill: skill['endurance_cost'] = skill.pop('stamina_cost') # 将JSON数据转换为易读的文本格式 char_text = f"角色：{char_data['name']}\n" char_text += f"攻击力：{char_data['attack']}\n" char_text += f"防御力：{char_data['defense']}\n" char_text += f"体力：{char_data['stamina']}\n" char_text += f"耐力：{char_data['endurance']}\n" char_text += f"法力：{char_data['mana']}\n" char_text += f"闪避：{char_data['dodge']}\n" char_text += f"速度：{char_data['speed']}\n" char_text += "技能：\n" for skill in char_data['skills']: char_text += f"- {skill['name']}：{skill['effect']}\n" if 'endurance_cost' in skill and 'mana_cost' in skill: char_text += f" 耐力消耗：{skill['endurance_cost']}，法力消耗：{skill['mana_cost']}\n" chunks.append(char_text) except json.JSONDecodeError as e: print(f"JSON解析错误: {e}") print(f"问题数据: {match.group()[:100]}...") continue except KeyError as e: print(f"缺少字段: {e}") print(f"问题数据: {match.group()[:100]}...") continue return chunks def split_documents(self, documents: list[Document]) -> list[Document]: """分割文档列表""" texts = [] metadatas = [] for doc in documents: texts.extend(self.split_text(doc.page_content)) metadatas.extend([doc.metadata] * len(self.split_text(doc.page_content))) return [Document(page_content=t, metadata=m) for t, m in zip(texts, metadatas)] def get_files(dir_path): file_list = [] for filepath, dirnames, filenames in os.walk(dir_path): for filename in filenames: file_list.append(os.path.join(filepath, filename)) return file_list def file_loader(file, loaders): if isinstance(file, tempfile._TemporaryFileWrapper): file = file.name if not os.path.isfile(file): [file_loader(os.path.join(file, f), loaders) for f in os.listdir(file)] return file_type = file.split('.')[-1].lower() if file_type == 'pdf': loaders.append(PyMuPDFLoader(file)) elif file_type == 'md': pattern = r"不存在|风控" match = re.search(pattern, file) if not match: loaders.append(UnstructuredMarkdownLoader(file)) elif file_type == 'txt': loaders.append(UnstructuredFileLoader(file)) elif file_type == 'docx': loaders.append(UnstructuredWordDocumentLoader(file)) return def create_db_info(files=DEFAULT_DB_PATH, embeddings="openai", persist_directory=DEFAULT_PERSIST_PATH): if embeddings == 'openai' or embeddings == 'm3e' or embeddings =='zhipuai': vectordb = create_db(files, persist_directory, embeddings) return "" def create_db(files=DEFAULT_DB_PATH, persist_directory=DEFAULT_PERSIST_PATH, embeddings="openai"): """ 该函数用于加载文件，切分文档，生成文档的嵌入向量，创建向量数据库。 参数: file: 存放文件的路径。 embeddings: 用于生产 Embedding 的模型 返回: vectordb: 创建的数据库。 """ if files == None: return "can't load empty file" if type(files) != list: files = [files] print(f"正在处理文件路径: {files}") loaders = [] [file_loader(file, loaders) for file in files] print(f"找到的加载器数量: {len(loaders)}") docs = [] for loader in loaders: if loader is not None: loaded_docs = loader.load() print(f"\n加载的文档数量: {len(loaded_docs)}") # 打印第一个文档的内容示例 if loaded_docs: print("\n文档内容示例:") print("-" * 50) print(loaded_docs[0].page_content[:500]) # 只打印前500个字符 print("-" * 50) print("\n文档元数据:") print(loaded_docs[0].metadata) print("-" * 50) docs.extend(loaded_docs) print(f"\n总文档数量: {len(docs)}") if len(docs) == 0: print("警告：没有找到任何文档！") return None # 使用自定义的角色文本分割器 text_splitter = CharacterTextSplitter() split_docs = text_splitter.split_documents(docs) print(f"\n分割后的文档数量: {len(split_docs)}") if len(split_docs) == 0: print("警告：分割后没有文档！") return None # 保存分割后的文档到文件 split_docs_dir = os.path.join(os.path.dirname(persist_directory), "split_docs") os.makedirs(split_docs_dir, exist_ok=True) split_docs_file = os.path.join(split_docs_dir, "split_documents.txt") with open(split_docs_file, "w", encoding="utf-8") as f: for i, doc in enumerate(split_docs, 1): f.write(f"\n文档 {i}:\n") f.write("-" * 50 + "\n") f.write(doc.page_content) f.write("\n" + "-" * 50 + "\n") print(f"\n分割后的文档已保存到: {split_docs_file}") if type(embeddings) == str: embeddings = get_embedding(embedding=embeddings) # 修正参数名称和初始化方式 vectordb = Chroma.from_documents( documents=split_docs, embedding=embeddings, persist_directory=persist_directory, collection_metadata={"hnsw:space": "cosine"} # 新增元数据配置 ) vectordb.persist() return vectordb def presit_knowledge_db(vectordb): """ 该函数用于持久化向量数据库。 参数: vectordb: 要持久化的向量数据库。 """ vectordb.persist() def load_knowledge_db(path, embeddings): """ 该函数用于加载向量数据库。 参数: path: 要加载的向量数据库路径。 embeddings: 向量数据库使用的 embedding 模型。 返回: vectordb: 加载的数据库。 """ vectordb = Chroma( persist_directory=path, embedding_function=embeddings ) return vectordb if __name__ == "__main__": create_db(embeddings="m3e")