from langchain.tools import tool from sqlalchemy.orm import Session from typing import Dict, Any, Tuple import logging from datetime import datetime import json import os import sys import re # 添加项目根目录到系统路径 sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) from database.db import SessionLocal, get_db from database.models import KnowledgeParsingContent, KnowledgeExtractionContent from gemini import GeminiProcessor # 配置日志 logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # 配置常量 BATCH_SIZE = 5 # 分批处理大小 SCORE_THRESHOLD = 70 # 评分阈值 # Define tools # evaluation_extraction_tool = Tool( # func=lambda request_id, query_word: _evaluation_extraction_tool(request_id, query_word), # name="evaluation_extraction_tool", # description="知识评估与抽取工具，用于处理数据库中的数据，执行评估并抽取内容" # ) @tool def evaluation_extraction_tool(request_id: str, query_word: str) -> str: """ 知识评估与抽取工具。持续处理数据库中的数据，分批执行评估并创建KnowledgeExtractionContent对象。 对于评分大于70分的内容，会进行抽取并更新KnowledgeExtractionContent对象。 Args: request_id: 请求ID，如果不提供则处理所有未处理的数据 query_word: 查询词，用于评估和抽取内容 Returns: str: "success" 表示处理完成，"no data" 表示没有数据需要处理 """ # 使用上下文管理器自动管理数据库连接的生命周期 with SessionLocal() as db: try: # 使用新的批量处理函数 result = execute_continuous_evaluation_extraction(request_id, db, query_word) return result except Exception as e: # 确保发生异常时回滚事务 db.rollback() logger.error(f"评估抽取过程中出错: {e}") return f"no data - 错误: {str(e)}" def execute_continuous_evaluation_extraction(request_id: str, db: Session, query_word: str) -> str: """持续执行评估循环，直到数据库没有数据""" logger.info(f"开始处理，request_id: {request_id}, query_word: {query_word}") total_processed = 0 offset = 0 try: while True: # 分批获取待评估的内容，使用offset实现分页 contents = get_batch_contents_for_evaluation(request_id, db, BATCH_SIZE, offset) logger.info(f"获取到 {len(contents)} 条待评估内容") if not contents: if total_processed > 0: logger.info(f"处理完成，共处理 {total_processed} 条内容") db.commit() # 确保最后一批数据被提交 return "success" return "no data" try: # 批量评估内容并创建KnowledgeExtractionContent对象 evaluation_results = batch_evaluate_content(contents, db, request_id, query_word) print(f"""evaluation_results: {evaluation_results}""") # 对评分大于阈值的内容进行抽取 high_score_results = [result for result in evaluation_results if result["score"] >= SCORE_THRESHOLD] if high_score_results: logger.info(f"发现 {len(high_score_results)} 条高分内容，进行抽取") batch_extract_and_save_content(high_score_results, db, request_id, query_word) total_processed += len(contents) offset += len(contents) # 更新offset值，以便下次获取下一批数据 db.commit() # 每批次处理完成后提交事务 except Exception as e: # 当前批次处理失败时回滚事务 db.rollback() logger.error(f"处理批次数据时出错: {e}") # 继续处理下一批数据 offset += len(contents) except Exception as e: # 发生严重异常时回滚事务并抛出异常 db.rollback() logger.error(f"执行评估抽取循环时出错: {e}") raise # 这里的代码永远不会被执行到，因为在while循环中，当contents为空时会返回 def get_batch_contents_for_evaluation(request_id: str, db: Session, batch_size: int, offset: int = 0) -> list: query = db.query(KnowledgeParsingContent).outerjoin( KnowledgeExtractionContent, KnowledgeParsingContent.id == KnowledgeExtractionContent.parsing_id ).filter( KnowledgeParsingContent.status == 2, # 已完成提取的数据 KnowledgeParsingContent.request_id == request_id, KnowledgeExtractionContent.parsing_id == None ) return query.offset(offset).limit(batch_size).all() def batch_evaluate_content(contents: list, db: Session, request_id: str, query_word: str) -> list: if not contents: return [] try: # 批量调用大模型进行评估 evaluation_results_raw = batch_call_llm_for_evaluation(contents, query_word) # 处理评估结果 evaluation_results = [] for i, (parsing_id, score, score_reason, parsing_data) in enumerate(evaluation_results_raw): # 创建KnowledgeExtractionContent对象 extraction_content = KnowledgeExtractionContent( request_id=request_id, parsing_id=parsing_id, score=score, score_reason=score_reason, create_at=datetime.now() ) db.add(extraction_content) evaluation_results.append({ "parsing_id": parsing_id, "score": score, "score_reason": score_reason, "parsing_data": parsing_data, "extraction_content": extraction_content }) return evaluation_results except Exception as e: logger.error(f"批量评估内容时出错: {e}") # 将所有内容标记为处理失败 for content in contents: content.status = 3 # 处理失败 return [] def batch_extract_and_save_content(evaluation_results: list, db: Session, request_id: str, query_word: str) -> list: if not evaluation_results: return [] try: # 批量调用大模型进行抽取 extraction_data_list = batch_call_llm_for_extraction(evaluation_results, query_word) # 保存抽取结果到数据库 success_ids = [] failed_ids = [] for i, (extracted_data, clean_reason) in enumerate(extraction_data_list): try: evaluation_result = evaluation_results[i] parsing_id = evaluation_result.get("parsing_id") if "extraction_content" in evaluation_result and parsing_id: # 更新已有对象的data字段和状态 extraction_content = evaluation_result["extraction_content"] extraction_content.data = extracted_data extraction_content.clean_reason = clean_reason extraction_content.status = 2 # 处理完成 success_ids.append(parsing_id) except Exception as e: logger.error(f"处理抽取结果 {i} 时出错: {e}") if i < len(evaluation_results): failed_ids.append(evaluation_results[i].get("parsing_id")) # 如果有失败的内容，将其标记为处理失败 if failed_ids: logger.warning(f"有 {len(failed_ids)} 条内容抽取失败") for result in evaluation_results: if result.get("parsing_id") in failed_ids and "extraction_content" in result: result["extraction_content"].status = 3 # 处理失败 return success_ids except Exception as e: logger.error(f"批量抽取和保存内容时出错: {e}") db.rollback() # 确保发生异常时回滚事务 return [] # 读取提示词文件 def read_prompt_file(file_path): """从文件中读取提示词""" try: with open(file_path, 'r', encoding='utf-8') as file: return file.read() except Exception as e: logger.error(f"读取提示词文件 {file_path} 失败: {str(e)}") return "" # 初始化 Gemini 处理器和提示词 gemini_processor = GeminiProcessor() # 加载评估和抽取提示词 project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) evaluation_prompt_path = os.path.join(project_root, 'prompt', 'evaluation.md') extraction_prompt_path = os.path.join(project_root, 'prompt', 'extraction.md') # 打印路径信息，用于调试 EVALUATION_PROMPT = read_prompt_file(evaluation_prompt_path) EXTRACTION_PROMPT = read_prompt_file(extraction_prompt_path) def batch_call_llm_for_evaluation(contents: list, query_word: str) -> list: """批量调用大模型进行内容评估 """ # 准备批量评估内容 evaluation_contents = [] for content in contents: evaluation_contents.append({ "query_word": query_word, "content": content.parsing_data }) try: # 批量调用 Gemini 进行评估 results = gemini_processor.batch_process(evaluation_contents, EVALUATION_PROMPT) # 处理返回结果 evaluation_results = [] for i, result in enumerate(results): result = re.sub(r'^\s*```json|\s*```\s*$', '', result, flags=re.MULTILINE).strip() result = json.loads(result) parsing_id = contents[i].id parsing_data = contents[i].parsing_data score = result.get("score", -2) score_reason = result.get("reason", "") evaluation_results.append((parsing_id, score, score_reason, parsing_data)) return evaluation_results except Exception as e: logger.error(f"批量评估过程异常: {str(e)}") # 返回默认结果 return [(content.id, 0, "评估过程异常", content.data if hasattr(content, 'data') else (content.parsing_data or "")) for content in contents] def batch_call_llm_for_extraction(evaluation_results: list, query_word: str) -> list: # 准备批量抽取内容 extraction_contents = [] for result in evaluation_results: parsing_data = result.get("parsing_data", "") extraction_contents.append({ "query_word": query_word, "content": parsing_data }) try: # 批量调用 Gemini 进行抽取 results = gemini_processor.batch_process(extraction_contents, EXTRACTION_PROMPT) # 处理返回结果 extraction_results = [] for i, result in enumerate(results): result = re.sub(r'^\s*```json|\s*```\s*$', '', result, flags=re.MULTILINE).strip() result = json.loads(result) extracted_data = result.get("extracted_content", "未提取到内容") clean_reason = result.get("analysis_reason", "未返回原因") extraction_results.append((extracted_data, clean_reason)) return extraction_results except Exception as e: logger.error(f"批量抽取过程异常: {str(e)}") # 返回空结果 return ["{}"] * len(evaluation_results)