从零构建生产级 AI 写作辅助系统：基于 LLM 的自动完成与风格迁移架构实践

从零构建生产级 AI 写作辅助系统：基于 LLM 的自动完成与风格迁移架构实践

当用户敲下第三个字时，模型如何在 300ms 内预测下一段情节？本文将带您深入 AI 写作辅助的工程内核。

1. 引言：AI 写作不止是“聊天框”

ChatGPT 可以写诗，Claude 可以写小说，但当我们试图将它们集成到 本地写作软件（如 Ulysses、Notion 或 VS Code 插件） 中时，会遇到巨大的工程鸿沟。

生产级 AI 写作辅助系统需要解决三个核心痛点：

1. 实时性（Latency）：用户敲击键盘后，补全建议必须在 200ms-500ms 内返回，否则体验断裂。
2. 连贯性（Coherence）：模型必须感知前 2000 字的故事线、人物关系和伏笔。
3. 风格可控（Style Control）：同一段情节，需要能输出“鲁迅风”、“古龙风”或“报告文学风”。

本文将分享一套基于 FastAPI + vLLM + Langfuse 的端到端架构，深入剖析我们是如何通过动态上下文剪枝和分阶段的采样策略来平衡速度与质量的。

2. 系统架构概览

我们采用 “检索增强 + 状态机” 的混合架构，而非简单的“Chat Completions 套壳”。

2.1 核心组件

• 前端 Editor：监听 input 事件，防抖（Debounce）500ms 后发送光标前 100 个字符。
• 网关层（Gateway）：负责用户鉴权、限流（Rate Limiting）和 多模型路由（小型模型处理续写，大型模型处理重构）。
• 状态存储（Redis）：维护用户会话的 KV Cache（键值缓存），避免重复计算历史 token。
• 推理引擎（Inference Engine）：基于 vLLM 部署 Qwen-2-7B-Instruct 或 Mistral-7B（兼顾显存与中文能力）。
• 风格迁移模块：一个独立的 LoRA 适配器热加载服务，可根据用户选择的风格动态切换 Adapter。

[Editor] --(SSE)--> [Gateway] --(gRPC)--> [Scheduler] --(Batch)--> [vLLM Workers]
     ^                                                    |
     |--------------------(Stream)-------------------------|

3. 深度技术拆解：上下文管理（Context Management）

这是 AI 写作工程中最难的一环。大模型的上下文窗口虽大（如 128k），但输入越长，推理越慢。对于写作场景，若每次都将整本书（10 万字）塞入 Prompt，TTFT（首 token 延迟）将高达数十秒，不可接受。

3.1 动态滑动窗口 + 记忆压缩

我们设计了一套 三重上下文结构：

1. 即时上下文（Immediate Context）：光标前的 512 tokens。包含用户刚刚输入的段落，用于保证行文流畅。
2. 场景记忆（Scene Memory）：当前章节的摘要（由模型每小时异步生成一次），约 200 tokens。
3. 全局风格向量（Global Style Vector）：用户选定风格后的 Embedding 向量，用于在 Prompt 中植入“规则”（例如：“不要使用‘然后’、‘接着’这类连接词”）。

代码实现（检索与拼接逻辑）：

from typing import List
from langchain.memory import ConversationSummaryBufferMemory

class WritingContextBuilder:
    def __init__(self, model_tokenizer):
        self.tokenizer = model_tokenizer
        self.max_input_len = 4096  # 硬限制
        self.immediate_tokens = 512
        self.summary_tokens = 200

    def build_prompt(self, current_text: str, scene_summary: str, style_rules: str):
        # 1. 截取即时文本
        immediate_chunk = self.tokenizer.decode(
            self.tokenizer.encode(current_text)[-self.immediate_tokens:]
        )
        
        # 2. 组装 System Prompt
        system_prompt = f"""
        你是一位资深作家。请遵循以下风格约束：
        {style_rules}
        
        当前故事情节背景：
        {scene_summary}
        """
        
        # 3. 计算总 Token，若超长则对 Summary 进行二次裁剪
        total_tokens = len(self.tokenizer.encode(system_prompt + immediate_chunk))
        if total_tokens > self.max_input_len:
            # 极端情况：暴力截断 Summary 至 50 tokens
            scene_summary = self.tokenizer.decode(
                self.tokenizer.encode(scene_summary)[:50]
            )
            
        return {
            "system": system_prompt,
            "user": f"继续接续以下内容，不要重复：\n{immediate_chunk}"
        }

4. 推理加速：KV Cache 复用与 Continuous Batching

对于多用户并发场景，如果每个请求都要重新计算前面 4096 个 tokens 的 Key-Value 矩阵，GPU 算力将被严重浪费。

4.1 基于 Redis 的 Prefix Cache

