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 # 添加项目根目录到系统路径 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 = 10 # 分批处理大小 SCORE_THRESHOLD = 70 # 评分阈值 # Define tools @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" 表示没有数据需要处理 """ try: db = SessionLocal() try: # 使用新的批量处理函数 result = execute_continuous_evaluation_extraction(request_id, db, query_word) return result finally: db.close() except Exception as e: logger.error(f"评估抽取过程中出错: {e}") return f"no data - 错误: {str(e)}" def execute_continuous_evaluation_extraction(request_id: str, db: Session, query_word: str) -> str: """持续执行评估循环，直到数据库没有数据""" total_processed = 0 while True: # 分批获取待评估的内容 contents = get_batch_contents_for_evaluation(request_id, db, BATCH_SIZE) if not contents: if total_processed > 0: logger.info(f"处理完成，共处理 {total_processed} 条内容") return "success" return "no data" # 批量评估内容并创建KnowledgeExtractionContent对象 evaluation_results = batch_evaluate_content(contents, db, request_id, query_word) # 对评分大于阈值的内容进行抽取 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) db.commit() # 这里的代码永远不会被执行到，因为在while循环中，当contents为空时会返回 def get_batch_contents_for_evaluation(request_id: str, db: Session, batch_size: int) -> list: """分批获取待评估的内容""" query = db.query(KnowledgeParsingContent).filter( KnowledgeParsingContent.status == 2 # 已完成提取的数据 ) # 如果指定了request_id，则只处理该request_id的数据 if request_id: query = query.filter(KnowledgeParsingContent.request_id == request_id) return query.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, reason, parsing_data) in enumerate(evaluation_results_raw): # 创建KnowledgeExtractionContent对象 extraction_content = KnowledgeExtractionContent( request_id=request_id, parsing_id=parsing_id, score=score, reason=reason, create_at=datetime.now() ) db.add(extraction_content) evaluation_results.append({ "parsing_id": parsing_id, "score": score, "reason": 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 [] # 批量调用大模型进行抽取 extraction_data_list = batch_call_llm_for_extraction(evaluation_results, query_word) # 保存抽取结果到数据库 success_ids = [] failed_ids = [] for i, extraction_data in enumerate(extraction_data_list): try: evaluation_result = evaluation_results[i] # 更新已有对象的data字段和状态 existing_extraction.data = evaluation_result["extraction_content"] existing_extraction.status = 2 # 处理完成 success_ids.append(parsing_id) except Exception as e: logger.error(f"处理抽取结果 {i} 时出错: {e}") 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 # 读取提示词文件 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') # 打印路径信息，用于调试 logger.info(f"评估提示词路径: {evaluation_prompt_path}") logger.info(f"抽取提示词路径: {extraction_prompt_path}") 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): parsing_id = contents[i].id parsing_data = contents[i].parsing_data if isinstance(result, dict) and "score" in result: # 正常结果 score = result.get("score", -2) reason = result.get("reason", "") else: # 异常结果 score = -2 reason = "评估失败" evaluation_results.append((parsing_id, 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): # 确保结果包含必要的字段 if not isinstance(result, dict): result = {"extracted_data": str(result)} extraction_results.append(json.dumps(result, ensure_ascii=False)) return extraction_results except Exception as e: logger.error(f"批量抽取过程异常: {str(e)}") # 返回空结果 return ["{}"] * len(evaluation_results)