import json import asyncio from typing import List, Dict, Optional from datetime import datetime from sqlalchemy.orm import Session from app.models.models import ( OptimizationSession, OptimizationSegment, SessionHistory, ChangeLog ) from app.services.ai_service import ( AIService, split_text_into_segments, count_chinese_characters, count_text_length, get_default_polish_prompt, get_default_enhance_prompt, get_emotion_polish_prompt, get_compression_prompt ) from app.services.concurrency import concurrency_manager from app.services.stream_manager import stream_manager from app.config import settings # 错误信息最大长度，避免数据库字段溢出 MAX_ERROR_MESSAGE_LENGTH = 500 class OptimizationService: """优化处理服务""" def __init__(self, db: Session, session_obj: OptimizationSession): self.db = db self.session_obj = session_obj self.polish_service: Optional[AIService] = None self.enhance_service: Optional[AIService] = None self.emotion_service: Optional[AIService] = None self.compression_service: Optional[AIService] = None def _init_ai_services(self): """初始化AI服务 改进的初始化逻辑： 1. 验证必需的配置项 2. 提供更详细的错误信息 3. 确保所有服务都正确初始化 """ try: # 润色服务 self.polish_service = AIService( model=self.session_obj.polish_model or settings.POLISH_MODEL, api_key=self.session_obj.polish_api_key or settings.POLISH_API_KEY, base_url=self.session_obj.polish_base_url or settings.POLISH_BASE_URL ) # 增强服务 self.enhance_service = AIService( model=self.session_obj.enhance_model or settings.ENHANCE_MODEL, api_key=self.session_obj.enhance_api_key or settings.ENHANCE_API_KEY, base_url=self.session_obj.enhance_base_url or settings.ENHANCE_BASE_URL ) # 感情文章润色服务 self.emotion_service = AIService( model=self.session_obj.emotion_model or settings.POLISH_MODEL, api_key=self.session_obj.emotion_api_key or settings.POLISH_API_KEY, base_url=self.session_obj.emotion_base_url or settings.POLISH_BASE_URL ) # 压缩服务 self.compression_service = AIService( model=settings.COMPRESSION_MODEL, api_key=settings.COMPRESSION_API_KEY or settings.OPENAI_API_KEY, base_url=settings.COMPRESSION_BASE_URL or settings.OPENAI_BASE_URL ) print(f"[INFO] 所有 AI 服务初始化成功，会话: {self.session_obj.session_id}") except Exception as e: error_msg = f"AI 服务初始化失败: {str(e)}" print(f"[ERROR] {error_msg}") raise Exception(error_msg) async def start_optimization(self): """开始优化流程""" try: # 初始化AI服务 self._init_ai_services() # 重置错误状态 self.session_obj.error_message = None self.session_obj.failed_segment_index = None self.db.commit() # 获取并发权限 acquired = await concurrency_manager.acquire(self.session_obj.session_id) if not acquired: self.session_obj.status = "queued" self.db.commit() # 等待获取权限 - acquire 方法内部已包含等待逻辑 acquired = await concurrency_manager.acquire(self.session_obj.session_id) if not acquired: raise Exception("等待并发权限超时") # 更新状态为处理中 self.session_obj.status = "processing" self.db.commit() # 检查是否已存在段落,避免重复创建 # 在每次循环前检查会话状态，如果被停止则中断执行 self.db.refresh(self.session_obj) if self.session_obj.status == "stopped": raise Exception("会话已被用户停止") existing_segments = self.db.query(OptimizationSegment).filter( OptimizationSegment.session_id == self.session_obj.id ).order_by(OptimizationSegment.segment_index).all() if not existing_segments: # 首次运行: 分割文本并创建段落记录 segments = split_text_into_segments(self.session_obj.original_text) self.session_obj.total_segments = len(segments) self.db.commit() for idx, segment_text in enumerate(segments): segment = OptimizationSegment( session_id=self.session_obj.id, segment_index=idx, stage="polish", original_text=segment_text, status="pending" ) self.db.add(segment) self.db.commit() else: # 继续运行: 同步总段落数 self.session_obj.total_segments = len(existing_segments) self.db.commit() # 根据处理模式执行不同的阶段 processing_mode = self.session_obj.processing_mode or 'paper_polish_enhance' if processing_mode == 'paper_polish': # 只进行论文润色 await self._process_stage("polish") elif processing_mode == 'paper_enhance': # 只进行论文增强（直接增强原文） await self._process_stage("enhance") elif processing_mode == 'emotion_polish': # 只进行感情文章润色 await self._process_stage("emotion_polish") elif processing_mode == 'paper_polish_enhance': # 论文润色 + 论文增强 await self._process_stage("polish") await self._process_stage("enhance") else: raise ValueError(f"不支持的处理模式: {processing_mode}") # 完成 self.session_obj.status = "completed" self.session_obj.completed_at = datetime.utcnow() self.session_obj.progress = 100.0 self.session_obj.failed_segment_index = None self.db.commit() except Exception as e: self.session_obj.status = "failed" # 安全地截断错误信息 error_msg = str(e) if len(error_msg) > MAX_ERROR_MESSAGE_LENGTH: error_msg = error_msg[:MAX_ERROR_MESSAGE_LENGTH - 50] + "... [错误信息已截断]" self.session_obj.error_message = error_msg self.db.commit() raise finally: # 释放并发权限 await concurrency_manager.release(self.session_obj.session_id) # 清理 AI 服务资源 self._cleanup_ai_services() def _cleanup_ai_services(self): """清理 AI 服务资源""" # 将服务引用设置为 None，让 Python 的垃圾回收处理 # AsyncOpenAI 客户端会自动清理连接 self.polish_service = None self.enhance_service = None self.emotion_service = None self.compression_service = None async def _process_stage(self, stage: str): """处理单个阶段""" print(f"\n[STAGE START] Stage: {stage}, Session: {self.session_obj.session_id}", flush=True) self.session_obj.current_stage = stage self.db.commit() # 获取该阶段的提示词 prompt = self._get_prompt(stage) # 获取AI服务 if stage == "emotion_polish": ai_service = self.emotion_service elif stage == "polish": ai_service = self.polish_service else: # enhance ai_service = self.enhance_service # 获取所有段落 segments = self.db.query(OptimizationSegment).filter( OptimizationSegment.session_id == self.session_obj.id ).order_by(OptimizationSegment.segment_index).all() # 如果存在失败段落，跳过已完成的段落 start_index = 0 if self.session_obj.failed_segment_index is not None: start_index = max(self.session_obj.failed_segment_index, 0) # 历史会话 - 只包含AI的回复内容 # 只加载 start_index 之前的段落到历史，避免重试时历史与当前处理位置不一致 history: List[Dict[str, str]] = [] total_chars = 0 for segment in segments[:start_index]: if segment.is_title: # 标题段落不参与历史上下文 continue if stage == "polish" and segment.polished_text: history.append({"role": "assistant", "content": segment.polished_text}) total_chars += count_chinese_characters(segment.polished_text) elif stage == "emotion_polish" and segment.polished_text: history.append({"role": "assistant", "content": segment.polished_text}) total_chars += count_chinese_characters(segment.polished_text) elif stage == "enhance" and segment.enhanced_text: history.append({"role": "assistant", "content": segment.enhanced_text}) total_chars += count_chinese_characters(segment.enhanced_text) print(f"[STAGE] Loaded {len(history)} history messages from segments[:start_index={start_index}]", flush=True) skip_threshold = max(settings.SEGMENT_SKIP_THRESHOLD, 0) # 获取处理模式，用于正确计算进度 processing_mode = self.session_obj.processing_mode or 'paper_polish_enhance' for idx, segment in enumerate(segments[start_index:], start=start_index): # 每次处理段落前检查会话状态 self.db.refresh(self.session_obj) if self.session_obj.status == "stopped": raise Exception("会话已被用户停止") # 更新进度（无论是否跳过都更新） self.session_obj.current_position = idx # 根据处理模式正确计算进度 if processing_mode == 'paper_polish_enhance': if stage == "polish": # 第一阶段占 0-50% progress = (idx / len(segments)) * 50 else: # enhance # 第二阶段占 50-100% progress = 50 + (idx / len(segments)) * 50 else: # 其他模式占 0-100% progress = (idx / len(segments)) * 100 self.session_obj.progress = min(progress, 100.0) self.db.commit() # 先判断标题和短段落（提前到这里） if count_text_length(segment.original_text) < skip_threshold: if not segment.is_title: segment.is_title = True segment.status = "completed" segment.polished_text = segment.original_text segment.enhanced_text = segment.original_text segment.completed_at = datetime.utcnow() segment.stage = stage self.db.commit() continue # 然后检查是否已处理 if stage in ["polish", "emotion_polish"] and segment.polished_text: continue if stage == "enhance": if segment.enhanced_text: continue if segment.is_title and not segment.enhanced_text: segment.enhanced_text = segment.polished_text or segment.original_text segment.status = "completed" segment.completed_at = segment.completed_at or datetime.utcnow() self.db.commit() continue try: print(f"\n[SEGMENT {idx}] Processing segment {idx+1}/{len(segments)}, Stage: {stage}", flush=True) print(f"[SEGMENT {idx}] Input Length: {count_text_length(segment.original_text)}", flush=True) segment.status = "processing" segment.stage = stage self.db.commit() # 准备输入文本 # 对于 enhance 阶段：如果有润色结果则使用，否则使用原文（适用于 paper_enhance 模式） if stage == "enhance": input_text = segment.polished_text if segment.polished_text else segment.original_text else: input_text = segment.original_text # 调用AI async def execute_call(): # 使用配置中的流式设置，默认非流式（False）以避免API阻止 use_stream = settings.USE_STREAMING if stage == "polish": response = await ai_service.polish_text(input_text, prompt, history, stream=use_stream) elif stage == "emotion_polish": response = await ai_service.polish_emotion_text(input_text, prompt, history, stream=use_stream) else: # enhance response = await ai_service.enhance_text(input_text, prompt, history, stream=use_stream) if use_stream: full_text = "" async for chunk in response: if chunk: full_text += chunk # 推送流式更新 await stream_manager.broadcast(self.session_obj.session_id, { "type": "content", "segment_index": idx, "stage": stage, "content": chunk, "full_text": full_text # 可选:发送全量或增量，这里发送增量chunk，全量用于恢复 }) return full_text else: return response output_text = await self._run_with_retry(idx, stage, execute_call) if stage in ["polish", "emotion_polish"]: segment.polished_text = output_text else: # enhance segment.enhanced_text = output_text segment.status = "completed" segment.completed_at = datetime.utcnow() self.db.commit() # 记录变更 await self._record_change(segment, input_text, output_text, stage) # 更新历史会话 - 只添加AI的回复内容 history.append({"role": "assistant", "content": output_text}) total_chars += count_chinese_characters(output_text) # 检查是否需要压缩历史 - 基于字符数阈值 if total_chars > settings.HISTORY_COMPRESSION_THRESHOLD: print(f"\n[HISTORY COMPRESS] Triggering compression, Stage: {stage}", flush=True) print(f"[HISTORY COMPRESS] Before: {total_chars} chars, {len(history)} messages", flush=True) compressed_history = await self._compress_history(history, stage) # 压缩后的历史替换原历史，用于后续处理 history = compressed_history # 重新计算字符数 total_chars = sum(count_chinese_characters(msg.get("content", "")) for msg in history) print(f"[HISTORY COMPRESS] After: {total_chars} chars, {len(history)} messages", flush=True) # 推送压缩通知给前端 await stream_manager.broadcast(self.session_obj.session_id, { "type": "history_compressed", "stage": stage, "message": f"历史会话已压缩（{stage} 阶段），节省上下文空间", "new_char_count": total_chars }) # 只在压缩后保存历史，减少数据库写入 await self._save_history(history, stage, total_chars) except Exception as e: import traceback error_trace = traceback.format_exc() print(f"[ERROR] Segment {idx} processing failed:", flush=True) print(error_trace, flush=True) segment.status = "failed" self.session_obj.failed_segment_index = idx # 安全地截断错误信息，避免数据库字段溢出 error_msg = str(e) if len(error_msg) > MAX_ERROR_MESSAGE_LENGTH: # 保留前面的主要错误信息和末尾的部分 prefix_len = MAX_ERROR_MESSAGE_LENGTH - 50 error_msg = error_msg[:prefix_len] + "... [错误信息已截断]" self.session_obj.error_message = error_msg self.db.commit() # 直接抛出原异常，保留堆栈 raise async def _run_with_retry(self, segment_index: int, stage: str, task): """执行单次任务，不自动重试""" try: return await task() except Exception as exc: raise Exception( f"段落 {segment_index + 1} 在 {stage} 阶段失败: {str(exc)}" ) def _get_prompt(self, stage: str) -> str: """获取提示词""" if stage == "polish": return get_default_polish_prompt() elif stage == "emotion_polish": return get_emotion_polish_prompt() else: # enhance return get_default_enhance_prompt() async def _compress_history( self, history: List[Dict[str, str]], stage: str ) -> List[Dict[str, str]]: """压缩历史会话 - 智能提取关键信息 压缩历史会话以减少token使用，但保留处理风格的关键特征。 压缩后的内容单独保存，不影响已完成的润色和增强文本。 如果压缩失败，返回最近的几条消息而不是抛出异常。 """ try: # 如果历史已经是压缩格式（system消息），直接返回 if len(history) == 1 and history[0].get("role") == "system": return history # 保留最近的2-3条消息作为风格参考 recent_messages = history[-3:] if len(history) > 3 else history # 选择合适的压缩提示词 if stage == "emotion_polish": compression_prompt = """你是一个专业的文本摘要助手。请压缩以下历史处理内容，提取关键风格特征： 1. 总结文本的表达风格和语言特点 2. 提取关键的修改方向和处理模式 3. 保留重要的词汇使用倾向 4. 删除重复的内容和冗余表述 要求： - 压缩后内容不超过原内容的30% - 只输出压缩后的摘要，不要添加任何解释和注释 历史处理内容：""" else: compression_prompt = """你是一个专业的学术文本摘要助手。请压缩以下历史处理内容，提取关键信息： 1. 保留论文的主要术语、核心概念和关键数据 2. 总结已处理段落的主题和要点 3. 提取处理风格和改进方向的关键特征 4. 删除重复内容和冗余表述 要求： - 压缩后内容不超过原内容的30% - 保持学术性和专业性 - 只输出压缩后的摘要文本，不要添加任何解释和注释 历史处理内容：""" compressed_summary = await self.compression_service.compress_history( recent_messages, compression_prompt ) # 返回压缩后的历史作为系统消息，用于后续段落的上下文参考 return [ { "role": "system", "content": f"之前处理的段落摘要：\n{compressed_summary}" } ] except Exception as e: # 压缩失败时，不抛出异常，而是返回最近的几条消息 print(f"[WARNING] 历史压缩失败: {str(e)}, 将使用最近的消息代替", flush=True) # 返回最近的2条消息，避免上下文过长 return history[-2:] if len(history) > 2 else history async def _save_history(self, history: List[Dict[str, str]], stage: str, char_count: int): """保存历史会话 - 只在压缩后保存 只有压缩后的历史才保存到数据库，以避免频繁写入导致数据库膨胀。 压缩后的内容单独保存，不影响已完成的润色和增强文本。 注意：未压缩的历史不会保存，因为： 1. 润色/增强后的文本已经保存在 segments 表中 2. 压缩只在字符数超过阈值时触发 3. 压缩后的历史用于后续段落的上下文参考 """ # 检测是否为压缩后的历史：压缩后只有一条 system 消息，包含之前处理的摘要 # 这种检测方式与 _compress_history 的返回格式保持一致 is_compressed = len(history) == 1 and history[0].get("role") == "system" if not is_compressed: return # 非压缩状态不保存，减少数据库写入 # 检查是否已存在该阶段的压缩记录 existing = self.db.query(SessionHistory).filter( SessionHistory.session_id == self.session_obj.id, SessionHistory.stage == stage, SessionHistory.is_compressed.is_(True) ).first() if existing: # 更新现有记录 existing.history_data = json.dumps(history, ensure_ascii=False) existing.character_count = char_count existing.created_at = datetime.utcnow() else: # 创建新记录 history_obj = SessionHistory( session_id=self.session_obj.id, stage=stage, history_data=json.dumps(history, ensure_ascii=False), is_compressed=True, character_count=char_count ) self.db.add(history_obj) self.db.commit() async def _record_change( self, segment: OptimizationSegment, before: str, after: str, stage: str ): """记录变更""" # 简单的变更检测 changes = { "before_length": len(before), "after_length": len(after), "changed": before != after } existing_log = self.db.query(ChangeLog).filter( ChangeLog.session_id == self.session_obj.id, ChangeLog.segment_index == segment.segment_index, ChangeLog.stage == stage ).order_by(ChangeLog.created_at.desc()).first() serialized_detail = json.dumps(changes, ensure_ascii=False) if existing_log: # 如果之前已经生成过同一段落同一阶段的记录，直接更新内容避免重复条目 existing_log.before_text = before existing_log.after_text = after existing_log.changes_detail = serialized_detail else: change_log = ChangeLog( session_id=self.session_obj.id, segment_index=segment.segment_index, stage=stage, before_text=before, after_text=after, changes_detail=serialized_detail ) self.db.add(change_log) self.db.commit()