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