Merge remote-tracking branch 'origin/master' into ji
# Conflicts: # frontend-web/src/pages/AgentChatPage.tsx
This commit is contained in:
吉鑫升 committed
commit
643f9b2c76
51 files changed
+2654
-2164
No files matched your search
@@ -108,25 +108,31 @@ async def _delete_checkpoint_thread(checkpointer, thread_id: str) -> None:
|
||||
)
|
||||
|
||||
|
||||
def _get_active_session_ids() -> set[str]:
|
||||
"""返回内存里**当前正在跑**的 session_id 集合。
|
||||
async def _get_active_session_ids_langgraph(thread_store) -> set[str]:
|
||||
"""新路径下「正在跑 / 等用户干预」的 thread_id 集合。
|
||||
|
||||
这些会话的写作流程还活着 —— 任何清理都不能动它们，否则用户「正在写」
|
||||
的稿子会被秒删，体验灾难。延迟到下一轮清理（24 小时后）再处理也比误
|
||||
删强。
|
||||
``threads_meta.status`` 由 ``start_run`` / ``run_agent`` 维护：
|
||||
- ``running``：graph 正在跑（节点中）
|
||||
- ``interrupted``：graph 暂停在 interrupt 点等用户回头填卡
|
||||
这两种都视作活跃，**绝不能删**。``idle / success / error`` 才是真正终止状态。
|
||||
|
||||
LangGraph Server 路径 + 老路径的 lead_agent 都用同一份 ``thread_store``，所以
|
||||
这里拿的 active 集合是全部 graph 的；筛 ai_writing 是按 assistant_id 过滤。
|
||||
"""
|
||||
try:
|
||||
# 延迟 import 避免循环依赖：cleanup 模块在 lifespan 启动时被 import，
|
||||
# routers/ai_writing 也被 include_router 主动 import，两者拓扑无环但
|
||||
# 还是放函数里更稳。
|
||||
from app.gateway.routers.ai_writing import _event_queues
|
||||
|
||||
return set(_event_queues.keys())
|
||||
except Exception:
|
||||
# 兜底：拿不到就当作没有活动会话（最差就是清掉一个刚跑完的会话，
|
||||
# 比误删正在写的要好得多 —— 用户至少能从 DB 历史里看到结果）
|
||||
logger.warning("ai_writing_cleanup: 读取活动会话列表失败", exc_info=True)
|
||||
if thread_store is None:
|
||||
return set()
|
||||
try:
|
||||
running = await thread_store.search(status="running", user_id=None, limit=1000)
|
||||
interrupted = await thread_store.search(status="interrupted", user_id=None, limit=1000)
|
||||
except Exception:
|
||||
logger.warning("ai_writing_cleanup: 读取 LangGraph 活动 thread 列表失败", exc_info=True)
|
||||
return set()
|
||||
# 仅关心 assistant_id='ai_writing' 的 thread —— 别的 graph 不归这套清理管
|
||||
active: set[str] = set()
|
||||
for row in running + interrupted:
|
||||
if row.get("assistant_id") == "ai_writing":
|
||||
active.add(row["thread_id"])
|
||||
return active
|
||||
|
||||
|
||||
async def run_cleanup_once(
|
||||
@@ -134,18 +140,27 @@ async def run_cleanup_once(
|
||||
retention_days: int,
|
||||
only_finished: bool = False,
|
||||
batch_size: int = 500,
|
||||
checkpointer=None,
|
||||
thread_store=None,
|
||||
) -> dict:
|
||||
"""执行一轮清理。返回统计信息。
|
||||
|
||||
**安全保证**：清理永远跳过内存里仍在跑的会话（``_event_queues`` 里有
|
||||
对应 key 的）—— 这条规则即使在 admin 手动触发清理时也强制生效，防止
|
||||
把用户「正在写」的稿子误删。
|
||||
Args:
|
||||
checkpointer: 主 LangGraph checkpointer（``app.state.checkpointer``）。
|
||||
AI 写作的 thread 状态由 LangGraph Server 注入，存在这里 —— admin 路由
|
||||
+ scheduler loop 都会从 ``app.state`` 取到引用传进来；测试场景下传 None
|
||||
会让 checkpoint 删除变成 no-op，但 DB 行仍会被清。
|
||||
thread_store: ``app.state.thread_store``。用来按 ``status in (running,
|
||||
interrupted)`` 查活跃 thread 列表做活跃保护。
|
||||
|
||||
**安全保证**：清理永远跳过仍在跑 / 等用户干预的会话（``threads_meta.status
|
||||
in (running, interrupted)`` + ``assistant_id='ai_writing'``）。即使 admin
|
||||
手动触发 ``retention_days=0`` 也强制生效。
|
||||
|
||||
幂等 + 可重入 —— 多次调用对结果只会越扫越干净。即便清理途中被中断，下次
|
||||
再跑会重新把 DB 里仍然过期的行扫到（checkpointer 删一个已不存在的 thread
|
||||
是 no-op）。
|
||||
"""
|
||||
from deerflow.agents.ai_writing.graph import get_ai_writing_checkpointer
|
||||
from deerflow.persistence.ai_writing_sessions import AIWritingSessionRepository
|
||||
from deerflow.persistence.engine import get_session_factory
|
||||
|
||||
@@ -155,8 +170,7 @@ async def run_cleanup_once(
|
||||
return {"deleted_db": 0, "deleted_checkpoints": 0, "skipped_active": 0}
|
||||
|
||||
repo = AIWritingSessionRepository(sf)
|
||||
checkpointer = get_ai_writing_checkpointer()
|
||||
active_ids = _get_active_session_ids()
|
||||
active_ids = await _get_active_session_ids_langgraph(thread_store)
|
||||
|
||||
cutoff = datetime.now(UTC) - timedelta(days=retention_days)
|
||||
logger.info(
|
||||
@@ -238,11 +252,16 @@ def _next_cron_delay_seconds(cron_expr: str, base: datetime | None = None) -> fl
|
||||
return max(1.0, (nxt - base).total_seconds())
|
||||
|
||||
|
||||
async def cleanup_scheduler_loop() -> None:
|
||||
async def cleanup_scheduler_loop(app=None) -> None:
|
||||
"""常驻后台任务：按 cron 调度跑清理。
|
||||
|
||||
多 worker 部署下文件锁做 leader election，没拿到锁的 worker 直接退出。
|
||||
被 cancel 时 CancelledError 不会被 except Exception 吞掉，能干净退出。
|
||||
|
||||
Args:
|
||||
app: FastAPI 实例，用来从 ``app.state`` 拿主 checkpointer / thread_store
|
||||
引用，喂给 ``run_cleanup_once`` —— 新路径下 ai_writing 的 thread 状态
|
||||
在主 checkpointer 里。``app=None`` 时回退到老 checkpointer（仅老路径）。
|
||||
"""
|
||||
from deerflow.config.app_config import get_app_config
|
||||
from deerflow.config.runtime_paths import runtime_home
|
||||
@@ -289,10 +308,17 @@ async def cleanup_scheduler_loop() -> None:
|
||||
logger.info("ai_writing_cleanup: 下一次清理 %.0f 秒后", delay)
|
||||
await asyncio.sleep(delay)
|
||||
try:
|
||||
# 每次循环重新读 app.state —— 不缓存，避免 lifespan 中重启 / 重置
|
||||
# checkpointer 时拿到死引用。app=None 时 main_checkpointer / store
|
||||
# 也为 None，run_cleanup_once 内部走老 fallback。
|
||||
main_checkpointer = getattr(getattr(app, "state", None), "checkpointer", None)
|
||||
main_thread_store = getattr(getattr(app, "state", None), "thread_store", None)
|
||||
await run_cleanup_once(
|
||||
retention_days=retention_days,
|
||||
only_finished=only_finished,
|
||||
batch_size=batch_size,
|
||||
checkpointer=main_checkpointer,
|
||||
thread_store=main_thread_store,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("ai_writing_cleanup: 本轮清理失败（不影响下一轮）")
|
||||
|
||||
@@ -71,12 +71,6 @@ logger = logging.getLogger(__name__)
|
||||
# firing signals into a worker that is stuck waiting for shutdown cleanup.
|
||||
_SHUTDOWN_HOOK_TIMEOUT_SECONDS = 5.0
|
||||
|
||||
# AI 写作 in-flight task 收尾的最大等待时间。一个 graph 片段通常在几秒内
|
||||
# 跑到下一个 interrupt 或结束，但单段草稿生成走完整个 LLM 可能耗时较久，
|
||||
# 给 30 秒兜底；超时则 cancel 强制结束，已落盘的 checkpoint 仍可用于重启
|
||||
# 后恢复，丢的只是最后一段 in-flight 但未 checkpoint 的进度。
|
||||
_AI_WRITING_DRAIN_TIMEOUT_SECONDS = 30.0
|
||||
|
||||
|
||||
async def _ensure_admin_user(app: FastAPI) -> None:
|
||||
"""Startup hook: handle first boot and migrate orphan threads otherwise.
|
||||
@@ -371,11 +365,10 @@ async def _curator_scheduler_loop() -> None:
|
||||
def _warn_if_multi_worker() -> None:
|
||||
"""检测 uvicorn / gunicorn 多 worker 部署并打 ERROR 日志。
|
||||
|
||||
AI 写作 + 主 agent 流式响应都依赖**进程内**的内存状态字典
|
||||
（``_event_queues``、``_session_pending_pause`` 等）。多 worker 下，
|
||||
用户的 /start /stream /resume 会被 OS 随机分配到不同 worker，
|
||||
per-worker 的内存字典互相看不见 → SSE 事件丢失、会话「凭空消失」、
|
||||
resume 走不通。
|
||||
主 agent / canvas / AI 写作流式响应在 Gateway 进程内通过 ``StreamBridge``
|
||||
维护 run_id → SSE 订阅者的内存 map。多 worker 下，用户在 worker A 发起
|
||||
``/runs/stream``、稍后在 worker B 调 join/resume，会因 in-process map 互相
|
||||
看不见而 404 / 卡死。
|
||||
|
||||
所以这里在 startup 打 ERROR 强提醒。**不抛异常 fail-fast** —— 万一
|
||||
运维确实知道自己在干嘛（比如配了 sticky session），不应该阻止启动。
|
||||
@@ -449,6 +442,18 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
# Run Alembic migrations before anything else touches the schema.
|
||||
await _run_db_migrations()
|
||||
|
||||
# 预热 recall_budget cache（从 DB 批量加载,使 get_effective_memory_config 同步可读）
|
||||
try:
|
||||
from deerflow.config.memory_config import load_recall_budget_cache
|
||||
|
||||
_mc_store = getattr(app.state, "memory_config_store", None)
|
||||
if _mc_store is not None:
|
||||
_recall_budgets = await _mc_store.load_all()
|
||||
load_recall_budget_cache(_recall_budgets)
|
||||
logger.info("Memory config cache warmed: %d user override(s)", len(_recall_budgets))
|
||||
except Exception:
|
||||
logger.exception("Memory config cache warm-up failed (non-fatal)")
|
||||
|
||||
# Ensure admin user exists (auto-create on first boot)
|
||||
# Must run AFTER langgraph_runtime so app.state.store is available for thread migration
|
||||
await _ensure_admin_user(app)
|
||||
@@ -474,7 +479,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
try:
|
||||
from app.gateway.ai_writing_cleanup import cleanup_scheduler_loop
|
||||
|
||||
app.state.ai_writing_cleanup_task = asyncio.create_task(cleanup_scheduler_loop())
|
||||
app.state.ai_writing_cleanup_task = asyncio.create_task(cleanup_scheduler_loop(app))
|
||||
except Exception:
|
||||
logger.exception("AI writing cleanup scheduler failed to start (non-fatal)")
|
||||
|
||||
@@ -526,32 +531,9 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
except Exception:
|
||||
logger.exception("AI writing cleanup scheduler shutdown error")
|
||||
|
||||
# 等待所有 in-flight 写作 graph task 跑到下一个 interrupt / 结束。
|
||||
# 这是「写作流程能完整继续」的关键 —— 如果直接关 checkpointer
|
||||
# connection，正在跑 astream 的 task 会在写 checkpoint 时抛错，
|
||||
# 已经生成但还没落盘的进度全丢。
|
||||
try:
|
||||
from app.gateway.routers.ai_writing import wait_for_inflight_graph_tasks
|
||||
|
||||
await wait_for_inflight_graph_tasks(timeout=_AI_WRITING_DRAIN_TIMEOUT_SECONDS)
|
||||
except Exception:
|
||||
logger.exception("Wait for in-flight AI writing tasks failed (non-fatal)")
|
||||
|
||||
# 关闭 AI 写作 graph 持有的 checkpointer 连接（aiosqlite Connection / pg pool）
|
||||
try:
|
||||
from deerflow.agents.ai_writing.graph import shutdown_ai_writing_graph
|
||||
|
||||
await asyncio.wait_for(
|
||||
shutdown_ai_writing_graph(),
|
||||
timeout=_SHUTDOWN_HOOK_TIMEOUT_SECONDS,
|
||||
)
|
||||
except TimeoutError:
|
||||
logger.warning(
|
||||
"AI writing checkpointer shutdown exceeded %.1fs; proceeding with worker exit.",
|
||||
_SHUTDOWN_HOOK_TIMEOUT_SECONDS,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to shutdown AI writing checkpointer")
|
||||
# 阶段 7 删冗余：AI 写作 graph 不再自管 checkpointer / in-flight task；
|
||||
# thread 状态由 LangGraph Server 主 checkpointer 接管，lifespan 出口由
|
||||
# ``langgraph_runtime()`` 的 contextmanager 统一关闭，无需在这里特殊处理。
|
||||
|
||||
try:
|
||||
scheduler_service = getattr(app.state, "scheduled_task_service", None)
|
||||
|
||||
@@ -130,14 +130,24 @@ async def langgraph_runtime(app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
|
||||
app.state.thread_share_store = make_thread_share_store(sf)
|
||||
|
||||
from deerflow.persistence.ai_writing_sessions import make_ai_writing_session_store
|
||||
from deerflow.persistence.ai_writing_sessions import (
|
||||
make_ai_writing_session_store,
|
||||
register_default_store as _register_ai_writing_default_store,
|
||||
)
|
||||
|
||||
app.state.ai_writing_session_store = make_ai_writing_session_store(sf)
|
||||
# 让 LangGraph Server 路径下的 AI 写作节点能跨 app/harness 边界拿到 repo。
|
||||
# 老 /start /stream /resume 路由仍走 request.app.state，互不影响。
|
||||
_register_ai_writing_default_store(app.state.ai_writing_session_store)
|
||||
|
||||
from deerflow.persistence.article_types import make_article_type_store
|
||||
|
||||
app.state.article_type_store = make_article_type_store(sf)
|
||||
|
||||
from deerflow.persistence.memory_config import make_user_memory_config_store
|
||||
|
||||
app.state.memory_config_store = make_user_memory_config_store(sf)
|
||||
|
||||
# Run event store (has its own factory with config-driven backend selection)
|
||||
run_events_config = getattr(config, "run_events", None)
|
||||
app.state.run_event_store = make_run_event_store(run_events_config)
|
||||
@@ -261,6 +271,14 @@ def get_article_type_store(request: Request):
|
||||
return getattr(request.app.state, "article_type_store", None)
|
||||
|
||||
|
||||
def get_memory_config_store(request: Request):
|
||||
"""Return the per-user memory config store."""
|
||||
val = getattr(request.app.state, "memory_config_store", None)
|
||||
if val is None:
|
||||
raise HTTPException(status_code=503, detail="Memory config store not available")
|
||||
return val
|
||||
|
||||
|
||||
def get_run_context(request: Request) -> RunContext:
|
||||
"""Build a :class:`RunContext` from ``app.state`` singletons.
|
||||
|
||||
|
||||
@@ -1,31 +1,26 @@
|
||||
"""AI 写作路由 —— 独立于原有路由，前缀 /api/ai-writing/
|
||||
"""AI 写作路由 —— 业务表 CRUD + transcript 持久化 + admin cleanup。
|
||||
|
||||
写作主流程（``/start /stream /resume /state``）已迁到 LangGraph Server 的
|
||||
``ai_writing`` assistant（前端走 ``useAIWritingStream`` hook + SDK），路由层
|
||||
不再承载 graph 推进，只保留以下「不依赖 graph 状态」的辅助接口：
|
||||
|
||||
POST /api/ai-writing/start 启动写作流程，返回 session_id
|
||||
GET /api/ai-writing/{id}/stream SSE 流（进度/干预/草稿事件）
|
||||
POST /api/ai-writing/{id}/resume 用户干预响应，恢复流程
|
||||
GET /api/ai-writing/{id}/state 查询当前流程状态
|
||||
GET /api/ai-writing/sessions 列出用户的会话历史
|
||||
GET /api/ai-writing/sessions/{id} 获取单个历史会话
|
||||
PATCH /api/ai-writing/sessions/{id} 重命名会话标题
|
||||
DELETE /api/ai-writing/sessions/{id} 删除会话
|
||||
PUT /api/ai-writing/sessions/{id}/transcript 保存完整对话时间线
|
||||
GET /api/ai-writing/article-types 列出所有文章类型
|
||||
POST /api/ai-writing/article-types 新增文章类型
|
||||
PATCH /api/ai-writing/article-types/{id} 更新文章类型
|
||||
DELETE /api/ai-writing/article-types/{id} 删除文章类型
|
||||
POST /api/ai-writing/admin/cleanup （admin）手动清理会话
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from collections.abc import AsyncGenerator
|
||||
from datetime import UTC, datetime
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
from langgraph.types import Command
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -34,38 +29,6 @@ router = APIRouter(prefix="/api/ai-writing", tags=["ai-writing"])
|
||||
|
||||
# ── 请求/响应模型 ─────────────────────────────────────────────────────────────
|
||||
|
||||
class StartWritingRequest(BaseModel):
|
||||
user_intent: str
|
||||
article_type: str = "analysis"
|
||||
target_audience: str = "普通读者"
|
||||
word_count_target: int = 800
|
||||
model_name: str = ""
|
||||
# 素材检索关键词数量：用户指定用几个关键词检索（1~10，默认 6）
|
||||
keyword_count: int = 6
|
||||
# 写作模式：strict=严格依据素材 / loose=可结合通用知识（默认 loose）
|
||||
writing_mode: str = "loose"
|
||||
|
||||
|
||||
class StartWritingResponse(BaseModel):
|
||||
session_id: str
|
||||
message: str = "写作流程已启动"
|
||||
|
||||
|
||||
class ResumeRequest(BaseModel):
|
||||
action: str
|
||||
pause_point: str = ""
|
||||
approved_material_ids: list[str] = Field(default_factory=list)
|
||||
extra_keywords: list[str] = Field(default_factory=list)
|
||||
edited_outline: Optional[dict] = None
|
||||
outline_feedback: str = ""
|
||||
user_revision_notes: str = ""
|
||||
override_verdict: bool = False
|
||||
selected_issue_indices: Optional[list[int]] = None
|
||||
# 素材不足求助点：每个被阻塞章节的处理决策
|
||||
# [{"sectionTitle": "...", "decision": "supplement|delete|loose"}]
|
||||
section_decisions: list[dict] = Field(default_factory=list)
|
||||
|
||||
|
||||
class RenameSessionRequest(BaseModel):
|
||||
title: str
|
||||
|
||||
@@ -85,8 +48,8 @@ class AdminCleanupRequest(BaseModel):
|
||||
class AdminCleanupResponse(BaseModel):
|
||||
deleted_db: int
|
||||
deleted_checkpoints: int
|
||||
# 因「仍在跑」被保护掉的会话数。永远跳过 _event_queues 里活跃的会话，
|
||||
# 即使 retention_days=0 也不会秒杀用户正在写的稿。
|
||||
# 因「仍在跑 / 等用户干预」被保护掉的会话数（threads_meta.status in
|
||||
# running/interrupted）。即使 retention_days=0 也不会秒杀用户正在写的稿。
|
||||
skipped_active: int = 0
|
||||
elapsed_seconds: float
|
||||
params: dict
|
||||
@@ -142,244 +105,7 @@ class SessionResponse(BaseModel):
|
||||
updated_at: str
|
||||
|
||||
|
||||
# ── 内部状态 ──────────────────────────────────────────────────────────────────
|
||||
|
||||
# session_id → asyncio.Queue
|
||||
_event_queues: dict[str, asyncio.Queue] = {}
|
||||
|
||||
# session_id → model_name（resume 时复用）
|
||||
_session_model_names: dict[str, str] = {}
|
||||
|
||||
# session_id → AIWritingSessionRepository（for DB saves）
|
||||
_session_repos: dict[str, object] = {}
|
||||
|
||||
# session_id → user_id (str | None, already resolved)
|
||||
_session_user_ids: dict[str, str | None] = {}
|
||||
|
||||
# session_id → accumulated intervention records
|
||||
_session_interventions: dict[str, list] = {}
|
||||
|
||||
# session_id → True if session ended with error
|
||||
_session_error_flags: dict[str, bool] = {}
|
||||
|
||||
# session_id → valid_article_types（list[dict] with key/label）
|
||||
_session_article_types: dict[str, list[dict]] = {}
|
||||
|
||||
# session_id → pending await_user event tuple for SSE reconnect re-delivery
|
||||
# asyncio.Queue is consume-once; if the old SSE generator dequeued the await_user event
|
||||
# before the client disconnected, the event is lost from the queue. We save a copy here
|
||||
# so that a new SSE connection (e.g., after user clicks "继续") can re-inject it.
|
||||
_session_pending_pause: dict[str, tuple] = {}
|
||||
|
||||
# Serializes session recovery so that a concurrent /stream + /resume on the same
|
||||
# lost session don't each rebuild a different queue object.
|
||||
_recovery_lock = asyncio.Lock()
|
||||
|
||||
# 在飞的 graph 推进任务（每次 /start 和 /resume 都会塞一个进来）。
|
||||
# Lifespan shutdown 时会等这些 task 跑到下一个 interrupt / 结束，再关闭
|
||||
# checkpointer 连接 —— 否则中途关 connection，task 内部的 sqlite 写入会抛
|
||||
# 错，导致 LLM 已经生成但还没落 checkpoint 的进度丢失。
|
||||
_inflight_graph_tasks: set[asyncio.Task] = set()
|
||||
|
||||
|
||||
# ── 工具函数 ──────────────────────────────────────────────────────────────────
|
||||
|
||||
def _sse(event: str, data: dict) -> str:
|
||||
"""组装一帧 SSE。
|
||||
|
||||
JSON 序列化对 progress_events 来说理论上不会失败（节点里塞的都是 str/dict
|
||||
/list/int 这类基础类型），但一旦塞进了 datetime / bytes / 自定义对象，
|
||||
``json.dumps`` 会抛 TypeError → SSE 生成器爆掉 → 客户端再也收不到任何
|
||||
事件，包括关键的 await_user 和 done。
|
||||
|
||||
这里用 ``default=str`` 兜底任何「带 __str__」的对象；如果连 ``str()``
|
||||
都炸了，最后还兜一次手写一条 error 帧推出去 —— 写作流程「能继续」比
|
||||
「漂亮地序列化」更重要。
|
||||
"""
|
||||
try:
|
||||
payload = json.dumps(data, ensure_ascii=False, default=str)
|
||||
except Exception as exc:
|
||||
logger.exception("ai_writing: SSE payload 序列化失败 event=%s", event)
|
||||
payload = json.dumps(
|
||||
{
|
||||
"session_id": data.get("session_id") if isinstance(data, dict) else None,
|
||||
"message": f"事件序列化失败（已降级）: {type(exc).__name__}",
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
event = "notice" # 不要硬塞回原 event，万一是 error 会污染语义
|
||||
return f"event: {event}\ndata: {payload}\n\n"
|
||||
|
||||
|
||||
def _ai_writing_config() -> dict:
|
||||
"""读取 config.yaml 的 ai_writing 段（含素材检索 provider 配置）。"""
|
||||
try:
|
||||
from deerflow.config import get_app_config
|
||||
|
||||
aw = getattr(get_app_config(), "ai_writing", None)
|
||||
if isinstance(aw, dict):
|
||||
return aw
|
||||
return dict(aw) if aw else {}
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
|
||||
def _make_config(session_id: str, model_name: str = "", valid_article_types: list[dict] | None = None) -> dict:
|
||||
return {
|
||||
"configurable": {
|
||||
"thread_id": session_id,
|
||||
"model_name": model_name or _session_model_names.get(session_id, ""),
|
||||
"ai_writing_config": _ai_writing_config(),
|
||||
# 把 SSE 队列注入，供草稿/大纲节点直接推流式 chunk
|
||||
"sse_queue": _event_queues.get(session_id),
|
||||
"sse_session_id": session_id,
|
||||
# 动态文章类型列表，intent_parser 从此处读取有效枚举
|
||||
"valid_article_types": valid_article_types or _session_article_types.get(session_id),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def _truncate_title(text: str, max_len: int = 50) -> str:
|
||||
text = text.strip()
|
||||
return text[:max_len] + ("…" if len(text) > max_len else "")
|
||||
|
||||
|
||||
def _spawn_graph_task(coro) -> asyncio.Task:
|
||||
"""启一个 graph 推进 task 并登记到 _inflight_graph_tasks。
|
||||
|
||||
用 ``add_done_callback`` 在完成时自动反登记 —— 不管 task 正常结束、抛
|
||||
异常、还是被 cancel，都会清理；这样 _inflight_graph_tasks 不会泄漏。
|
||||
"""
|
||||
task = asyncio.create_task(coro)
|
||||
_inflight_graph_tasks.add(task)
|
||||
task.add_done_callback(_inflight_graph_tasks.discard)
|
||||
return task
|
||||
|
||||
|
||||
async def wait_for_inflight_graph_tasks(timeout: float) -> tuple[int, int]:
|
||||
"""Lifespan shutdown 入口：阻塞等待所有 in-flight graph task 收尾。
|
||||
|
||||
每个 ``_stream_graph_segment`` 跑到下一个 ``interrupt()`` 或 ``END`` 就
|
||||
会自然返回；通常只需要等几秒到几十秒。超时后剩下的会被 ``cancel()``，
|
||||
但 cancel 完仍要等它们 ack 退出（带极短的 grace），否则下一步关
|
||||
checkpointer 时还会撞上活的协程。
|
||||
|
||||
Returns:
|
||||
(drained, cancelled): 自然结束的数量、被 cancel 强制结束的数量
|
||||
"""
|
||||
if not _inflight_graph_tasks:
|
||||
return 0, 0
|
||||
|
||||
pending = set(_inflight_graph_tasks)
|
||||
initial = len(pending)
|
||||
logger.info("ai_writing: 等待 %d 个 in-flight graph task 收尾（timeout=%.1fs）", initial, timeout)
|
||||
|
||||
done, still_running = await asyncio.wait(pending, timeout=timeout)
|
||||
drained = len(done)
|
||||
|
||||
cancelled = 0
|
||||
if still_running:
|
||||
logger.warning(
|
||||
"ai_writing: 仍有 %d 个 task 超时未结束，强制 cancel", len(still_running)
|
||||
)
|
||||
for t in still_running:
|
||||
t.cancel()
|
||||
cancelled += 1
|
||||
# 给 cancel 一点 ack 时间，避免下游关连接时撞活协程
|
||||
await asyncio.wait(still_running, timeout=2.0)
|
||||
|
||||
logger.info("ai_writing: in-flight 收尾完成 drained=%d cancelled=%d", drained, cancelled)
|
||||
return drained, cancelled
|
||||
|
||||
|
||||
# ── 路由 ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
@router.post("/start", response_model=StartWritingResponse)
|
||||
async def start_writing(body: StartWritingRequest, request: Request):
|
||||
"""启动写作流程，返回 session_id。
|
||||
|
||||
**失败兜底契约**：本接口除非 Pydantic 校验 body 失败（422），其他任何
|
||||
后端组件（DB / checkpointer / 文章类型仓库 / 用户解析）失败都不会导致
|
||||
/start 返回 5xx。哪怕兜底走的是「极简内存模式」也要先把 session_id
|
||||
交给前端，让它能起 SSE 接到 error 事件，避免 UI 卡死在 loading 状态。
|
||||
"""
|
||||
from deerflow.agents.ai_writing.graph import get_ai_writing_graph
|
||||
from app.gateway.deps import get_ai_writing_session_store, get_article_type_store, get_current_user
|
||||
|
||||
# get_ai_writing_graph() 内部已对 sqlite/postgres 失败做了 InMemorySaver
|
||||
# 降级，理论上永不抛。但 build_ai_writing_graph 拼图结构也可能因下游
|
||||
# 节点 import 失败抛错 —— 这种情况 graph 完全不可用，必须早 fail，前端
|
||||
# 弹错处理。
|
||||
try:
|
||||
graph = await get_ai_writing_graph()
|
||||
except Exception as e:
|
||||
logger.exception("ai_writing: graph 构建彻底失败，无法启动")
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail=f"AI 写作服务暂不可用: {type(e).__name__}",
|
||||
) from e
|
||||
|
||||
session_id = str(uuid.uuid4())
|
||||
queue: asyncio.Queue = asyncio.Queue()
|
||||
_event_queues[session_id] = queue
|
||||
_session_model_names[session_id] = body.model_name
|
||||
_session_interventions[session_id] = []
|
||||
|
||||
# 加载动态文章类型列表，注入给 intent_parser 节点
|
||||
article_type_store = get_article_type_store(request)
|
||||
valid_article_types: list[dict] | None = None
|
||||
if article_type_store is not None:
|
||||
try:
|
||||
records = await article_type_store.list()
|
||||
if records:
|
||||
valid_article_types = [{"key": r["key"], "label": r["label"]} for r in records]
|
||||
except Exception:
|
||||
logger.exception("ai_writing: failed to load article types for session=%s", session_id)
|
||||
_session_article_types[session_id] = valid_article_types or []
|
||||
|
||||
# 持久化会话记录（DB 可能未配置，则跳过）
|
||||
repo = get_ai_writing_session_store(request)
|
||||
user_id = await get_current_user(request)
|
||||
if repo is not None:
|
||||
try:
|
||||
title = _truncate_title(body.user_intent)
|
||||
await repo.create(
|
||||
id=session_id,
|
||||
title=title,
|
||||
user_intent=body.user_intent,
|
||||
user_id=user_id,
|
||||
)
|
||||
_session_repos[session_id] = repo
|
||||
_session_user_ids[session_id] = user_id
|
||||
except Exception:
|
||||
logger.exception("ai_writing: failed to create session record: session=%s", session_id)
|
||||
|
||||
initial_state = {
|
||||
"user_intent": body.user_intent,
|
||||
"article_type": body.article_type,
|
||||
"target_audience": body.target_audience,
|
||||
"word_count_target": body.word_count_target,
|
||||
# 规整到 [1, 10]，非法值回退默认 6
|
||||
"keyword_count": max(1, min(10, body.keyword_count)),
|
||||
# 写作模式，非法值回退 loose
|
||||
"writing_mode": body.writing_mode if body.writing_mode in ("strict", "loose") else "loose",
|
||||
"revision_count": 0,
|
||||
"max_revisions": 3,
|
||||
"status": "researching",
|
||||
"drafts": [],
|
||||
"progress_events": [],
|
||||
"material_package": None,
|
||||
"current_outline": None,
|
||||
"current_draft": None,
|
||||
"review_result": None,
|
||||
"last_user_intervention": None,
|
||||
}
|
||||
config = _make_config(session_id, body.model_name, valid_article_types=valid_article_types)
|
||||
|
||||
_spawn_graph_task(_stream_graph_segment(session_id, initial_state, config, graph))
|
||||
return StartWritingResponse(session_id=session_id)
|
||||
|
||||
|
||||
@router.get("/sessions", response_model=list[SessionResponse])
|
||||
async def list_sessions(request: Request):
|
||||
"""列出当前用户的 AI 写作历史会话。"""
|
||||
@@ -460,8 +186,6 @@ async def save_transcript(session_id: str, body: SaveTranscriptRequest, request:
|
||||
# 会话不存在或不属于当前用户：同样不抛 500，返回 success=False 即可
|
||||
return {"success": False, "reason": "session_not_found"}
|
||||
return {"success": True}
|
||||
|
||||
|
||||
@router.get("/article-types", response_model=list[ArticleTypeResponse])
|
||||
async def list_article_types(request: Request):
|
||||
"""列出所有文章类型，按 sort_order 升序排列。"""
|
||||
@@ -535,91 +259,6 @@ async def delete_article_type(article_type_id: str, request: Request):
|
||||
return {"success": True}
|
||||
|
||||
|
||||
@router.get("/{session_id}/stream")
|
||||
async def stream_progress(session_id: str):
|
||||
"""SSE 流：推送进度事件、干预请求、草稿内容。"""
|
||||
return StreamingResponse(
|
||||
_sse_generator(session_id),
|
||||
media_type="text/event-stream",
|
||||
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/{session_id}/resume")
|
||||
async def resume_writing(session_id: str, body: ResumeRequest, request: Request):
|
||||
"""用户提交干预决策，恢复暂停的流程。"""
|
||||
from deerflow.agents.ai_writing.graph import get_ai_writing_graph
|
||||
from app.gateway.deps import get_current_user
|
||||
|
||||
graph = await get_ai_writing_graph()
|
||||
|
||||
if session_id not in _event_queues:
|
||||
# 内存会话注册表丢失（进程内清理后重连等）。先尝试从 checkpointer 恢复，
|
||||
# 让用户能在干预点继续，而不是直接 404 中断整个写作流程。
|
||||
recovery = await _recover_session(session_id, request)
|
||||
if recovery != "interrupted":
|
||||
logger.warning(
|
||||
"ai_writing: resume on unrecoverable session=%s (recovery=%s)", session_id, recovery
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail="会话状态已丢失或流程已结束，无法继续，请重新发起写作。",
|
||||
)
|
||||
|
||||
config = _make_config(session_id)
|
||||
intervention = {
|
||||
"action": body.action,
|
||||
"approvedMaterialIds": body.approved_material_ids,
|
||||
"extraKeywords": body.extra_keywords,
|
||||
"editedOutline": body.edited_outline,
|
||||
"outlineFeedback": body.outline_feedback,
|
||||
"userRevisionNotes": body.user_revision_notes,
|
||||
"overrideVerdict": body.override_verdict,
|
||||
"selectedIssueIndices": body.selected_issue_indices,
|
||||
"sectionDecisions": body.section_decisions,
|
||||
}
|
||||
|
||||
# Record this intervention for session history
|
||||
record = {
|
||||
"pausePoint": body.pause_point,
|
||||
"action": body.action,
|
||||
"timestamp": datetime.now(UTC).isoformat(),
|
||||
}
|
||||
if session_id in _session_interventions:
|
||||
_session_interventions[session_id].append(record)
|
||||
repo = _session_repos.get(session_id)
|
||||
user_id = _session_user_ids.get(session_id)
|
||||
if repo is not None:
|
||||
try:
|
||||
await repo.update(
|
||||
session_id,
|
||||
user_id=user_id,
|
||||
completed_interventions=_session_interventions[session_id],
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("ai_writing: failed to save intervention: session=%s", session_id)
|
||||
|
||||
# Clear pending pause so SSE reconnect won't re-inject the stale await_user event
|
||||
_session_pending_pause.pop(session_id, None)
|
||||
|
||||
_spawn_graph_task(
|
||||
_stream_graph_segment(session_id, Command(resume=intervention), config, graph)
|
||||
)
|
||||
return {"message": "已提交，流程恢复中"}
|
||||
|
||||
|
||||
@router.get("/{session_id}/state")
|
||||
async def get_state(session_id: str):
|
||||
"""查询当前流程状态快照。"""
|
||||
from deerflow.agents.ai_writing.graph import get_ai_writing_graph
|
||||
|
||||
graph = await get_ai_writing_graph()
|
||||
config = _make_config(session_id)
|
||||
snapshot = await graph.aget_state(config)
|
||||
if not snapshot or not snapshot.values:
|
||||
raise HTTPException(status_code=404, detail="Session not found")
|
||||
return snapshot.values
|
||||
|
||||
|
||||
# ── 管理接口 ──────────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -688,9 +327,20 @@ async def admin_cleanup(
|
||||
await _require_admin(request)
|
||||
params = _resolve_cleanup_params(body)
|
||||
|
||||
# 透传主 LangGraph checkpointer / thread_store —— 新路径下 ai_writing 的
|
||||
# thread 状态在主 checkpointer 里，thread_store 用来按 status 过滤活跃 thread。
|
||||
# getattr 嵌套防御：测试用 stub Request 时 ``request.app`` 可能不存在。
|
||||
_app_state = getattr(getattr(request, "app", None), "state", None)
|
||||
main_checkpointer = getattr(_app_state, "checkpointer", None)
|
||||
main_thread_store = getattr(_app_state, "thread_store", None)
|
||||
|
||||
started = time.monotonic()
|
||||
try:
|
||||
result = await run_cleanup_once(**params)
|
||||
result = await run_cleanup_once(
|
||||
**params,
|
||||
checkpointer=main_checkpointer,
|
||||
thread_store=main_thread_store,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception("ai_writing admin cleanup failed: params=%s", params)
|
||||
# 抛 500 是合理的 —— 管理员触发清理失败需要看到失败原因，跟「业务
|
||||
@@ -715,292 +365,3 @@ async def admin_cleanup(
|
||||
elapsed_seconds=round(elapsed, 3),
|
||||
params=params,
|
||||
)
|
||||
|
||||
|
||||
# ── 内部辅助 ──────────────────────────────────────────────────────────────────
|
||||
|
||||
async def _stream_graph_segment(
|
||||
session_id: str,
|
||||
input_or_command,
|
||||
config: dict,
|
||||
graph,
|
||||
):
|
||||
"""运行一个 Graph 片段（初始运行或 resume），将输出推入队列。
|
||||
遇到 interrupt 时停止（不关闭队列），等待下次 resume。
|
||||
Graph 结束时关闭队列（放入 None 哨兵）。
|
||||
"""
|
||||
queue = _event_queues.get(session_id)
|
||||
if queue is None:
|
||||
return
|
||||
|
||||
reached_end = False
|
||||
is_error = False
|
||||
try:
|
||||
async for chunk in graph.astream(input_or_command, config=config, stream_mode="updates"):
|
||||
node_name = next(iter(chunk))
|
||||
node_output = chunk[node_name]
|
||||
|
||||
# interrupt() 触发
|
||||
if node_name == "__interrupt__":
|
||||
# node_output 是 (Interrupt(...),) 元组
|
||||
interrupt_list = node_output if isinstance(node_output, (list, tuple)) else [node_output]
|
||||
for intr in interrupt_list:
|
||||
interrupt_data = intr.value if hasattr(intr, "value") else intr
|
||||
event_tuple = ("await_user", interrupt_data)
|
||||
await queue.put(event_tuple)
|
||||
# Save for SSE reconnect re-delivery: if the old generator consumed this
|
||||
# event before the client disconnected, a new connection would block forever.
|
||||
_session_pending_pause[session_id] = event_tuple
|
||||
return # 等待 resume，不关闭队列
|
||||
|
||||
if not isinstance(node_output, dict):
|
||||
continue
|
||||
|
||||
# 推送 progress_events
|
||||
for evt in node_output.get("progress_events") or []:
|
||||
await queue.put((_event_type(evt), {**evt, "session_id": session_id}))
|
||||
|
||||
# 素材就绪
|
||||
if pkg := node_output.get("material_package"):
|
||||
await queue.put(("materials_ready", {
|
||||
"session_id": session_id,
|
||||
"keywords": pkg.get("keywords", []),
|
||||
"materials": pkg.get("materials", []),
|
||||
"summary": pkg.get("summary", ""),
|
||||
}))
|
||||
|
||||
# 大纲就绪
|
||||
if outline := node_output.get("current_outline"):
|
||||
await queue.put(("outline_ready", {"session_id": session_id, "outline": outline}))
|
||||
|
||||
# 草稿就绪 — 保存到 DB
|
||||
if draft := node_output.get("current_draft"):
|
||||
draft_title = draft.get("title", "")
|
||||
draft_markdown = draft.get("full_markdown", "")
|
||||
await queue.put(("draft_ready", {
|
||||
"session_id": session_id,
|
||||
"title": draft_title,
|
||||
"markdown": draft_markdown,
|
||||
"revision_count": draft.get("revision_count", 0),
|
||||
}))
|
||||
repo = _session_repos.get(session_id)
|
||||
user_id = _session_user_ids.get(session_id)
|
||||
if repo is not None:
|
||||
try:
|
||||
await repo.update(
|
||||
session_id,
|
||||
user_id=user_id,
|
||||
status="draft_ready",
|
||||
title=draft_title or _truncate_title(
|
||||
next(iter(_session_interventions.get(session_id, [{}])), {}).get("action", "草稿"),
|
||||
),
|
||||
draft_title=draft_title,
|
||||
draft_markdown=draft_markdown,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("ai_writing: failed to save draft: session=%s", session_id)
|
||||
|
||||
# 编辑审核就绪 — 保存到 DB
|
||||
if review := node_output.get("review_result"):
|
||||
await queue.put(("review_ready", {"session_id": session_id, "review_result": review}))
|
||||
repo = _session_repos.get(session_id)
|
||||
user_id = _session_user_ids.get(session_id)
|
||||
if repo is not None:
|
||||
try:
|
||||
await repo.update(
|
||||
session_id,
|
||||
user_id=user_id,
|
||||
status="review_ready",
|
||||
review_result=review,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("ai_writing: failed to save review: session=%s", session_id)
|
||||
|
||||
# 状态 done：Graph 结束
|
||||
if node_output.get("status") == "done":
|
||||
reached_end = True
|
||||
|
||||
# astream 正常结束 = Graph 执行到 END
|
||||
reached_end = True
|
||||
await queue.put(("done", {"session_id": session_id, "message": "写作完成"}))
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("ai_writing graph segment failed: session=%s", session_id)
|
||||
await queue.put(("error", {"session_id": session_id, "message": str(e)}))
|
||||
is_error = True
|
||||
reached_end = True
|
||||
finally:
|
||||
if reached_end:
|
||||
# 更新最终状态到 DB
|
||||
repo = _session_repos.get(session_id)
|
||||
user_id = _session_user_ids.get(session_id)
|
||||
if repo is not None:
|
||||
try:
|
||||
await repo.update(
|
||||
session_id,
|
||||
user_id=user_id,
|
||||
status="error" if is_error else "done",
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("ai_writing: failed to update final status: session=%s", session_id)
|
||||
|
||||
await queue.put(None) # 哨兵，通知 SSE 生成器关闭
|
||||
_event_queues.pop(session_id, None)
|
||||
_session_model_names.pop(session_id, None)
|
||||
_session_repos.pop(session_id, None)
|
||||
_session_user_ids.pop(session_id, None)
|
||||
_session_interventions.pop(session_id, None)
|
||||
_session_error_flags.pop(session_id, None)
|
||||
_session_article_types.pop(session_id, None)
|
||||
_session_pending_pause.pop(session_id, None)
|
||||
|
||||
|
||||
def _event_type(evt: dict) -> str:
|
||||
t = evt.get("type", "progress")
|
||||
# 过滤掉已在上方单独推送的事件类型，避免重复
|
||||
if t in ("materials_ready", "outline_ready", "draft_ready", "review_ready", "await_user",
|
||||
"draft_chunk", "outline_chunk"):
|
||||
return "step_done"
|
||||
return t
|
||||
|
||||
|
||||
async def _load_graph_snapshot(session_id: str):
|
||||
"""读取 AI 写作 Graph 在该 session 上的状态快照（失败/无状态返回 None）。
|
||||
|
||||
Graph 的 checkpointer 现在按 config.yaml 持久化到磁盘（默认 SQLite），
|
||||
所以容器重启后只要数据卷还在，仍然能恢复出来。
|
||||
"""
|
||||
try:
|
||||
from deerflow.agents.ai_writing.graph import get_ai_writing_graph
|
||||
|
||||
graph = await get_ai_writing_graph()
|
||||
snapshot = await graph.aget_state(_make_config(session_id))
|
||||
except Exception:
|
||||
logger.exception("ai_writing: aget_state failed during recovery: session=%s", session_id)
|
||||
return None
|
||||
if not snapshot or not getattr(snapshot, "values", None):
|
||||
return None
|
||||
return snapshot
|
||||
|
||||
|
||||
def _extract_interrupt(snapshot) -> dict | None:
|
||||
"""从快照里取出当前挂起的 interrupt 负载（即 await_user 的数据）。"""
|
||||
for task in getattr(snapshot, "tasks", None) or ():
|
||||
for intr in getattr(task, "interrupts", None) or ():
|
||||
value = getattr(intr, "value", None)
|
||||
if isinstance(value, dict):
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
async def _recover_session(session_id: str, request: Request | None = None) -> str:
|
||||
"""内存会话注册表丢失时，尝试从 checkpointer 重建会话。
|
||||
|
||||
返回:
|
||||
- "interrupted" —— 图正停在某个干预点，已重建队列并预置 await_user 事件，可继续 resume
|
||||
- "finished" —— 图已完成，已重建队列并预置 done 事件（+草稿）
|
||||
- "unrecoverable" —— 无可恢复状态（进程/容器重启过，或 session 从未存在）
|
||||
|
||||
幂等：内存里已有队列时直接当作 "interrupted" 返回。整段在 `_recovery_lock`
|
||||
下执行，避免并发的 /stream 与 /resume 各自重建出不同的队列对象。
|
||||
"""
|
||||
async with _recovery_lock:
|
||||
if session_id in _event_queues:
|
||||
return "interrupted"
|
||||
|
||||
snapshot = await _load_graph_snapshot(session_id)
|
||||
if snapshot is None:
|
||||
return "unrecoverable"
|
||||
|
||||
queue: asyncio.Queue = asyncio.Queue()
|
||||
_event_queues[session_id] = queue
|
||||
_session_interventions.setdefault(session_id, [])
|
||||
_session_model_names.setdefault(session_id, "")
|
||||
_session_article_types.setdefault(session_id, [])
|
||||
|
||||
# 尽量恢复 DB 仓库引用，使后续 resume 仍能持久化草稿/审核结果
|
||||
if request is not None and session_id not in _session_repos:
|
||||
try:
|
||||
from app.gateway.deps import get_ai_writing_session_store, get_current_user
|
||||
|
||||
repo = get_ai_writing_session_store(request)
|
||||
if repo is not None:
|
||||
_session_repos[session_id] = repo
|
||||
_session_user_ids[session_id] = await get_current_user(request)
|
||||
except Exception:
|
||||
logger.exception("ai_writing: failed to restore repo during recovery: session=%s", session_id)
|
||||
|
||||
interrupt_val = _extract_interrupt(snapshot)
|
||||
if interrupt_val is not None:
|
||||
event = ("await_user", {**interrupt_val, "session_id": session_id})
|
||||
await queue.put(event)
|
||||
_session_pending_pause[session_id] = event
|
||||
logger.info("ai_writing: recovered interrupted session from checkpointer: session=%s", session_id)
|
||||
return "interrupted"
|
||||
|
||||
values = snapshot.values or {}
|
||||
if not getattr(snapshot, "next", None) or values.get("status") == "done":
|
||||
draft = values.get("current_draft") or {}
|
||||
if draft.get("full_markdown"):
|
||||
await queue.put(("draft_ready", {
|
||||
"session_id": session_id,
|
||||
"title": draft.get("title", ""),
|
||||
"markdown": draft.get("full_markdown", ""),
|
||||
"revision_count": draft.get("revision_count", 0),
|
||||
}))
|
||||
await queue.put(("done", {"session_id": session_id, "message": "写作已完成"}))
|
||||
await queue.put(None) # 哨兵：让 SSE 生成器读完即关闭
|
||||
logger.info("ai_writing: recovered finished session from checkpointer: session=%s", session_id)
|
||||
return "finished"
|
||||
|
||||
# 既不在干预点、也未结束：上次运行多半被中断在某个节点，无法续推
|
||||
await queue.put(("error", {
|
||||
"session_id": session_id,
|
||||
"message": "写作流程已中断，请重新发起写作。",
|
||||
}))
|
||||
await queue.put(None)
|
||||
logger.warning("ai_writing: session checkpoint stuck mid-flow, not resumable: session=%s", session_id)
|
||||
return "finished"
|
||||
|
||||
|
||||
async def _sse_generator(session_id: str) -> AsyncGenerator[str, None]:
|
||||
"""从队列消费事件，推送为 SSE 帧。"""
|
||||
# 等待队列创建（最多 5 秒）
|
||||
for _ in range(50):
|
||||
if session_id in _event_queues:
|
||||
break
|
||||
await asyncio.sleep(0.1)
|
||||
else:
|
||||
# 内存里没有这个会话队列。不要直接报「Session not found」——
|
||||
# 先尝试从 checkpointer 恢复（图可能仍停在某个干预点 / 已完成）。
|
||||
recovery = await _recover_session(session_id)
|
||||
if recovery == "unrecoverable":
|
||||
yield _sse("connected", {"session_id": session_id})
|
||||
yield _sse("error", {
|
||||
"session_id": session_id,
|
||||
"message": "会话状态已丢失（服务可能已重启），请重新发起写作。",
|
||||
})
|
||||
return
|
||||
# 恢复成功：_event_queues[session_id] 已重建并预置事件，继续走正常消费逻辑。
|
||||
|
||||
queue = _event_queues.get(session_id)
|
||||
if queue is None:
|
||||
# 极端兜底：恢复后又被并发清理
|
||||
yield _sse("connected", {"session_id": session_id})
|
||||
yield _sse("error", {"session_id": session_id, "message": "会话状态已丢失，请重新发起写作。"})
|
||||
return
|
||||
yield _sse("connected", {"session_id": session_id})
|
||||
|
||||
# Re-inject the saved await_user event if the queue is empty and a pending pause exists.
|
||||
# This covers the case where the old SSE generator dequeued the await_user event and
|
||||
# yielded it to a disconnecting client — the event is then lost from the consume-once queue.
|
||||
if queue.empty() and session_id in _session_pending_pause:
|
||||
await queue.put(_session_pending_pause[session_id])
|
||||
|
||||
while True:
|
||||
item = await queue.get()
|
||||
if item is None:
|
||||
break
|
||||
event_type, data = item
|
||||
yield _sse(event_type, data)
|
||||
await asyncio.sleep(0)
|
||||
@@ -21,15 +21,16 @@ from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Query
|
||||
from fastapi import APIRouter, HTTPException, Query, Request
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.gateway.deps import get_memory_config_store
|
||||
from deerflow.config.agents_config import AGENT_NAME_PATTERN
|
||||
from deerflow.config.memory_config import (
|
||||
USER_WRITABLE_FIELDS,
|
||||
delete_user_memory_override,
|
||||
delete_user_memory_override_db,
|
||||
get_effective_memory_config,
|
||||
save_user_memory_override,
|
||||
save_user_memory_override_db,
|
||||
)
|
||||
from deerflow.config.paths import get_paths
|
||||
from deerflow.runtime.user_context import get_effective_user_id
|
||||
@@ -295,14 +296,17 @@ async def get_memory_v2_config() -> MemoryV2ConfigResponse:
|
||||
"/memory/v2/config",
|
||||
response_model=MemoryV2ConfigResponse,
|
||||
summary="Patch Memory V2 Config",
|
||||
description="把白名单字段写入当前用户的 per-user 配置覆盖。非白名单字段被忽略。",
|
||||
description="把 hindsight.recall_budget 写入当前用户的 per-user 配置覆盖 (DB 存储)。",
|
||||
)
|
||||
async def patch_memory_v2_config(request: MemoryConfigPatchRequest) -> MemoryV2ConfigResponse:
|
||||
"""更新当前用户的 per-user 记忆配置覆盖。"""
|
||||
async def patch_memory_v2_config(request: MemoryConfigPatchRequest, req: Request) -> MemoryV2ConfigResponse:
|
||||
"""更新当前用户的 per-user 记忆配置覆盖 (仅 hindsight.recall_budget)。"""
|
||||
user_id = get_effective_user_id()
|
||||
store = get_memory_config_store(req)
|
||||
try:
|
||||
save_user_memory_override(user_id, request.overrides)
|
||||
except OSError as exc:
|
||||
await save_user_memory_override_db(user_id, request.overrides, store)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
except Exception as exc: # noqa: BLE001
|
||||
raise HTTPException(status_code=500, detail="Failed to save memory config override.") from exc
|
||||
|
||||
config = get_effective_memory_config(user_id)
|
||||
@@ -317,18 +321,20 @@ async def patch_memory_v2_config(request: MemoryConfigPatchRequest) -> MemoryV2C
|
||||
response_model=MemoryV2ConfigResponse,
|
||||
summary="Delete Memory V2 Config Override",
|
||||
description=(
|
||||
"撤销当前用户的 per-user 配置覆盖。传 ?field=injection.enabled 只撤销该字段;"
|
||||
"不传 field 则清空整个 per-user 覆盖。"
|
||||
"撤销当前用户的 per-user 配置覆盖。传 ?field=hindsight.recall_budget 只撤销该字段;"
|
||||
"不传 field 则清空全部覆盖。"
|
||||
),
|
||||
)
|
||||
async def delete_memory_v2_config(
|
||||
req: Request,
|
||||
field: str | None = Query(None, description="要撤销的点号路径字段;省略则清空全部覆盖"),
|
||||
) -> MemoryV2ConfigResponse:
|
||||
"""撤销当前用户的 per-user 记忆配置覆盖。"""
|
||||
user_id = get_effective_user_id()
|
||||
store = get_memory_config_store(req)
|
||||
try:
|
||||
delete_user_memory_override(user_id, field)
|
||||
except OSError as exc:
|
||||
await delete_user_memory_override_db(user_id, field, store)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
raise HTTPException(status_code=500, detail="Failed to delete memory config override.") from exc
|
||||
|
||||
config = get_effective_memory_config(user_id)
|
||||
|
||||
@@ -100,7 +100,7 @@ _DEFAULT_ASSISTANT_ID = "lead_agent"
|
||||
|
||||
# Built-in LangGraph graph IDs registered in langgraph.json.
|
||||
# These are NOT stored in the agents DB and must bypass the custom-agent access check.
|
||||
_BUILTIN_GRAPH_IDS: frozenset[str] = frozenset({"lead_agent", "canvas_agent"})
|
||||
_BUILTIN_GRAPH_IDS: frozenset[str] = frozenset({"lead_agent", "canvas_agent", "ai_writing"})
|
||||
|
||||
|
||||
def _prepend_prompt_prefix(graph_input: dict[str, Any], prefix: str) -> dict[str, Any]:
|
||||
@@ -186,15 +186,20 @@ def merge_run_context_overrides(config: dict[str, Any], context: Mapping[str, An
|
||||
def resolve_agent_factory(assistant_id: str | None):
|
||||
"""Resolve the agent factory callable from config.
|
||||
|
||||
Built-in graphs (``lead_agent``, ``canvas_agent``) are dispatched directly
|
||||
to their own factory. Custom agents share the ``lead_agent`` factory with
|
||||
``agent_id`` injected into configurable/context by :func:`build_run_config`.
|
||||
Built-in graphs (``lead_agent``, ``canvas_agent``, ``ai_writing``) are dispatched
|
||||
directly to their own factory. Custom agents share the ``lead_agent`` factory
|
||||
with ``agent_id`` injected into configurable/context by :func:`build_run_config`.
|
||||
"""
|
||||
if assistant_id == "canvas_agent":
|
||||
from deerflow.agents.canvas_agent.agent import make_canvas_agent
|
||||
|
||||
return make_canvas_agent
|
||||
|
||||
if assistant_id == "ai_writing":
|
||||
from deerflow.agents.ai_writing.graph import make_ai_writing_graph
|
||||
|
||||
return make_ai_writing_graph
|
||||
|
||||
from deerflow.agents.lead_agent.agent import make_lead_agent
|
||||
|
||||
return make_lead_agent
|
||||
@@ -400,6 +405,11 @@ async def start_run(
|
||||
)
|
||||
else:
|
||||
await run_ctx.thread_store.update_status(thread_id, "running")
|
||||
# SDK 的 ``client.threads.create()`` 不传 assistant_id —— 先建出来空的 thread，
|
||||
# 真正提交 run 时才知道用哪个 assistant。这里回填一次，避免历史列表 / 按
|
||||
# assistant 筛选时拿不到关联。
|
||||
if not existing.get("assistant_id") and body.assistant_id:
|
||||
await run_ctx.thread_store.update_assistant_id(thread_id, body.assistant_id)
|
||||
except Exception:
|
||||
logger.warning("Failed to upsert thread_meta for %s (non-fatal)", sanitize_log_param(thread_id))
|
||||
|
||||
@@ -441,6 +451,9 @@ async def start_run(
|
||||
stream_subgraphs=body.stream_subgraphs,
|
||||
interrupt_before=body.interrupt_before,
|
||||
interrupt_after=body.interrupt_after,
|
||||
# 透传 SDK ``Command``：resume / update / goto。Pregel 会据此从最近
|
||||
# checkpoint 处恢复，而不是把 graph_input 当新一轮输入跑。
|
||||
command=getattr(body, "command", None),
|
||||
)
|
||||
)
|
||||
record.task = task
|
||||
|
||||
@@ -202,4 +202,4 @@ memory:
|
||||
ai_writing:
|
||||
keyword_model: deepseek-chat # 检索词生成专用模型，留空回退写作主模型
|
||||
rank_model: deepseek-chat # 素材排序专用模型，留空回退写作主模型
|
||||
max_materials: 50 # 合并去重后的素材总量上限，范围 [1, 100]，默认 20
|
||||
max_materials: 20 # 合并去重后的素材总量上限，范围 [1, 100]，默认 20
|
||||
@@ -7,7 +7,8 @@
|
||||
"env": ".env",
|
||||
"graphs": {
|
||||
"lead_agent": "deerflow.agents:make_lead_agent",
|
||||
"canvas_agent": "deerflow.agents.canvas_agent:make_canvas_agent"
|
||||
"canvas_agent": "deerflow.agents.canvas_agent:make_canvas_agent",
|
||||
"ai_writing": "deerflow.agents.ai_writing.graph:make_ai_writing_graph"
|
||||
},
|
||||
"auth": {
|
||||
"path": "./app/gateway/langgraph_auth.py:auth"
|
||||
|
||||
+2
-14
@@ -1,15 +1,3 @@
|
||||
from .graph import (
|
||||
AI_WRITING_GRAPH,
|
||||
build_ai_writing_graph,
|
||||
get_ai_writing_checkpointer,
|
||||
get_ai_writing_graph,
|
||||
shutdown_ai_writing_graph,
|
||||
)
|
||||
from .graph import make_ai_writing_graph
|
||||
|
||||
__all__ = [
|
||||
"AI_WRITING_GRAPH",
|
||||
"build_ai_writing_graph",
|
||||
"get_ai_writing_graph",
|
||||
"get_ai_writing_checkpointer",
|
||||
"shutdown_ai_writing_graph",
|
||||
]
|
||||
__all__ = ["make_ai_writing_graph"]
|
||||
+53
-216
@@ -3,6 +3,7 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.graph import END, StateGraph
|
||||
from langgraph.types import interrupt
|
||||
|
||||
@@ -15,6 +16,30 @@ from .nodes.editor import editor_node
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def _mark_session_done(config: RunnableConfig | None) -> None:
|
||||
"""把业务表行 status 写成 ``done``。
|
||||
|
||||
pause_draft_confirm action=finalize / pause_review_confirm force_finalize 时调，
|
||||
让 LangGraph Server 路径下新会话有正确终态（老路径靠 ``_stream_graph_segment``
|
||||
finally 块写终态，新路径完全不走它）。失败不影响 graph 流程。
|
||||
"""
|
||||
if not config:
|
||||
return
|
||||
thread_id = (config.get("configurable") or {}).get("thread_id")
|
||||
if not thread_id:
|
||||
return
|
||||
from deerflow.persistence.ai_writing_sessions import get_default_store
|
||||
from deerflow.runtime.user_context import AUTO
|
||||
|
||||
repo = get_default_store()
|
||||
if repo is None:
|
||||
return
|
||||
try:
|
||||
await repo.update(thread_id, user_id=AUTO, status="done")
|
||||
except Exception:
|
||||
logger.exception("ai_writing: mark_session_done failed (non-fatal) thread_id=%s", thread_id)
|
||||
|
||||
# ── 干预节点 ──────────────────────────────────────────────────────────────────
|
||||
|
||||
async def pause_material_confirm(state: AIWritingState) -> dict:
|
||||
@@ -127,7 +152,7 @@ async def pause_outline_confirm(state: AIWritingState) -> dict:
|
||||
return updates
|
||||
|
||||
|
||||
async def pause_draft_confirm(state: AIWritingState) -> dict:
|
||||
async def pause_draft_confirm(state: AIWritingState, config: RunnableConfig) -> dict:
|
||||
"""PAUSE-3：草稿确认"""
|
||||
draft = state["current_draft"]
|
||||
user_input: dict = interrupt({
|
||||
@@ -141,6 +166,7 @@ async def pause_draft_confirm(state: AIWritingState) -> dict:
|
||||
action = user_input.get("action", "to_editor")
|
||||
|
||||
if action == "finalize":
|
||||
await _mark_session_done(config)
|
||||
return {
|
||||
"last_user_intervention": user_input,
|
||||
"status": "done",
|
||||
@@ -223,7 +249,7 @@ async def pause_section_help(state: AIWritingState) -> dict:
|
||||
}
|
||||
|
||||
|
||||
async def pause_review_confirm(state: AIWritingState) -> dict:
|
||||
async def pause_review_confirm(state: AIWritingState, config: RunnableConfig) -> dict:
|
||||
"""PAUSE-4：编辑打回确认"""
|
||||
user_input: dict = interrupt({
|
||||
"pause_point": "review_confirm",
|
||||
@@ -234,6 +260,7 @@ async def pause_review_confirm(state: AIWritingState) -> dict:
|
||||
})
|
||||
|
||||
if user_input.get("overrideVerdict") or user_input.get("action") == "force_finalize":
|
||||
await _mark_session_done(config)
|
||||
return {
|
||||
"last_user_intervention": user_input,
|
||||
"status": "done",
|
||||
@@ -299,25 +326,26 @@ def after_review_pause(state: AIWritingState) -> str:
|
||||
return "end" if state["status"] == "done" else "writer_draft"
|
||||
|
||||
|
||||
# ── 组装 Graph ────────────────────────────────────────────────────────────────
|
||||
|
||||
def build_ai_writing_graph(checkpointer=None):
|
||||
"""构建并 compile AI 写作 graph。
|
||||
# ── LangGraph Server 入口 ────────────────────────────────────────────────────
|
||||
#
|
||||
# compile 出裸 graph（不带 checkpointer），由 LangGraph Server 注入主 checkpointer
|
||||
# （langgraph.json 里 ``make_checkpointer`` 配的那个）。AI 写作的 thread 状态与
|
||||
# 主聊天共享同一份 saver —— LangGraph schema 用 ``(thread_id, checkpoint_ns)`` 做
|
||||
# 主键，UUID 不会撞。
|
||||
def make_ai_writing_graph(config=None):
|
||||
"""LangGraph Server graph factory：不接收 checkpointer，由 Server 注入。
|
||||
|
||||
Args:
|
||||
checkpointer: 可选的 LangGraph checkpointer。
|
||||
- ``None`` → 退回 ``MemorySaver`` —— 仅供测试 / 开发使用。
|
||||
进程或容器重启后会话状态全部丢失，前端拿着旧 session_id 调
|
||||
``/resume`` 会得到 404。
|
||||
- 其他 → 传入的 (Async)BaseCheckpointSaver 实例，用于持久化。
|
||||
config: LangGraph Server 调用工厂时注入的 RunnableConfig（暂未消费）。
|
||||
"""
|
||||
g = StateGraph(AIWritingState)
|
||||
|
||||
g.add_node("intent_parser", intent_parser_node)
|
||||
g.add_node("researcher", researcher_node)
|
||||
g.add_node("pause_material", pause_material_confirm)
|
||||
g.add_node("writer_outline", writer_outline_node)
|
||||
g.add_node("pause_outline", pause_outline_confirm)
|
||||
g.add_node("intent_parser", intent_parser_node)
|
||||
g.add_node("researcher", researcher_node)
|
||||
g.add_node("pause_material", pause_material_confirm)
|
||||
g.add_node("writer_outline", writer_outline_node)
|
||||
g.add_node("pause_outline", pause_outline_confirm)
|
||||
g.add_node("writer_draft", writer_draft_node)
|
||||
g.add_node("pause_section_help", pause_section_help)
|
||||
g.add_node("pause_draft", pause_draft_confirm)
|
||||
@@ -325,214 +353,23 @@ def build_ai_writing_graph(checkpointer=None):
|
||||
g.add_node("pause_review", pause_review_confirm)
|
||||
|
||||
g.set_entry_point("intent_parser")
|
||||
g.add_edge("intent_parser", "researcher")
|
||||
g.add_edge("intent_parser", "researcher")
|
||||
g.add_conditional_edges("researcher", after_researcher,
|
||||
{"pause_material": "pause_material", "writer_draft": "writer_draft"})
|
||||
|
||||
{"pause_material": "pause_material", "writer_draft": "writer_draft"})
|
||||
g.add_conditional_edges("pause_material", after_material_pause,
|
||||
{"researcher": "researcher", "writer_outline": "writer_outline"})
|
||||
|
||||
{"researcher": "researcher", "writer_outline": "writer_outline"})
|
||||
g.add_edge("writer_outline", "pause_outline")
|
||||
g.add_conditional_edges("pause_outline", after_outline_pause,
|
||||
{"writer_outline": "writer_outline", "writer_draft": "writer_draft"})
|
||||
|
||||
{"writer_outline": "writer_outline", "writer_draft": "writer_draft"})
|
||||
g.add_conditional_edges("writer_draft", after_writer_draft,
|
||||
{"pause_section_help": "pause_section_help", "pause_draft": "pause_draft"})
|
||||
|
||||
{"pause_section_help": "pause_section_help", "pause_draft": "pause_draft"})
|
||||
g.add_conditional_edges("pause_section_help", after_section_help_pause,
|
||||
{"researcher": "researcher", "writer_draft": "writer_draft"})
|
||||
|
||||
{"researcher": "researcher", "writer_draft": "writer_draft"})
|
||||
g.add_conditional_edges("pause_draft", after_draft_pause,
|
||||
{"end": END, "writer_draft": "writer_draft", "editor": "editor"})
|
||||
|
||||
{"end": END, "writer_draft": "writer_draft", "editor": "editor"})
|
||||
g.add_conditional_edges("editor", after_review,
|
||||
{"pause_draft": "pause_draft", "pause_review": "pause_review"})
|
||||
|
||||
{"pause_draft": "pause_draft", "pause_review": "pause_review"})
|
||||
g.add_conditional_edges("pause_review", after_review_pause,
|
||||
{"end": END, "writer_draft": "writer_draft"})
|
||||
{"end": END, "writer_draft": "writer_draft"})
|
||||
|
||||
if checkpointer is None:
|
||||
# MemorySaver 仅作为 fallback —— 提示一下，正式部署应通过
|
||||
# get_ai_writing_graph() 拿到带持久化 checkpointer 的 graph。
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
|
||||
checkpointer = MemorySaver()
|
||||
return g.compile(checkpointer=checkpointer)
|
||||
|
||||
|
||||
# ── 持久化 checkpointer 单例（懒加载）─────────────────────────────────────────
|
||||
|
||||
# 模块级 graph 槽位。import 时为 None；首次 ``get_ai_writing_graph()`` 调用时
|
||||
# 按 config.yaml 的 checkpointer 段懒构建出真正的 graph。
|
||||
# 测试可通过 ``monkeypatch.setattr(graph_mod, "AI_WRITING_GRAPH", stub)``
|
||||
# 直接注入桩；本模块所有访问点都走 ``get_ai_writing_graph()`` 这一入口，会优先
|
||||
# 返回 module attribute 的值，保留测试桩注入路径。
|
||||
AI_WRITING_GRAPH = None
|
||||
|
||||
# 持有 async checkpointer（底层 aiosqlite Connection 或 psycopg pool 的所有者）。
|
||||
# 进程生命周期内常驻；shutdown 时由 ``shutdown_ai_writing_graph`` 显式关闭。
|
||||
_async_checkpointer_singleton = None
|
||||
|
||||
|
||||
def _fallback_to_memory_saver(reason: str):
|
||||
"""退回 InMemorySaver 的统一兜底入口。
|
||||
|
||||
把 reason 打到 warning 日志里方便运维定位。InMemorySaver 进程重启会丢
|
||||
状态，但**当次写作能跑完** —— 我们的首要目标是「写作流程能继续」，
|
||||
持久化是锦上添花。
|
||||
"""
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
logger.warning(
|
||||
"ai_writing: 持久化 checkpointer 不可用，退回 InMemorySaver。原因: %s",
|
||||
reason,
|
||||
)
|
||||
return InMemorySaver()
|
||||
|
||||
|
||||
async def _get_async_checkpointer():
|
||||
"""按 config.yaml 的 checkpointer 段返回一个 async checkpointer 实例。
|
||||
|
||||
**绝不抛异常** —— 任何失败一律退回 InMemorySaver。这是为了保证「写作
|
||||
流程能继续」这个首要目标：持久化坏了用户至少还能写完当次稿，重启会丢
|
||||
会话但不影响当下创作。
|
||||
|
||||
优先级：
|
||||
1. ``checkpointer:`` 段（``type=sqlite|postgres|memory``）
|
||||
2. fallback 到 ``InMemorySaver`` 并打 warning
|
||||
|
||||
返回的 saver 与主 agent 公用同一 sqlite 文件 —— LangGraph 的 schema
|
||||
用 ``(thread_id, checkpoint_ns)`` 做主键，AI 写作的 thread_id 是
|
||||
会话 uuid，跟主 agent 的 LangGraph thread_id 不会撞。
|
||||
"""
|
||||
global _async_checkpointer_singleton
|
||||
if _async_checkpointer_singleton is not None:
|
||||
return _async_checkpointer_singleton
|
||||
|
||||
try:
|
||||
from deerflow.config.checkpointer_config import get_checkpointer_config
|
||||
|
||||
cp_cfg = get_checkpointer_config()
|
||||
except Exception:
|
||||
logger.exception("ai_writing: 加载 checkpointer 配置失败")
|
||||
_async_checkpointer_singleton = _fallback_to_memory_saver("加载 config 抛异常")
|
||||
return _async_checkpointer_singleton
|
||||
|
||||
if cp_cfg is None:
|
||||
_async_checkpointer_singleton = _fallback_to_memory_saver(
|
||||
"config.yaml 未配置 checkpointer 段"
|
||||
)
|
||||
return _async_checkpointer_singleton
|
||||
|
||||
if cp_cfg.type == "memory":
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
_async_checkpointer_singleton = InMemorySaver()
|
||||
logger.info("ai_writing: 配置 type=memory，使用 InMemorySaver")
|
||||
return _async_checkpointer_singleton
|
||||
|
||||
if cp_cfg.type == "sqlite":
|
||||
try:
|
||||
import aiosqlite
|
||||
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
||||
from deerflow.runtime.store._sqlite_utils import (
|
||||
ensure_sqlite_parent_dir,
|
||||
resolve_sqlite_conn_str,
|
||||
)
|
||||
|
||||
conn_str = resolve_sqlite_conn_str(
|
||||
cp_cfg.connection_string or "ai_writing_checkpoints.db"
|
||||
)
|
||||
ensure_sqlite_parent_dir(conn_str)
|
||||
# 直接持有 aiosqlite.Connection —— 进程生命周期，shutdown 时显式关。
|
||||
conn = await aiosqlite.connect(conn_str)
|
||||
saver = AsyncSqliteSaver(conn)
|
||||
await saver.setup()
|
||||
_async_checkpointer_singleton = saver
|
||||
logger.info("ai_writing: 使用 AsyncSqliteSaver (%s)", conn_str)
|
||||
return saver
|
||||
except Exception as exc:
|
||||
# sqlite 失败原因可能是：依赖未装 / 磁盘满 / 文件权限 / 路径不存在 /
|
||||
# 文件被外部进程独占锁。任何一个都不能让 AI 写作整个不可用。
|
||||
logger.exception(
|
||||
"ai_writing: AsyncSqliteSaver 初始化失败，conn_str=%r",
|
||||
getattr(cp_cfg, "connection_string", None),
|
||||
)
|
||||
_async_checkpointer_singleton = _fallback_to_memory_saver(
|
||||
f"sqlite 初始化失败: {type(exc).__name__}: {exc}"
|
||||
)
|
||||
return _async_checkpointer_singleton
|
||||
|
||||
if cp_cfg.type == "postgres":
|
||||
try:
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
|
||||
if not cp_cfg.connection_string:
|
||||
raise ValueError("checkpointer.connection_string 是 postgres 必填项")
|
||||
pool = AsyncConnectionPool(cp_cfg.connection_string, open=False)
|
||||
await pool.open()
|
||||
saver = AsyncPostgresSaver(pool)
|
||||
await saver.setup()
|
||||
_async_checkpointer_singleton = saver
|
||||
logger.info("ai_writing: 使用 AsyncPostgresSaver")
|
||||
return saver
|
||||
except Exception as exc:
|
||||
logger.exception("ai_writing: AsyncPostgresSaver 初始化失败")
|
||||
_async_checkpointer_singleton = _fallback_to_memory_saver(
|
||||
f"postgres 初始化失败: {type(exc).__name__}: {exc}"
|
||||
)
|
||||
return _async_checkpointer_singleton
|
||||
|
||||
# 未知 type —— config 配错也不能让服务挂
|
||||
_async_checkpointer_singleton = _fallback_to_memory_saver(
|
||||
f"未知 checkpointer.type={cp_cfg.type!r}"
|
||||
)
|
||||
return _async_checkpointer_singleton
|
||||
|
||||
|
||||
async def get_ai_writing_graph():
|
||||
"""懒加载并返回 AI 写作 graph 单例。
|
||||
|
||||
第一次调用按 config 选定的 checkpointer 构建 graph 并缓存；后续调用直接
|
||||
返回缓存。测试可在调用前通过 ``monkeypatch.setattr`` 直接替换模块级
|
||||
``AI_WRITING_GRAPH`` —— 本函数会优先返回已设的值。
|
||||
"""
|
||||
global AI_WRITING_GRAPH
|
||||
if AI_WRITING_GRAPH is not None:
|
||||
return AI_WRITING_GRAPH
|
||||
checkpointer = await _get_async_checkpointer()
|
||||
AI_WRITING_GRAPH = build_ai_writing_graph(checkpointer=checkpointer)
|
||||
return AI_WRITING_GRAPH
|
||||
|
||||
|
||||
def get_ai_writing_checkpointer():
|
||||
"""返回当前已实例化的 checkpointer（清理任务用它来删 thread）。
|
||||
|
||||
若尚未懒构建过则返回 None。
|
||||
"""
|
||||
return _async_checkpointer_singleton
|
||||
|
||||
|
||||
async def shutdown_ai_writing_graph() -> None:
|
||||
"""关闭 saver 持有的底层连接，并清空 graph 单例。Lifespan shutdown 调用。"""
|
||||
global AI_WRITING_GRAPH, _async_checkpointer_singleton
|
||||
saver = _async_checkpointer_singleton
|
||||
_async_checkpointer_singleton = None
|
||||
AI_WRITING_GRAPH = None
|
||||
if saver is None:
|
||||
return
|
||||
# aiosqlite saver: saver.conn 是 aiosqlite.Connection
|
||||
conn = getattr(saver, "conn", None)
|
||||
if conn is not None and hasattr(conn, "close"):
|
||||
try:
|
||||
await conn.close()
|
||||
except Exception:
|
||||
logger.warning("ai_writing: 关闭 checkpointer connection 失败", exc_info=True)
|
||||
# postgres saver: saver.pool 是 AsyncConnectionPool
|
||||
pool = getattr(saver, "pool", None)
|
||||
if pool is not None and hasattr(pool, "close"):
|
||||
try:
|
||||
await pool.close()
|
||||
except Exception:
|
||||
logger.warning("ai_writing: 关闭 checkpointer pool 失败", exc_info=True)
|
||||
return g.compile()
|
||||
+48
@@ -15,6 +15,54 @@ from langchain_core.runnables import RunnableConfig
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_ai_writing_config() -> dict:
|
||||
"""读 ``config.yaml`` 的 ``ai_writing`` 段。
|
||||
|
||||
阶段 7 之前由路由层 ``_make_config`` 把这份 dict 注入到 ``configurable.ai_writing_config``
|
||||
给节点使用。阶段 7 删冗余后 ``_make_config`` 和注入路径全部消失（前端 ``stream.submit``
|
||||
的 config 只传 ``model_name``）—— 节点改成自己直接从 ``get_app_config()`` 读，跨边界
|
||||
更干净，``config.yaml`` 改完也不用重启前端 / 重发起会话。
|
||||
|
||||
返回字段（按使用方）：
|
||||
- researcher：``keyword_model`` / ``rank_model`` / ``max_materials`` /
|
||||
``search_provider`` 等
|
||||
- editor：``pass_threshold``
|
||||
"""
|
||||
try:
|
||||
from deerflow.config import get_app_config
|
||||
|
||||
aw = getattr(get_app_config(), "ai_writing", None)
|
||||
if isinstance(aw, dict):
|
||||
return aw
|
||||
return dict(aw) if aw else {}
|
||||
except Exception:
|
||||
logger.exception("ai_writing: 读取 app config 失败，使用空配置")
|
||||
return {}
|
||||
|
||||
|
||||
def safe_get_stream_writer():
|
||||
"""返回一个 stream writer callable；在没有 LangGraph runtime 上下文时退化成 no-op。
|
||||
|
||||
AI 写作节点用它把流式 chunk 推到 LangGraph 原生 ``custom`` 通道。
|
||||
场景：
|
||||
|
||||
- LangGraph Server 路径（``client.runs.stream`` 订阅了 ``stream_mode=["custom", ...]``）
|
||||
→ 真实 writer，事件透传给 SDK 消费方。
|
||||
- 老 ``/start /stream /resume`` 路径（``graph.astream(stream_mode="updates")``）
|
||||
→ 真实 writer 仍可调用，只是没人订阅 custom 通道，事件被静默丢弃。
|
||||
这是阶段 2「硬切到 get_stream_writer」的预期行为：老前端 EventSource 不再收到
|
||||
token-level chunk，但整体进度卡（draft_ready / review_ready 这些 state 级事件）
|
||||
仍正常 —— 阶段 3 前端切到 useStream 后立刻补齐。
|
||||
- 单测里直接调节点函数（无 LangGraph runtime） → 返回 no-op，调用方零分支。
|
||||
"""
|
||||
try:
|
||||
from langgraph.config import get_stream_writer # 延迟 import，便于单测打补丁
|
||||
|
||||
return get_stream_writer()
|
||||
except Exception:
|
||||
return lambda _payload: None
|
||||
|
||||
|
||||
def get_model(config: RunnableConfig, model_name: str | None = None):
|
||||
"""创建对话模型。
|
||||
|
||||
|
||||
+51
-25
@@ -1,8 +1,8 @@
|
||||
"""编辑节点：三项独立 LLM 审查（事实/逻辑/语言），逐 token 流式推送以支持打字机效果。"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Callable
|
||||
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
@@ -11,7 +11,7 @@ from ..state import AIWritingState, ReviewResult, ReviewIssue
|
||||
from ..prompts.editor_prompts import (
|
||||
EDITOR_SYSTEM, FACT_REVIEW_USER, LOGIC_REVIEW_USER, LANGUAGE_REVIEW_USER,
|
||||
)
|
||||
from ._utils import get_model, parse_json, _content_to_text, ThinkTagParser
|
||||
from ._utils import get_ai_writing_config, get_model, parse_json, _content_to_text, ThinkTagParser, safe_get_stream_writer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -40,9 +40,9 @@ def _parse_issues(raw_issues: list, category: str) -> list[ReviewIssue]:
|
||||
|
||||
async def _stream_one_review(
|
||||
model, user_prompt: str, category: str,
|
||||
sse_queue: asyncio.Queue | None, session_id: str,
|
||||
writer: Callable[[dict], None],
|
||||
) -> str:
|
||||
"""流式执行单项审查，逐 chunk 推送 review_chunk 事件，返回累积的正文文本（用于 JSON 解析）。"""
|
||||
"""流式执行单项审查，逐 chunk 经 LangGraph custom 通道推送 review_chunk 事件，返回累积的正文文本（用于 JSON 解析）。"""
|
||||
full_text = "" # 仅累积正文（非思考），供后续 JSON 解析
|
||||
parser = ThinkTagParser()
|
||||
async for chunk in model.astream([
|
||||
@@ -55,16 +55,42 @@ async def _stream_one_review(
|
||||
for is_thinking, text in parser.feed(delta):
|
||||
if not is_thinking:
|
||||
full_text += text
|
||||
if sse_queue is not None:
|
||||
await sse_queue.put(("review_chunk", {
|
||||
"session_id": session_id,
|
||||
"category": category,
|
||||
"text": text,
|
||||
"is_thinking": is_thinking,
|
||||
}))
|
||||
writer({
|
||||
"type": "review_chunk",
|
||||
"category": category,
|
||||
"text": text,
|
||||
"is_thinking": is_thinking,
|
||||
})
|
||||
return full_text
|
||||
|
||||
|
||||
async def _save_review_to_repo(config: RunnableConfig, review: ReviewResult) -> None:
|
||||
"""把 review_result 异步落到业务表（best-effort，写失败不影响流程）。
|
||||
|
||||
``user_id`` 走 ``resolve_user_id(AUTO)`` 读 contextvar（与 intent_parser 建行、
|
||||
writer_draft 写草稿同套规则）。
|
||||
"""
|
||||
cfg = config.get("configurable", {}) or {}
|
||||
thread_id = cfg.get("thread_id")
|
||||
if not thread_id:
|
||||
return
|
||||
from deerflow.persistence.ai_writing_sessions import get_default_store
|
||||
from deerflow.runtime.user_context import AUTO
|
||||
|
||||
repo = get_default_store()
|
||||
if repo is None:
|
||||
return
|
||||
try:
|
||||
await repo.update(
|
||||
thread_id,
|
||||
user_id=AUTO,
|
||||
status="review_ready",
|
||||
review_result=dict(review),
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("editor: save_review failed (non-fatal) thread_id=%s", thread_id)
|
||||
|
||||
|
||||
async def editor_node(state: AIWritingState, config: RunnableConfig) -> dict:
|
||||
draft = state.get("current_draft")
|
||||
if not draft:
|
||||
@@ -78,9 +104,9 @@ async def editor_node(state: AIWritingState, config: RunnableConfig) -> dict:
|
||||
mats = _materials_summary(pkg)
|
||||
|
||||
cfg = config.get("configurable", {})
|
||||
sse_queue: asyncio.Queue | None = cfg.get("sse_queue")
|
||||
session_id: str = cfg.get("sse_session_id", "")
|
||||
ai_writing_config = cfg.get("ai_writing_config", {}) or {}
|
||||
writer = safe_get_stream_writer()
|
||||
# 阶段 7 后从 app config 直接读，不走 configurable 注入
|
||||
ai_writing_config = get_ai_writing_config()
|
||||
pass_threshold = int(ai_writing_config.get("pass_threshold", _DEFAULT_PASS_THRESHOLD))
|
||||
|
||||
REVIEW_TASKS = [
|
||||
@@ -110,7 +136,7 @@ async def editor_node(state: AIWritingState, config: RunnableConfig) -> dict:
|
||||
for category, user_prompt in REVIEW_TASKS:
|
||||
fallback = {"score": 80, "issues": []}
|
||||
try:
|
||||
full_text = await _stream_one_review(model, user_prompt, category, sse_queue, session_id)
|
||||
full_text = await _stream_one_review(model, user_prompt, category, writer)
|
||||
res = parse_json(full_text, fallback)
|
||||
except Exception as e:
|
||||
logger.warning("editor review [%s] failed: %s", category, e)
|
||||
@@ -120,16 +146,13 @@ async def editor_node(state: AIWritingState, config: RunnableConfig) -> dict:
|
||||
results[category] = res
|
||||
|
||||
# 这一项审查完成 —— 推送结构化结果（前端据此把流式文本替换为评分/意见）
|
||||
if sse_queue is not None:
|
||||
issues = _parse_issues(res.get("issues", []), category)
|
||||
await sse_queue.put(("review_item_done", {
|
||||
"session_id": session_id,
|
||||
"category": category,
|
||||
"score": int(res.get("score", 80)),
|
||||
"issues": [dict(iss) for iss in issues],
|
||||
}))
|
||||
else:
|
||||
logger.debug("editor [%s] done (no sse_queue)", category)
|
||||
issues = _parse_issues(res.get("issues", []), category)
|
||||
writer({
|
||||
"type": "review_item_done",
|
||||
"category": category,
|
||||
"score": int(res.get("score", 80)),
|
||||
"issues": [dict(iss) for iss in issues],
|
||||
})
|
||||
|
||||
# 汇总
|
||||
fact_score = int(results["fact"].get("score", 80))
|
||||
@@ -156,6 +179,9 @@ async def editor_node(state: AIWritingState, config: RunnableConfig) -> dict:
|
||||
revision_notes=all_issues,
|
||||
)
|
||||
|
||||
# 阶段 2：节点内并存 DB 写入（路由层同样也写一遍，阶段 7 删冗余时移除路由那份）
|
||||
await _save_review_to_repo(config, review)
|
||||
|
||||
# 分项审查出现模型异常时，以非致命 notice 提示用户，审核结论仍照常产出
|
||||
notices: list[dict] = []
|
||||
if review_errors:
|
||||
|
||||
+48
@@ -1,6 +1,7 @@
|
||||
"""意图解析节点：将用户自然语言意图结构化为写作配置。"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import re
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
@@ -8,6 +9,49 @@ from langchain_core.runnables import RunnableConfig
|
||||
from ..state import AIWritingState
|
||||
from ._utils import llm_json
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _truncate_title(text: str, max_len: int = 50) -> str:
|
||||
"""与老 /start 路由里 _truncate_title 行为一致：UTF 字符截断 + 省略号。"""
|
||||
text = (text or "").strip()
|
||||
if len(text) <= max_len:
|
||||
return text
|
||||
return text[:max_len].rstrip() + "…"
|
||||
|
||||
|
||||
async def _upsert_business_row(configurable: dict, state: AIWritingState) -> None:
|
||||
"""在 intent_parser 入口幂等创建 ai_writing_sessions 业务表行。
|
||||
|
||||
LangGraph Server 路径下没有 /start 路由建行，必须节点自建，否则列表页查不到、
|
||||
transcript / draft / review 后续写入都会因 update 找不到行而静默丢数据。
|
||||
resume 路径下 thread_id 已经有行，upsert_session 只刷新 updated_at。
|
||||
|
||||
user_id 走 ``resolve_user_id(AUTO)``：FastAPI 请求线程把它写到 contextvar，
|
||||
asyncio.create_task 继承父任务的 context，所以这里能拿到当前用户；
|
||||
contextvar 没设时（测试 / 离线场景）抛 RuntimeError，try/except 兜底跳过。
|
||||
"""
|
||||
from deerflow.persistence.ai_writing_sessions import get_default_store
|
||||
from deerflow.runtime.user_context import AUTO
|
||||
|
||||
thread_id = configurable.get("thread_id")
|
||||
if not thread_id:
|
||||
return
|
||||
repo = get_default_store()
|
||||
if repo is None:
|
||||
return
|
||||
user_intent = (state.get("user_intent") or "").strip()
|
||||
try:
|
||||
await repo.upsert_session(
|
||||
id=thread_id,
|
||||
title=_truncate_title(user_intent) or "AI 写作",
|
||||
user_intent=user_intent,
|
||||
user_id=AUTO,
|
||||
status="in_progress",
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("intent_parser: upsert_session failed (non-fatal) thread_id=%s", thread_id)
|
||||
|
||||
# 兜底类型列表，当 config 中未注入 valid_article_types 时使用
|
||||
_DEFAULT_ARTICLE_TYPES: list[dict] = [
|
||||
{"key": "news", "label": "新闻资讯"},
|
||||
@@ -79,6 +123,10 @@ async def intent_parser_node(state: AIWritingState, config: RunnableConfig) -> d
|
||||
valid_keys: list[str] = [t["key"] for t in valid_article_types]
|
||||
default_key: str = valid_keys[0] if valid_keys else "analysis"
|
||||
|
||||
# 业务表 upsert（LangGraph Server 路径下没有 /start 路由建行，节点入口自建；
|
||||
# resume 时幂等不重复建。失败不阻断写作流程）。
|
||||
await _upsert_business_row(configurable, state)
|
||||
|
||||
type_options = _build_type_options(valid_article_types)
|
||||
system = _SYSTEM_TEMPLATE.format(type_options=type_options)
|
||||
user = _USER_TEMPLATE.format(
|
||||
|
||||
+73
-45
@@ -15,12 +15,12 @@ from ..prompts.researcher_prompts import (
|
||||
EXTRACT_KEYWORDS_SYSTEM, EXTRACT_KEYWORDS_USER,
|
||||
RANK_MATERIALS_SYSTEM, RANK_MATERIALS_USER,
|
||||
)
|
||||
from ._utils import llm_json
|
||||
from ._utils import get_ai_writing_config, llm_json, safe_get_stream_writer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 素材总量上限边界：config.yaml 的 ai_writing.max_materials，非法值回退默认 20
|
||||
_DEFAULT_MAX_MATERIALS = 20
|
||||
# 素材总量上限边界：config.yaml 的 ai_writing.max_materials，非法值回退默认 50
|
||||
_DEFAULT_MAX_MATERIALS = 50
|
||||
_MIN_MAX_MATERIALS = 1
|
||||
_MAX_MAX_MATERIALS = 100
|
||||
|
||||
@@ -93,68 +93,80 @@ async def researcher_node(state: AIWritingState, config: RunnableConfig) -> dict
|
||||
article_type = state.get("article_type", "analysis")
|
||||
target_audience = state.get("target_audience", "普通读者")
|
||||
|
||||
# 检查是否是补充搜索（PAUSE-1 re_search action）
|
||||
# 检查干预模式（PAUSE-1 素材确认 / 素材不足求助）
|
||||
intervention = state.get("last_user_intervention") or {}
|
||||
extra_keywords: list[str] = intervention.get("_extra_keywords") or intervention.get("extra_keywords") or []
|
||||
action = intervention.get("action")
|
||||
|
||||
# 素材不足求助点触发的「补充检索」：保留原大纲，检索完直接回写作、不走素材确认/大纲
|
||||
is_supplement = intervention.get("action") == "supplement_search"
|
||||
is_supplement = action == "supplement_search"
|
||||
# PAUSE-1 素材确认卡用户点「追加检索」并填了新关键词：
|
||||
# - 跳过 LLM 重新生成检索词（节省一次调用 + 尊重用户意图）
|
||||
# - 只检索 extra_keywords 这几个新词，不重检索之前 LLM 出过的词
|
||||
# - 保留 existing material_package，按 URL 去重合并
|
||||
is_append_only = action == "re_search" and bool(extra_keywords)
|
||||
|
||||
# 用户指定的检索关键词数量（1~10），决定生成多少个检索词、检索多少个方向
|
||||
keyword_count = _resolve_keyword_count(state.get("keyword_count"))
|
||||
|
||||
# SSE 实时进度推送（与 outline/draft 节点共用 sse_queue 机制）
|
||||
# 实时进度推送：阶段 2 起统一走 LangGraph 原生 custom 通道。
|
||||
cfg = config.get("configurable", {})
|
||||
sse_queue = cfg.get("sse_queue")
|
||||
sse_session_id: str = cfg.get("sse_session_id", "")
|
||||
writer = safe_get_stream_writer()
|
||||
|
||||
# 素材收集专家专用模型：检索词生成 / 素材排序 可在 config.yaml 的
|
||||
# ai_writing 段单独配置（keyword_model / rank_model），留空则回退写作主模型。
|
||||
ai_writing_config = cfg.get("ai_writing_config", {}) or {}
|
||||
# 注：阶段 7 后这份配置由节点直接读 app config，不再走 configurable 注入。
|
||||
ai_writing_config = get_ai_writing_config()
|
||||
keyword_model: str | None = ai_writing_config.get("keyword_model") or None
|
||||
rank_model: str | None = ai_writing_config.get("rank_model") or None
|
||||
# 素材总量上限（config.yaml: ai_writing.max_materials），按检索词数均分到每个词
|
||||
max_materials = _resolve_max_materials(ai_writing_config.get("max_materials"))
|
||||
|
||||
async def _emit(message: str) -> None:
|
||||
if sse_queue is None:
|
||||
return
|
||||
await sse_queue.put(("research_progress", {
|
||||
"session_id": sse_session_id,
|
||||
"message": message,
|
||||
}))
|
||||
def _emit(message: str) -> None:
|
||||
writer({"type": "research_progress", "message": message})
|
||||
|
||||
# Step 1: 提取关键词
|
||||
await _emit("正在生成检索词…")
|
||||
kw_result = await llm_json(
|
||||
EXTRACT_KEYWORDS_SYSTEM,
|
||||
EXTRACT_KEYWORDS_USER.format(
|
||||
user_intent=intent,
|
||||
article_type=article_type,
|
||||
target_audience=target_audience,
|
||||
keyword_count=keyword_count,
|
||||
),
|
||||
fallback={"keywords": [intent]},
|
||||
config=config,
|
||||
model_name=keyword_model,
|
||||
)
|
||||
# 清洗关键词：剔除占位符 / 空值，全空时回退写作意图
|
||||
keywords: list[str] = _clean_keywords(kw_result.get("keywords"), intent)
|
||||
if extra_keywords:
|
||||
keywords = list(dict.fromkeys(keywords + extra_keywords))
|
||||
kw_result: dict = {} # 走追加分支时跳过 LLM 调用，但下方 notice 拼接需要兜底值
|
||||
if is_append_only:
|
||||
# 追加模式：直接用用户填的关键词，不再调 LLM
|
||||
keywords: list[str] = _clean_keywords(extra_keywords, intent)
|
||||
_emit(f"追加检索 {len(keywords)} 个关键词…")
|
||||
else:
|
||||
_emit("正在生成检索词…")
|
||||
kw_result = await llm_json(
|
||||
EXTRACT_KEYWORDS_SYSTEM,
|
||||
EXTRACT_KEYWORDS_USER.format(
|
||||
user_intent=intent,
|
||||
article_type=article_type,
|
||||
target_audience=target_audience,
|
||||
keyword_count=keyword_count,
|
||||
),
|
||||
fallback={"keywords": [intent]},
|
||||
config=config,
|
||||
model_name=keyword_model,
|
||||
)
|
||||
# 清洗关键词：剔除占位符 / 空值，全空时回退写作意图
|
||||
keywords = _clean_keywords(kw_result.get("keywords"), intent)
|
||||
if extra_keywords:
|
||||
keywords = list(dict.fromkeys(keywords + extra_keywords))
|
||||
|
||||
# Step 2: 并发搜索
|
||||
provider = get_search_provider(ai_writing_config)
|
||||
|
||||
# 按用户指定的关键词数量截断；补充搜索时不截断用户额外添加的关键词
|
||||
search_limit = keyword_count
|
||||
if extra_keywords:
|
||||
search_limit = max(search_limit, len(keywords))
|
||||
search_keywords = keywords[:search_limit]
|
||||
# 按用户指定的关键词数量截断。
|
||||
# 追加模式：keywords 本身就是用户填的，全部检索不截断（用户填 N 个就搜 N 个）。
|
||||
# 其他模式：受 keyword_count 限制，但 extra_keywords 在内时上限放宽到 keywords 总数。
|
||||
if is_append_only:
|
||||
search_keywords = keywords
|
||||
else:
|
||||
search_limit = keyword_count
|
||||
if extra_keywords:
|
||||
search_limit = max(search_limit, len(keywords))
|
||||
search_keywords = keywords[:search_limit]
|
||||
total = len(search_keywords)
|
||||
# 每个检索词分到的检索条数 = 总量上限 / 检索词数（向上取整）
|
||||
per_keyword = _per_keyword_limit(max_materials, total)
|
||||
await _emit(f"检索词已生成，开始检索 {total} 个方向…")
|
||||
_emit(f"检索词已生成，开始检索 {total} 个方向…")
|
||||
|
||||
# 并发搜索，但每完成一个就推送进度（计数在单线程 asyncio 下无需加锁）
|
||||
completed = 0
|
||||
@@ -166,18 +178,22 @@ async def researcher_node(state: AIWritingState, config: RunnableConfig) -> dict
|
||||
finally:
|
||||
completed += 1
|
||||
if completed < total:
|
||||
await _emit(f"已完成 {completed}/{total} 次检索，正在继续检索相关内容…")
|
||||
_emit(f"已完成 {completed}/{total} 次检索，正在继续检索相关内容…")
|
||||
else:
|
||||
await _emit(f"已完成全部 {total} 次检索")
|
||||
_emit(f"已完成全部 {total} 次检索")
|
||||
|
||||
search_tasks = [_search_one(kw) for kw in search_keywords]
|
||||
all_results_nested: list[list[SearchResult]] = await asyncio.gather(*search_tasks, return_exceptions=True)
|
||||
|
||||
# 合并去重
|
||||
# - 首次检索 / 普通 re_search（不带 extra_keywords，用户只想 LLM 重出）：旧素材丢弃
|
||||
# - 素材不足求助 supplement_search：保留旧素材 + 追加新检索结果
|
||||
# - PAUSE-1 追加检索 is_append_only：保留旧素材 + 仅检索新关键词
|
||||
seen_urls: set[str] = set()
|
||||
raw_materials: list[Material] = []
|
||||
existing_pkg = state.get("material_package")
|
||||
if existing_pkg and intervention.get("action") != "re_search":
|
||||
preserve_existing = existing_pkg and (is_supplement or is_append_only)
|
||||
if preserve_existing:
|
||||
for m in existing_pkg.get("materials", []):
|
||||
seen_urls.add(m.get("url", ""))
|
||||
raw_materials.append(m)
|
||||
@@ -208,7 +224,7 @@ async def researcher_node(state: AIWritingState, config: RunnableConfig) -> dict
|
||||
search_errors.extend(getattr(provider, "errors", []))
|
||||
|
||||
# Step 3: 相关度排序（如素材较多，让 LLM 排序）
|
||||
await _emit("检索完成，正在筛选排序素材…")
|
||||
_emit("检索完成，正在筛选排序素材…")
|
||||
rank_error = ""
|
||||
if len(raw_materials) > max_materials:
|
||||
materials_list = "\n".join(
|
||||
@@ -241,7 +257,13 @@ async def researcher_node(state: AIWritingState, config: RunnableConfig) -> dict
|
||||
if not materials:
|
||||
summary = "未检索到相关素材，可重新搜索或在无素材的情况下继续写作"
|
||||
|
||||
pkg = MaterialPackage(keywords=keywords, materials=materials, summary=summary)
|
||||
# 追加模式下，pkg.keywords 体现到目前为止累积的检索词（旧 + 新追加），
|
||||
# 让前端列表 / 历史回显能看到完整的搜索路径。
|
||||
pkg_keywords = keywords
|
||||
if is_append_only and existing_pkg:
|
||||
old_kws = list(existing_pkg.get("keywords") or [])
|
||||
pkg_keywords = list(dict.fromkeys(old_kws + keywords))
|
||||
pkg = MaterialPackage(keywords=pkg_keywords, materials=materials, summary=summary)
|
||||
|
||||
# 容错告警：检索接口 / 模型调用出错时以非致命 notice 提示用户，不打断写作流程
|
||||
notices: list[dict] = []
|
||||
@@ -259,12 +281,18 @@ async def researcher_node(state: AIWritingState, config: RunnableConfig) -> dict
|
||||
raw="\n\n----------\n\n".join(raw_dumps[:3]),
|
||||
))
|
||||
|
||||
# 进度事件文案：追加模式下让用户看到只搜了哪几个新词
|
||||
step_start_msg = (
|
||||
f"追加检索关键词：{', '.join(keywords[:3])}"
|
||||
if is_append_only
|
||||
else f"搜索关键词：{', '.join(keywords[:3])}"
|
||||
)
|
||||
return {
|
||||
"material_package": pkg,
|
||||
# 补充检索：直接回写作（沿用原大纲）；常规检索：进入素材确认暂停点
|
||||
"status": "writing_draft" if is_supplement else "awaiting_material_confirm",
|
||||
"progress_events": [
|
||||
{"type": "step_start", "agent_name": "素材收集专家", "message": f"搜索关键词：{', '.join(keywords[:3])}"},
|
||||
{"type": "step_start", "agent_name": "素材收集专家", "message": step_start_msg},
|
||||
*notices,
|
||||
{"type": "materials_ready", "agent_name": "素材收集专家", "message": summary},
|
||||
],
|
||||
|
||||
+53
-22
@@ -1,7 +1,6 @@
|
||||
"""作家节点 Phase 2：按大纲逐段写正文，支持流式 chunk 输出。"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import re
|
||||
|
||||
@@ -13,7 +12,7 @@ from ..prompts.writer_prompts import (
|
||||
DRAFT_SECTION_SYSTEM, DRAFT_SECTION_USER, DRAFT_EXPAND_USER,
|
||||
WRITER_MODE_LOOSE, WRITER_MODE_STRICT, SECTION_BLOCKED_MARKER,
|
||||
)
|
||||
from ._utils import get_model, _content_to_text, count_words, ThinkTagParser
|
||||
from ._utils import get_model, _content_to_text, count_words, ThinkTagParser, safe_get_stream_writer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -64,6 +63,35 @@ def _get_section_materials(material_ids: list[str], pkg: dict) -> str:
|
||||
return "\n".join(lines) if lines else "（无相关素材）"
|
||||
|
||||
|
||||
async def _save_draft_to_repo(config: RunnableConfig, draft: DraftArticle) -> None:
|
||||
"""把 current_draft 异步落到业务表（best-effort，写失败不影响流程）。
|
||||
|
||||
``user_id`` 走 ``resolve_user_id(AUTO)`` 读 contextvar（FastAPI 请求线程设置，
|
||||
asyncio.create_task 继承）。intent_parser 在阶段 4 已用同样规则建行；这里 update
|
||||
时也用 AUTO 才能命中 owner filter。
|
||||
"""
|
||||
cfg = config.get("configurable", {}) or {}
|
||||
thread_id = cfg.get("thread_id")
|
||||
if not thread_id:
|
||||
return
|
||||
from deerflow.persistence.ai_writing_sessions import get_default_store
|
||||
from deerflow.runtime.user_context import AUTO
|
||||
|
||||
repo = get_default_store()
|
||||
if repo is None:
|
||||
return
|
||||
try:
|
||||
await repo.update(
|
||||
thread_id,
|
||||
user_id=AUTO,
|
||||
status="draft_ready",
|
||||
draft_title=draft.get("title") or "",
|
||||
draft_markdown=draft.get("full_markdown") or "",
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("writer_draft: save_draft failed (non-fatal) thread_id=%s", thread_id)
|
||||
|
||||
|
||||
async def writer_draft_node(state: AIWritingState, config: RunnableConfig) -> dict:
|
||||
outline = state.get("current_outline")
|
||||
if not outline:
|
||||
@@ -80,10 +108,10 @@ async def writer_draft_node(state: AIWritingState, config: RunnableConfig) -> di
|
||||
writing_mode = state.get("writing_mode", "loose")
|
||||
loose_titles: set[str] = set(intervention.get("_loose_sections") or [])
|
||||
|
||||
# SSE 流式推送所需
|
||||
# 流式推送：阶段 2 起统一走 LangGraph 原生 custom 通道（get_stream_writer）。
|
||||
# cfg 仍保留以便阶段 4 取 user_id / thread_id 用。
|
||||
cfg = config.get("configurable", {})
|
||||
sse_queue: asyncio.Queue | None = cfg.get("sse_queue")
|
||||
session_id: str = cfg.get("sse_session_id", "")
|
||||
writer = safe_get_stream_writer()
|
||||
|
||||
# 修改意见（来自用户或编辑）
|
||||
revision_notes = ""
|
||||
@@ -174,25 +202,23 @@ async def writer_draft_node(state: AIWritingState, config: RunnableConfig) -> di
|
||||
for is_thinking, text in parser.feed(delta):
|
||||
if not is_thinking:
|
||||
section_buf += text
|
||||
if sse_queue is not None:
|
||||
await sse_queue.put(("draft_chunk", {
|
||||
"session_id": session_id,
|
||||
"section_index": i,
|
||||
"section_title": sec_title,
|
||||
"text": text,
|
||||
"is_thinking": is_thinking,
|
||||
"is_section_done": False,
|
||||
}))
|
||||
writer({
|
||||
"type": "draft_chunk",
|
||||
"section_index": i,
|
||||
"section_title": sec_title,
|
||||
"text": text,
|
||||
"is_thinking": is_thinking,
|
||||
"is_section_done": False,
|
||||
})
|
||||
|
||||
# 章节完成
|
||||
if sse_queue is not None:
|
||||
await sse_queue.put(("draft_chunk", {
|
||||
"session_id": session_id,
|
||||
"section_index": i,
|
||||
"section_title": sec_title,
|
||||
"text": "",
|
||||
"is_section_done": True,
|
||||
}))
|
||||
writer({
|
||||
"type": "draft_chunk",
|
||||
"section_index": i,
|
||||
"section_title": sec_title,
|
||||
"text": "",
|
||||
"is_section_done": True,
|
||||
})
|
||||
|
||||
cleaned_buf = section_buf.strip()
|
||||
blocked_reason = ""
|
||||
@@ -289,6 +315,11 @@ async def writer_draft_node(state: AIWritingState, config: RunnableConfig) -> di
|
||||
revision_count=revision_count,
|
||||
)
|
||||
|
||||
# 阶段 2：节点内并存 DB 写入（路由层 _stream_graph_segment 同样也会写一遍，
|
||||
# 阶段 7 删冗余时移除路由层那份）。new path (LangGraph Server) 下路由层
|
||||
# 不再介入，节点写就是唯一来源。
|
||||
await _save_draft_to_repo(config, draft)
|
||||
|
||||
# 分节写作出现模型异常时，以非致命 notice 提示用户，草稿仍照常产出
|
||||
notices: list[dict] = []
|
||||
if section_errors:
|
||||
|
||||
+8
-12
@@ -1,7 +1,6 @@
|
||||
"""作家节点 Phase 1：生成文章大纲，支持流式 chunk 输出。"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
|
||||
@@ -12,7 +11,7 @@ from ..state import AIWritingState, Outline, OutlineSection
|
||||
from ..prompts.writer_prompts import (
|
||||
OUTLINE_SYSTEM, OUTLINE_USER, OUTLINE_REVISE_SYSTEM, OUTLINE_REVISE_USER,
|
||||
)
|
||||
from ._utils import get_model, parse_json, _content_to_text, ThinkTagParser
|
||||
from ._utils import get_model, parse_json, _content_to_text, ThinkTagParser, safe_get_stream_writer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -79,10 +78,8 @@ async def writer_outline_node(state: AIWritingState, config: RunnableConfig) ->
|
||||
feedback_section=feedback_section,
|
||||
)
|
||||
|
||||
# SSE 流式推送
|
||||
cfg = config.get("configurable", {})
|
||||
sse_queue: asyncio.Queue | None = cfg.get("sse_queue")
|
||||
session_id: str = cfg.get("sse_session_id", "")
|
||||
# 流式推送：阶段 2 起统一走 LangGraph 原生 custom 通道（get_stream_writer）。
|
||||
writer = safe_get_stream_writer()
|
||||
|
||||
full_text = "" # 只累积正文（非思考）用于 JSON 解析
|
||||
parser = ThinkTagParser()
|
||||
@@ -101,12 +98,11 @@ async def writer_outline_node(state: AIWritingState, config: RunnableConfig) ->
|
||||
for is_thinking, text in parser.feed(delta):
|
||||
if not is_thinking:
|
||||
full_text += text
|
||||
if sse_queue is not None:
|
||||
await sse_queue.put(("outline_chunk", {
|
||||
"session_id": session_id,
|
||||
"text": text,
|
||||
"is_thinking": is_thinking,
|
||||
}))
|
||||
writer({
|
||||
"type": "outline_chunk",
|
||||
"text": text,
|
||||
"is_thinking": is_thinking,
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error("outline streaming failed: %s", e, exc_info=True)
|
||||
return {
|
||||
|
||||
+84
-94
@@ -12,21 +12,21 @@
|
||||
injection: {...} # 系统提示词注入
|
||||
security: {...} # 内容安全扫描
|
||||
|
||||
三层配置覆盖
|
||||
两层配置覆盖
|
||||
============
|
||||
``get_effective_memory_config(user_id)`` 按优先级合并 (高 -> 低):
|
||||
1. per-user 覆盖 ``{base_dir}/users/{user_id}/memory_config.json``
|
||||
2. 全局覆盖 ``{base_dir}/memory_user_overrides.json``
|
||||
3. config.yaml ``memory:`` 段 (即 ``get_memory_config()``)
|
||||
4. Pydantic 默认值
|
||||
1. per-user 覆盖 数据库 ``user_memory_configs`` 表 (经进程内 cache 读取)
|
||||
(有 per-user 记录时自动强制 enabled=True)
|
||||
2. config.yaml ``memory:`` 段 (即 ``get_memory_config()``)
|
||||
3. Pydantic 默认值
|
||||
|
||||
第 1、2 层只允许覆盖 :data:`USER_WRITABLE_FIELDS` 白名单内的字段。
|
||||
第 1 层只允许覆盖 :data:`USER_WRITABLE_FIELDS` 白名单内的字段。
|
||||
cache 由 Gateway lifespan 启动时预热,写入时同步更新,保持同步读取路径。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
@@ -129,31 +129,26 @@ class MemoryConfig(BaseModel):
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# per-user 运行时覆盖
|
||||
# per-user 运行时覆盖 — 白名单
|
||||
# ===========================================================================
|
||||
|
||||
# 用户可在运行时覆盖的字段白名单 (点号路径)。
|
||||
# 不在此表中的字段必须改 config.yaml 并重启 (例如 provider / api_key / 字符配额)。
|
||||
# 用户唯一可覆盖的字段:召回深度。
|
||||
# injection.* / builtin.* / enabled 始终由 config.yaml 控制,用户不可关闭。
|
||||
USER_WRITABLE_FIELDS: frozenset[str] = frozenset(
|
||||
{
|
||||
# injection.enabled and injection.include_builtin are always ON;
|
||||
# removing them from the whitelist makes _filter_to_whitelist drop
|
||||
# any per-user override for these fields, so stored "false" values
|
||||
# in existing memory_config.json files are silently ignored.
|
||||
"injection.max_tokens",
|
||||
"injection.include_hindsight",
|
||||
# hindsight.memory_mode is fixed to "hybrid"; not user-writable.
|
||||
# hindsight.bank_mission is fixed; not user-writable.
|
||||
"hindsight.recall_budget",
|
||||
}
|
||||
)
|
||||
|
||||
_VALID_RECALL_BUDGETS: frozenset[str] = frozenset({"low", "mid", "high"})
|
||||
|
||||
|
||||
def _filter_to_whitelist(overrides: dict[str, Any]) -> dict[str, Any]:
|
||||
"""把覆盖 dict 过滤为只含白名单字段的嵌套 dict。
|
||||
|
||||
输入既支持嵌套形式 ``{"injection": {"enabled": false}}``,也支持
|
||||
点号扁平形式 ``{"injection.enabled": false}``。非白名单字段被丢弃并 log。
|
||||
输入支持嵌套形式 ``{"hindsight": {"recall_budget": "high"}}`` 或
|
||||
点号扁平形式 ``{"hindsight.recall_budget": "high"}``。
|
||||
非白名单字段被丢弃并 log。
|
||||
"""
|
||||
result: dict[str, Any] = {}
|
||||
|
||||
@@ -184,21 +179,6 @@ def _deep_merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any
|
||||
return merged
|
||||
|
||||
|
||||
def _load_json_file(path: Any) -> dict[str, Any]:
|
||||
"""读取一个 JSON 覆盖文件;不存在或解析失败时返回空 dict。"""
|
||||
try:
|
||||
from pathlib import Path
|
||||
|
||||
p = Path(path)
|
||||
if not p.exists():
|
||||
return {}
|
||||
data = json.loads(p.read_text(encoding="utf-8"))
|
||||
return data if isinstance(data, dict) else {}
|
||||
except (OSError, json.JSONDecodeError) as exc:
|
||||
logger.warning("读取记忆配置覆盖文件失败 %s: %s", path, exc)
|
||||
return {}
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# 全局单例 + 加载入口
|
||||
# ===========================================================================
|
||||
@@ -223,82 +203,88 @@ def load_memory_config_from_dict(config_dict: dict) -> None:
|
||||
_memory_config = MemoryConfig(**config_dict)
|
||||
|
||||
|
||||
# ===========================================================================
|
||||
# per-user recall_budget 进程内 cache
|
||||
#
|
||||
# Gateway lifespan 启动时调用 load_recall_budget_cache() 从 DB 批量加载。
|
||||
# API 写入时同步更新 cache,使 get_effective_memory_config() 保持同步读取。
|
||||
# ===========================================================================
|
||||
|
||||
_recall_budget_cache: dict[str, str] = {}
|
||||
|
||||
|
||||
def load_recall_budget_cache(data: dict[str, str]) -> None:
|
||||
"""用 DB 批量加载结果替换整个 cache (启动时调用一次)。"""
|
||||
global _recall_budget_cache
|
||||
_recall_budget_cache = dict(data)
|
||||
logger.info("recall_budget cache loaded: %d user(s)", len(_recall_budget_cache))
|
||||
|
||||
|
||||
def set_recall_budget_cache(user_id: str, value: str) -> None:
|
||||
"""写入单个用户的 recall_budget cache。"""
|
||||
_recall_budget_cache[user_id] = value
|
||||
|
||||
|
||||
def delete_recall_budget_cache(user_id: str) -> None:
|
||||
"""从 cache 中移除某用户的 recall_budget。"""
|
||||
_recall_budget_cache.pop(user_id, None)
|
||||
|
||||
|
||||
def get_recall_budget_from_cache(user_id: str) -> str | None:
|
||||
"""返回 cache 中某用户的 recall_budget,未设置时返回 None。"""
|
||||
return _recall_budget_cache.get(user_id)
|
||||
|
||||
|
||||
def get_effective_memory_config(user_id: str) -> MemoryConfig:
|
||||
"""按三层覆盖优先级,返回某用户的生效记忆配置。
|
||||
"""返回某用户的生效记忆配置 (同步,从进程内 cache 读取 per-user 覆盖)。
|
||||
|
||||
优先级 (高 -> 低):per-user 覆盖 > 全局覆盖 > config.yaml > 默认值。
|
||||
第 1、2 层只应用 :data:`USER_WRITABLE_FIELDS` 白名单内的字段。
|
||||
|
||||
覆盖文件读取失败不会抛异常 —— 退化为只用 config.yaml 级别配置。
|
||||
优先级 (高 -> 低):per-user DB 覆盖 (经 cache) > config.yaml > 默认值。
|
||||
有 per-user 记录时自动强制 enabled=True。
|
||||
"""
|
||||
base = _memory_config.model_dump()
|
||||
|
||||
recall_budget = _recall_budget_cache.get(user_id)
|
||||
if recall_budget is not None:
|
||||
base.setdefault("hindsight", {})["recall_budget"] = recall_budget
|
||||
base["enabled"] = True
|
||||
try:
|
||||
from deerflow.config.paths import get_paths
|
||||
|
||||
paths = get_paths()
|
||||
global_overrides = _filter_to_whitelist(
|
||||
_load_json_file(paths.memory_global_overrides_file)
|
||||
)
|
||||
user_overrides = _filter_to_whitelist(
|
||||
_load_json_file(paths.user_memory_config_file(user_id))
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 - 配置覆盖失败必须降级,不能阻断对话
|
||||
logger.warning("加载 per-user 记忆配置覆盖失败,退回 config.yaml: %s", exc)
|
||||
return _memory_config
|
||||
|
||||
merged = _deep_merge(_deep_merge(base, global_overrides), user_overrides)
|
||||
try:
|
||||
return MemoryConfig(**merged)
|
||||
return MemoryConfig(**base)
|
||||
except Exception as exc: # noqa: BLE001 - 覆盖产生非法配置时降级
|
||||
logger.warning("合并后的记忆配置非法,退回 config.yaml: %s", exc)
|
||||
return _memory_config
|
||||
|
||||
|
||||
def save_user_memory_override(user_id: str, patch: dict[str, Any]) -> dict[str, Any]:
|
||||
"""把白名单字段写入某用户的 per-user 覆盖文件,返回写入后的完整覆盖 dict。
|
||||
# ===========================================================================
|
||||
# Gateway 调用的异步 DB 对接函数
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
async def save_user_memory_override_db(user_id: str, patch: dict[str, Any], store: Any) -> dict[str, Any]:
|
||||
"""把白名单字段写入 DB store 并更新 cache,返回生效配置 dict。
|
||||
|
||||
``patch`` 支持嵌套或点号扁平形式;非白名单字段被忽略。
|
||||
``store`` 为 ``UserMemoryConfigStore`` 实例 (避免循环 import,用 Any)。
|
||||
"""
|
||||
from deerflow.config.paths import get_paths
|
||||
|
||||
path = get_paths().user_memory_config_file(user_id)
|
||||
current = _load_json_file(path)
|
||||
filtered = _filter_to_whitelist(patch)
|
||||
merged = _deep_merge(current, filtered)
|
||||
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(merged, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
return merged
|
||||
recall_budget = (filtered.get("hindsight") or {}).get("recall_budget")
|
||||
if recall_budget is not None:
|
||||
if recall_budget not in _VALID_RECALL_BUDGETS:
|
||||
raise ValueError(f"recall_budget 必须为 {sorted(_VALID_RECALL_BUDGETS)} 之一,收到:{recall_budget!r}")
|
||||
await store.set_recall_budget(user_id, recall_budget)
|
||||
set_recall_budget_cache(user_id, recall_budget)
|
||||
return get_effective_memory_config(user_id).model_dump()
|
||||
|
||||
|
||||
def delete_user_memory_override(user_id: str, field: str | None = None) -> dict[str, Any]:
|
||||
"""撤销某用户的 per-user 覆盖。
|
||||
async def delete_user_memory_override_db(user_id: str, field: str | None, store: Any) -> dict[str, Any]:
|
||||
"""撤销某用户的 per-user 覆盖并更新 cache,返回生效配置 dict。
|
||||
|
||||
``field`` 为点号路径 (如 ``injection.enabled``) 时只删该字段;
|
||||
为 ``None`` 时清空整个 per-user 覆盖文件。返回剩余的覆盖 dict。
|
||||
``field=None`` 清空全部;``field="hindsight.recall_budget"`` 只清该字段。
|
||||
"""
|
||||
from deerflow.config.paths import get_paths
|
||||
|
||||
path = get_paths().user_memory_config_file(user_id)
|
||||
if field is None:
|
||||
if path.exists():
|
||||
path.unlink()
|
||||
return {}
|
||||
|
||||
current = _load_json_file(path)
|
||||
if "." in field:
|
||||
section, sub_field = field.split(".", 1)
|
||||
if isinstance(current.get(section), dict):
|
||||
current[section].pop(sub_field, None)
|
||||
if not current[section]:
|
||||
current.pop(section, None)
|
||||
if field is None or field == "hindsight.recall_budget":
|
||||
await store.delete(user_id)
|
||||
delete_recall_budget_cache(user_id)
|
||||
else:
|
||||
current.pop(field, None)
|
||||
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(current, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
return current
|
||||
logger.warning("忽略删除非白名单的记忆配置覆盖字段: %s", field)
|
||||
return get_effective_memory_config(user_id).model_dump()
|
||||
|
||||
|
||||
__all__ = [
|
||||
@@ -312,6 +298,10 @@ __all__ = [
|
||||
"set_memory_config",
|
||||
"load_memory_config_from_dict",
|
||||
"get_effective_memory_config",
|
||||
"save_user_memory_override",
|
||||
"delete_user_memory_override",
|
||||
"load_recall_budget_cache",
|
||||
"set_recall_budget_cache",
|
||||
"delete_recall_budget_cache",
|
||||
"get_recall_budget_from_cache",
|
||||
"save_user_memory_override_db",
|
||||
"delete_user_memory_override_db",
|
||||
]
|
||||
+36
-1
@@ -8,4 +8,39 @@ def make_ai_writing_session_store(session_factory) -> AIWritingSessionRepository
|
||||
return AIWritingSessionRepository(session_factory)
|
||||
|
||||
|
||||
__all__ = ["AIWritingSessionRow", "AIWritingSessionRepository", "make_ai_writing_session_store"]
|
||||
# ── 进程级默认 repo 单例 ──────────────────────────────────────────────────────
|
||||
#
|
||||
# AI 写作 graph 注册成 LangGraph Server assistant 后，节点跑在 harness 进程里，
|
||||
# 拿不到 FastAPI 的 ``request.app.state``；harness → app 又是禁止的 import 方向。
|
||||
# 这里提供一个 app 启动时注册、harness 节点运行时读取的「跨边界单例」通道。
|
||||
#
|
||||
# 写入：app/gateway/deps.py 在创建完 session_factory 之后调
|
||||
# ``register_default_store(app.state.ai_writing_session_store)``
|
||||
# 读取：节点里 ``get_default_store()``，None 表示未配置 DB（跳过写入即可）。
|
||||
_default_store: AIWritingSessionRepository | None = None
|
||||
|
||||
|
||||
def register_default_store(store: AIWritingSessionRepository | None) -> None:
|
||||
"""登记进程级默认 AI 写作 session store。
|
||||
|
||||
幂等：可多次调用，覆盖前值。传 ``None`` 用于显式清理（主要给测试）。
|
||||
"""
|
||||
global _default_store
|
||||
_default_store = store
|
||||
|
||||
|
||||
def get_default_store() -> AIWritingSessionRepository | None:
|
||||
"""返回进程级默认 AI 写作 session store；未配置时为 ``None``。
|
||||
|
||||
节点调用方应在拿到 ``None`` 时跳过 DB 写入（写入是「锦上添花」非致命）。
|
||||
"""
|
||||
return _default_store
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AIWritingSessionRow",
|
||||
"AIWritingSessionRepository",
|
||||
"make_ai_writing_session_store",
|
||||
"register_default_store",
|
||||
"get_default_store",
|
||||
]
|
||||
+12
-4
@@ -4,11 +4,11 @@ from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from sqlalchemy import DateTime, String, Text
|
||||
from sqlalchemy import String, Text
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from deerflow.persistence.base import Base
|
||||
from deerflow.persistence.types import PortableLongText
|
||||
from deerflow.persistence.types import BeijingDateTime, PortableLongText
|
||||
|
||||
|
||||
class AIWritingSessionRow(Base):
|
||||
@@ -26,5 +26,13 @@ class AIWritingSessionRow(Base):
|
||||
# 完整对话时间线（前端组装的 progressEvents + completedInterventions），用于历史回显。
|
||||
# 用 PortableLongText：MySQL 上映射为 LONGTEXT，避免 TEXT 64KB 上限导致保存 500。
|
||||
transcript: Mapped[str | None] = mapped_column(PortableLongText, nullable=True) # JSON object
|
||||
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
|
||||
updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
|
||||
# 时间列必须用 BeijingDateTime（项目规范，与 threads_meta 等其它表对齐）—— MySQL
|
||||
# DATETIME 不存时区，raw ``DateTime(timezone=True)`` 会把 UTC wall clock 写入
|
||||
# 列，读出来比真实北京时间晚 8 小时。``BeijingDateTime`` 在 bind 时 UTC→Beijing、
|
||||
# 在 result 时 attach +08:00，前后端拿到的就是北京时间且 tz-aware。
|
||||
created_at: Mapped[datetime] = mapped_column(BeijingDateTime(), default=lambda: datetime.now(UTC))
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
BeijingDateTime(),
|
||||
default=lambda: datetime.now(UTC),
|
||||
onupdate=lambda: datetime.now(UTC),
|
||||
)
|
||||
+46
@@ -63,6 +63,52 @@ class AIWritingSessionRepository:
|
||||
await session.refresh(row)
|
||||
return self._row_to_dict(row)
|
||||
|
||||
async def upsert_session(
|
||||
self,
|
||||
*,
|
||||
id: str,
|
||||
title: str,
|
||||
user_intent: str | None = None,
|
||||
user_id: str | None | _AutoSentinel = AUTO,
|
||||
status: str = "in_progress",
|
||||
) -> dict:
|
||||
"""幂等：行不存在则创建，行存在则保留原数据（不覆盖 draft / review 等已有写入）。
|
||||
|
||||
AI 写作 graph 的 ``intent_parser_node`` 在每次进入时调一次（首启 + resume 时
|
||||
都会跑），所以必须幂等。resume 时若直接 ``create`` 会因主键冲突抛错；用
|
||||
``update`` 又找不到行而 no-op。这个方法把两条路合一：行不存在新建，行存在
|
||||
只刷新 ``updated_at`` 以便列表页排序前置。
|
||||
"""
|
||||
resolved_user_id = resolve_user_id(user_id, method_name="AIWritingSessionRepository.upsert_session")
|
||||
now = datetime.now(UTC)
|
||||
async with self._sf() as session:
|
||||
row = await session.get(AIWritingSessionRow, id)
|
||||
if row is None:
|
||||
row = AIWritingSessionRow(
|
||||
id=id,
|
||||
user_id=resolved_user_id,
|
||||
title=title,
|
||||
user_intent=user_intent,
|
||||
status=status,
|
||||
draft_title=None,
|
||||
draft_markdown=None,
|
||||
completed_interventions=None,
|
||||
review_result=None,
|
||||
transcript=None,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
session.add(row)
|
||||
else:
|
||||
# 行已存在（resume 路径）：仅刷新 updated_at，保留 user_id / title / draft / review 等。
|
||||
# 若 user_id 之前为空（migration 兼容场景），用解析出来的填上。
|
||||
if row.user_id is None and resolved_user_id is not None:
|
||||
row.user_id = resolved_user_id
|
||||
row.updated_at = now
|
||||
await session.commit()
|
||||
await session.refresh(row)
|
||||
return self._row_to_dict(row)
|
||||
|
||||
async def update(
|
||||
self,
|
||||
session_id: str,
|
||||
|
||||
+23
@@ -0,0 +1,23 @@
|
||||
"""Per-user memory configuration override persistence."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from deerflow.persistence.memory_config.base import UserMemoryConfigStore
|
||||
from deerflow.persistence.memory_config.memory import MemoryUserMemoryConfigStore
|
||||
from deerflow.persistence.memory_config.model import UserMemoryConfigRow
|
||||
from deerflow.persistence.memory_config.sql import SqlUserMemoryConfigRepository
|
||||
|
||||
|
||||
def make_user_memory_config_store(session_factory) -> UserMemoryConfigStore:
|
||||
if session_factory is None:
|
||||
return MemoryUserMemoryConfigStore()
|
||||
return SqlUserMemoryConfigRepository(session_factory)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"UserMemoryConfigStore",
|
||||
"UserMemoryConfigRow",
|
||||
"MemoryUserMemoryConfigStore",
|
||||
"SqlUserMemoryConfigRepository",
|
||||
"make_user_memory_config_store",
|
||||
]
|
||||
+32
@@ -0,0 +1,32 @@
|
||||
"""Abstract persistence API for per-user memory config overrides."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
class UserMemoryConfigStore(ABC):
|
||||
|
||||
@abstractmethod
|
||||
async def get_recall_budget(self, user_id: str) -> str | None:
|
||||
"""Return the user's recall_budget override, or None if not set."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
async def set_recall_budget(self, user_id: str, value: str) -> None:
|
||||
"""Upsert the user's recall_budget override."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
async def delete(self, user_id: str) -> None:
|
||||
"""Remove all per-user overrides for the given user."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
async def load_all(self) -> dict[str, str]:
|
||||
"""Bulk-load all user recall_budget overrides.
|
||||
|
||||
Returns {user_id: recall_budget} for all rows where recall_budget IS NOT NULL.
|
||||
Called once at startup to warm the in-memory cache.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
+23
@@ -0,0 +1,23 @@
|
||||
"""In-memory user memory config store (used when database.backend=memory)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from deerflow.persistence.memory_config.base import UserMemoryConfigStore
|
||||
|
||||
|
||||
class MemoryUserMemoryConfigStore(UserMemoryConfigStore):
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._data: dict[str, str] = {}
|
||||
|
||||
async def get_recall_budget(self, user_id: str) -> str | None:
|
||||
return self._data.get(user_id)
|
||||
|
||||
async def set_recall_budget(self, user_id: str, value: str) -> None:
|
||||
self._data[user_id] = value
|
||||
|
||||
async def delete(self, user_id: str) -> None:
|
||||
self._data.pop(user_id, None)
|
||||
|
||||
async def load_all(self) -> dict[str, str]:
|
||||
return dict(self._data)
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
"""ORM model for per-user memory configuration overrides."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from sqlalchemy import String
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from deerflow.persistence.base import Base
|
||||
from deerflow.persistence.types import BeijingDateTime
|
||||
|
||||
|
||||
class UserMemoryConfigRow(Base):
|
||||
__tablename__ = "user_memory_configs"
|
||||
|
||||
# One row per user; upserted on write.
|
||||
user_id: Mapped[str] = mapped_column(String(64), primary_key=True)
|
||||
|
||||
# The only user-writable field. NULL means "use global default from config.yaml".
|
||||
# Allowed values: "low" | "mid" | "high".
|
||||
recall_budget: Mapped[str | None] = mapped_column(String(16), nullable=True)
|
||||
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
BeijingDateTime(),
|
||||
nullable=False,
|
||||
default=lambda: datetime.now(UTC),
|
||||
onupdate=lambda: datetime.now(UTC),
|
||||
)
|
||||
+50
@@ -0,0 +1,50 @@
|
||||
"""SQLAlchemy-backed user memory config repository."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
|
||||
from deerflow.persistence.memory_config.base import UserMemoryConfigStore
|
||||
from deerflow.persistence.memory_config.model import UserMemoryConfigRow
|
||||
|
||||
|
||||
class SqlUserMemoryConfigRepository(UserMemoryConfigStore):
|
||||
|
||||
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
||||
self._sf = session_factory
|
||||
|
||||
async def get_recall_budget(self, user_id: str) -> str | None:
|
||||
async with self._sf() as session:
|
||||
row = await session.get(UserMemoryConfigRow, user_id)
|
||||
return row.recall_budget if row is not None else None
|
||||
|
||||
async def set_recall_budget(self, user_id: str, value: str) -> None:
|
||||
async with self._sf() as session:
|
||||
row = await session.get(UserMemoryConfigRow, user_id)
|
||||
if row is None:
|
||||
row = UserMemoryConfigRow(
|
||||
user_id=user_id,
|
||||
recall_budget=value,
|
||||
updated_at=datetime.now(UTC),
|
||||
)
|
||||
session.add(row)
|
||||
else:
|
||||
row.recall_budget = value
|
||||
row.updated_at = datetime.now(UTC)
|
||||
await session.commit()
|
||||
|
||||
async def delete(self, user_id: str) -> None:
|
||||
async with self._sf() as session:
|
||||
row = await session.get(UserMemoryConfigRow, user_id)
|
||||
if row is not None:
|
||||
await session.delete(row)
|
||||
await session.commit()
|
||||
|
||||
async def load_all(self) -> dict[str, str]:
|
||||
stmt = select(UserMemoryConfigRow).where(UserMemoryConfigRow.recall_budget.is_not(None))
|
||||
async with self._sf() as session:
|
||||
result = await session.execute(stmt)
|
||||
return {row.user_id: row.recall_budget for row in result.scalars() if row.recall_budget is not None}
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
"""Create user_memory_configs table.
|
||||
|
||||
Replaces the per-user file-based memory config override
|
||||
({base_dir}/users/{user_id}/memory_config.json) with a DB row.
|
||||
Stores only hindsight.recall_budget — the single user-writable field
|
||||
after the memory config simplification.
|
||||
|
||||
Revision ID: 20260521_01
|
||||
Revises: 20260519_02
|
||||
Create Date: 2026-05-21
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from sqlalchemy import inspect
|
||||
|
||||
|
||||
revision = "20260521_01"
|
||||
down_revision = "20260519_02"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def _has_table(table_name: str) -> bool:
|
||||
inspector = inspect(op.get_bind())
|
||||
return table_name in inspector.get_table_names()
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
if _has_table("user_memory_configs"):
|
||||
return
|
||||
op.create_table(
|
||||
"user_memory_configs",
|
||||
sa.Column("user_id", sa.String(64), primary_key=True),
|
||||
sa.Column("recall_budget", sa.String(16), nullable=True),
|
||||
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
if _has_table("user_memory_configs"):
|
||||
op.drop_table("user_memory_configs")
|
||||
+2
@@ -16,6 +16,7 @@ there is no matching entity directory.
|
||||
"""
|
||||
|
||||
from deerflow.persistence.agents.model import AgentRow
|
||||
from deerflow.persistence.memory_config.model import UserMemoryConfigRow
|
||||
from deerflow.persistence.feedback.model import FeedbackRow
|
||||
from deerflow.persistence.llm_metrics.model import LlmCallMetricRow
|
||||
from deerflow.persistence.models.run_event import RunEventRow
|
||||
@@ -55,5 +56,6 @@ __all__ = [
|
||||
"ThreadMetaRow",
|
||||
"ThreadShareRow",
|
||||
"ToolCallMetricRow",
|
||||
"UserMemoryConfigRow",
|
||||
"UserRow",
|
||||
]
|
||||
+11
@@ -56,6 +56,17 @@ class ThreadMetaStore(abc.ABC):
|
||||
async def update_status(self, thread_id: str, status: str, *, user_id: str | None | _AutoSentinel = AUTO) -> None:
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
async def update_assistant_id(self, thread_id: str, assistant_id: str, *, user_id: str | None | _AutoSentinel = AUTO) -> None:
|
||||
"""Backfill ``assistant_id`` for a thread row.
|
||||
|
||||
Used by ``start_run`` when the row was pre-created via ``POST /threads``
|
||||
without an assistant_id (the SDK's ``client.threads.create()`` doesn't
|
||||
pass it), and the first run on the thread reveals the assistant. No-op
|
||||
if the thread does not exist or the owner check fails.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
async def update_metadata(self, thread_id: str, metadata: dict, *, user_id: str | None | _AutoSentinel = AUTO) -> None:
|
||||
"""Merge ``metadata`` into the thread's metadata field.
|
||||
|
||||
+8
@@ -117,6 +117,14 @@ class MemoryThreadMetaStore(ThreadMetaStore):
|
||||
record["updated_at"] = now_iso()
|
||||
await self._store.aput(THREADS_NS, thread_id, record)
|
||||
|
||||
async def update_assistant_id(self, thread_id: str, assistant_id: str, *, user_id: str | None | _AutoSentinel = AUTO) -> None:
|
||||
record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.update_assistant_id")
|
||||
if record is None:
|
||||
return
|
||||
record["assistant_id"] = assistant_id
|
||||
record["updated_at"] = now_iso()
|
||||
await self._store.aput(THREADS_NS, thread_id, record)
|
||||
|
||||
async def update_metadata(self, thread_id: str, metadata: dict, *, user_id: str | None | _AutoSentinel = AUTO) -> None:
|
||||
record = await self._get_owned_record(thread_id, user_id, "MemoryThreadMetaStore.update_metadata")
|
||||
if record is None:
|
||||
|
||||
+14
@@ -174,6 +174,20 @@ class ThreadMetaRepository(ThreadMetaStore):
|
||||
await session.execute(update(ThreadMetaRow).where(ThreadMetaRow.thread_id == thread_id).values(status=status, updated_at=datetime.now(UTC)))
|
||||
await session.commit()
|
||||
|
||||
async def update_assistant_id(
|
||||
self,
|
||||
thread_id: str,
|
||||
assistant_id: str,
|
||||
*,
|
||||
user_id: str | None | _AutoSentinel = AUTO,
|
||||
) -> None:
|
||||
resolved_user_id = resolve_user_id(user_id, method_name="ThreadMetaRepository.update_assistant_id")
|
||||
async with self._sf() as session:
|
||||
if not await self._check_ownership(session, thread_id, resolved_user_id):
|
||||
return
|
||||
await session.execute(update(ThreadMetaRow).where(ThreadMetaRow.thread_id == thread_id).values(assistant_id=assistant_id, updated_at=datetime.now(UTC)))
|
||||
await session.commit()
|
||||
|
||||
async def update_metadata(
|
||||
self,
|
||||
thread_id: str,
|
||||
|
||||
@@ -130,8 +130,16 @@ async def run_agent(
|
||||
stream_subgraphs: bool = False,
|
||||
interrupt_before: list[str] | Literal["*"] | None = None,
|
||||
interrupt_after: list[str] | Literal["*"] | None = None,
|
||||
command: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""Execute an agent in the background, publishing events to *bridge*."""
|
||||
"""Execute an agent in the background, publishing events to *bridge*.
|
||||
|
||||
``command`` 来自 LangGraph SDK 的 ``Command`` 协议（``{resume, update, goto}``）。
|
||||
存在时把它翻译成 ``langgraph.types.Command`` 并作为 ``astream`` 的入参 ——
|
||||
Pregel 检查到 Command 输入会从最近 checkpoint 处恢复（典型场景：interrupt 之后
|
||||
的 ``Command(resume=user_input)``）；不存在时仍按 ``graph_input`` dict 起新 run。
|
||||
AI 写作的 4 个干预卡就靠这个把 user_input 传回 ``pause_*_confirm`` 节点。
|
||||
"""
|
||||
|
||||
# Unpack infrastructure dependencies from RunContext.
|
||||
checkpointer = ctx.checkpointer
|
||||
@@ -269,10 +277,23 @@ async def run_agent(
|
||||
logger.info("Run %s: streaming with modes %s (requested: %s)", run_id, lg_modes, requested_modes)
|
||||
|
||||
# 7. Stream using graph.astream
|
||||
# 如果客户端传了 Command（典型：interrupt 后续提的 ``Command(resume=...)``），
|
||||
# 把它翻译成 langgraph.types.Command 作为 astream 输入；否则按 graph_input dict
|
||||
# 起新一轮。Pregel 检查到 Command 会从最近 checkpoint 处恢复，而不是重头跑。
|
||||
stream_input: Any = graph_input
|
||||
if command is not None:
|
||||
from langgraph.types import Command as _LGCommand
|
||||
|
||||
stream_input = _LGCommand(
|
||||
resume=command.get("resume"),
|
||||
update=command.get("update"),
|
||||
goto=command.get("goto"),
|
||||
)
|
||||
|
||||
if len(lg_modes) == 1 and not stream_subgraphs:
|
||||
# Single mode, no subgraphs: astream yields raw chunks
|
||||
single_mode = lg_modes[0]
|
||||
async for chunk in agent.astream(graph_input, config=runnable_config, stream_mode=single_mode):
|
||||
async for chunk in agent.astream(stream_input, config=runnable_config, stream_mode=single_mode):
|
||||
if record.abort_event.is_set():
|
||||
logger.info("Run %s abort requested — stopping", run_id)
|
||||
break
|
||||
@@ -281,7 +302,7 @@ async def run_agent(
|
||||
else:
|
||||
# Multiple modes or subgraphs: astream yields tuples
|
||||
async for item in agent.astream(
|
||||
graph_input,
|
||||
stream_input,
|
||||
config=runnable_config,
|
||||
stream_mode=lg_modes,
|
||||
subgraphs=stream_subgraphs,
|
||||
|
||||
+23
-3
@@ -10,6 +10,7 @@ Consumers: ``deerflow.runtime.runs.worker`` (SSE publishing) and
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
from typing import Any
|
||||
|
||||
|
||||
@@ -35,6 +36,16 @@ def serialize_lc_object(obj: Any) -> Any:
|
||||
return obj.dict()
|
||||
except Exception:
|
||||
pass
|
||||
# Dataclass instances —— LangGraph 的 ``Interrupt`` 就是 dataclass 没有
|
||||
# model_dump 也没有 dict 方法；之前 fallback 到 ``str(obj)`` 会把
|
||||
# ``Interrupt(value=..., id='...')`` 序列化成字符串，下游 SDK
|
||||
# (``@langchain/langgraph-sdk``) 期待 dict 结构 → 直接丢弃 → AI 写作
|
||||
# 的 4 个干预卡永远不弹（参 ai-writing-优化.md §阶段 3 排查记录）。
|
||||
if dataclasses.is_dataclass(obj) and not isinstance(obj, type):
|
||||
try:
|
||||
return {k: serialize_lc_object(v) for k, v in dataclasses.asdict(obj).items()}
|
||||
except Exception:
|
||||
pass
|
||||
# Last resort
|
||||
try:
|
||||
return str(obj)
|
||||
@@ -45,12 +56,21 @@ def serialize_lc_object(obj: Any) -> Any:
|
||||
def serialize_channel_values(channel_values: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Serialize channel values, stripping internal LangGraph keys.
|
||||
|
||||
Internal keys like ``__pregel_*`` and ``__interrupt__`` are removed
|
||||
to match what the LangGraph Platform API returns.
|
||||
``__pregel_*`` 键始终剥离。``__interrupt__`` 仅在**为空时**剥离 —— 因为
|
||||
``@langchain/langgraph-sdk`` 的 ``useStream`` 只在 ``event === "values"``
|
||||
通道里读 ``data.__interrupt__`` 并通过 ``stream.interrupt`` 暴露给上层。
|
||||
AI 写作的 4 个干预卡（material/outline/draft/review）就是靠它触发的；
|
||||
继续过滤会让前端永远拿不到 interrupt。
|
||||
|
||||
主聊天不消费 ``stream.interrupt``，但若我们让一个空数组溜过去，SDK 会
|
||||
fallback 成 ``{when: "breakpoint"}`` 的伪 interrupt —— 所以仅在数组非空
|
||||
时保留，两侧都不破坏。
|
||||
"""
|
||||
result: dict[str, Any] = {}
|
||||
for key, value in channel_values.items():
|
||||
if key.startswith("__pregel_") or key == "__interrupt__":
|
||||
if key.startswith("__pregel_"):
|
||||
continue
|
||||
if key == "__interrupt__" and not value:
|
||||
continue
|
||||
result[key] = serialize_lc_object(value)
|
||||
return result
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
"""AI 写作 assistant 冒烟脚本：验证 ai_writing graph 注册到 LangGraph Server 后，
|
||||
interrupt/Command(resume=...)/state 持久化/custom stream 全链路是否跑通。
|
||||
|
||||
阶段 0 用名为 ai_writing_probe；阶段 1 起改为正式 assistant 名 ai_writing，本脚本随之切换。
|
||||
|
||||
前置条件：
|
||||
1. langgraph.json 已注册 ai_writing（见 ai-writing-优化.md 阶段 1）。
|
||||
2. 后端已通过 `make dev` 启动，Gateway + LangGraph Server 监听 8001。
|
||||
3. config.yaml 配好至少一个可用 model，researcher 用到的 search_provider 也已配置
|
||||
—— 否则 researcher 节点会失败（这是业务依赖，不是迁移问题，仍然算「跑通」）。
|
||||
|
||||
运行：
|
||||
cd offline-backend-20260512/backend
|
||||
PYTHONPATH=. uv run python scripts/probe_ai_writing_assistant.py
|
||||
|
||||
通过条件（与 ai-writing-优化.md §阶段 0 §通过条件 对齐）：
|
||||
- assistants 列表里能看到 ai_writing_probe
|
||||
- 跑到 pause_material_confirm 能拿到 interrupt 事件
|
||||
- Command(resume={"action": "confirm", ...}) 能成功续跑
|
||||
- 重建客户端后 client.threads.get_state(thread_id) 仍能拿到上一阶段 state
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
from langgraph_sdk import get_client
|
||||
except ImportError:
|
||||
print("ERROR: langgraph-sdk 未安装。run: uv pip install langgraph-sdk", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
ASSISTANT_ID = "ai_writing"
|
||||
DEFAULT_URL = os.environ.get("DEERFLOW_LANGGRAPH_URL", "http://localhost:8001/api/langgraph")
|
||||
|
||||
|
||||
def _banner(text: str) -> None:
|
||||
print(f"\n{'=' * 8} {text} {'=' * 8}")
|
||||
|
||||
|
||||
def _summarize_event(event: Any) -> str:
|
||||
"""裁剪日志，避免一屏全是大段 markdown / 素材列表。"""
|
||||
s = repr(event)
|
||||
return s if len(s) <= 240 else s[:240] + "…"
|
||||
|
||||
|
||||
async def _stream_until_interrupt(
|
||||
client, thread_id: str, *, input_payload: dict | None = None, command: dict | None = None
|
||||
) -> dict | None:
|
||||
"""跑一次 runs.stream，捕获第一个 interrupt（含 pause payload）后返回，否则返回 None。
|
||||
|
||||
收到 interrupt 时把 payload 一并打出来，便于人工核对 4 个 pause 点的 schema。
|
||||
"""
|
||||
interrupt_payload: dict | None = None
|
||||
kwargs: dict[str, Any] = {
|
||||
"stream_mode": ["values", "messages-tuple", "custom", "updates"],
|
||||
}
|
||||
if input_payload is not None:
|
||||
kwargs["input"] = input_payload
|
||||
if command is not None:
|
||||
kwargs["command"] = command
|
||||
|
||||
async for chunk in client.runs.stream(thread_id, ASSISTANT_ID, **kwargs):
|
||||
# langgraph-sdk 把不同 mode 的事件混在同一 stream 里，按 chunk.event 分类。
|
||||
event_name = getattr(chunk, "event", None) or (
|
||||
chunk.get("event") if isinstance(chunk, dict) else None
|
||||
)
|
||||
data = getattr(chunk, "data", None) or (
|
||||
chunk.get("data") if isinstance(chunk, dict) else None
|
||||
)
|
||||
print(f"[stream] event={event_name} data={_summarize_event(data)}")
|
||||
|
||||
# SDK 在遇到 interrupt() 时会发一个特殊事件 —— 不同版本字段名不同，做下兼容。
|
||||
if event_name == "updates" and isinstance(data, dict):
|
||||
for _node, node_update in data.items():
|
||||
if isinstance(node_update, dict) and "__interrupt__" in node_update:
|
||||
interrupt_payload = node_update["__interrupt__"]
|
||||
if event_name in {"interrupt", "__interrupt__"}:
|
||||
interrupt_payload = data
|
||||
break
|
||||
|
||||
return interrupt_payload
|
||||
|
||||
|
||||
async def main() -> int:
|
||||
print(f"LangGraph URL = {DEFAULT_URL}")
|
||||
client = get_client(url=DEFAULT_URL)
|
||||
|
||||
# ── Step 1: 确认 assistant 已注册 ──────────────────────────────────────────
|
||||
_banner("Step 1: 列 assistant，确认 probe 已注册")
|
||||
assistants = await client.assistants.search()
|
||||
names = [a.get("name") or a.get("graph_id") for a in assistants]
|
||||
print(f"已注册 assistants: {names}")
|
||||
if not any(ASSISTANT_ID in (n or "") for n in names):
|
||||
print(f"FAIL: 未找到 {ASSISTANT_ID}，请检查 langgraph.json 注册是否生效。")
|
||||
return 1
|
||||
|
||||
# ── Step 2: 创建 thread 并触发第一次 run ──────────────────────────────────
|
||||
_banner("Step 2: 创建 thread + 启动写作（应停在 pause_material）")
|
||||
thread = await client.threads.create()
|
||||
thread_id = thread["thread_id"]
|
||||
print(f"thread_id = {thread_id}")
|
||||
|
||||
start_input = {
|
||||
"user_intent": "写一篇关于 LangGraph 多 worker 部署最佳实践的分析报道，目标读者是后端工程师，约 800 字",
|
||||
"article_type": "analysis",
|
||||
"target_audience": "后端工程师",
|
||||
"word_count_target": 800,
|
||||
"keyword_count": 4,
|
||||
"writing_mode": "loose",
|
||||
"max_revisions": 2,
|
||||
"status": "researching",
|
||||
"drafts": [],
|
||||
"progress_events": [],
|
||||
}
|
||||
payload = await _stream_until_interrupt(client, thread_id, input_payload=start_input)
|
||||
if not payload:
|
||||
print("FAIL: 未收到 interrupt 事件 —— pause_material 可能没停住，或 researcher 抛错。")
|
||||
return 2
|
||||
print(f"\nPAUSE-1 payload = {_summarize_event(payload)}")
|
||||
|
||||
# ── Step 3: 确认素材，应停在 pause_outline ────────────────────────────────
|
||||
_banner("Step 3: Command(resume) 确认素材（应停在 pause_outline）")
|
||||
payload = await _stream_until_interrupt(
|
||||
client, thread_id,
|
||||
command={"resume": {"action": "confirm"}}, # 不指定 approvedMaterialIds → 默认全部通过
|
||||
)
|
||||
if not payload:
|
||||
print("FAIL: 未停在 pause_outline。")
|
||||
return 3
|
||||
print(f"\nPAUSE-2 payload = {_summarize_event(payload)}")
|
||||
|
||||
# ── Step 4: 确认大纲，应停在 pause_draft 或 pause_section_help ───────────
|
||||
_banner("Step 4: Command(resume) 确认大纲（应停在 pause_draft 或 pause_section_help）")
|
||||
payload = await _stream_until_interrupt(
|
||||
client, thread_id,
|
||||
command={"resume": {"action": "confirm"}},
|
||||
)
|
||||
if not payload:
|
||||
print("FAIL: 未停在 pause_draft / pause_section_help。")
|
||||
return 4
|
||||
print(f"\nPAUSE-3 payload = {_summarize_event(payload)}")
|
||||
|
||||
# ── Step 5: 定稿 ─────────────────────────────────────────────────────────
|
||||
_banner("Step 5: Command(resume) finalize 草稿（应跑到 END）")
|
||||
payload = await _stream_until_interrupt(
|
||||
client, thread_id,
|
||||
command={"resume": {"action": "finalize"}},
|
||||
)
|
||||
if payload:
|
||||
print(f"NOTE: 收到额外 interrupt（可能编辑审核打回） payload = {_summarize_event(payload)}")
|
||||
print("尝试强制定稿…")
|
||||
payload = await _stream_until_interrupt(
|
||||
client, thread_id,
|
||||
command={"resume": {"action": "force_finalize", "overrideVerdict": True}},
|
||||
)
|
||||
if payload:
|
||||
print("WARN: 强制定稿后仍有 interrupt，graph 路由可能有意外分支。")
|
||||
|
||||
# ── Step 6: 验证 state 持久化 ────────────────────────────────────────────
|
||||
_banner("Step 6: 重新建客户端 client.threads.get_state(thread_id) 应能拿到最终 state")
|
||||
client2 = get_client(url=DEFAULT_URL)
|
||||
state = await client2.threads.get_state(thread_id)
|
||||
values = state.get("values") or {}
|
||||
final_status = values.get("status")
|
||||
has_draft = bool(values.get("current_draft"))
|
||||
print(f"final status = {final_status}, has_draft = {has_draft}, "
|
||||
f"next nodes = {state.get('next')}")
|
||||
if final_status not in {"done", "writing_draft", "revising"}:
|
||||
print("WARN: 最终 status 不是 done，请人工确认上面流程是否正常。")
|
||||
print("\nPROBE OK —— 阶段 0 验证通过的核心条件已成立（interrupt + resume + state 持久化）。")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(asyncio.run(main()))
|
||||
Reference in new issue
Block a user