import os import pathlib from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker, joinedload from contextlib import contextmanager from sqlalchemy.orm.attributes import instance_state from src import config from src.utils import logger from server.models.kb_models import Base, KnowledgeDatabase, KnowledgeFile, KnowledgeNode class KBDBManager: """知识库数据库管理器""" def __init__(self): self.db_path = os.path.join(config.save_dir, "data", "knowledge.db") self.ensure_db_dir() # 创建SQLAlchemy引擎 self.engine = create_engine(f"sqlite:///{self.db_path}") # 创建会话工厂 self.Session = sessionmaker(bind=self.engine) # 确保表存在 self.create_tables() def ensure_db_dir(self): """确保数据库目录存在""" db_dir = os.path.dirname(self.db_path) pathlib.Path(db_dir).mkdir(parents=True, exist_ok=True) def create_tables(self): """创建数据库表""" Base.metadata.create_all(self.engine) @contextmanager def get_session(self): """获取数据库会话的上下文管理器""" session = self.Session() try: yield session session.commit() except Exception as e: session.rollback() logger.error(f"数据库操作失败: {e}") raise finally: session.close() def _detach_safely(self, obj): """安全地分离对象，确保属性已加载""" if obj is None: return None # 确保主键已加载 if hasattr(obj, 'id'): _ = obj.id if hasattr(obj, 'db_id'): _ = obj.db_id # 根据需要添加其他必须预加载的属性 return obj # 知识库操作方法 def get_all_databases(self): """获取所有知识库""" with self.get_session() as session: # 使用eager loading加载关联的files databases = session.query(KnowledgeDatabase).options( joinedload(KnowledgeDatabase.files) ).all() # 转换为字典并返回，避免后续延迟加载 return [self._to_dict_safely(db) for db in databases] def get_database_by_id(self, db_id): """根据ID获取知识库""" with self.get_session() as session: # 使用eager loading加载关联的files db = session.query(KnowledgeDatabase).options( joinedload(KnowledgeDatabase.files).joinedload(KnowledgeFile.nodes) ).filter_by(db_id=db_id).first() # 转换为字典并返回，避免后续延迟加载 return self._to_dict_safely(db) if db else None def _to_dict_safely(self, obj): """安全地将对象转换为字典，避免延迟加载问题""" if hasattr(obj, 'to_dict'): return obj.to_dict() return obj def create_database(self, db_id, name, description, embed_model=None, dimension=None, metadata=None): """创建知识库""" with self.get_session() as session: db = KnowledgeDatabase( db_id=db_id, name=name, description=description, embed_model=embed_model, dimension=dimension, meta_info=metadata or {} # 存储到meta_info字段 ) session.add(db) session.flush() # 立即写入数据库，获取ID # 手动将必要的数据加载到内存中 db_dict = { "db_id": db_id, "name": name, "description": description, "embed_model": embed_model, "dimension": dimension, "metadata": metadata or {}, # 返回时使用metadata键 "files": {} } return db_dict def delete_database(self, db_id): """删除知识库""" with self.get_session() as session: db = session.query(KnowledgeDatabase).filter_by(db_id=db_id).first() if db: session.delete(db) return True return False def update_database(self, db_id, name, description): """更新知识库信息""" with self.get_session() as session: db = session.query(KnowledgeDatabase).filter_by(db_id=db_id).first() if not db: raise ValueError(f"数据库 {db_id} 不存在") # 更新字段 db.name = name db.description = description session.commit() # 返回更新后的数据库信息 return self._to_dict_safely(db) # 文件操作方法 def add_file(self, db_id, file_id, filename, path, file_type, status="waiting"): """添加文件""" with self.get_session() as session: file = KnowledgeFile( file_id=file_id, database_id=db_id, filename=filename, path=path, file_type=file_type, status=status ) session.add(file) session.flush() # 返回字典而非对象，避免会话关闭后的延迟加载问题 return { "file_id": file_id, "filename": filename, "path": path, "type": file_type, "status": status, "created_at": file.created_at.timestamp() if file.created_at else None, "nodes": [] } def update_file_status(self, file_id, status): """更新文件状态""" with self.get_session() as session: file = session.query(KnowledgeFile).filter_by(file_id=file_id).first() if file: file.status = status return True return False def delete_file(self, file_id): """删除文件""" with self.get_session() as session: file = session.query(KnowledgeFile).filter_by(file_id=file_id).first() if file: session.delete(file) return True return False def get_files_by_database(self, db_id): """获取知识库下的所有文件""" with self.get_session() as session: files = session.query(KnowledgeFile).options( joinedload(KnowledgeFile.nodes) ).filter_by(database_id=db_id).all() return [self._to_dict_safely(file) for file in files] def get_file_by_id(self, file_id): """根据ID获取文件""" with self.get_session() as session: file = session.query(KnowledgeFile).options( joinedload(KnowledgeFile.nodes) ).filter_by(file_id=file_id).first() return self._to_dict_safely(file) if file else None # 知识块操作方法 def add_node(self, file_id, text, hash_value=None, start_char_idx=None, end_char_idx=None, metadata=None): """添加知识块""" with self.get_session() as session: node = KnowledgeNode( file_id=file_id, text=text, hash=hash_value, start_char_idx=start_char_idx, end_char_idx=end_char_idx, meta_info=metadata or {} ) session.add(node) session.flush() # 返回字典而非对象，避免会话关闭后的延迟加载问题 return { "id": node.id, "file_id": file_id, "text": text, "hash": hash_value, "start_char_idx": start_char_idx, "end_char_idx": end_char_idx, "metadata": metadata or {} } def get_nodes_by_file(self, file_id): """获取文件下的所有知识块""" with self.get_session() as session: nodes = session.query(KnowledgeNode).filter_by(file_id=file_id).all() return [self._to_dict_safely(node) for node in nodes] def get_nodes_by_filter(self, file_id=None, search_text=None, limit=100): """根据条件筛选知识块""" with self.get_session() as session: query = session.query(KnowledgeNode) if file_id: query = query.filter_by(file_id=file_id) if search_text: query = query.filter(KnowledgeNode.text.like(f"%{search_text}%")) nodes = query.limit(limit).all() return [self._to_dict_safely(node) for node in nodes] # 创建全局知识库数据库管理器实例 kb_db_manager = KBDBManager()