diff --git a/README.md b/README.md index e3456f3..3838e98 100644 --- a/README.md +++ b/README.md @@ -63,6 +63,18 @@ Chat Completions API，`EXTRA_BODY` 中的关闭推理等扩展参数照常透 模型调用继承图的 callbacks、metadata 和 tags，并标记 `llm.complete` / `llm.stream` 及 `llm_attempt`。这是追踪集成的基础，日志、存储和整轮 SSE 的观测需分别接入。 +事故采集使用 LangGraph 原生并行分支：`collect` 生成回复， +`extract_accident_fields` 提取字段，二者经汇合边进入 `merge_collect`。 +两个分支读取同一轮的参考时间和状态，各写独立通道，合并后统一提交。 +人员核实的抽取依赖本轮回复，保持串行依赖。 + +增量输出由 LangChain 原生 `astream_events(version="v2")` 消费： +`RunnableGenerator` 输出通过阶段前缀校验的文本；`validated_reply` 的 +`on_chain_end` 表示整条回复已通过校验，立即刷新分句器最后一句。 +抽取模型的事件不转成聊天正文；`formUpdate` 和 `done` 仍等待业务提交成功。 +客户端关闭流时由原生事件流取消并等待图任务，公开 SSE 事件类型不变。 +节点职责、事件映射和提交时序见 [docs/agent-streaming.md](docs/agent-streaming.md)。 + ## HTTP 接口 - `GET /`：存活检查。 diff --git a/docs/agent-streaming.md b/docs/agent-streaming.md new file mode 100644 index 0000000..75bc0fe --- /dev/null +++ b/docs/agent-streaming.md @@ -0,0 +1,82 @@ +# LangGraph 原生并行与流式输出 + +文本模型使用 ChatOpenAI，业务图嵌入 FastAPI。对外仍使用原来的 +stage_code / text_delta / formUpdate / done / error SSE 契约。 + +## 图内并行 + +事故采集阶段： + +```mermaid +flowchart LR + P[normalize_input] --> R[collect：生成并校验回复] + P --> X[extract_accident_fields：提取候选字段] + R --> M[merge_collect] + X --> M + M --> C[应用服务事务提交] +``` + +条件边只在采集阶段选中两个分支；welcome、拍照、人员核实走自己的分支。 +通过 `add_edge(["collect", "extract_accident_fields"], "merge_collect")` +保证两个分支都完成后，汇合节点只运行一次。 + +`normalize_input` 固定服务器参考时间并清空本轮临时通道。回复节点只更新 +回复、阶段和对话簿记；提取节点只写 `extracted_accident_fields`；汇合节点 +更新 accident_fields / turn_patch，再清空候选通道。两条分支都读取本轮开始时 +的快照，不需要为并发写同一字段增加合并 reducer。 + +人员核实的提取提示词当前依赖本轮模型回复，并负责声明人员切换。 +本次保留它的串行依赖与提示词，仍能在抽取完成前发送完整回复。 + +## 原生流与校验 + +后端使用 LangChain `astream_events(version="v2")`，其生命周期管理负责 +消费期间的任务调度，以及流关闭后的取消与等待。 + +| 原生事件 | 条件 | 应用服务适配 | +|---|---|---| +| on_chain_stream | name=validated_reply_tokens | TextDelta | +| on_chain_end | name=validated_reply | ReplyFinished，刷新分句器 | +| on_chain_end | 根 accident_turn 正常结束 | 已提交结果，可发送 FormUpdate | + +`validated_reply_tokens` 是原生 RunnableGenerator，封装现有模型文本迭代器， +等待完整合法的 state 前缀后才产出第一段文本。原始 ChatOpenAI 的 model 事件 +仍可被 callbacks 追踪，但公开文本只消费受控回复 Runnable 的事件。 +原生 messages 流包含校验前 token，单纯按模型或节点过滤不能替代前缀校验。 + +`validated_reply` 是原生 RunnableLambda。完整回复通过格式、阶段转换和可选 +时间依据校验后，其正常 on_chain_end 才表示回复结束。模型结束事件本身 +不能证明这些业务校验通过，因此不作为刷新最后一句的依据。 + +字段提取保留 complete / ainvoke，显式 stream=False，避免事件追踪 callback +把该调用隐式切换为流式。辅助模型输出不会映射为 TextDelta。 + +## 回复结束与事务完成 + +ReplyFinished 仅是后端中立适配层的内部标记，不新增公开 SSE 事件。 +它立即刷新 useTextChunk 的剩余完整回复，覆盖有/无字段差量和 +needFormUpdate=false 的情况。 + +文本可以早于提取和业务事务完成到达客户端；FormUpdate 与 done 仍在事务 +成功后出现。完整回复失败、图失败、业务提交失败时，发送 error 而不发送 done。 +已经输出的文字不可撤回，客户端仍以 done 判断本轮成功。 + +取消流时必须等图分支与模型关闭后才释放会话锁。持久化模式继续从业务提交 +快照恢复；并行分支产生的未提交 checkpoint 不成为业务成功状态。 + +## 验证 + +`test/backends/test_native_parallel_streaming.py` 使用事件屏障验证： + +- 回复和抽取同时运行，汇合与提交等待两者完成。 +- 抽取阻塞时已经收到完整正文和 ReplyFinished。 +- 分句器最后一句及时刷新；公开 SSE 类型、文本顺序和终止事件保持兼容。 +- 抽取 JSON 不泄漏，原生节点和汇合在事件流中可见。 +- 断开取消剩余抽取，SQLite 未提交状态可恢复并重试。 +- 完整回复校验失败和业务提交失败均不产生虚假成功。 +- ChatOpenAI 的真实 HTTP mock 路径保留普通提取与流式回复的请求参数。 + +官方接口参考： + +- https://docs.langchain.com/oss/python/langgraph/streaming +- https://docs.langchain.com/oss/python/langgraph/use-graph-api diff --git a/docs/assignment6-compat.md b/docs/assignment6-compat.md index 706e43d..2951b9c 100644 --- a/docs/assignment6-compat.md +++ b/docs/assignment6-compat.md @@ -90,6 +90,18 @@ ## 流式、并发与字段一致性约定 +- 事故采集的回复与抽取为独立 LangGraph 节点，通过原生条件分支并行， + `add_edge(["collect", "extract_accident_fields"], "merge_collect")` 汇合； + 分支各写独立通道，参考时间在 normalize_input 固定，合并后统一提交。 + 人员抽取依赖本轮模型回复，暂保留串行，不改变其提示词和人员切换语义。 +- 模型文本与生命周期由原生 `astream_events` / `RunnableGenerator` 驱动， + 不再自行维护图后台任务、队列或 custom event bus。原始模型 token 先经过 + 业务前缀校验；只有受控回复 Runnable 的文本流能映射为 TextDelta。 + `validated_reply` 的正常 chain-end 映射为内部 ReplyFinished，立即刷新 + 分句缓冲；它不会对外新增 SSE 事件，也不表示字段已提交。 + 即使本轮无 formUpdate 或 needFormUpdate=false，最后一句也不等待抽取结束。 + 业务提交失败仍发 error，不能发 formUpdate/done 冒充成功。 + - 阶段前缀先缓冲并按当前节点的转换规则校验，非法前缀不推送； 未输出时仍可使用一次修复预算。已开始输出后，全文格式或时间依据 校验失败会以错误终止，不在同一流中拼接修复答案，不提交本轮业务字段。 diff --git a/docs/langgraph-backend-migration-plan.md b/docs/langgraph-backend-migration-plan.md index a7810b4..7eaf37e 100644 --- a/docs/langgraph-backend-migration-plan.md +++ b/docs/langgraph-backend-migration-plan.md @@ -370,7 +370,7 @@ docs/langgraph/ | 5. 拍照与核实 | ✅ 完成 | 真实模型 4/4 + 12/12 + 2/2（20261021 基准） | | 6. 四接口业务状态 | 已实现 | 应用服务、异步 repository、字段别名、会话串行与接口回归 | | 7. 持久化 | SQLite/PostgreSQL 初版已实现 | `src/storage/`；真实两库集成、进程退出恢复、幂等和版本冲突；见 [agent-storage.md](agent-storage.md) | -| 8. 真流式 | 已实现核心路径 | 模型增量、前缀校验、取消与提交回调；流式终止事件回归 | +| 8. 真流式 | 已实现核心路径 | 原生 astream_events / Runnable 流、采集并行节点、回复校验后提前刷新；见 [agent-streaming.md](agent-streaming.md) | | 9—10 | 未完成 | 生产访问边界、容量验证与双后端等价性验收继续保留 | 数据库实现阶段全量测试 315 通过，包含真实 SQLite/PostgreSQL 集成测试。 @@ -380,6 +380,11 @@ docs/langgraph/ 提示词来源和逐字比对见 `test/agent/test_workflow_prompts.py`。下方真实模型 黄金评测结果是此前轮次，数据库开发和本次模型适配迁移均未重新运行 46 条评测。 +原生并行/流式更新：采集回复与字段提取通过图条件边并行并原生汇合， +人员核实保留依赖回复的串行提取。validated_reply 的原生结束事件触发 +最后一句刷新，业务提交成功后才发送字段差量与 done。取消、失败和 +SQLite 快照恢复有独立回归；详见 [agent-streaming.md](agent-streaming.md)。 + ### 已落地的架构决策 1. **LLM 提议、代码校验（propose-validate）**：所有 stage_code 由节点内 LLM 产出（提示词即程序），代码只做前缀解析、白名单校验（含 clientMode 出口）和落库；曾尝试的代码强制推进/短促词守卫已按此分工移除。 @@ -420,6 +425,7 @@ src/agent/ confirm.py 核实节点（字段提取按 person_index 隔离） llm.py LLMService Protocol / Fake / ChatOpenAI(超时重试温度) / workflow fake chat_model.py ChatOpenAI 流式迭代器的显式关闭与并发上下文隔离 + reply_stream.py 原生 RunnableGenerator 的前缀校验流 fields.py 字段三态校验；clock.py 可控时钟 prompts/ 七份版本化提示词（来源 nodeId + 逐字原文） ``` diff --git a/src/agent/collect.py b/src/agent/collect.py index 572f282..74f69f1 100644 --- a/src/agent/collect.py +++ b/src/agent/collect.py @@ -92,6 +92,23 @@ async def validated_reply( 或关键字替模型选择动作。非法输出不写回图状态；修复仍失败则抛出 不含模型原文的错误。 """ + async def invoke(payload): + return await _validated_reply(**payload) + + # 原生 chain-end 事件表示整条回复已经通过校验；它早于字段提取/业务提交。 + from langchain_core.runnables import RunnableLambda + + return await RunnableLambda(invoke).ainvoke( + dict(llm=llm, system_prompt=system_prompt, user_message=user_message, + current=current, validator=validator, history=history, + reply_validator=reply_validator), + config={"run_name": "validated_reply", "tags": ["reply"]}, + ) + + +async def _validated_reply( + llm, system_prompt, user_message, current, *, validator, history, reply_validator, +): for attempt in range(2): prompt = system_prompt if attempt: @@ -273,15 +290,15 @@ def _bookkeeping(state: AgentState) -> tuple[int, int]: return no_reply_count, state.get("clarify_count") or 0 -async def run_collect(state: AgentState, llm: LLMService, *, now: datetime | None = None) -> dict: - """采集节点核心：一轮对话的 LLM 接入（参考实现）。 +def collect_time(state: AgentState, now: datetime | None = None) -> datetime: + if now is None and state.get("turn_observed_at"): + return datetime.fromisoformat(state["turn_observed_at"]) + return observation_time(now) - 流程：一次模型调用 → 前缀校验（不合规修复一次）→ - 白名单守卫（非法提议修复一次）→ 字段候选落库 → 返回状态增量。 - 任何一步修复后仍不合规都抛 CollectReplyError：不提交业务状态、 - 错误信息不携带原始模型内容。 - """ - moment = observation_time(now) + +async def collect_reply(state: AgentState, llm: LLMService, *, now: datetime | None = None) -> dict: + """只生成并校验回复；不等待独立的字段抽取分支。""" + moment = collect_time(state, now) evidence = time_evidence(state, moment) system_prompt = render_prompt(COLLECTION, state, now=moment) + context_prompt(state, time_context=evidence) user_message = _build_user_message(state, now=moment) @@ -290,11 +307,6 @@ async def run_collect(state: AgentState, llm: LLMService, *, now: datetime | Non llm, system_prompt, user_message, current, history=conversation_history(state), reply_validator=lambda code, body: valid_time_reply(code, body, evidence), ) - - - # 4) 字段提取与表单差量（提取调用在对话调用之外，见 _extract_turn_fields） - extracted = await _extract_turn_fields(state, llm, now=moment, time_context=evidence) - accident_fields, patch = _merge_form(state, extracted) no_reply_count, clarify_count = _bookkeeping(state) # 5) 状态增量：response_text 保留完整前缀文本，API 层负责解析； @@ -311,8 +323,30 @@ async def run_collect(state: AgentState, llm: LLMService, *, now: datetime | Non "response_text": raw, "stage_code": stage, "raw_status_code": stage, - "accident_fields": accident_fields, - "turn_patch": patch, "dialogue": dialogue, "dialogue_timestamps": timestamps, } + + +async def collect_extraction(state: AgentState, llm: LLMService) -> dict: + """与回复读取同一份轮次快照，返回独立的候选结果通道。""" + moment = collect_time(state) + extracted = await _extract_turn_fields( + state, llm, now=moment, time_context=time_evidence(state, moment), + ) + return {"extracted_accident_fields": extracted} + + +def merge_collect(state: AgentState) -> dict: + form, patch = _merge_form(state, state.get("extracted_accident_fields") or {}) + return {"accident_fields": form, "turn_patch": patch, "extracted_accident_fields": {}} + + +async def run_collect(state: AgentState, llm: LLMService, *, now: datetime | None = None) -> dict: + """独立调用入口；图内通过两个原生节点并行执行回复和提取。""" + moment = observation_time(now) + snapshot = {**state, "turn_observed_at": moment.isoformat()} + reply = await collect_reply(snapshot, llm) + extraction = await collect_extraction(snapshot, llm) + merged = merge_collect({**snapshot, **reply, **extraction}) + return {**reply, **merged} diff --git a/src/agent/graph.py b/src/agent/graph.py index 060151d..7a49f56 100644 --- a/src/agent/graph.py +++ b/src/agent/graph.py @@ -8,6 +8,8 @@ - 节点间大切换由 ``route_by_stage`` 条件边按 ``stage_code`` 分发 （对照 Condition#2/#7/#9 等）；一轮 HTTP 请求 = 一次图执行， 跨轮记忆由 checkpoint 承载； +- 采集阶段回复与字段提取由原生条件分支同时调度，汇合后合并差量； + 两个分支读取同一轮时间与状态，各自写独立通道； - 节点内阶段码一律由 LLM 提议（提示词即程序），代码只校验 （前缀格式、白名单、clientMode 出口）与落库； - 依赖注入走 LangGraph 原生机制：``context_schema=GraphDeps``， @@ -18,7 +20,8 @@ 图结构:: START → normalize_input → route ─┬→ welcome（首轮固定问候）→ END - ├→ collect ───────────→ END + ├→ collect ──────────┐ + │ extract_accident_fields ─┴→ merge_collect → END ├→ single_photo ──────→ END ├→ dual_photo ────────→ END ├→ single_confirm ────→ END @@ -40,7 +43,7 @@ from langgraph.checkpoint.memory import InMemorySaver from langgraph.graph import END, START, StateGraph from langgraph.runtime import Runtime -from .collect import run_collect +from .collect import collect_reply, collect_extraction, collect_time, merge_collect from .confirm import run_dual_confirm, run_single_confirm from .events import parse_client_event from .llm import LLMService, workflow_fake_llm @@ -74,6 +77,10 @@ def normalize_input(state: AgentState) -> dict: "input_text": cleaned, "client_event": event.value if event else "", "turn_patch": {}, + "extracted_accident_fields": {}, + "turn_observed_at": ( + collect_time({}) if route_by_stage(state) == "collect" else observation_time() + ).isoformat(), } @@ -135,7 +142,16 @@ def welcome(state: AgentState) -> dict: async def collect_node(state: AgentState, runtime: Runtime[GraphDeps]) -> dict: """事故信息收集节点（工作流「事故信息收集」chatNode）。""" - return await run_collect(state, runtime.context.llm) + return await collect_reply(state, runtime.context.llm) + + +async def extract_accident_fields(state: AgentState, runtime: Runtime[GraphDeps]) -> dict: + return await collect_extraction(state, runtime.context.llm) + + +def route_turn(state: AgentState) -> str | list[str]: + target = route_by_stage(state) + return ["collect", "extract_accident_fields"] if target == "collect" else target async def single_photo_node(state: AgentState, runtime: Runtime[GraphDeps]) -> dict: @@ -163,7 +179,7 @@ def build_graph(checkpointer=None): 参数: checkpointer: 自定义 checkpointer；缺省 ``InMemorySaver``， - 作业 7 替换为 ``AsyncPostgresSaver``。 + 持久化模式由运行时注入 SQLite/PostgreSQL saver。 LLM 依赖不在此处绑定——调用时通过 ``config={"configurable": {"thread_id": <会话>}}`` 区分会话， @@ -173,6 +189,8 @@ def build_graph(checkpointer=None): builder.add_node("normalize_input", normalize_input) builder.add_node("welcome", welcome) builder.add_node("collect", collect_node) + builder.add_node("extract_accident_fields", extract_accident_fields) + builder.add_node("merge_collect", merge_collect) builder.add_node("single_photo", single_photo_node) builder.add_node("dual_photo", dual_photo_node) builder.add_node("single_confirm", single_confirm_node) @@ -181,10 +199,11 @@ def build_graph(checkpointer=None): builder.add_edge(START, "normalize_input") builder.add_conditional_edges( "normalize_input", - route_by_stage, + route_turn, { "welcome": "welcome", "collect": "collect", + "extract_accident_fields": "extract_accident_fields", "single_photo": "single_photo", "dual_photo": "dual_photo", "single_confirm": "single_confirm", @@ -192,8 +211,10 @@ def build_graph(checkpointer=None): END: END, }, ) + builder.add_edge(["collect", "extract_accident_fields"], "merge_collect") + builder.add_edge("merge_collect", END) for name in ( - "welcome", "collect", "single_photo", "dual_photo", + "welcome", "single_photo", "dual_photo", "single_confirm", "dual_confirm", ): builder.add_edge(name, END) diff --git a/src/agent/llm.py b/src/agent/llm.py index 22daf80..7567635 100644 --- a/src/agent/llm.py +++ b/src/agent/llm.py @@ -201,6 +201,8 @@ class OpenAILLMService: response = await client.ainvoke( messages, config={"run_name": "llm.complete", "metadata": {"llm_attempt": attempt + 1}}, + # astream_events 的 callback 不能把抽取调用隐式改成流式。 + stream=False, ) content = self._text(response) if not content: diff --git a/src/agent/reply_stream.py b/src/agent/reply_stream.py new file mode 100644 index 0000000..cac942d --- /dev/null +++ b/src/agent/reply_stream.py @@ -0,0 +1,76 @@ +"""LangChain 原生 Runnable 文本流：在输出前验证阶段前缀。""" + +from contextlib import aclosing + +from langchain_core.runnables import RunnableGenerator + +from .collect import STATE_PREFIX_RE, parse_state_reply + +REPLY_TOKENS_RUN = "validated_reply_tokens" +REPLY_RUN = "validated_reply" + + +class ValidatedStreamingLLM: + """为现有 LLMService 添加受控流；抽取调用仍走 complete。 + + 流事件由 RunnableGenerator 原生生成，不自行创建任务或事件队列。 + 每个请求独立实例，避免并行请求共享前缀/重试状态。 + """ + + def __init__(self, inner): + self._inner = inner + self.response_started = False + + async def complete(self, system_prompt, user_message, *, history=None): + return await self._inner.complete(system_prompt, user_message, history=history) + + async def complete_validated(self, system_prompt, user_message, *, history=None, stage_validator): + raw = "" + + async def generate(inputs): + nonlocal raw + # RunnableGenerator 接收上游输入流；每次调用恰好一个 payload。 + async for _ in inputs: + stream_tokens = getattr(self._inner, "stream_tokens", None) + if stream_tokens is None: + raw = await self.complete(system_prompt, user_message, history=history) + try: + code, _ = parse_state_reply(raw) + except ValueError: + return + if stage_validator(code): + self.response_started = True + yield raw + return + + chunks = [] + prefix = "" + accepted = False + try: + async with aclosing(stream_tokens(system_prompt, user_message, history=history)) as source: + async for token in source: + chunks.append(token) + if not accepted: + prefix += token + match = STATE_PREFIX_RE.match(prefix) + if match is None: + if len(prefix) > 1024 or "" in prefix: + return + continue + if not stage_validator(match.group(1)): + return + accepted = True + self.response_started = True + yield prefix + else: + yield token + finally: + raw = "".join(chunks) + + runnable = RunnableGenerator(generate).with_config( + run_name=REPLY_TOKENS_RUN, tags=["reply_tokens"], + ) + async with aclosing(runnable.astream(None)) as source: + async for _ in source: + pass # 原生 astream_events 消费这些 chain-stream 事件。 + return raw diff --git a/src/agent/state.py b/src/agent/state.py index a95c41c..02106d5 100644 --- a/src/agent/state.py +++ b/src/agent/state.py @@ -40,6 +40,8 @@ class AgentState(TypedDict): clarify_count: int # 澄清计数（阈值按阶段配置，作业 4-9） no_reply_count: int # 连续无回复计数（两次 → 0004） turn_patch: dict # 本轮输出字段差量；每轮开始清空（作业 3-7） + turn_observed_at: str # 并行分支共用的服务器参考时间 + extracted_accident_fields: dict # 抽取分支临时通道；汇合后清空 state_version: int # 业务版本（作业 7 并发冲突检测） # ---- 对话上下文（作业 4/5） ---- diff --git a/src/api/endpoints.py b/src/api/endpoints.py index c0cc50e..f624e69 100644 --- a/src/api/endpoints.py +++ b/src/api/endpoints.py @@ -11,6 +11,7 @@ from ..backends.chat import ( ChatBackendRateLimitError, ChatInput, FormUpdate, + ReplyFinished, TextDelta, ) from ..core.fastgpt_client import ( @@ -309,6 +310,17 @@ async def chat( f"ttfb_ms={(time.perf_counter() - stream_started_at) * 1000:.1f}" ) + if isinstance(event, ReplyFinished): + # 回复已独立通过校验；抽取/提交可能仍在进行。 + # 最后一句立即刷新，对外不新增事件，done 仍等待提交。 + if state_code_found: + text = clean_later_state_tags("", final=True) + for text_event in build_text_delta_events(text): + yield text_event + for text_event in flush_text_chunker_events(): + yield text_event + continue + if isinstance(event, FormUpdate) and not module_form_sent: if event.data: form_update_payload = event.data diff --git a/src/backends/chat.py b/src/backends/chat.py index a93a27c..715de2c 100644 --- a/src/backends/chat.py +++ b/src/backends/chat.py @@ -40,7 +40,12 @@ class FormUpdate: data: Any -ChatStreamEvent = TextDelta | FormUpdate +@dataclass(frozen=True) +class ReplyFinished: + """完整回复已经校验；用于刷新文本缓冲，不表示业务提交成功。""" + + +ChatStreamEvent = TextDelta | FormUpdate | ReplyFinished @dataclass(frozen=True) diff --git a/src/backends/langgraph.py b/src/backends/langgraph.py index 22c3868..5a29067 100644 --- a/src/backends/langgraph.py +++ b/src/backends/langgraph.py @@ -1,27 +1,19 @@ -"""LangGraph implementation of the backend-neutral chat contract（作业 2）。 +"""LangGraph 后端：原生 astream_events 驱动增量回复与结束信号。 -实现 ``ChatBackend`` Protocol，把图执行结果翻译为公共契约： -- ``complete()`` 返回 ``ChatResult``，``content`` 保留 - ``XXXX`` 前缀（API 层负责解析）； -- ``stream()`` 产生 ``TextDelta`` / ``FormUpdate``。 - 作业 8 起为真实增量：图内对话调用的模型 token 逐条透出， - 字段提取等辅助调用经标记旁路（见 _TokenPassThroughLLM）。 - -通过 ``build_graph()`` 编译的图构造；会话用 -``configurable.thread_id = chat_input.session_id`` 绑定。 - -对应测试：test/backends/test_langgraph_backend.py、 -test/backends/test_langgraph_streaming.py +图调度、并行分支和事件生命周期交给 LangGraph/LangChain。 +只转发经过前缀校验的回复 Runnable 文本，抽取模型的原生事件不作为正文。 +业务提交成功后才公开 FormUpdate；公开 SSE done 仍由 API 层产生。 """ -import asyncio import inspect -from contextlib import suppress -from collections.abc import AsyncIterator, Awaitable, Callable +from collections.abc import AsyncIterator +from contextlib import aclosing + +from langchain_core.runnables import RunnableLambda -from ..agent.collect import STATE_PREFIX_RE, parse_state_reply from ..agent.graph import GraphDeps, build_graph, workflow_fake_llm from ..agent.llm import LLMAuthError, LLMRateLimitError, LLMService +from ..agent.reply_stream import REPLY_RUN, REPLY_TOKENS_RUN, ValidatedStreamingLLM from .chat import ( ChatBackend, ChatBackendAuthenticationError, @@ -30,77 +22,15 @@ from .chat import ( ChatResult, ChatStreamEvent, FormUpdate, + ReplyFinished, TextDelta, ) -# 流式桥接的队列哨兵：图任务完成后入队，消费端据此收尾 -_DONE = object() - - -class _TokenPassThroughLLM: - """先验证阶段前缀；已经输出的回复不能重生成后拼接。""" - - def __init__(self, inner: LLMService, sink: Callable[[str], Awaitable[None]]): - self._inner = inner - self._sink = sink - self.response_started = False - - async def complete(self, system_prompt, user_message, *, history=None): - # 提取等未经过 validated_reply 的辅助调用始终静默。 - return await self._inner.complete(system_prompt, user_message, history=history) - - async def complete_validated( - self, system_prompt, user_message, *, history=None, stage_validator, - ) -> str: - stream_tokens = getattr(self._inner, "stream_tokens", None) - if stream_tokens is None: - raw = await self.complete(system_prompt, user_message, history=history) - try: - code, _ = parse_state_reply(raw) - except ValueError: - return raw - if stage_validator(code): - self.response_started = True - await self._sink(raw) - return raw - - chunks: list[str] = [] - prefix = "" - accepted = False - source = stream_tokens(system_prompt, user_message, history=history) - try: - async for token in source: - chunks.append(token) - if not accepted: - prefix += token - match = STATE_PREFIX_RE.match(prefix) - if match is None: - if len(prefix) > 1024 or "" in prefix: - return "".join(chunks) - continue - if not stage_validator(match.group(1)): - return "".join(chunks) - accepted = True - self.response_started = True - await self._sink(prefix) - else: - await self._sink(token) - return "".join(chunks) - finally: - close = getattr(source, "aclose", None) - if close is not None: - await close() +TURN_RUN = "accident_turn" class LangGraphBackend(ChatBackend): - """Adapt a compiled LangGraph to the neutral chat contract.""" - def __init__(self, graph=None, llm=None): - """参数: - graph: 编译后的图；缺省调用 ``build_graph()``。 - llm: LLM 服务（运行时依赖）；缺省为离线 workflow fake。 - 生产环境由 lifespan 注入 ``build_llm_service()`` 的实例。 - """ self._graph = graph if graph is not None else build_graph() self._llm = llm if llm is not None else workflow_fake_llm() @@ -109,17 +39,6 @@ class LangGraphBackend(ChatBackend): await delete_checkpoint(self._graph, session_id) async def complete(self, chat_input: ChatInput) -> ChatResult: - """非流式：执行一轮图并翻译结果。 - - TODO(作业 2): 实现该方法。要求： - - 用 ``ainvoke`` 执行图，``thread_id`` 取 - ``chat_input.session_id``； - - 初始 state 至少包含 ``session_id``、``input_text``； - - 返回 ``ChatResult(content=<响应文本（含 state 前缀）>, - status_code=None, form_update={})``； - - 本阶段不处理 ``need_form_update`` 的真实表单差量 - （作业 6 起实现）。 - """ result = await self._complete_turn(chat_input) patch = result.get("turn_patch") or {} return ChatResult( @@ -131,17 +50,7 @@ class LangGraphBackend(ChatBackend): async def _complete_turn( self, chat_input: ChatInput, llm: LLMService | None = None, ) -> dict: - """执行一轮图，把 LLM 层错误翻译为后端中立错误。 - - llm：运行时 LLM 依赖；缺省用构造注入的实例。stream 路径传入 - token 透传包装器（_TokenPassThroughLLM）。 - - 错误码等价（对照 FastGPT 后端）： - LLMAuthError -> ChatBackendAuthenticationError（API 层 401）； - LLMRateLimitError -> ChatBackendRateLimitError（API 层 429）。 - """ - graph_input: dict = dict(chat_input.injected_fields or {}) - # 恢复快照中的上一轮 input_text 不得覆盖本次请求。 + graph_input = dict(chat_input.injected_fields or {}) graph_input.update({ "session_id": chat_input.session_id, "input_text": chat_input.text, @@ -153,7 +62,6 @@ class LangGraphBackend(ChatBackend): config={"configurable": {"thread_id": chat_input.session_id}}, context=GraphDeps(llm=llm or self._llm), ) - # 即使消费者在图完成后断开，业务 repository 仍与 checkpoint 一致。 if chat_input.on_commit is not None: pending = chat_input.on_commit(result.get("turn_patch") or {}, state=result) if inspect.isawaitable(pending): @@ -165,58 +73,43 @@ class LangGraphBackend(ChatBackend): raise ChatBackendRateLimitError(str(exc)) from exc async def stream(self, chat_input: ChatInput) -> AsyncIterator[ChatStreamEvent]: - """流式（作业 8 真实增量）：模型 token 逐条透出。 + llm = ValidatedStreamingLLM(self._llm) - 结构：后台任务执行一轮图（LLM 经 _TokenPassThroughLLM 把 - 对话 token 推入队列），本生成器从队列消费并产出 - ``TextDelta``；图完成后产出 ``FormUpdate``（差量）。 - - API 层的缓冲解析（state 标签跨片段拼接、stage 事件优先） - 本就支持任意分片，无需改动； - - 生成器被关闭（客户端断开）时取消图任务，停止 LLM 调用； - - 图任务异常在消费端重新抛出（错误映射沿用 _complete_turn）。 - """ - queue: asyncio.Queue = asyncio.Queue(maxsize=32) - token_count = 0 + async def run_turn(request): + return await self._complete_turn(request, llm=llm) - async def sink(token: str) -> None: - nonlocal token_count - token_count += 1 - await queue.put(token) + turn = RunnableLambda(run_turn).with_config(run_name=TURN_RUN) + text_started = False + reply_finished = False + result = None + # 底层负责关闭时取消并等待任务，避免业务锁释放后残留图执行。 + source = turn.astream_events( + chat_input, version="v2", + include_names=[TURN_RUN, REPLY_RUN, REPLY_TOKENS_RUN], + ) + async with aclosing(source): + async for event in source: + name = event["name"] + kind = event["event"] + if name == REPLY_TOKENS_RUN and kind == "on_chain_stream": + text = event["data"]["chunk"] + if text: + text_started = True + yield TextDelta(text) + elif name == REPLY_RUN and kind == "on_chain_end": + # validated_reply 正常结束，完整回复才算通过校验。 + if not reply_finished: + reply_finished = True + yield ReplyFinished() + elif name == TURN_RUN and kind == "on_chain_end" and not event["parent_ids"]: + result = event["data"]["output"] - async def run_turn() -> dict: - try: - result = await self._complete_turn( - chat_input, llm=_TokenPassThroughLLM(self._llm, sink) - ) - except asyncio.CancelledError: - raise - except Exception as exc: - await queue.put(exc) - return {} - await queue.put(_DONE) - return result - - task = asyncio.ensure_future(run_turn()) - try: - while True: - item = await queue.get() - if item is _DONE: - break - if isinstance(item, BaseException): - raise item - yield TextDelta(item) - result = await task - if token_count == 0: - # 本回合无 LLM 调用（welcome 等固定应答轮）： - # sink 无 token 可推，回退为完整文本切片（旧行为） - text = result.get("response_text") or "" - if text: - yield TextDelta(text) - form_update = result.get("turn_patch") or {} - if form_update and chat_input.need_form_update: - yield FormUpdate(form_update) - finally: - if not task.done(): - task.cancel() - with suppress(asyncio.CancelledError): - await task + if result is None: + raise RuntimeError("Graph stream ended without a committed result") + if not text_started: + text = result.get("response_text") or "" + if text: + yield TextDelta(text) + patch = result.get("turn_patch") or {} + if patch and chat_input.need_form_update: + yield FormUpdate(patch) diff --git a/test/backends/test_native_parallel_streaming.py b/test/backends/test_native_parallel_streaming.py new file mode 100644 index 0000000..a871521 --- /dev/null +++ b/test/backends/test_native_parallel_streaming.py @@ -0,0 +1,298 @@ +"""并行节点和原生流事件的时序、提交与取消契约。""" + +import asyncio +from contextlib import aclosing +import json + +import httpx +import pytest + +from src.agent.collect import CollectReplyError +from src.agent.graph import build_graph +from src.agent.llm import OpenAILLMService +from src.api import endpoints +from src.backends.chat import ChatInput, FormUpdate, ReplyFinished, TextDelta +from src.backends.langgraph import LangGraphBackend +from src.core.config import Config +from src.schemas.models import ProcessRequest_chat +from src.services.app_service import AccidentAppService +from src.storage.runtime import StorageSettings, open_storage +from src.services.repository import InMemorySessionRepository + + +class ParallelLLM: + def __init__(self, patch=None): + self.patch = patch if patch is not None else {'has_injury': False} + self.extraction_started = asyncio.Event() + self.reply_started = asyncio.Event() + self.release_extraction = asyncio.Event() + self.extraction_cancelled = asyncio.Event() + self.reply_done = asyncio.Event() + + async def complete(self, system, user, **kwargs): + if '【字段提取】' in system: + self.extraction_started.set() + try: + await self.release_extraction.wait() + except asyncio.CancelledError: + self.extraction_cancelled.set() + raise + return json.dumps(self.patch) + self.reply_started.set() + # 若是串行实现，会卡在这里；两个原生图节点必须同时运行。 + await self.extraction_started.wait() + self.reply_done.set() + return '1002请问您还在事故现场吗？' + + async def stream_tokens(self, system, user, **kwargs): + self.reply_started.set() + await self.extraction_started.wait() + yield '1002' + yield '请问您还在事故现场吗？' + self.reply_done.set() + + +async def service_for(llm): + repository = InMemorySessionRepository() + backend = LangGraphBackend(llm=llm) + service = AccidentAppService(backend, repository) + await service.chat(ChatInput('session', 'start')) + return service, repository + + +@pytest.mark.asyncio +async def test_non_stream_nodes_fan_out_before_join_and_commit(): + llm = ParallelLLM() + service, repo = await service_for(llm) + task = asyncio.create_task(service.chat(ChatInput('session', '继续'))) + try: + await asyncio.wait_for(llm.reply_done.wait(), 3) + assert llm.extraction_started.is_set() + assert not task.done() + assert await repo.load('session') == {} + llm.release_extraction.set() + await asyncio.wait_for(task, 3) + assert (await repo.load('session'))['ywrysw'] is False + finally: + llm.release_extraction.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_reply_finished_precedes_slow_extraction_and_commit(): + llm = ParallelLLM() + service, repo = await service_for(llm) + events = [] + async with aclosing(service.stream(ChatInput('session', '继续', need_form_update=True))) as source: + while True: + event = await asyncio.wait_for(anext(source), 3) + events.append(event) + if isinstance(event, ReplyFinished): + break + assert llm.reply_done.is_set() + assert not llm.release_extraction.is_set() + assert all(not isinstance(event, FormUpdate) for event in events) + assert await repo.load('session') == {} + llm.release_extraction.set() + events.extend([event async for event in source]) + assert isinstance(events[-1], FormUpdate) + assert (await repo.load('session'))['ywrysw'] is False + assert not service._locks._entries + + +@pytest.mark.asyncio +@pytest.mark.parametrize('patch', [{}, {'has_injury': False}]) +@pytest.mark.parametrize('visible', [False, True]) +async def test_sse_flushes_final_sentence_before_extraction_with_unchanged_public_events(monkeypatch, patch, visible): + llm = ParallelLLM(patch) + service, repo = await service_for(llm) + monkeypatch.setattr(Config, 'AGENT_BACKEND', 'langgraph') + monkeypatch.setattr(endpoints, 'get_app_service', lambda: service) + response = await endpoints.chat( + ProcessRequest_chat(sessionId='session', timeStamp='probe', text='继续', + needFormUpdate=visible, useTextChunk=True), + stream=True, backend=service._backend, idempotency_key=None, + ) + received = [] + async with aclosing(response.body_iterator) as source: + while True: + event = await asyncio.wait_for(anext(source), 3) + received.append(event) + if 'event: text_delta' in event: + break + assert '请问您还在事故现场吗？' in received[-1] + assert llm.extraction_started.is_set() + assert not llm.release_extraction.is_set() + assert not any('event: done' in event for event in received) + assert await repo.load('session') == {} + llm.release_extraction.set() + received.extend([event async for event in source]) + body = ''.join(received) + assert body.count('event: text_delta') == 1 + assert body.count('event: done') == 1 + assert ('event: formUpdate' in body) is (visible and bool(patch)) + assert 'reply_finished' not in body + if not visible: + assert 'has_injury' not in body + + +@pytest.mark.asyncio +async def test_closing_after_reply_finished_cancels_extraction_before_session_unlock(): + llm = ParallelLLM() + service, repo = await service_for(llm) + async with aclosing(service.stream(ChatInput('session', '继续'))) as source: + while not isinstance(await asyncio.wait_for(anext(source), 3), ReplyFinished): + pass + assert llm.extraction_cancelled.is_set() + assert not service._locks._entries + assert await repo.load('session') == {} + + +@pytest.mark.asyncio +async def test_invalid_full_reply_never_emits_finished_or_commits_and_cancels_parallel_extraction(): + class Invalid(ParallelLLM): + async def stream_tokens(self, *args, **kwargs): + await self.extraction_started.wait() + yield '1002partial' + yield '9999' + llm = Invalid() + service, repo = await service_for(llm) + events = [] + with pytest.raises(CollectReplyError): + async for event in service.stream(ChatInput('session', '继续')): + events.append(event) + assert any(isinstance(event, TextDelta) for event in events) + assert not any(isinstance(event, ReplyFinished) for event in events) + assert llm.extraction_cancelled.is_set() + assert await repo.load('session') == {} + assert not service._locks._entries + + +@pytest.mark.asyncio +async def test_native_graph_nodes_are_visible_and_join_once(): + llm = ParallelLLM() + llm.release_extraction.set() + graph = build_graph() + from src.agent.graph import GraphDeps + names = [] + async for event in graph.astream_events( + {'session_id': 'native', 'input_text': '继续', 'stage_code': '1002'}, + config={'configurable': {'thread_id': 'native'}}, context=GraphDeps(llm=llm), version='v2', + ): + if event['event'] == 'on_chain_end': + names.append(event['name']) + assert names.count('collect') == names.count('extract_accident_fields') == names.count('merge_collect') == 1 + + +@pytest.mark.asyncio +async def test_confirmation_reply_finishes_before_its_dependent_extraction(): + class LLM: + def __init__(self): + self.release = asyncio.Event() + async def complete(self, *args, **kwargs): + await self.release.wait() + return '{}' + async def stream_tokens(self, *args, **kwargs): + yield '3002请问您的姓名是什么？' + llm = LLM() + backend = LangGraphBackend(llm=llm) + async with aclosing(backend.stream(ChatInput( + 'confirm', '是的', injected_fields={'stage_code': '3002'}, + ))) as source: + received = [] + while True: + event = await asyncio.wait_for(anext(source), 3) + received.append(event) + if isinstance(event, ReplyFinished): + break + assert ''.join(event.text for event in received if isinstance(event, TextDelta)) == '3002请问您的姓名是什么？' + assert not llm.release.is_set() + llm.release.set() + assert [event async for event in source] == [] + + +@pytest.mark.asyncio +async def test_sqlite_disconnect_after_reply_finishes_preserves_committed_snapshot(tmp_path): + llm = ParallelLLM() + settings = StorageSettings(kind='sqlite', sqlite_path=str(tmp_path/'parallel.sqlite3'), auto_migrate=True) + async with open_storage(settings) as storage: + backend = LangGraphBackend(build_graph(storage.checkpointer), llm=llm) + service = AccidentAppService(backend, storage.repository, locks=storage.locks) + await service.chat(ChatInput('session', 'start')) + async with aclosing(service.stream(ChatInput('session', '继续', request_key='interrupted'))) as source: + while not isinstance(await asyncio.wait_for(anext(source), 3), ReplyFinished): + pass + assert llm.extraction_cancelled.is_set() + assert not storage.repository.database.connection.in_transaction + assert not storage.checkpointer.conn.in_transaction + async with storage.repository.database.transaction() as conn: + row = await conn.one('SELECT version,graph_state FROM agent_sessions WHERE session_id=?', ('session',)) + assert row['version'] == 1 + assert json.loads(row['graph_state'])['turn_count'] == 1 + llm.release_extraction.set() + await service.chat(ChatInput('session', '继续', request_key='interrupted')) + assert (await storage.repository.load('session'))['ywrysw'] is False + assert not service._locks._entries + + +@pytest.mark.asyncio +async def test_commit_failure_after_early_reply_still_emits_error_without_form_update_or_done(monkeypatch): + llm = ParallelLLM() + service, repo = await service_for(llm) + monkeypatch.setattr(Config, 'AGENT_BACKEND', 'langgraph') + monkeypatch.setattr(endpoints, 'get_app_service', lambda: service) + async def fail(*args): + raise OSError('commit unavailable') + monkeypatch.setattr(repo, 'apply_update', fail) + response = await endpoints.chat( + ProcessRequest_chat(sessionId='session', timeStamp='probe', text='继续', + needFormUpdate=True, useTextChunk=True), + stream=True, backend=service._backend, idempotency_key=None, + ) + received = [] + async with aclosing(response.body_iterator) as source: + while True: + event = await asyncio.wait_for(anext(source), 3) + received.append(event) + if 'event: text_delta' in event: + break + llm.release_extraction.set() + received.extend([event async for event in source]) + body = ''.join(received) + assert 'event: text_delta' in body and 'event: error' in body + assert 'event: done' not in body and 'event: formUpdate' not in body + assert await repo.load('session') == {} + + +@pytest.mark.asyncio +async def test_real_chatopenai_native_events_do_not_leak_extraction_or_force_streaming(monkeypatch): + requests = [] + def handler(request): + payload = json.loads(request.content) + requests.append(payload) + extraction = '【字段提取】' in payload['messages'][0]['content'] + if extraction: + assert payload['stream'] is False + return httpx.Response(200, json={ + 'id': 'extract', 'object': 'chat.completion', 'created': 0, 'model': 'probe', + 'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': '{"has_injury":false}'}, 'finish_reason': 'stop'}], + }) + assert payload['stream'] is True + data = {'id': 'reply', 'object': 'chat.completion.chunk', 'created': 0, 'model': 'probe', + 'choices': [{'index': 0, 'delta': {'content': '1002请问您还在事故现场吗？'}, 'finish_reason': None}]} + return httpx.Response(200, headers={'content-type': 'text/event-stream'}, + content='data: '+json.dumps(data)+'\n\ndata: [DONE]\n\n') + monkeypatch.setattr(Config, 'LLM_HTTP_TRUST_ENV', False) + llm = OpenAILLMService('probe', 'offline', base_url='http://probe.local/v1') + llm._get_client().http_async_client._transport = httpx.MockTransport(handler) + try: + service, repo = await service_for(llm) + events = [event async for event in service.stream(ChatInput('session', '继续', need_form_update=True))] + text = ''.join(event.text for event in events if isinstance(event, TextDelta)) + assert text == '1002请问您还在事故现场吗？' + assert len(requests) == 2 + assert sum(isinstance(event, ReplyFinished) for event in events) == 1 + assert isinstance(events[-1], FormUpdate) + finally: + await llm.aclose() diff --git a/test/services/test_stream_retry_cleanup.py b/test/services/test_stream_retry_cleanup.py index 2de87a6..991ec4a 100644 --- a/test/services/test_stream_retry_cleanup.py +++ b/test/services/test_stream_retry_cleanup.py @@ -37,7 +37,7 @@ async def test_valid_prefix_reaches_client_before_model_finishes(): assert not finished.is_set() release.set() remaining = [e async for e in events] - assert ''.join(e.text for e in remaining) == 'tail' + assert ''.join(e.text for e in remaining if isinstance(e, TextDelta)) == 'tail' assert finished.is_set()