deerflow-code/offline-backend-20260512/backend/app/gateway/routers/writing.py

"""Writing rewrite API — SSE streaming endpoint.

POST /api/writing/rewrite
  Returns text/event-stream with frames:
    event: thinking  data: {"chunk": "..."}
    event: text      data: {"chunk": "..."}
    event: done      data: {}
    event: error     data: {"message": "..."}
"""

from __future__ import annotations

import json
import logging

from fastapi import APIRouter, Depends
from fastapi.responses import StreamingResponse
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from pydantic import BaseModel, Field

from app.gateway.deps import get_config
from deerflow.config.app_config import AppConfig
from deerflow.models import create_chat_model

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/writing", tags=["writing"])

# ── Prompts ───────────────────────────────────────────────────────────────────

_SYSTEM_PROMPT = """你是一位专业的中文写作助手，擅长润色、扩写、缩写和调整文风。

规则：
- 只输出改写后的正文，不加任何前缀、解释或括号说明。
- 保持原文的语言（中文输入 → 中文输出，英文输入 → 英文输出）。
- 改写结果应与原文长度相近，除非任务要求扩写或缩写。
- 不要添加 markdown 标题或特殊格式，除非原文本身包含。
"""

_ACTION_INSTRUCTIONS: dict[str, str] = {
    "polish": "请对以下文字进行润色，使其表达更流畅、自然，同时保留原意。",
    "expand": "请对以下文字进行扩写，补充细节和说明，使内容更丰富饱满。",
    "shorten": "请对以下文字进行缩写，保留核心意思，去除冗余表达，使其更简洁。",
    "tone": "请调整以下文字的语气风格。",
    "imitate": "请模仿给定范文的写作风格、句式与语气，对以下文字进行仿写改写，在保留原文核心信息的前提下尽量贴近范文的表达方式。",
}

# 仿写参考范文的最大长度（防止超长 prompt）
_REFERENCE_MAX_CHARS = 2000

_TONE_SUFFIX: dict[str, str] = {
    "professional": "目标风格：专业、正式。",
    "casual": "目标风格：轻松、口语化。",
    "readable": "目标风格：通俗易懂、易于阅读。",
    "academic": "目标风格：学术严谨、逻辑清晰。",
    "expressive": "目标风格：生动有文采、富有感染力。",
}


# ── Request model ─────────────────────────────────────────────────────────────


class SelectionInfo(BaseModel):
    from_: int = Field(..., alias="from", description="ProseMirror start offset")
    to: int = Field(..., description="ProseMirror end offset")
    text: str = Field(..., description="Selected plain text")

    model_config = {"populate_by_name": True}


class SurroundingContext(BaseModel):
    before: str = Field(default="", description="Text before the selection")
    after: str = Field(default="", description="Text after the selection")


class RewriteRequest(BaseModel):
    thread_id: str | None = Field(default=None, alias="threadId")
    selection: SelectionInfo
    action: str = Field(..., description="polish | expand | shorten | tone | imitate")
    tone: str | None = Field(default=None)
    model_name: str | None = Field(default=None, alias="modelName")
    surrounding_context: SurroundingContext = Field(default_factory=SurroundingContext, alias="surroundingContext")
    custom_instruction: str | None = Field(default=None, alias="customInstruction")
    reference_text: str | None = Field(default=None, alias="referenceText", description="仿写参考范文（action=imitate）")
    follow_up_instruction: str | None = Field(default=None, alias="followUpInstruction")
    previous_result: str | None = Field(default=None, alias="previousResult")

    model_config = {"populate_by_name": True}


# ── Helpers ───────────────────────────────────────────────────────────────────


def _sse(event: str, data: dict) -> str:
    return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"


