# rag_manager.py from __future__ import annotations import os import asyncio import uuid from pkg.rag.knowledge.services.parser import FileParser from pkg.rag.knowledge.services.chunker import Chunker from pkg.rag.knowledge.services.database import create_db_and_tables, SessionLocal, KnowledgeBase, File, Chunk from pkg.core import app from pkg.rag.knowledge.services.embedder import Embedder from pkg.rag.knowledge.services.retriever import Retriever from pkg.rag.knowledge.services.chroma_manager import ChromaIndexManager class RAGManager: ap: app.Application def __init__(self, ap: app.Application): self.ap = ap self.chroma_manager = ChromaIndexManager() self.parser = FileParser() self.chunker = Chunker() # Initialize Embedder with targeted model type and name self.embedder = Embedder( model_type='third_party_api', model_name_key='bge-m3', chroma_manager=self.chroma_manager ) self.retriever = Retriever( model_type='third_party_api', model_name_key='bge-m3', chroma_manager=self.chroma_manager ) async def initialize_rag_system(self): """Initializes the RAG system by creating database tables.""" await asyncio.to_thread(create_db_and_tables) async def create_knowledge_base( self, kb_name: str, kb_description: str, embedding_model_uuid: str = '', top_k: int = 5 ): """ Creates a new knowledge base if it doesn't already exist. """ try: if not kb_name: raise ValueError('Knowledge base name must be set while creating.') def _create_kb_sync(): session = SessionLocal() try: kb = session.query(KnowledgeBase).filter_by(name=kb_name).first() if not kb: id = str(uuid.uuid4()) new_kb = KnowledgeBase( name=kb_name, description=kb_description, embedding_model_uuid=embedding_model_uuid, top_k=top_k, id=id, ) session.add(new_kb) session.commit() session.refresh(new_kb) self.ap.logger.info(f"Knowledge Base '{kb_name}' created.") return new_kb.id else: self.ap.logger.info(f"Knowledge Base '{kb_name}' already exists.") except Exception as e: session.rollback() self.ap.logger.error(f"Error in _create_kb_sync for '{kb_name}': {str(e)}", exc_info=True) raise finally: session.close() return await asyncio.to_thread(_create_kb_sync) except Exception as e: self.ap.logger.error(f"Error creating knowledge base '{kb_name}': {str(e)}", exc_info=True) raise async def get_all_knowledge_bases(self): """ Retrieves all knowledge bases from the database. """ try: def _get_all_kbs_sync(): session = SessionLocal() try: return session.query(KnowledgeBase).all() finally: session.close() return await asyncio.to_thread(_get_all_kbs_sync) except Exception as e: self.ap.logger.error(f'Error retrieving knowledge bases: {str(e)}', exc_info=True) return [] async def get_knowledge_base_by_id(self, kb_id: str): """ Retrieves a specific knowledge base by its ID. """ try: def _get_kb_sync(kb_id_param): session = SessionLocal() try: return session.query(KnowledgeBase).filter_by(id=kb_id_param).first() finally: session.close() return await asyncio.to_thread(_get_kb_sync, kb_id) except Exception as e: self.ap.logger.error(f'Error retrieving knowledge base with ID {kb_id}: {str(e)}', exc_info=True) return None async def get_files_by_knowledge_base(self, kb_id: str): """ Retrieves files associated with a specific knowledge base by querying the File table directly. """ try: def _get_files_sync(kb_id_param): session = SessionLocal() try: return session.query(File).filter_by(kb_id=kb_id_param).all() finally: session.close() return await asyncio.to_thread(_get_files_sync, kb_id) except Exception as e: self.ap.logger.error(f'Error retrieving files for knowledge base ID {kb_id}: {str(e)}', exc_info=True) return [] async def get_all_files(self): """ Retrieves all files stored in the database, regardless of their association with any specific knowledge base. """ try: def _get_all_files_sync(): session = SessionLocal() try: return session.query(File).all() finally: session.close() return await asyncio.to_thread(_get_all_files_sync) except Exception as e: self.ap.logger.error(f'Error retrieving all files: {str(e)}', exc_info=True) return [] async def store_data(self, file_path: str, kb_id: str, file_type: str, file_id: str = None): """ Parses, chunks, embeds, and stores data from a given file into the RAG system. Associates the file with a knowledge base using kb_id in the File table. """ self.ap.logger.info(f'Starting data storage process for file: {file_path}') session = SessionLocal() file_obj = None try: kb = session.query(KnowledgeBase).filter_by(id=kb_id).first() if not kb: self.ap.logger.info(f'Knowledge Base "{kb_id}" does not exist. ') self.ap.logger.info(f'Created Knowledge Base with ID: {kb_id}') else: self.ap.logger.info(f"Knowledge Base '{kb_id}' already exists.") file_name = os.path.basename(file_path) text = await self.parser.parse(file_path) if not text: self.ap.logger.warning(f'No text extracted from file {file_path}. ') return chunks_texts = await self.chunker.chunk(text) self.ap.logger.info(f"Chunked file '{file_name}' into {len(chunks_texts)} chunks.") await self.embedder.embed_and_store(file_id=file_id, chunks=chunks_texts) self.ap.logger.info(f'Data storage process completed for file: {file_path}') except Exception as e: session.rollback() self.ap.logger.error(f'Error in store_data for file {file_path}: {str(e)}', exc_info=True) raise finally: if file_id: file_obj = session.query(File).filter_by(id=file_id).first() if file_obj: file_obj.status = 1 session.close() async def retrieve_data(self, query: str): """ Retrieves relevant data chunks based on a given query using the configured retriever. """ self.ap.logger.info(f"Starting data retrieval process for query: '{query}'") try: retrieved_chunks = await self.retriever.retrieve(query) self.ap.logger.info(f'Successfully retrieved {len(retrieved_chunks)} chunks for query.') return retrieved_chunks except Exception as e: self.ap.logger.error(f"Error in retrieve_data for query '{query}': {str(e)}", exc_info=True) return [] async def delete_data_by_file_id(self, file_id: str): """ Deletes all data associated with a specific file ID, including its chunks and vectors, and the file record itself. """ self.ap.logger.info(f'Starting data deletion process for file_id: {file_id}') session = SessionLocal() try: # delete vectors await asyncio.to_thread(self.chroma_manager.delete_by_file_id_sync, file_id) self.ap.logger.info(f'Deleted embeddings from ChromaDB for file_id: {file_id}') chunks_to_delete = session.query(Chunk).filter_by(file_id=file_id).all() for chunk in chunks_to_delete: session.delete(chunk) self.ap.logger.info(f'Deleted {len(chunks_to_delete)} chunk records for file_id: {file_id}') file_to_delete = session.query(File).filter_by(id=file_id).first() if file_to_delete: session.delete(file_to_delete) try: await self.ap.storage_mgr.storage_provider.delete(file_id) except Exception as e: self.ap.logger.error( f'Error deleting file from storage for file_id {file_id}: {str(e)}', exc_info=True ) self.ap.logger.info(f'Deleted file record for file_id: {file_id}') else: self.ap.logger.warning( f'File with ID {file_id} not found in database. Skipping deletion of file record.' ) session.commit() self.ap.logger.info(f'Successfully completed data deletion for file_id: {file_id}') except Exception as e: session.rollback() self.ap.logger.error(f'Error deleting data for file_id {file_id}: {str(e)}', exc_info=True) raise finally: session.close() async def delete_kb_by_id(self, kb_id: str): """ Deletes a knowledge base and all associated files, chunks, and vectors. This involves querying for associated files and then deleting them. """ self.ap.logger.info(f'Starting deletion of knowledge base with ID: {kb_id}') session = SessionLocal() try: kb_to_delete = session.query(KnowledgeBase).filter_by(id=kb_id).first() if not kb_to_delete: self.ap.logger.warning(f'Knowledge Base with ID {kb_id} not found.') return files_to_delete = session.query(File).filter_by(kb_id=kb_id).all() session.close() for file_obj in files_to_delete: try: await self.delete_data_by_file_id(file_obj.id) except Exception as file_del_e: self.ap.logger.error(f'Failed to delete file ID {file_obj.id} during KB deletion: {file_del_e}') session = SessionLocal() try: kb_final_delete = session.query(KnowledgeBase).filter_by(id=kb_id).first() if kb_final_delete: session.delete(kb_final_delete) session.commit() self.ap.logger.info(f'Successfully deleted knowledge base with ID: {kb_id}') else: self.ap.logger.warning( f'Knowledge Base with ID {kb_id} not found after file deletion, skipping KB deletion.' ) except Exception as kb_del_e: session.rollback() self.ap.logger.error(f'Error deleting KnowledgeBase record for ID {kb_id}: {kb_del_e}', exc_info=True) raise finally: session.close() except Exception as e: # 如果在最初获取 KB 或文件列表时出错 if session.is_active: session.rollback() self.ap.logger.error( f'Error during overall knowledge base deletion for ID {kb_id}: {str(e)}', exc_info=True ) raise finally: if session.is_active: session.close() async def get_file_content_by_file_id(self, file_id: str) -> str: file_bytes = await self.ap.storage_mgr.storage_provider.load(file_id) _, ext = os.path.splitext(file_id.lower()) ext = ext.lstrip('.') try: text = file_bytes.decode('utf-8') except UnicodeDecodeError: return '[非文本文件或编码无法识别]' if ext in ['txt', 'md', 'csv', 'log', 'py', 'html']: return text else: return f'[未知类型: .{ext}]' async def relate_file_id_with_kb(self, knowledge_base_uuid: str, file_id: str) -> None: """ Associates a file with a knowledge base by updating the kb_id in the File table. """ self.ap.logger.info(f'Associating file ID {file_id} with knowledge base UUID {knowledge_base_uuid}') session = SessionLocal() try: # 查询知识库是否存在 kb = session.query(KnowledgeBase).filter_by(id=knowledge_base_uuid).first() if not kb: self.ap.logger.error(f'Knowledge Base with UUID {knowledge_base_uuid} not found.') return if not await self.ap.storage_mgr.storage_provider.exists(file_id): self.ap.logger.error(f'File with ID {file_id} does not exist.') return self.ap.logger.info(f'File with ID {file_id} exists, proceeding with association.') # add new file record file_to_update = File( id=file_id, kb_id=kb.id, file_name=file_id, path=os.path.join('data', 'storage', file_id), file_type=os.path.splitext(file_id)[1].lstrip('.'), status=0, ) session.add(file_to_update) session.commit() self.ap.logger.info( f'Successfully associated file ID {file_id} with knowledge base UUID {knowledge_base_uuid}' ) except Exception as e: session.rollback() self.ap.logger.error( f'Error associating file ID {file_id} with knowledge base UUID {knowledge_base_uuid}: {str(e)}', exc_info=True, ) finally: # 进行文件解析 try: await self.store_data( file_path=os.path.join('data', 'storage', file_id), kb_id=knowledge_base_uuid, file_type=os.path.splitext(file_id)[1].lstrip('.'), file_id=file_id, ) except Exception: # 如果存储数据时出错，更新文件状态为失败 file_obj = session.query(File).filter_by(id=file_id).first() if file_obj: file_obj.status = 2 session.commit() self.ap.logger.error(f'Error storing data for file ID {file_id}', exc_info=True) session.close()