董宏凯andCursor
5de74e90dc
AI 写作页默认展示智能体 Tab，在页内复用 ai-writing 模块配置与 AgentChatPage；确认类型后侧栏扁平展示类型名并支持整块收起动画。同步历史回看、流式续跑、改写菜单与 researcher 检索相关改动。 Co-authored-by: Cursor <cursoragent@cursor.com>
214 lines
9.3 KiB
Python
214 lines
9.3 KiB
Python
"""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"},
|
||
)
|