from configs import ( EMBEDDING_MODEL, DEFAULT_VS_TYPE, ZH_TITLE_ENHANCE, CHUNK_SIZE, OVERLAP_SIZE, logger, log_verbose ) from server.knowledge_base.utils import ( get_file_path, list_kbs_from_folder, list_files_from_folder, files2docs_in_thread, KnowledgeFile ) from server.knowledge_base.kb_service.base import KBServiceFactory from server.db.models.conversation_model import ConversationModel from server.db.models.message_model import MessageModel from server.db.repository.knowledge_file_repository import add_file_to_db # ensure Models are imported from server.db.repository.knowledge_metadata_repository import add_summary_to_db from server.db.base import Base, engine from server.db.session import session_scope import os from dateutil.parser import parse from typing import Literal, List def create_tables(): Base.metadata.create_all(bind=engine) def reset_tables(): Base.metadata.drop_all(bind=engine) create_tables() def import_from_db( sqlite_path: str = None, # csv_path: str = None, ) -> bool: """ 在知识库与向量库无变化的情况下，从备份数据库中导入数据到 info.db。 适用于版本升级时，info.db 结构变化，但无需重新向量化的情况。 请确保两边数据库表名一致，需要导入的字段名一致 当前仅支持 sqlite """ import sqlite3 as sql from pprint import pprint models = list(Base.registry.mappers) try: con = sql.connect(sqlite_path) con.row_factory = sql.Row cur = con.cursor() tables = [x["name"] for x in cur.execute("select name from sqlite_master where type='table'").fetchall()] for model in models: table = model.local_table.fullname if table not in tables: continue print(f"processing table: {table}") with session_scope() as session: for row in cur.execute(f"select * from {table}").fetchall(): data = {k: row[k] for k in row.keys() if k in model.columns} if "create_time" in data: data["create_time"] = parse(data["create_time"]) pprint(data) session.add(model.class_(**data)) con.close() return True except Exception as e: print(f"无法读取备份数据库：{sqlite_path}。错误信息：{e}") return False def file_to_kbfile(kb_name: str, files: List[str]) -> List[KnowledgeFile]: kb_files = [] for file in files: try: kb_file = KnowledgeFile(filename=file, knowledge_base_name=kb_name) kb_files.append(kb_file) except Exception as e: msg = f"{e}，已跳过" logger.error(f'{e.__class__.__name__}: {msg}', exc_info=e if log_verbose else None) return kb_files def folder2db( kb_names: List[str], mode: Literal["recreate_vs", "update_in_db", "increment"], vs_type: Literal["faiss", "milvus", "pg", "chromadb"] = DEFAULT_VS_TYPE, embed_model: str = EMBEDDING_MODEL, chunk_size: int = CHUNK_SIZE, chunk_overlap: int = OVERLAP_SIZE, zh_title_enhance: bool = ZH_TITLE_ENHANCE, ): """ use existed files in local folder to populate database and/or vector store. set parameter `mode` to: recreate_vs: recreate all vector store and fill info to database using existed files in local folder fill_info_only(disabled): do not create vector store, fill info to db using existed files only update_in_db: update vector store and database info using local files that existed in database only increment: create vector store and database info for local files that not existed in database only """ def files2vs(kb_name: str, kb_files: List[KnowledgeFile]): for success, result in files2docs_in_thread(kb_files, chunk_size=chunk_size, chunk_overlap=chunk_overlap, zh_title_enhance=zh_title_enhance): if success: _, filename, docs = result print(f"正在将 {kb_name}/{filename} 添加到向量库，共包含{len(docs)}条文档") kb_file = KnowledgeFile(filename=filename, knowledge_base_name=kb_name) kb_file.splited_docs = docs kb.add_doc(kb_file=kb_file, not_refresh_vs_cache=True) else: print(result) kb_names = kb_names or list_kbs_from_folder() for kb_name in kb_names: kb = KBServiceFactory.get_service(kb_name, vs_type, embed_model) if not kb.exists(): kb.create_kb() # 清除向量库，从本地文件重建 if mode == "recreate_vs": kb.clear_vs() kb.create_kb() kb_files = file_to_kbfile(kb_name, list_files_from_folder(kb_name)) files2vs(kb_name, kb_files) kb.save_vector_store() # # 不做文件内容的向量化，仅将文件元信息存到数据库 # # 由于现在数据库存了很多与文本切分相关的信息，单纯存储文件信息意义不大，该功能取消。 # elif mode == "fill_info_only": # files = list_files_from_folder(kb_name) # kb_files = file_to_kbfile(kb_name, files) # for kb_file in kb_files: # add_file_to_db(kb_file) # print(f"已将 {kb_name}/{kb_file.filename} 添加到数据库") # 以数据库中文件列表为基准，利用本地文件更新向量库 elif mode == "update_in_db": files = kb.list_files() kb_files = file_to_kbfile(kb_name, files) files2vs(kb_name, kb_files) kb.save_vector_store() # 对比本地目录与数据库中的文件列表，进行增量向量化 elif mode == "increment": db_files = kb.list_files() folder_files = list_files_from_folder(kb_name) files = list(set(folder_files) - set(db_files)) kb_files = file_to_kbfile(kb_name, files) files2vs(kb_name, kb_files) kb.save_vector_store() else: print(f"unsupported migrate mode: {mode}") def prune_db_docs(kb_names: List[str]): """ delete docs in database that not existed in local folder. it is used to delete database docs after user deleted some doc files in file browser """ for kb_name in kb_names: kb = KBServiceFactory.get_service_by_name(kb_name) if kb is not None: files_in_db = kb.list_files() files_in_folder = list_files_from_folder(kb_name) files = list(set(files_in_db) - set(files_in_folder)) kb_files = file_to_kbfile(kb_name, files) for kb_file in kb_files: kb.delete_doc(kb_file, not_refresh_vs_cache=True) print(f"success to delete docs for file: {kb_name}/{kb_file.filename}") kb.save_vector_store() def prune_folder_files(kb_names: List[str]): """ delete doc files in local folder that not existed in database. it is used to free local disk space by delete unused doc files. """ for kb_name in kb_names: kb = KBServiceFactory.get_service_by_name(kb_name) if kb is not None: files_in_db = kb.list_files() files_in_folder = list_files_from_folder(kb_name) files = list(set(files_in_folder) - set(files_in_db)) for file in files: os.remove(get_file_path(kb_name, file)) print(f"success to delete file: {kb_name}/{file}")