feat(word-count): scene budget check with expand/compress rewrite (#15)

backend/app/api/ai_schemas.py

@@ -1,5 +1,7 @@
 """AI generation schemas: SceneCard, SceneDraft, Context Pack."""
 
+from typing import Literal
+
 from pydantic import BaseModel, Field
 
 
@@ -47,3 +49,38 @@ class ChapterSummaryModel(BaseModel):
     plot_threads: list[str] = Field(
         default_factory=list, description="情节线索"
     )
+
+
+class RewriteRequest(BaseModel):
+    scene_id: int
+    text: str = Field(
+        description="当前场景正文", max_length=100000
+    )
+    target_chars: int = Field(
+        default=1500, ge=100, le=50000,
+        description="目标字数",
+    )
+    mode: Literal["expand", "compress"] = Field(
+        description="重写模式: expand 或 compress",
+    )
+
+
+class WordCountCheckRequest(BaseModel):
+    text: str = Field(
+        description="当前场景正文", max_length=100000
+    )
+    target_chars: int = Field(
+        default=1500, ge=100, le=50000,
+        description="目标字数",
+    )
+
+
+class WordCountCheck(BaseModel):
+    status: str = Field(description="within / over / under")
+    actual_chars: int
+    target_chars: int
+    delta: int
+    deviation: float
+    suggestion: str | None = Field(
+        default=None, description="compress / expand / None"
+    )

backend/app/api/generation.py

@@ -6,7 +6,14 @@
 from fastapi.responses import StreamingResponse
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from app.api.ai_schemas import SceneCard, SceneCardRequest, SceneDraftRequest
+from app.api.ai_schemas import (
+    RewriteRequest,
+    SceneCard,
+    SceneCardRequest,
+    SceneDraftRequest,
+    WordCountCheck,
+    WordCountCheckRequest,
+)
 from app.core.config import settings
 from app.core.database import get_db
 from app.core.llm import call_llm_stream, instructor_client
@@ -15,6 +22,7 @@
     assemble_context_pack,
     get_scene_project_id,
 )
+from app.services.word_count import build_rewrite_prompt, check_word_budget
 
 router = APIRouter(prefix="/api/generate", tags=["generation"])
 
@@ -115,3 +123,63 @@ async def event_stream():
     return StreamingResponse(
         event_stream(), media_type="text/event-stream"
     )
+
+
+@router.post("/word-count-check", response_model=WordCountCheck)
+async def word_count_check(req: WordCountCheckRequest):
+    """Check if scene text fits within the target char budget."""
+    return check_word_budget(req.text, req.target_chars)
+
+
+@router.post("/rewrite")
+async def rewrite_scene(
+    req: RewriteRequest, db: AsyncSession = Depends(get_db)
+):
+    """Expand or compress scene text to fit target char budget via SSE."""
+    scene = await db.get(Scene, req.scene_id)
+    if not scene:
+        raise HTTPException(404, "Scene not found")
+
+    budget = check_word_budget(req.text, req.target_chars)
+    if budget["status"] == "within":
+        raise HTTPException(
+            400,
+            f"Text already within budget "
+            f"(deviation={budget['deviation']:.1%})",
+        )
+
+    expected_mode = budget["suggestion"]
+    if req.mode != expected_mode:
+        raise HTTPException(
+            400,
+            f"Mode '{req.mode}' conflicts with budget status "
+            f"'{budget['status']}': expected '{expected_mode}'",
+        )
+
+    prompt = build_rewrite_prompt(
+        req.text, req.target_chars, req.mode
+    )
+
+    async def event_stream():
+        total_text = ""
+        async for chunk in call_llm_stream(
+            messages=[{"role": "user", "content": prompt}]
+        ):
+            total_text += chunk
+            payload = json.dumps(
+                {"text": chunk}, ensure_ascii=False
+            )
+            yield f"data: {payload}\n\n"
+
+        result = check_word_budget(total_text, req.target_chars)
+        done_data = {
+            "done": True,
+            "char_count": len(total_text),
+            "budget": result,
+        }
+        payload = json.dumps(done_data, ensure_ascii=False)
+        yield f"data: {payload}\n\n"
+
+    return StreamingResponse(
+        event_stream(), media_type="text/event-stream"
+    )

backend/app/services/word_count.py

@@ -0,0 +1,72 @@
+"""Word count budget checking for scene generation."""
+
+from typing import Literal
+
+
+def check_word_budget(
+    text: str,
+    target_chars: int,
+    tolerance: float = 0.15,
+) -> dict:
+    """Check if text length is within the target budget.
+
+    Returns dict with:
+        status: 'within' | 'over' | 'under'
+        actual_chars: int
+        target_chars: int
+        delta: int (actual - target)
+        deviation: float (abs ratio)
+        suggestion: 'compress' | 'expand' | None
+    """
+    actual = len(text)
+    delta = actual - target_chars
+    if target_chars <= 0:
+        deviation = float(actual) if actual > 0 else 0.0
+    else:
+        deviation = abs(delta) / target_chars
+
+    if deviation <= tolerance:
+        status: Literal["within", "over", "under"] = "within"
+        suggestion = None
+    elif delta > 0:
+        status = "over"
+        suggestion = "compress"
+    else:
+        status = "under"
+        suggestion = "expand"
+
+    return {
+        "status": status,
+        "actual_chars": actual,
+        "target_chars": target_chars,
+        "delta": delta,
+        "deviation": round(deviation, 3),
+        "suggestion": suggestion,
+    }
+
+
+def build_rewrite_prompt(
+    text: str,
+    target_chars: int,
+    mode: str,
+) -> str:
+    """Build a prompt for expanding or compressing scene text."""
+    actual = len(text)
+    if mode == "expand":
+        return (
+            f"你是一位专业的中文小说作家。以下场景正文目前有 {actual} 字，"
+            f"目标是约 {target_chars} 字。"
+            "请在保持原有情节和人物不变的前提下，扩写以下内容，"
+            "增加细节描写、对话或环境描写。\n\n"
+            f"{text}\n\n"
+            "请直接输出扩写后的完整场景正文，不要输出任何标记或说明。"
+        )
+    else:
+        return (
+            f"你是一位专业的中文小说编辑。以下场景正文目前有 {actual} 字，"
+            f"目标是约 {target_chars} 字。"
+            "请在保持核心情节和关键对话不变的前提下，精简以下内容，"
+            "删除冗余描写和不必要的细节。\n\n"
+            f"{text}\n\n"
+            "请直接输出精简后的完整场景正文，不要输出任何标记或说明。"
+        )