"""Document optimization & evaluation — SSE streaming API.""" import json import sys import time from pathlib import Path from fastapi import APIRouter, HTTPException, Query, Request from fastapi.responses import StreamingResponse # Ensure app dir and root are on path APP_DIR = Path(__file__).resolve().parent.parent ROOT_DIR = APP_DIR.parent for p in (str(APP_DIR), str(ROOT_DIR)): if p not in sys.path: sys.path.insert(0, p) from db import get_distractor_set, save_record from config_ui import get_missing_key_message router = APIRouter(prefix="/api/optimize", tags=["optimize"]) def _pct(before: float, after: float) -> str: if before == 0: return "N/A" change = (after - before) / before * 100 return f"{change:+.1f}%" def _optimize_stream( query: str, document: str, set_id: int, engine_llm: str, max_rounds: int, ): """Synchronous generator yielding SSE events for the full optimization flow.""" from engine import run_evaluate, run_rewrite_iterative # Validate key msg = get_missing_key_message(engine_llm) if msg: yield f"data: {json.dumps({'type': 'error', 'message': msg})}\n\n" return # Load distractor set ds = get_distractor_set(set_id) if not ds: yield f"data: {json.dumps({'type': 'error', 'message': '干扰文档集合不存在'})}\n\n" return distractors = json.loads(ds["distractors"]) text_list = [document.strip()] + [d for d in distractors if d.strip()] target_id = 0 # ---- Step 1: Evaluate original (baseline) ---- yield f"data: {json.dumps({'type': 'progress', 'stage': 'baseline', 'text': '评估原始文档...'})}\n\n" try: vanilla_scores = run_evaluate( query=query.strip(), text_list=text_list, target_id=target_id, engine_llm=engine_llm, ) yield f"data: {json.dumps({'type': 'baseline_complete', 'scores': vanilla_scores})}\n\n" except ValueError as e: yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n" return except Exception as e: yield f"data: {json.dumps({'type': 'error', 'message': f'原始文档评估失败: {e}'})}\n\n" return # ---- Step 2: Iterative rewrite + evaluate ---- rounds_data = [] def progress_callback(round_num: int, status: str): """Called by run_rewrite_iterative for each stage.""" stage_text = "改写中" if status == "rewriting" else ("优化中" if status == "refining" else "评估中") # This callback is synchronous and called from inside the engine. # We can't yield from here directly in a generator pattern, # so we store the state and yield from the wrapper. progress_callback.last = {"round": round_num, "status": status, "text": stage_text} progress_callback.last = None try: # We need to run the iterative process and capture results. # The progress_callback can't yield SSE events directly, # so we wrap each round manually to emit SSE. from engine import run_rewrite, refine_document total_rounds = max_rounds rewritten = None # Round 1: Standard rewrite yield f"data: {json.dumps({'type': 'progress', 'round': 1, 'total_rounds': total_rounds, 'stage': 'rewriting', 'text': f'第 1/{total_rounds} 轮：改写中...'})}\n\n" rewritten = run_rewrite(document=document, engine_llm=engine_llm) yield f"data: {json.dumps({'type': 'progress', 'round': 1, 'total_rounds': total_rounds, 'stage': 'evaluating', 'text': f'第 1/{total_rounds} 轮：评估中...'})}\n\n" text_list_eval = [rewritten] + text_list[1:] scores = run_evaluate( query=query.strip(), text_list=text_list_eval, target_id=target_id, engine_llm=engine_llm, ) rounds_data.append({"round": 1, "text": rewritten, "scores": scores}) yield f"data: {json.dumps({'type': 'round_complete', 'round': 1, 'scores': scores, 'text': rewritten})}\n\n" # Rounds 2+: Refine with feedback for rnd in range(2, total_rounds + 1): prev = rounds_data[-1] # Check if already excellent if all(v >= 0.9 for v in prev["scores"].values()): break yield f"data: {json.dumps({'type': 'progress', 'round': rnd, 'total_rounds': total_rounds, 'stage': 'refining', 'text': f'第 {rnd}/{total_rounds} 轮：优化中...'})}\n\n" refined = refine_document( original_doc=document, previous_rewrite=prev["text"], scores=prev["scores"], query=query, engine_llm=engine_llm, ) yield f"data: {json.dumps({'type': 'progress', 'round': rnd, 'total_rounds': total_rounds, 'stage': 'evaluating', 'text': f'第 {rnd}/{total_rounds} 轮：评估中...'})}\n\n" text_list_eval = [refined] + text_list[1:] scores = run_evaluate( query=query.strip(), text_list=text_list_eval, target_id=target_id, engine_llm=engine_llm, ) rounds_data.append({"round": rnd, "text": refined, "scores": scores}) yield f"data: {json.dumps({'type': 'round_complete', 'round': rnd, 'scores': scores, 'text': refined})}\n\n" # Early stop if rnd >= 3: prev_prev = rounds_data[-2] prev_cur = rounds_data[-1] if prev_cur["scores"]["wordpos"] <= prev_prev["scores"]["wordpos"] * 1.02: break # ---- Step 3: Best result ---- best = max(rounds_data, key=lambda x: x["scores"]["wordpos"]) rewritten_text = best["text"] optimized_scores = best["scores"] # Save to history save_record( model=engine_llm, query=query.strip(), original_doc=document.strip(), rewritten_doc=rewritten_text, vanilla_scores=vanilla_scores, optimized_scores=optimized_scores, distractors=distractors, ) yield f"data: {json.dumps({ 'type': 'complete', 'vanilla_scores': vanilla_scores, 'optimized_scores': optimized_scores, 'rewritten_text': rewritten_text, 'rounds': rounds_data, 'improvements': { 'pos': _pct(vanilla_scores['pos'], optimized_scores['pos']), 'word': _pct(vanilla_scores['word'], optimized_scores['word']), 'wordpos': _pct(vanilla_scores['wordpos'], optimized_scores['wordpos']), }, })}\n\n" except ValueError as e: yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n" except Exception as e: yield f"data: {json.dumps({'type': 'error', 'message': f'优化失败: {e}'})}\n\n" @router.get("/stream") async def optimize_stream( query: str = Query(..., description="用户查询"), document: str = Query(..., description="原始文档"), set_id: int = Query(..., description="干扰文档集合 ID"), engine_llm: str = Query("doubao", description="引擎"), max_rounds: int = Query(2, ge=1, le=5, description="最大迭代轮数"), ): """SSE endpoint — stream optimization + evaluation progress.""" # Validate if not query.strip(): raise HTTPException(status_code=400, detail="查询不能为空") if not document.strip(): raise HTTPException(status_code=400, detail="文档不能为空") if max_rounds < 1 or max_rounds > 5: raise HTTPException(status_code=400, detail="迭代轮数需要 1-5") return StreamingResponse( _optimize_stream(query.strip(), document.strip(), set_id, engine_llm, max_rounds), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", }, )