我们将每个会话（Session ID）的前缀 KV Cache 存储在共享内存中。当用户仅新增 50 个字符时，系统直接复用之前的 Cache，仅增量计算新 Token。

# 在 vLLM 启动参数中启用 prefix caching
# python -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2-7B-Instruct --enable-prefix-caching --max-num-seqs 256

实测数据：开启 Prefix Caching 后，在连续写作场景下，首 Token 延迟从 1800ms 降低至 320ms，降幅达 82%。

4.2 投机采样（Speculative Decoding）在写作中的应用

由于写作不像代码有明确答案，我们使用了 Draft Model（草稿模型） 策略。利用一个 1.5B 的“快思考”模型生成 Top-5 候选词，再由 7B 的“慢思考”模型进行验证。

此项优化使得生成吞吐量（Throughput）从 35 tokens/s 提升至 52 tokens/s。

5. 风格迁移：LoRA 插拔式架构

要让 AI 写出“张爱玲的风格”或“刘慈欣的风格”，全量微调（Full Fine-tuning）是不现实的。我们采用 LoRA（低秩适配） 技术。

5.1 动态 Adapter 加载

我们预先训练了 10 个风格 LoRA 权重（每个仅 50MB）。在 API 请求中传入 style_id，系统通过 peft 库动态添加或卸载 adapter。

from peft import PeftModel, PeftConfig

class StyleManager:
    def __init__(self, base_model):
        self.base_model = base_model
        self.current_adapter = None

    def switch_style(self, adapter_path: str):
        if self.current_adapter:
            # 卸载旧 Adapter
            self.base_model.unload_adapter()
        
        # 加载新 Adapter，只需几毫秒
        self.base_model.load_adapter(adapter_path, adapter_name="style_adapter")
        self.current_adapter = adapter_path
        print(f"Switched to style: {adapter_path}")

5.2 风格向量 Prompt 增强

除了 LoRA，我们在 Prompt 层面还注入了 风格负向提示（Negative Prompt），这在写作中至关重要。

// 古龙风示例
Positive Constraint: 对话极简，多使用短句，氛围萧瑟。
Negative Constraint: 严禁使用“因为…所以…”，严禁使用现代网络用语。

6. 评估体系：如何量化“写得好不好”？

这是 AI 写作领域最棘手的问题。我们建立了一套 自动化 + 人工反馈（RLHF） 的双轨评估机制：

1. 困惑度（Perplexity）：衡量生成文本的“惊喜度”。过低的困惑度意味着文本平庸，过高的困惑度意味着逻辑混乱。
2. Repetition Score（重复度评分）：统计 N-gram 重复率。重复率 > 0.3 则触发惩罚机制（Penalty）。
3. 用户行为埋点：采纳率（Acceptance Rate） 是最关键的北极星指标。如果模型生成的建议被用户直接采纳（按下 Tab 键），记为正反馈；若被立即删除，记为负反馈。

# 在 FastAPI 中间件中记录反馈
@app.post("/feedback")
async def collect_feedback(session_id: str, suggestion_id: str, accepted: bool):
    # 若拒绝，将当前生成结果作为 Negative Example 存入数据库
    # 用于未来模型的 DPO（直接偏好优化）训练
    if not accepted:
        store_negative_sample(session_id, suggestion_id)

7. 血泪踩坑实录

7.1 中文标点导致的 Token 爆炸

坑：一个中文标点（如“。”）在某些 tokenizer 下占用 1 个 token，但在 Qwen 的 tokenizer 中可能是 1-2 个。这导致用户端显示字数与后端 Token 计算严重不符。 解：前端显示“剩余配额”时，强制使用 tiktoken 或 transformers 库的 encode 函数进行精确计算，而非简单计算 len(text)。

7.2 SSE 流式传输中的格式错乱

坑：当模型生成包含换行符 \n 或 Markdown 表格时，前端 SSE 解析 data: 行会报错。 解：强制将输出进行 JSON 序列化转义，并在后端使用 json.dumps 确保多行文本被安全包裹。

# 正确示例
async def stream_generator():
    async for chunk in model.generate():
        yield f"data: {json.dumps({'text': chunk}, ensure_ascii=False)}\n\n"
    yield f"data: [DONE]\n\n"

8. 总结与展望

AI 写作辅助并非简单的“套壳 GPT”，它涉及 实时系统设计、KV Cache 调度、风格注入与用户体验心理学的交叉领域。

通过本文的架构优化，我们成功将系统的 采纳率从初期的 18% 提升至 37%。未来，随着 MoE（混合专家模型） 的普及，我们或许可以让“情节构思”、“润色词藻”和“校对纠错”分别由不同的专家模型并行处理，实现真正的“人机合著”。

互动讨论：你在开发 AI 写作或 Copilot 类应用时，遇到了哪些特有的 Bad Case？是“角色错乱”还是“情节死循环”？评论区欢迎交流！