192 lines
8.0 KiB
Python
192 lines
8.0 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": "请调整以下文字的语气风格。",
|
|
}
|
|
|
|
_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")
|
|
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")
|
|
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
|
|
|
|
|
|
# ── 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 = 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}"
|
|
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)]
|
|
|
|
# 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"},
|
|
)
|