import os import json import time from src.utils import hashstr, setup_logger, is_text_pdf from src.models.embedding import get_embedding_model logger = setup_logger("DataBaseManager") class DataBaseManager: def __init__(self, config=None) -> None: self.config = config self.database_path = os.path.join(config.save_dir, "data", "database.json") self.embed_model = get_embedding_model(config) if self.config.enable_knowledge_base: from src.core.knowledgebase import KnowledgeBase self.knowledge_base = KnowledgeBase(config, self.embed_model) if self.config.enable_knowledge_graph: from src.core.graphbase import GraphDatabase self.graph_base = GraphDatabase(self.config, self.embed_model) self.graph_base.start() else: self.graph_base = None self.data = {"databases": [], "graph": {}} self._load_databases() self._update_database() def _load_databases(self): """将数据库的信息保存到本地的文件里面""" if not os.path.exists(self.database_path): return with open(self.database_path, "r") as f: data = json.load(f) self.data = { "databases": [DataBaseLite(**db) for db in data["databases"]], "graph": data["graph"] } # 检查所有文件，如果出现状态是 processing 的，那么设置为 failed for db in self.data["databases"]: for file in db.files: if file["status"] == "processing" or file["status"] == "waiting": file["status"] = "failed" def _save_databases(self): """将数据库的信息保存到本地的文件里面""" self._update_database() with open(self.database_path, "w+") as f: json.dump({ "databases": [db.to_dict() for db in self.data["databases"]], "graph": self.data["graph"] }, f, ensure_ascii=False, indent=4) def _update_database(self): self.id2db = {db.db_id: db for db in self.data["databases"]} self.name2db = {db.name: db for db in self.data["databases"]} self.metaname2db = {db.metaname: db for db in self.data["databases"]} def get_databases(self): self._update_database() knowledge_base_collections = self.knowledge_base.get_collection_names() if len(self.data["databases"]) != len(knowledge_base_collections): logger.warning(f"Database number not match, {knowledge_base_collections}") for db in self.data["databases"]: db.update(self.knowledge_base.get_collection_info(db.metaname)) return {"databases": [db.to_dict() for db in self.data["databases"]]} def get_graph(self): if self.config.enable_knowledge_graph: self.data["graph"].update(self.graph_base.get_database_info("neo4j")) return {"graph": self.data["graph"]} else: return {"message": "Graph base not enabled", "graph": {}} def create_database(self, database_name, description, db_type, dimension): from src.config import EMBED_MODEL_INFO dimension = dimension or EMBED_MODEL_INFO[self.config.embed_model]["dimension"] new_database = DataBaseLite(database_name, description, db_type, embed_model=self.config.embed_model, dimension=dimension) self.knowledge_base.add_collection(new_database.metaname, dimension) self.data["databases"].append(new_database) self._save_databases() return self.get_databases() def add_files(self, db_id, files, params=None): db = self.get_kb_by_id(db_id) if db.embed_model != self.config.embed_model: logger.error(f"Embed model not match, {db.embed_model} != {self.config.embed_model}") return {"message": f"Embed model not match, cur: {self.config.embed_model}", "status": "failed"} # Preprocessing the files to the queue new_files = [] for file in files: new_file = { "file_id": "file_" + hashstr(file + str(time.time())), "filename": os.path.basename(file), "path": file, "type": file.split(".")[-1], "status": "waiting", "created_at": time.time() } db.files.append(new_file) new_files.append(new_file) from src.core.indexing import chunk for new_file in new_files: file_id = new_file["file_id"] idx = self.get_idx_by_fileid(db, file_id) db.files[idx]["status"] = "processing" try: nodes = chunk(new_file["path"], params=params) self.knowledge_base.add_documents( docs=[node.text for node in nodes], collection_name=db.metaname, file_id=file_id) idx = self.get_idx_by_fileid(db, file_id) db.files[idx]["status"] = "done" except Exception as e: logger.error(f"Failed to add documents to collection {db.metaname}, {e}") idx = self.get_idx_by_fileid(db, file_id) db.files[idx]["status"] = "failed" self._save_databases() return {"message": "全部解析完成", "status": "success"} def get_database_info(self, db_id): db = self.get_kb_by_id(db_id) if db is None: return None else: db.update(self.knowledge_base.get_collection_info(db.metaname)) return db.to_dict() def read_text(self, file, params=None): support_format = [".pdf", ".txt", ".md"] assert os.path.exists(file), "File not found" logger.info(f"Try to read file {file}") if not os.path.isfile(file): logger.error(f"Directory not supported now!") raise NotImplementedError("Directory not supported now!") if file.endswith(".pdf"): if is_text_pdf(file): from src.core.filereader import pdfreader return pdfreader(file) else: from src.plugins import pdf2txt return pdf2txt(file, return_text=True) elif file.endswith(".txt") or file.endswith(".md"): from src.core.filereader import plainreader return plainreader(file) else: logger.error(f"File format not supported, only support {support_format}") raise Exception(f"File format not supported, only support {support_format}") def delete_file(self, db_id, file_id): db = self.get_kb_by_id(db_id) file_idx_to_delete = self.get_idx_by_fileid(db, file_id) self.knowledge_base.client.delete( collection_name=db.metaname, filter=f"file_id == '{file_id}'"), del db.files[file_idx_to_delete] self._save_databases() def get_file_info(self, db_id, file_id): db = self.get_kb_by_id(db_id) if db is None: return {"message": "database not found"}, 404 lines = self.knowledge_base.client.query( collection_name=db.metaname, filter=f"file_id == '{file_id}'", output_fields=["id", "text", "file_id", "hash"] ) return {"lines": lines} def chunking(self, text, params=None): chunk_method = params.get("chunk_method", "fixed") chunk_size = params.get("chunk_size", 500) """将文本切分成固定大小的块""" chunks = [] for i in range(0, len(text), chunk_size): chunks.append(text[i:i + chunk_size]) return chunks def delete_database(self, db_id): db = self.get_kb_by_id(db_id) if db is None: return {"message": "database not found"}, 404 self.knowledge_base.client.drop_collection(db.metaname) self.data["databases"] = [d for d in self.data["databases"] if d.db_id != db_id] self._save_databases() return {"message": "删除成功"} def get_kb_by_id(self, db_id): for db in self.data["databases"]: if db.db_id == db_id: return db return None def get_idx_by_fileid(self, db, file_id): for idx, f in enumerate(db.files): if f["file_id"] == file_id: return idx class DataBaseLite: def __init__(self, name, description, db_type, dimension=None, **kwargs) -> None: self.name = name self.description = description self.db_type = db_type self.dimension = dimension self.db_id = kwargs.get("db_id", hashstr(name)) self.metaname = kwargs.get("metaname", f"{db_type[:1]}{hashstr(name)}") self.metadata = kwargs.get("metaname", {}) self.files = kwargs.get("files", []) self.embed_model = kwargs.get("embed_model", None) def id2file(self, file_id): for f in self.files: if f["file_id"] == file_id: return f return None def update(self, metadata): self.metadata = metadata def to_dict(self): return { "name": self.name, "description": self.description, "db_type": self.db_type, "db_id": self.db_id, "embed_model": self.embed_model, "metaname": self.metaname, "metadata": self.metadata, "files": self.files, "dimension": self.dimension } def to_json(self): return json.dumps(self.to_dict(), ensure_ascii=False) def __str__(self): return self.to_json()