def _extract_chunk_parts(chunk: object) -> tuple[str, str]:
    """Return (thinking_delta, text_delta) from a LangChain streaming chunk."""
    thinking = ""
    text = ""

    content = getattr(chunk, "content", "")
    if isinstance(content, str):
        text = content
    elif isinstance(content, list):
        for block in content:
            if isinstance(block, str):
                text += block
            elif isinstance(block, dict):
                btype = block.get("type", "")
                if btype in ("thinking", "thinking_delta"):
                    thinking += block.get("thinking", "")
                elif btype in ("text", "text_delta"):
                    text += block.get("text", "")
                elif btype == "reasoning":
                    thinking += block.get("content", "")

    # DeepSeek / OpenAI o-series: reasoning in additional_kwargs
    if not thinking:
        kwargs = getattr(chunk, "additional_kwargs", {}) or {}
        r = kwargs.get("reasoning_content") or ""
        if isinstance(r, str):
            thinking = r

    return thinking, text


# ── Prompt building (pure, testable) ──────────────────────────────────────────


def build_rewrite_messages(body: RewriteRequest) -> tuple[str, list]:
    """Build (resolved_action, messages) for the rewrite request.

    Pure function (no model / IO) so the prompt assembly — including 仿写
    (imitate) reference-sample injection — can be unit-tested directly.
    """
    action = body.action if body.action in _ACTION_INSTRUCTIONS else "polish"
    instruction = _ACTION_INSTRUCTIONS[action]
    if action == "tone" and body.tone and body.tone in _TONE_SUFFIX:
        instruction = f"{instruction}\n{_TONE_SUFFIX[body.tone]}"

    context_parts: list[str] = []
    before = body.surrounding_context.before.strip()
    after = body.surrounding_context.after.strip()
    if before or after:
        context_parts.append("文章上下文（仅供参考，不要改写）：")
        if before:
            context_parts.append(f"[前文] {before[-300:]}")
        if after:
            context_parts.append(f"[后文] {after[:300]}")
        context_parts.append("")

    context_block = "\n".join(context_parts)
    base_instruction = instruction
    if body.custom_instruction:
        base_instruction = f"{instruction}\n用户额外要求：{body.custom_instruction}"
    if action == "imitate" and body.reference_text and body.reference_text.strip():
        reference = body.reference_text.strip()[:_REFERENCE_MAX_CHARS]
        base_instruction = f"{base_instruction}\n\n【范文】（请模仿其文风、句式、语气与结构，但不要照抄其具体内容）：\n{reference}"
    user_message = f"{base_instruction}\n\n{context_block}需要改写的文字：\n{body.selection.text}"

    if body.follow_up_instruction and body.previous_result:
        messages = [
            SystemMessage(content=_SYSTEM_PROMPT),
            HumanMessage(content=user_message),
            AIMessage(content=body.previous_result),
            HumanMessage(content=body.follow_up_instruction),
        ]
    else:
        messages = [SystemMessage(content=_SYSTEM_PROMPT), HumanMessage(content=user_message)]

    return action, messages


# ── Endpoint ──────────────────────────────────────────────────────────────────


@router.post(
    "/rewrite",
    summary="AI Text Rewrite (SSE stream)",
    response_class=StreamingResponse,
)
async def rewrite_selection(
    body: RewriteRequest,
    config: AppConfig = Depends(get_config),
) -> StreamingResponse:
    action, messages = build_rewrite_messages(body)

    # Auto-enable thinking when the selected model supports it
    resolved_name = body.model_name or (config.models[0].name if config.models else None)
    model_cfg = config.get_model_config(resolved_name) if resolved_name else None
    thinking_enabled = bool(model_cfg and getattr(model_cfg, "supports_thinking", False))

    async def generate():
        try:
            model = create_chat_model(name=body.model_name, thinking_enabled=thinking_enabled, app_config=config)
            async for chunk in model.astream(messages, config={"run_name": "writing_rewrite"}):
                thinking_delta, text_delta = _extract_chunk_parts(chunk)
                if thinking_delta:
                    yield _sse("thinking", {"chunk": thinking_delta})
                if text_delta:
                    yield _sse("text", {"chunk": text_delta})
            yield _sse("done", {})
        except Exception as e:
            logger.exception("writing/rewrite failed: thread=%s action=%s", body.thread_id, action)
            yield _sse("error", {"message": str(e)})

    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={"X-Accel-Buffering": "no", "Cache-Control": "no-cache"},
    )