app/api/optimizer.py

"""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",
        },
    )