152 lines
6.2 KiB
Python
152 lines
6.2 KiB
Python
from configs import (EMBEDDING_MODEL, DEFAULT_VS_TYPE, ZH_TITLE_ENHANCE,
|
|
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, SupportedVSType
|
|
from server.db.repository.knowledge_file_repository import add_file_to_db
|
|
from server.db.base import Base, engine
|
|
import os
|
|
from typing import Literal, Any, List
|
|
|
|
|
|
def create_tables():
|
|
Base.metadata.create_all(bind=engine)
|
|
|
|
|
|
def reset_tables():
|
|
Base.metadata.drop_all(bind=engine)
|
|
create_tables()
|
|
|
|
|
|
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_name: str,
|
|
mode: Literal["recreate_vs", "fill_info_only", "update_in_db", "increament"],
|
|
vs_type: Literal["faiss", "milvus", "pg", "chromadb"] = DEFAULT_VS_TYPE,
|
|
embed_model: str = EMBEDDING_MODEL,
|
|
chunk_size: int = -1,
|
|
chunk_overlap: int = -1,
|
|
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: 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
|
|
increament: create vector store and database info for local files that not existed in database only
|
|
'''
|
|
kb = KBServiceFactory.get_service(kb_name, vs_type, embed_model)
|
|
kb.create_kb()
|
|
|
|
if mode == "recreate_vs":
|
|
files_count = kb.count_files()
|
|
print(f"知识库 {kb_name} 中共有 {files_count} 个文档。\n即将清除向量库。")
|
|
kb.clear_vs()
|
|
files_count = kb.count_files()
|
|
print(f"清理后，知识库 {kb_name} 中共有 {files_count} 个文档。")
|
|
|
|
kb_files = file_to_kbfile(kb_name, list_files_from_folder(kb_name))
|
|
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.add_doc(kb_file=kb_file, docs=docs, not_refresh_vs_cache=True)
|
|
else:
|
|
print(result)
|
|
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)
|
|
|
|
for kb_file in kb_files:
|
|
kb.update_doc(kb_file, not_refresh_vs_cache=True)
|
|
kb.save_vector_store()
|
|
elif mode == "increament":
|
|
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)
|
|
|
|
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} 添加到向量库")
|
|
kb_file = KnowledgeFile(filename=filename, knowledge_base_name=kb_name)
|
|
kb.add_doc(kb_file=kb_file, docs=docs, not_refresh_vs_cache=True)
|
|
else:
|
|
print(result)
|
|
kb.save_vector_store()
|
|
else:
|
|
print(f"unspported migrate mode: {mode}")
|
|
|
|
|
|
def recreate_all_vs(
|
|
vs_type: Literal["faiss", "milvus", "pg", "chromadb"] = DEFAULT_VS_TYPE,
|
|
embed_mode: str = EMBEDDING_MODEL,
|
|
**kwargs: Any,
|
|
):
|
|
'''
|
|
used to recreate a vector store or change current vector store to another type or embed_model
|
|
'''
|
|
for kb_name in list_kbs_from_folder():
|
|
folder2db(kb_name, "recreate_vs", vs_type, embed_mode, **kwargs)
|
|
|
|
|
|
def prune_db_files(kb_name: str):
|
|
'''
|
|
delete files in database that not existed in local folder.
|
|
it is used to delete database files after user deleted some doc files in file browser
|
|
'''
|
|
kb = KBServiceFactory.get_service_by_name(kb_name)
|
|
if kb.exists():
|
|
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)
|
|
kb.save_vector_store()
|
|
return kb_files
|
|
|
|
def prune_folder_files(kb_name: str):
|
|
'''
|
|
delete doc files in local folder that not existed in database.
|
|
is is used to free local disk space by delete unused doc files.
|
|
'''
|
|
kb = KBServiceFactory.get_service_by_name(kb_name)
|
|
if kb.exists():
|
|
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))
|
|
return files
|