xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-08-14 14:31:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: the fucking logger

Browse Source
This commit is contained in:
Junyan Qin
2025-07-11 21:37:31 +08:00
parent a79a22a74d
commit 6d788cadbc
1 changed files with 48 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/rag/knowledge/mgr.py
+48 -54
View File
@@ -1,6 +1,5 @@
# rag_manager.py # rag_manager.py
from __future__ import annotations from __future__ import annotations
import logging
import os import os
import asyncio import asyncio
import uuid import uuid
@@ -13,9 +12,8 @@ from pkg.core import app
class RAGManager: class RAGManager:
ap: app.Application ap: app.Application
def __init__(self, ap: app.Application, logger: logging.Logger = None): def __init__(self, ap: app.Application):
self.ap = ap self.ap = ap
self.logger = logger or logging.getLogger(__name__)
self.chroma_manager = None self.chroma_manager = None
self.parser = FileParser() self.parser = FileParser()
self.chunker = Chunker() self.chunker = Chunker()
@@ -52,20 +50,20 @@ class RAGManager:
session.add(new_kb) session.add(new_kb)
session.commit() session.commit()
session.refresh(new_kb) session.refresh(new_kb)
self.logger.info(f"Knowledge Base '{kb_name}' created.") self.ap.logger.info(f"Knowledge Base '{kb_name}' created.")
return new_kb.id return new_kb.id
else: else:
self.logger.info(f"Knowledge Base '{kb_name}' already exists.") self.ap.logger.info(f"Knowledge Base '{kb_name}' already exists.")
except Exception as e: except Exception as e:
session.rollback() session.rollback()
self.logger.error(f"Error in _create_kb_sync for '{kb_name}': {str(e)}", exc_info=True) self.ap.logger.error(f"Error in _create_kb_sync for '{kb_name}': {str(e)}", exc_info=True)
raise raise
finally: finally:
session.close() session.close()
return await asyncio.to_thread(_create_kb_sync) return await asyncio.to_thread(_create_kb_sync)
except Exception as e: except Exception as e:
self.logger.error(f"Error creating knowledge base '{kb_name}': {str(e)}", exc_info=True) self.ap.logger.error(f"Error creating knowledge base '{kb_name}': {str(e)}", exc_info=True)
raise raise
async def get_all_knowledge_bases(self): async def get_all_knowledge_bases(self):
@@ -83,7 +81,7 @@ class RAGManager:
return await asyncio.to_thread(_get_all_kbs_sync) return await asyncio.to_thread(_get_all_kbs_sync)
except Exception as e: except Exception as e:
self.logger.error(f'Error retrieving knowledge bases: {str(e)}', exc_info=True) self.ap.logger.error(f'Error retrieving knowledge bases: {str(e)}', exc_info=True)
return [] return []
async def get_knowledge_base_by_id(self, kb_id: str): async def get_knowledge_base_by_id(self, kb_id: str):
@@ -101,7 +99,7 @@ class RAGManager:
return await asyncio.to_thread(_get_kb_sync, kb_id) return await asyncio.to_thread(_get_kb_sync, kb_id)
except Exception as e: except Exception as e:
self.logger.error(f'Error retrieving knowledge base with ID {kb_id}: {str(e)}', exc_info=True) self.ap.logger.error(f'Error retrieving knowledge base with ID {kb_id}: {str(e)}', exc_info=True)
return None return None
async def get_files_by_knowledge_base(self, kb_id: str): async def get_files_by_knowledge_base(self, kb_id: str):
@@ -119,7 +117,7 @@ class RAGManager:
return await asyncio.to_thread(_get_files_sync, kb_id) return await asyncio.to_thread(_get_files_sync, kb_id)
except Exception as e: except Exception as e:
self.logger.error(f'Error retrieving files for knowledge base ID {kb_id}: {str(e)}', exc_info=True) self.ap.logger.error(f'Error retrieving files for knowledge base ID {kb_id}: {str(e)}', exc_info=True)
return [] return []
async def get_all_files(self): async def get_all_files(self):
@@ -138,7 +136,7 @@ class RAGManager:
return await asyncio.to_thread(_get_all_files_sync) return await asyncio.to_thread(_get_all_files_sync)
except Exception as e: except Exception as e:
self.logger.error(f'Error retrieving all files: {str(e)}', exc_info=True) self.ap.logger.error(f'Error retrieving all files: {str(e)}', exc_info=True)
return [] return []
async def store_data( async def store_data(
@@ -148,27 +146,25 @@ class RAGManager:
Parses, chunks, embeds, and stores data from a given file into the RAG system. 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. Associates the file with a knowledge base using kb_id in the File table.
""" """
self.logger.info(f'Starting data storage process for file: {file_path}') self.ap.logger.info(f'Starting data storage process for file: {file_path}')
session = SessionLocal() session = SessionLocal()
file_obj = None file_obj = None
try: try:
kb = session.query(KnowledgeBase).filter_by(name=kb_name).first() kb = session.query(KnowledgeBase).filter_by(name=kb_name).first()
if not kb: if not kb:
kb = KnowledgeBase(name=kb_name, description=kb_description) kb = KnowledgeBase(name=kb_name, description=kb_description)
session.add(kb) session.add(kb)
session.commit() session.commit()
session.refresh(kb) session.refresh(kb)
self.logger.info(f"Knowledge Base '{kb_name}' created during store_data.") self.ap.logger.info(f"Knowledge Base '{kb_name}' created during store_data.")
else: else:
self.logger.info(f"Knowledge Base '{kb_name}' already exists.") self.ap.logger.info(f"Knowledge Base '{kb_name}' already exists.")
file_name = os.path.basename(file_path) file_name = os.path.basename(file_path)
existing_file = session.query(File).filter_by(kb_id=kb.id, file_name=file_name).first() existing_file = session.query(File).filter_by(kb_id=kb.id, file_name=file_name).first()
if existing_file: if existing_file:
self.logger.warning( self.ap.logger.warning(
f"File '{file_name}' already exists in knowledge base '{kb_name}'. Skipping storage." f"File '{file_name}' already exists in knowledge base '{kb_name}'. Skipping storage."
) )
return return
@@ -177,32 +173,32 @@ class RAGManager:
session.add(file_obj) session.add(file_obj)
session.commit() session.commit()
session.refresh(file_obj) session.refresh(file_obj)
self.logger.info( self.ap.logger.info(
f"File record '{file_name}' added to database with ID: {file_obj.id}, associated with KB ID: {kb.id}" f"File record '{file_name}' added to database with ID: {file_obj.id}, associated with KB ID: {kb.id}"
) )
text = await self.parser.parse(file_path) text = await self.parser.parse(file_path)
if not text: if not text:
self.logger.warning(f'No text extracted from file {file_path}. Deleting file record ID: {file_obj.id}.') self.ap.logger.warning(
f'No text extracted from file {file_path}. Deleting file record ID: {file_obj.id}.'
)
session.delete(file_obj) session.delete(file_obj)
session.commit() session.commit()
return return
chunks_texts = await self.chunker.chunk(text) chunks_texts = await self.chunker.chunk(text)
self.logger.info(f"Chunked file '{file_name}' into {len(chunks_texts)} chunks.") self.ap.logger.info(f"Chunked file '{file_name}' into {len(chunks_texts)} chunks.")
await self.embedder.embed_and_store(file_id=file_obj.id, chunks=chunks_texts) await self.embedder.embed_and_store(file_id=file_obj.id, chunks=chunks_texts)
self.logger.info(f'Data storage process completed for file: {file_path}') self.ap.logger.info(f'Data storage process completed for file: {file_path}')
except Exception as e: except Exception as e:
session.rollback() session.rollback()
self.logger.error(f'Error in store_data for file {file_path}: {str(e)}', exc_info=True) self.ap.logger.error(f'Error in store_data for file {file_path}: {str(e)}', exc_info=True)
if file_obj and file_obj.id: if file_obj and file_obj.id:
try: try:
await asyncio.to_thread(self.chroma_manager.delete_by_file_id_sync, file_obj.id) await asyncio.to_thread(self.chroma_manager.delete_by_file_id_sync, file_obj.id)
except Exception as chroma_e: except Exception as chroma_e:
self.logger.warning( self.ap.logger.warning(
f'Could not clean up ChromaDB entries for file_id {file_obj.id} after store_data failure: {chroma_e}' f'Could not clean up ChromaDB entries for file_id {file_obj.id} after store_data failure: {chroma_e}'
) )
raise raise
@@ -213,13 +209,13 @@ class RAGManager:
""" """
Retrieves relevant data chunks based on a given query using the configured retriever. Retrieves relevant data chunks based on a given query using the configured retriever.
""" """
self.logger.info(f"Starting data retrieval process for query: '{query}'") self.ap.logger.info(f"Starting data retrieval process for query: '{query}'")
try: try:
retrieved_chunks = await self.retriever.retrieve(query) retrieved_chunks = await self.retriever.retrieve(query)
self.logger.info(f'Successfully retrieved {len(retrieved_chunks)} chunks for query.') self.ap.logger.info(f'Successfully retrieved {len(retrieved_chunks)} chunks for query.')
return retrieved_chunks return retrieved_chunks
except Exception as e: except Exception as e:
self.logger.error(f"Error in retrieve_data for query '{query}': {str(e)}", exc_info=True) self.ap.logger.error(f"Error in retrieve_data for query '{query}': {str(e)}", exc_info=True)
return [] return []
async def delete_data_by_file_id(self, file_id: str): async def delete_data_by_file_id(self, file_id: str):
@@ -227,32 +223,34 @@ class RAGManager:
Deletes all data associated with a specific file ID, including its chunks and vectors, Deletes all data associated with a specific file ID, including its chunks and vectors,
and the file record itself. and the file record itself.
""" """
self.logger.info(f'Starting data deletion process for file_id: {file_id}') self.ap.logger.info(f'Starting data deletion process for file_id: {file_id}')
session = SessionLocal() session = SessionLocal()
try: try:
# 1. 从 ChromaDB 删除 embeddings # 1. 从 ChromaDB 删除 embeddings
await asyncio.to_thread(self.chroma_manager.delete_by_file_id_sync, file_id) await asyncio.to_thread(self.chroma_manager.delete_by_file_id_sync, file_id)
self.logger.info(f'Deleted embeddings from ChromaDB for file_id: {file_id}') self.ap.logger.info(f'Deleted embeddings from ChromaDB for file_id: {file_id}')
# 2. 删除与文件关联的 chunks 记录 # 2. 删除与文件关联的 chunks 记录
chunks_to_delete = session.query(Chunk).filter_by(file_id=file_id).all() chunks_to_delete = session.query(Chunk).filter_by(file_id=file_id).all()
for chunk in chunks_to_delete: for chunk in chunks_to_delete:
session.delete(chunk) session.delete(chunk)
self.logger.info(f'Deleted {len(chunks_to_delete)} chunk records for file_id: {file_id}') self.ap.logger.info(f'Deleted {len(chunks_to_delete)} chunk records for file_id: {file_id}')
# 3. 删除文件记录本身 # 3. 删除文件记录本身
file_to_delete = session.query(File).filter_by(id=file_id).first() file_to_delete = session.query(File).filter_by(id=file_id).first()
if file_to_delete: if file_to_delete:
session.delete(file_to_delete) session.delete(file_to_delete)
self.logger.info(f'Deleted file record for file_id: {file_id}') self.ap.logger.info(f'Deleted file record for file_id: {file_id}')
else: else:
self.logger.warning(f'File with ID {file_id} not found in database. Skipping deletion of file record.') self.ap.logger.warning(
f'File with ID {file_id} not found in database. Skipping deletion of file record.'
)
session.commit() session.commit()
self.logger.info(f'Successfully completed data deletion for file_id: {file_id}') self.ap.logger.info(f'Successfully completed data deletion for file_id: {file_id}')
except Exception as e: except Exception as e:
session.rollback() session.rollback()
self.logger.error(f'Error deleting data for file_id {file_id}: {str(e)}', exc_info=True) self.ap.logger.error(f'Error deleting data for file_id {file_id}: {str(e)}', exc_info=True)
raise raise
finally: finally:
session.close() session.close()
@@ -262,45 +260,39 @@ class RAGManager:
Deletes a knowledge base and all associated files, chunks, and vectors. Deletes a knowledge base and all associated files, chunks, and vectors.
This involves querying for associated files and then deleting them. This involves querying for associated files and then deleting them.
""" """
self.logger.info(f'Starting deletion of knowledge base with ID: {kb_id}') self.ap.logger.info(f'Starting deletion of knowledge base with ID: {kb_id}')
session = SessionLocal() session = SessionLocal()
try: try:
kb_to_delete = session.query(KnowledgeBase).filter_by(id=kb_id).first() kb_to_delete = session.query(KnowledgeBase).filter_by(id=kb_id).first()
if not kb_to_delete: if not kb_to_delete:
self.logger.warning(f'Knowledge Base with ID {kb_id} not found.') self.ap.logger.warning(f'Knowledge Base with ID {kb_id} not found.')
return return
files_to_delete = session.query(File).filter_by(kb_id=kb_id).all() files_to_delete = session.query(File).filter_by(kb_id=kb_id).all()
session.close() session.close()
for file_obj in files_to_delete: for file_obj in files_to_delete:
try: try:
await self.delete_data_by_file_id(file_obj.id) await self.delete_data_by_file_id(file_obj.id)
except Exception as file_del_e: except Exception as file_del_e:
self.logger.error(f'Failed to delete file ID {file_obj.id} during KB deletion: {file_del_e}') self.ap.logger.error(f'Failed to delete file ID {file_obj.id} during KB deletion: {file_del_e}')
session = SessionLocal() session = SessionLocal()
try: try:
kb_final_delete = session.query(KnowledgeBase).filter_by(id=kb_id).first() kb_final_delete = session.query(KnowledgeBase).filter_by(id=kb_id).first()
if kb_final_delete: if kb_final_delete:
session.delete(kb_final_delete) session.delete(kb_final_delete)
session.commit() session.commit()
self.logger.info(f'Successfully deleted knowledge base with ID: {kb_id}') self.ap.logger.info(f'Successfully deleted knowledge base with ID: {kb_id}')
else: else:
self.logger.warning( self.ap.logger.warning(
f'Knowledge Base with ID {kb_id} not found after file deletion, skipping KB deletion.' f'Knowledge Base with ID {kb_id} not found after file deletion, skipping KB deletion.'
) )
except Exception as kb_del_e: except Exception as kb_del_e:
session.rollback() session.rollback()
self.logger.error(f'Error deleting KnowledgeBase record for ID {kb_id}: {kb_del_e}', exc_info=True) self.ap.logger.error(f'Error deleting KnowledgeBase record for ID {kb_id}: {kb_del_e}', exc_info=True)
raise raise
finally: finally:
session.close() session.close()
@@ -309,7 +301,9 @@ class RAGManager:
# 如果在最初获取 KB 或文件列表时出错 # 如果在最初获取 KB 或文件列表时出错
if session.is_active: if session.is_active:
session.rollback() session.rollback()
self.logger.error(f'Error during overall knowledge base deletion for ID {kb_id}: {str(e)}', exc_info=True) self.ap.logger.error(
f'Error during overall knowledge base deletion for ID {kb_id}: {str(e)}', exc_info=True
)
raise raise
finally: finally:
if session.is_active: if session.is_active:
@@ -335,29 +329,29 @@ class RAGManager:
""" """
Associates a file with a knowledge base by updating the kb_id in the File table. Associates a file with a knowledge base by updating the kb_id in the File table.
""" """
self.logger.info(f'Associating file ID {file_id} with knowledge base UUID {knowledge_base_uuid}') self.ap.logger.info(f'Associating file ID {file_id} with knowledge base UUID {knowledge_base_uuid}')
session = SessionLocal() session = SessionLocal()
try: try:
# 查询知识库是否存在 # 查询知识库是否存在
kb = session.query(KnowledgeBase).filter_by(id=knowledge_base_uuid).first() kb = session.query(KnowledgeBase).filter_by(id=knowledge_base_uuid).first()
if not kb: if not kb:
self.logger.error(f'Knowledge Base with UUID {knowledge_base_uuid} not found.') self.ap.logger.error(f'Knowledge Base with UUID {knowledge_base_uuid} not found.')
return return
# 更新文件的 kb_id # 更新文件的 kb_id
file_to_update = session.query(File).filter_by(id=file_id).first() file_to_update = session.query(File).filter_by(id=file_id).first()
if not file_to_update: if not file_to_update:
self.logger.error(f'File with ID {file_id} not found.') self.ap.logger.error(f'File with ID {file_id} not found.')
return return
file_to_update.kb_id = kb.id file_to_update.kb_id = kb.id
session.commit() session.commit()
self.logger.info( self.ap.logger.info(
f'Successfully associated file ID {file_id} with knowledge base UUID {knowledge_base_uuid}' f'Successfully associated file ID {file_id} with knowledge base UUID {knowledge_base_uuid}'
) )
except Exception as e: except Exception as e:
session.rollback() session.rollback()
self.logger.error( self.ap.logger.error(
f'Error associating file ID {file_id} with knowledge base UUID {knowledge_base_uuid}: {str(e)}', f'Error associating file ID {file_id} with knowledge base UUID {knowledge_base_uuid}: {str(e)}',
exc_info=True, exc_info=True,
) )