- git mv 保留历史:ui/(Vue3+Vite)、sdk/(BFF 中间层)
- 路径语义修正:db_path/figures_cache 改相对 sdk/(BFF cwd),
历史库 sdk/backend/data 迁移至 sdk/data/(会话/图片缓存无损)
- 引用修复 ~15 处:config.py、deploy/*.sh、vitest/playwright 配置、
单测导入路径、docs、README、gitignore、.env.example
- 骨架:agent/{qa,methodology,legacy-rag,vendor}、knowledge/manifest.yaml、
tooling/{compile,kb,crypto,assemble,audit,tools,scripts} + README
- 门禁验证:sdk pytest 125 通过、ui vitest 158 通过、
生产 BFF 脚本重启四上游全绿 + 静态托管 + 会话库无损
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
701 lines
30 KiB
Python
701 lines
30 KiB
Python
"""services/sse_relay.py —— 上游 QA SSE → 前端契约帧中继(A2 归属)。
|
||
|
||
上游:127.0.0.1:7920 POST /api/v1/qa-v1-1-73/chat/stream
|
||
(真实帧序列见 DISCOVERY.md §4.2:`event: <name>` + `data: <JSON>` 两行一组,
|
||
data 可能多行,期间有 `: keepalive` 注释帧,终端为 `event: complete` + 完整 JSON。)
|
||
|
||
下游(前端契约帧,API.md §4):
|
||
event: progress data: {stage, phase, label, details, at} ← phase 由 BFF 计算
|
||
event: complete data: {<同 /chat 的归一化响应>}
|
||
event: raw data: {<无法分类的上游 JSON>} ← 绝不丢数据
|
||
event: error data: {code, message, retryable}
|
||
event: timeout data: {message}
|
||
: ping (15s 保活注释帧)
|
||
|
||
规则:
|
||
- 上游 HTTP 头阶段失败且 detail.error_code=="invalid_llm_citation" → 自动重试一次,
|
||
仍失败 → `event: error`(INVALID_CITATION)。
|
||
- 看门狗:15s 无输出发 `: ping`;240s 无任何上游事件 → `event: timeout` 并关闭。
|
||
- 客户端断开(StreamingResponse 取消 → CancelledError 传播)→ `async with` 退出,
|
||
上游流随之关闭。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import time
|
||
from collections import deque
|
||
from datetime import datetime, timezone
|
||
from typing import Any, AsyncIterator, Callable, Optional
|
||
|
||
import anyio
|
||
import httpx
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 常量(TODO(A2→A1 对齐点):A1 的 app.config.settings 就绪后,
|
||
# 上游地址与「超时阶梯」改由 settings 提供,此处仅作兜底默认值)
|
||
# ---------------------------------------------------------------------------
|
||
QA_BASE_URL = "http://127.0.0.1:7920"
|
||
QA_UPSTREAM_PREFIX = "/api/v1/qa-v1-1-73"
|
||
CONNECT_TIMEOUT = 10.0
|
||
STREAM_IDLE_TIMEOUT = 1200.0 # 2026-08-28:长文(5000字)时代上游图检索静默期可超 240s(高负载下更久),
|
||
# 原 240s 会把正常流误判超时掐断;1200s 覆盖最坏情况(图检索 ~10 分钟)。
|
||
STREAM_KEEPALIVE_INTERVAL = 15.0
|
||
SEARCH_READ_TIMEOUT = 180.0 # 上游 /search 读超时(实测约 92s,见 DISCOVERY.md)
|
||
INVALID_CITATION_ERROR_CODE = "invalid_llm_citation"
|
||
|
||
# API.md §4:SSE 响应头
|
||
sse_headers = {
|
||
"Cache-Control": "no-cache",
|
||
"X-Accel-Buffering": "no",
|
||
}
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 阶段 → 相位映射(API.md §4 表;与 i18n/labels.js 同源,PM 归属)。
|
||
# 未知 stage → phase "retrieval",label = stage 原文。
|
||
# ---------------------------------------------------------------------------
|
||
STAGE_PHASE: dict[str, str] = {
|
||
"understanding_started": "planning",
|
||
"entities_identified": "planning",
|
||
"related_search_started": "planning",
|
||
"concepts_found": "retrieval",
|
||
"subquestion_retrieved": "retrieval",
|
||
"agent_review_started": "retrieval",
|
||
"subquestion_retry_started": "retrieval",
|
||
"evidence_accepted": "evidence",
|
||
"evidence_review_completed": "evidence",
|
||
"answer_generation_started": "generation",
|
||
"completed": "done",
|
||
}
|
||
|
||
STAGE_LABEL: dict[str, str] = {
|
||
"understanding_started": "开始理解问题",
|
||
"entities_identified": "已识别实体",
|
||
"related_search_started": "开始关联检索",
|
||
"concepts_found": "发现相关概念",
|
||
"subquestion_retrieved": "子问题召回完成",
|
||
"agent_review_started": "开始智能体审核",
|
||
"subquestion_retry_started": "子问题重试中",
|
||
"evidence_accepted": "证据已采纳",
|
||
"evidence_review_completed": "证据审核完成",
|
||
"answer_generation_started": "开始生成答案",
|
||
"completed": "检索完成",
|
||
}
|
||
|
||
UNKNOWN_STAGE_PHASE = "retrieval"
|
||
|
||
|
||
def phase_for(stage: str) -> tuple[str, str]:
|
||
"""stage → (phase, label);未知 stage 回落 retrieval / stage 原文。"""
|
||
return (
|
||
STAGE_PHASE.get(stage, UNKNOWN_STAGE_PHASE),
|
||
STAGE_LABEL.get(stage, stage),
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 错误文案与载荷(TODO(A2→A1 对齐点):HTTP JSON 信封换用 app.upstream.envelope())
|
||
# ---------------------------------------------------------------------------
|
||
ERROR_MESSAGES: dict[str, tuple[str, bool]] = {
|
||
"UPSTREAM_DOWN": ("知识问答服务当前不可用,请稍后重试", True),
|
||
"UPSTREAM_ERROR": ("知识问答服务返回错误", False),
|
||
"TIMEOUT": ("检索超时", False),
|
||
"INVALID_CITATION": ("引用校验失败,请重试", False),
|
||
"VALIDATION": ("请求参数不合法", False),
|
||
"INTERNAL": ("内部错误", False),
|
||
}
|
||
|
||
RETRYABLE_ERRORS = {"TIMEOUT", "INVALID_CITATION", "UPSTREAM_ERROR"}
|
||
|
||
|
||
def error_payload(code: str, message: Optional[str] = None) -> dict:
|
||
"""SSE `event: error` 的 data 载荷(API.md §4:{code, message, retryable})。"""
|
||
default_message, _degraded = ERROR_MESSAGES.get(code, ("内部错误", False))
|
||
return {
|
||
"code": code,
|
||
"message": message or default_message,
|
||
"retryable": code in RETRYABLE_ERRORS,
|
||
}
|
||
|
||
|
||
def is_invalid_citation(body: Any) -> bool:
|
||
"""上游错误体 `{"detail": {"error_code": "invalid_llm_citation"}}` 判定。"""
|
||
if not isinstance(body, dict):
|
||
return False
|
||
detail = body.get("detail")
|
||
if isinstance(detail, dict):
|
||
return detail.get("error_code") == INVALID_CITATION_ERROR_CODE
|
||
return body.get("error_code") == INVALID_CITATION_ERROR_CODE
|
||
|
||
|
||
def status_to_error(status: int, detail: Any) -> tuple[str, Optional[str]]:
|
||
"""HTTP≥400 头阶段失败 → (code, message)。"""
|
||
if is_invalid_citation(detail):
|
||
return "INVALID_CITATION", None
|
||
if isinstance(detail, dict):
|
||
d = detail.get("detail")
|
||
if isinstance(d, dict):
|
||
return "UPSTREAM_ERROR", d.get("message") or d.get("error")
|
||
return "UPSTREAM_ERROR", detail.get("message") or detail.get("error")
|
||
return "UPSTREAM_ERROR", None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 上游 client 与地址(TODO(A2→A1 对齐点):A1 的 app.upstream client 单例 /
|
||
# settings 上游地址就绪后,以下两函数直接改为其代理)
|
||
# ---------------------------------------------------------------------------
|
||
_qa_client: Optional[httpx.AsyncClient] = None
|
||
|
||
|
||
def get_qa_client() -> httpx.AsyncClient:
|
||
"""QA 上游 httpx client。
|
||
|
||
注意(2026-08-25 PM 集成修复):**不复用 A1 upstream 单例**。
|
||
A1 的 get_client() 是 async 函数且默认 read=300s,而 QA 流式路径
|
||
需要 read=None + 应用层看门狗(长静默 SSE 会被 300s 读超时误杀)。
|
||
两个 client 语义不同,各自独立持有。
|
||
"""
|
||
global _qa_client
|
||
if _qa_client is None or _qa_client.is_closed:
|
||
_qa_client = httpx.AsyncClient(
|
||
timeout=httpx.Timeout(connect=CONNECT_TIMEOUT, read=None, write=30.0, pool=CONNECT_TIMEOUT),
|
||
)
|
||
return _qa_client
|
||
|
||
|
||
def qa_base_url() -> str:
|
||
"""QA 上游 base URL。A1 settings 就绪后从配置读取。"""
|
||
try:
|
||
from ..config import settings # A1 归属
|
||
|
||
for attr in ("qa_base_url", "qa_url", "upstream_qa_url", "qa_upstream_base"):
|
||
value = getattr(settings, attr, None)
|
||
if value:
|
||
return str(value).rstrip("/")
|
||
except Exception:
|
||
pass
|
||
return QA_BASE_URL
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 归一化(TODO(A2→A1 对齐点):A1 的 app.services.normalize.normalize_qa_response
|
||
# 就绪后,get_qa_normalizer 直接返回之,以下 fallback 全部删除)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def get_qa_normalizer() -> Callable[[dict], dict]:
|
||
"""QA 响应归一化函数:优先 A1 实现,未就绪时用等价本地逻辑。"""
|
||
try:
|
||
from .normalize import normalize_qa_response # A1 归属
|
||
|
||
return normalize_qa_response
|
||
except (ImportError, AttributeError):
|
||
return _fallback_normalize_qa
|
||
|
||
|
||
def _fallback_normalize_qa(payload: dict) -> dict:
|
||
"""等价本地归一化(API.md §7 顶层别名 / §8 source / §9 图谱链条双形态容错)。"""
|
||
from .normalize import pick # A1 已提供
|
||
|
||
src = payload if isinstance(payload, dict) else {}
|
||
sources = [_fallback_source(i) for i in (pick(src, "sources", "citations", default=[]) or [])]
|
||
used = [_fallback_source(i) for i in (pick(src, "used_sources", "references", default=[]) or [])]
|
||
chains, graph = _fallback_graph_chains(src)
|
||
return {
|
||
"query_id": pick(src, "query_id"),
|
||
"version": pick(src, "version"),
|
||
"mode": "chat", # 调用方(qa.py)按端点覆盖 chat|search
|
||
"answer": pick(src, "answer", "response", "message", default=""),
|
||
"source_count": pick(src, "source_count", default=len(sources)),
|
||
"sources": sources,
|
||
"used_sources": used,
|
||
"chains": chains,
|
||
"graph": graph,
|
||
"retrieval_plan": pick(src, "retrieval_plan", "agentic_plan", "query_plan"),
|
||
"retrieval_notes": pick(src, "retrieval_notes", "notes"),
|
||
"timing": pick(src, "timing", "timings"),
|
||
"citation_map": pick(src, "citation_map", "citation_validation"),
|
||
"raw": src,
|
||
}
|
||
|
||
|
||
_SOURCE_PASSTHROUGH_KEYS = (
|
||
"output_layer", "claim_scope", "cluster_id", "support_count", "section", "species",
|
||
)
|
||
_SOURCE_METADATA_KEYS = (
|
||
"vector_score", "rrf_score", "rerank_score", "relation_id", "assertion_id",
|
||
"cluster_member_count", "cluster_member_assertion_ids", "cluster_member_paper_ids",
|
||
"cluster_member_pmids", "retrieval_methods", "retrieval_queries", "kb",
|
||
)
|
||
|
||
|
||
def _fallback_source(item: Any) -> dict:
|
||
"""§8 source 别名归一化;缺失字段键必须存在(值为 null)。"""
|
||
from .normalize import pick # A1 已提供
|
||
|
||
if not isinstance(item, dict):
|
||
item = {}
|
||
doi = pick(item, "doi")
|
||
url = pick(item, "url", "doi_url")
|
||
if url is None and doi:
|
||
url = f"https://doi.org/{doi}"
|
||
title = pick(item, "title", "paper_title")
|
||
metadata: dict[str, Any] = {}
|
||
for k in _SOURCE_PASSTHROUGH_KEYS:
|
||
if k in item:
|
||
metadata[k] = item[k]
|
||
for k in _SOURCE_METADATA_KEYS:
|
||
if k in item:
|
||
metadata[k] = item[k]
|
||
if "source" in item:
|
||
# 真实 item 的 source 是断言主体实体名而非论文标题
|
||
if title is None:
|
||
title = item["source"]
|
||
else:
|
||
metadata.setdefault("source_entity", item["source"])
|
||
return {
|
||
"id": pick(item, "evidence_id", "id"),
|
||
"title": title,
|
||
"paper_id": pick(item, "paper_id", "doc_id", "pubmed_id", "id", "chunk_id"),
|
||
"chunk_id": pick(item, "chunk_id"),
|
||
"pmid": pick(item, "pmid"),
|
||
"pmcid": pick(item, "pmcid"),
|
||
"doi": doi,
|
||
"url": url,
|
||
"snippet": pick(item, "content", "evidence", "text", "snippet"),
|
||
"score": pick(item, "final_score", "score", "confidence", "rerank_score"),
|
||
"locator": pick(item, "locator"),
|
||
"section": pick(item, "section"),
|
||
"support_count": pick(item, "support_count"),
|
||
"type": pick(item, "type"),
|
||
"metadata": metadata,
|
||
}
|
||
|
||
|
||
def _fallback_graph_chains(src: dict) -> tuple[list[dict], dict]:
|
||
"""§9 图谱与链条:native_graph_chains(list) / graph_knowledge_chains(dict) 双形态容错。
|
||
|
||
返回 (chains, graph);graph = {nodes, edges, chains}。
|
||
"""
|
||
from .normalize import pick # A1 已提供
|
||
|
||
raw_chains: list = pick(src, "native_graph_chains", default=[]) or []
|
||
if isinstance(raw_chains, dict): # 容错:意外为 dict
|
||
raw_chains = raw_chains.get("chains") or []
|
||
gkc = pick(src, "graph_knowledge_chains")
|
||
if isinstance(gkc, dict):
|
||
raw_chains = raw_chains + (gkc.get("chains") or [])
|
||
elif isinstance(gkc, list):
|
||
raw_chains = raw_chains + gkc
|
||
|
||
chains: list[dict] = []
|
||
nodes: dict[str, dict] = {}
|
||
edges: list[dict] = []
|
||
seen_chain: set = set()
|
||
|
||
for idx, chain in enumerate(raw_chains):
|
||
if not isinstance(chain, dict):
|
||
continue
|
||
chain_id = pick(chain, "chain_id", default=f"chain_{idx}")
|
||
if chain_id in seen_chain: # native 与 graph_knowledge 两形态可能镜像同数据
|
||
continue
|
||
seen_chain.add(chain_id)
|
||
hops_out: list[dict] = []
|
||
for hop in (pick(chain, "hops", default=[]) or []):
|
||
if not isinstance(hop, dict):
|
||
continue
|
||
source_name = pick(hop, "from_name", "source")
|
||
target_name = pick(hop, "to_name", "target")
|
||
src_id = pick(hop, "from_id", "source_canonical_id") or source_name
|
||
tgt_id = pick(hop, "to_id", "target_canonical_id") or target_name
|
||
hops_out.append({
|
||
"source": source_name,
|
||
"target": target_name,
|
||
"relation": pick(hop, "relation"),
|
||
"source_type": pick(hop, "source_type"),
|
||
"target_type": pick(hop, "target_type"),
|
||
"paper_id": pick(hop, "paper_id"),
|
||
"pmid": pick(hop, "pmid"),
|
||
"pmcid": pick(hop, "pmcid"),
|
||
"doi": pick(hop, "doi"),
|
||
})
|
||
nodes.setdefault(src_id, {"id": src_id, "label": source_name, "type": pick(hop, "source_type")})
|
||
nodes.setdefault(tgt_id, {"id": tgt_id, "label": target_name, "type": pick(hop, "target_type")})
|
||
edges.append({
|
||
"source": src_id,
|
||
"target": tgt_id,
|
||
"relation": pick(hop, "relation"),
|
||
"label": pick(hop, "relation"),
|
||
"chain_id": chain_id,
|
||
"hop": pick(hop, "hop_index"),
|
||
"paper_id": pick(hop, "paper_id"),
|
||
"pmid": pick(hop, "pmid"),
|
||
"pmcid": pick(hop, "pmcid"),
|
||
"doi": pick(hop, "doi"),
|
||
})
|
||
chains.append({"chain_id": chain_id, "hops": hops_out})
|
||
|
||
# knowledge_graph 直出形态(上游若给出现成 nodes/edges,直接用)
|
||
kg = pick(src, "knowledge_graph")
|
||
if isinstance(kg, dict):
|
||
for n in (kg.get("nodes") or []):
|
||
if isinstance(n, dict) and n.get("id") not in nodes:
|
||
nodes[n["id"]] = n
|
||
edges = (kg.get("edges") or []) + edges
|
||
elif not edges:
|
||
edges = pick(src, "graph_edges", default=[]) or []
|
||
|
||
graph = {"nodes": list(nodes.values()), "edges": edges, "chains": chains}
|
||
return chains, graph
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# SSE 帧工具
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def utc_now_iso() -> str:
|
||
return datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z")
|
||
|
||
|
||
def sse_frame(event: str, data: Any = None) -> str:
|
||
text = data if isinstance(data, str) else json.dumps(data, ensure_ascii=False)
|
||
return f"event: {event}\ndata: {text}\n\n"
|
||
|
||
|
||
def _progress_payload(payload: dict, at: str) -> dict:
|
||
stage = payload.get("stage")
|
||
phase, label = phase_for(stage)
|
||
details = payload.get("details")
|
||
if not isinstance(details, dict):
|
||
details = {}
|
||
return {"stage": stage, "phase": phase, "label": label, "details": details, "at": at}
|
||
|
||
|
||
def _classify(event_name: Optional[str], data_text: str) -> Optional[tuple[str, Any]]:
|
||
"""一条上游事件 → (kind, payload);kind ∈ progress|complete|error|raw。
|
||
|
||
None 表示无内容可发(空帧)。无 event 名时按 payload 形态分类。
|
||
"""
|
||
if not data_text.strip():
|
||
if not event_name:
|
||
return None
|
||
return ("raw", {"payload": data_text, "event_name": event_name})
|
||
try:
|
||
payload = json.loads(data_text)
|
||
except (json.JSONDecodeError, ValueError):
|
||
# 不是 JSON —— 透传不丢数据
|
||
return ("raw", {"payload": data_text, "event_name": event_name or None})
|
||
if not isinstance(payload, dict):
|
||
return ("raw", payload)
|
||
if "stage" in payload:
|
||
return ("progress", payload)
|
||
if event_name == "error":
|
||
return ("error", _error_payload_from_upstream(payload))
|
||
if event_name == "complete" or "answer" in payload or "success" in payload:
|
||
return ("complete", payload)
|
||
return ("raw", payload)
|
||
|
||
|
||
def _error_payload_from_upstream(payload: dict) -> dict:
|
||
if is_invalid_citation(payload):
|
||
return error_payload("INVALID_CITATION")
|
||
detail = payload.get("detail")
|
||
if isinstance(detail, dict):
|
||
return error_payload("UPSTREAM_ERROR", detail.get("message") or detail.get("error"))
|
||
return error_payload("UPSTREAM_ERROR", payload.get("message") or payload.get("error"))
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 中继核心
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
async def relay_stream(
|
||
body: dict[str, Any],
|
||
*,
|
||
client: Optional[httpx.AsyncClient] = None,
|
||
base_url: Optional[str] = None,
|
||
normalizer: Optional[Callable[[dict], dict]] = None,
|
||
idle_timeout: Optional[float] = None,
|
||
keepalive_interval: Optional[float] = None,
|
||
now_fn: Callable[[], str] = utc_now_iso,
|
||
evidence_first: bool = False,
|
||
) -> AsyncIterator[str]:
|
||
"""POST 上游 /chat/stream 并把帧中继为前端契约帧(API.md §4)。
|
||
|
||
evidence_first(2026-08-26 方案 3):为真时并发打上游 /search(无 LLM,
|
||
约 92s),归一化后作为非终态 `event: evidence` 帧在 chat 的 progress 与
|
||
complete 之间先推给前端;search 失败静默降级为纯 chat 流。
|
||
|
||
客户端断开 → StreamingResponse 取消本生成器 → CancelledError 传播 →
|
||
`async with client.stream(...)` 退出,上游流与并发 search 随之关闭。
|
||
"""
|
||
client = client or get_qa_client()
|
||
normalizer = normalizer or get_qa_normalizer()
|
||
idle_timeout = STREAM_IDLE_TIMEOUT if idle_timeout is None else idle_timeout
|
||
keepalive_interval = STREAM_KEEPALIVE_INTERVAL if keepalive_interval is None else keepalive_interval
|
||
url = f"{base_url or qa_base_url()}{QA_UPSTREAM_PREFIX}/chat/stream"
|
||
timeout = httpx.Timeout(connect=CONNECT_TIMEOUT, read=None, write=30.0, pool=CONNECT_TIMEOUT)
|
||
|
||
# 并发 evidence 任务:search 归一化结果或 None(失败/禁用)
|
||
async def _fetch_evidence() -> Optional[dict]:
|
||
if not evidence_first:
|
||
return None
|
||
# 2026-08-26 性能修正:evidence 检索关闭图链(native_chain_enabled=false),
|
||
# 只走 Qdrant 向量检索——否则与主 chat 管线并发抢 Neo4j 图遍历,
|
||
# 实测 chain_retrieval 从 109s 涨到 182s。证据帧因此更快(几十秒),
|
||
# 图链留给主管线在 complete 里完整呈现。
|
||
search_body = {**body, "use_llm": False, "native_chain_enabled": False}
|
||
search_url = f"{base_url or qa_base_url()}{QA_UPSTREAM_PREFIX}/search"
|
||
try:
|
||
async with client.stream(
|
||
"POST",
|
||
search_url,
|
||
json=search_body,
|
||
timeout=httpx.Timeout(connect=CONNECT_TIMEOUT, read=SEARCH_READ_TIMEOUT, write=30.0, pool=CONNECT_TIMEOUT),
|
||
headers={"Accept": "application/json"},
|
||
) as resp:
|
||
if resp.status_code >= 400:
|
||
logger.warning("evidence search 上游返回 %d,降级为纯 chat 流", resp.status_code)
|
||
return None
|
||
raw = await resp.aread()
|
||
payload = json.loads(raw.decode("utf-8", errors="replace"))
|
||
normalized = normalizer(payload if isinstance(payload, dict) else {})
|
||
normalized["mode"] = "search"
|
||
normalized["answer"] = ""
|
||
return normalized
|
||
except Exception as exc:
|
||
logger.warning("evidence search 失败,降级为纯 chat 流: %s", exc)
|
||
return None
|
||
|
||
evidence_task = asyncio.create_task(_fetch_evidence())
|
||
evidence_sent = False
|
||
try:
|
||
for attempt in (1, 2):
|
||
try:
|
||
async with client.stream(
|
||
"POST", url, json=body, timeout=timeout, headers={"Accept": "text/event-stream"}
|
||
) as resp:
|
||
if resp.status_code >= 400:
|
||
detail = await _safe_aread(resp)
|
||
if attempt == 1 and is_invalid_citation(detail):
|
||
logger.warning(
|
||
"上游 /chat/stream 返回 %d invalid_llm_citation,自动重试(第 1 次)",
|
||
resp.status_code,
|
||
)
|
||
continue
|
||
code, message = status_to_error(resp.status_code, detail)
|
||
yield sse_frame("error", error_payload(code, message))
|
||
return
|
||
async for frame_text in _pump_events(
|
||
resp,
|
||
normalizer,
|
||
idle_timeout,
|
||
keepalive_interval,
|
||
now_fn,
|
||
):
|
||
# evidence 就绪且未发出 → 在下一帧之前先推(progress 与 complete 之间)
|
||
if not evidence_sent and evidence_task.done():
|
||
evidence_sent = True
|
||
try:
|
||
evidence_result = evidence_task.result()
|
||
except Exception:
|
||
evidence_result = None
|
||
if evidence_result is not None:
|
||
yield sse_frame("evidence", evidence_result)
|
||
yield frame_text
|
||
# 流已终态关闭(正常路径由 _pump_events 在终态帧后 return)
|
||
return
|
||
except httpx.ReadTimeout:
|
||
yield sse_frame("error", error_payload("TIMEOUT"))
|
||
return
|
||
except (httpx.ConnectError, httpx.ConnectTimeout, httpx.PoolTimeout) as exc:
|
||
logger.warning("上游 /chat/stream 连接失败: %s", exc)
|
||
yield sse_frame("error", error_payload("UPSTREAM_DOWN"))
|
||
return
|
||
except httpx.HTTPError as exc:
|
||
logger.warning("上游 /chat/stream 请求失败: %s", exc)
|
||
yield sse_frame("error", error_payload("UPSTREAM_ERROR"))
|
||
return
|
||
# 理论上不可达(attempt 循环内必然 return);防御性兜底
|
||
yield sse_frame("error", error_payload("UPSTREAM_ERROR"))
|
||
finally:
|
||
if not evidence_task.done():
|
||
evidence_task.cancel()
|
||
|
||
|
||
async def _inject_figures_async(normalized: dict) -> dict:
|
||
"""图文并茂(2026-08-28):异步预取答案引用的论文图片,注入 markdown 图注块。
|
||
|
||
流程:从 normalized 的 sources 提取 PMCID 清单 → 逐个 resolve_figures(本地优先
|
||
PMC 兜底,缓存 1h)→ 组装 figures_map(含 pmid 供兜底路径)→ 调
|
||
normalize.inject_figure_markdown 纯同步注入。无图/失败静默降级。
|
||
"""
|
||
from .figures import resolve_figures
|
||
from .normalize import inject_figure_markdown
|
||
|
||
answer = normalized.get("answer") or ""
|
||
sources = normalized.get("sources") or []
|
||
if not answer or not sources:
|
||
return normalized
|
||
|
||
# 收集答案里引用的 PMCID([[FIG:]] 标记 + sources 的 pmcid)
|
||
import re
|
||
|
||
marker_pmcids = re.findall(r"\[\[FIG:([A-Za-z0-9]+)\]\]", answer)
|
||
figures_map: dict = {}
|
||
# 来源里 pmcid 优先(有 pmid 供兜底)
|
||
for s in sources:
|
||
pmcid = (s.get("pmcid") or "").strip().upper()
|
||
pmid = s.get("pmid")
|
||
if not pmcid:
|
||
continue
|
||
figures_map.setdefault(pmcid, {"pmid": pmid, "figures": [], "license": ""})
|
||
# 标记里的 PMCID 也加入(即使 sources 没有)
|
||
for pmcid in marker_pmcids:
|
||
figures_map.setdefault(pmcid.upper(), {"pmid": None, "figures": [], "license": ""})
|
||
|
||
# 预取图片(每个 PMCID 最多 2 张够注入用)
|
||
for pmcid, data in list(figures_map.items()):
|
||
try:
|
||
resolved = await resolve_figures(pmcid, limit=2)
|
||
data["figures"] = resolved.get("figures") or []
|
||
data["license"] = resolved.get("license") or ""
|
||
except Exception as exc:
|
||
logger.info("预取图片失败(%s): %s", pmcid, exc)
|
||
|
||
base_url = "" # figures.resolve 已返回代理端点绝对路径,无需拼接
|
||
normalized["answer"] = inject_figure_markdown(answer, figures_map, base_url)
|
||
return normalized
|
||
|
||
|
||
async def _safe_aread(resp: httpx.Response) -> Any:
|
||
try:
|
||
raw = await resp.aread()
|
||
if not raw:
|
||
return None
|
||
return json.loads(raw.decode("utf-8", errors="replace"))
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
async def _pump_events(
|
||
resp: httpx.Response,
|
||
normalizer: Callable[[dict], dict],
|
||
idle_timeout: float,
|
||
keepalive_interval: float,
|
||
now_fn: Callable[[], str],
|
||
) -> AsyncIterator[str]:
|
||
"""消费上游 SSE 字节流并产出契约帧。
|
||
|
||
设计要点:上游读取放在独立的 producer task(一次 `anext` 可能挂数分钟——
|
||
上游逐字节慢速输出是常态)。看门狗只在 `Event.wait()` 上施加 `fail_after`,
|
||
避免取消传播进 httpx 的字节生成器(asyncio 后端下那会破坏后续所有 await,
|
||
使 timeout/ping 帧发不出去)。
|
||
|
||
看门狗:
|
||
- 上游字节静默 `keepalive_interval` → `: ping`;
|
||
- 自上次完整上游事件起 `idle_timeout` 无事件 → `event: timeout` 并关闭。
|
||
"""
|
||
chunks = resp.aiter_raw()
|
||
queue: deque[tuple[str, bool]] = deque() # (frame_text, is_terminal)
|
||
wakeup = anyio.Event() # 持久事件:consumer 等待前先 clear,避免残留 set 造成忙循环
|
||
|
||
def enqueue(frame_text: str, terminal: bool) -> None:
|
||
queue.append((frame_text, terminal))
|
||
wakeup.set()
|
||
|
||
async def producer() -> None:
|
||
buf = b""
|
||
event_name: Optional[str] = None
|
||
data_lines: list[str] = []
|
||
|
||
async def dispatch_local() -> list[tuple[str, bool]]:
|
||
nonlocal event_name, data_lines
|
||
kind = _classify(event_name, "\n".join(data_lines))
|
||
event_name, data_lines = None, []
|
||
if kind is None:
|
||
return []
|
||
kind_name, payload = kind
|
||
if kind_name == "progress":
|
||
return [(sse_frame("progress", _progress_payload(payload, now_fn())), False)]
|
||
if kind_name == "complete":
|
||
try:
|
||
normalized = normalizer(payload)
|
||
except Exception:
|
||
logger.exception("归一化上游 complete 帧失败")
|
||
return [(sse_frame("error", error_payload("INTERNAL")), True)]
|
||
# 图文并茂(2026-08-28):异步预取图片并注入 markdown 图注块。
|
||
# 注入失败静默降级——图片不影响答案主链。
|
||
try:
|
||
normalized = await _inject_figures_async(normalized)
|
||
except Exception as exc:
|
||
logger.warning("图片注入失败(降级为纯文本答案): %s", exc)
|
||
return [(sse_frame("complete", normalized), True)]
|
||
if kind_name == "error":
|
||
return [(sse_frame("error", payload), True)]
|
||
return [(sse_frame("raw", payload), False)]
|
||
|
||
try:
|
||
async for chunk in chunks:
|
||
buf += chunk
|
||
while b"\n" in buf:
|
||
raw_line, buf = buf.split(b"\n", 1)
|
||
line = raw_line.rstrip(b"\r").decode("utf-8", errors="replace")
|
||
if line == "":
|
||
for frame_text, terminal in await dispatch_local():
|
||
enqueue(frame_text, terminal)
|
||
if queue and any(t for _, t in queue):
|
||
return # 终态已入队,停止读取,连接由 aclose 关闭
|
||
elif line.startswith(":"):
|
||
continue # 上游注释帧(keepalive)忽略;保活由本中继自己的 : ping 承担
|
||
elif line.startswith("event:"):
|
||
event_name = line[len("event:"):].strip() or None
|
||
elif line.startswith("data:"):
|
||
data_lines.append(line[len("data:"):].strip())
|
||
# 其他 SSE 字段(id:/retry:)忽略
|
||
# 上游流在终态帧之前关闭:残余帧照发(绝不丢数据),再报错
|
||
for frame_text, terminal in await dispatch_local():
|
||
enqueue(frame_text, terminal)
|
||
enqueue(sse_frame("error", error_payload("UPSTREAM_ERROR", "知识问答服务连接意外关闭")), True)
|
||
except httpx.ReadTimeout:
|
||
enqueue(sse_frame("error", error_payload("TIMEOUT")), True)
|
||
except (httpx.ConnectError, httpx.ConnectTimeout):
|
||
enqueue(sse_frame("error", error_payload("UPSTREAM_DOWN")), True)
|
||
except httpx.HTTPError as exc:
|
||
logger.warning("上游 SSE 流中断: %s", exc)
|
||
enqueue(sse_frame("error", error_payload("UPSTREAM_ERROR")), True)
|
||
|
||
async with anyio.create_task_group() as tg:
|
||
tg.start_soon(producer)
|
||
last_upstream_event = time.monotonic()
|
||
while True:
|
||
if queue:
|
||
last_upstream_event = time.monotonic()
|
||
frame_text, terminal = queue.popleft()
|
||
yield frame_text
|
||
if terminal:
|
||
tg.cancel_scope.cancel()
|
||
return
|
||
continue
|
||
now = time.monotonic()
|
||
if now - last_upstream_event >= idle_timeout:
|
||
yield sse_frame("timeout", {"message": "检索超时"})
|
||
tg.cancel_scope.cancel()
|
||
return
|
||
wakeup = anyio.Event()
|
||
try:
|
||
with anyio.fail_after(min(keepalive_interval, idle_timeout - (now - last_upstream_event))):
|
||
await wakeup.wait()
|
||
except TimeoutError:
|
||
now = time.monotonic()
|
||
if now - last_upstream_event >= idle_timeout:
|
||
yield sse_frame("timeout", {"message": "检索超时"})
|
||
tg.cancel_scope.cancel()
|
||
return
|
||
yield ": ping\n\n"
|