JOJO
e6d8f194d2
- 模型输出 【cite:src_xxx】/【file:相对路径】 标记，渲染为可交互来源胶囊（悬停延迟开/点击固定/滚动收起） - 后端：modules/citations.py 注册表与落库校验（剥离幻觉 id 与不存在的文件）；web_search/extract_webpage 注册来源并在工具结果带 citation 信息；assistant 消息落库时 finalize 并挂载 message.metadata.citations - 前端：marker 输出瞬间即渲染（工具结果查表 + 文件 token 自解析），task_complete 后权威裁决与富化；弹层支持网页摘要/文件内容片段/图片预览，文件头部点击打开右侧预览面板，宿主机模式可在文件管理器中打开 - prompt：三份系统提示加入引用格式与规则（提到文件名不等于引用；图片等二进制文件同样可引用）
284 lines
10 KiB
Python
284 lines
10 KiB
Python
"""Inline Citations：引用注册表、marker 解析与校验。
|
||
|
||
数据流：
|
||
工具执行（web_search / extract_webpage）→ CitationRegistry.register_url()
|
||
→ 工具结果文本带 [src_xxx]（模型可见）
|
||
→ 模型输出 【cite:src_xxx】/【file:相对路径】 marker
|
||
→ assistant 消息落库前 finalize_message_citations() 校验 + 挂载 message.citations
|
||
→ 前端渲染 citation chip
|
||
|
||
设计要点：
|
||
- 网页来源用 cite: 前缀 + ID（src_<sha1(规范化 url) 前 10 位>），确定性、同 URL 天然去重；
|
||
- 文件来源用 file: 前缀 + 工作区相对路径（可带 #L120-148 / #p7 locator），不经过 registry；
|
||
- registry 是会话运行期的临时查找表，不落盘；历史消息靠 message.citations 自包含恢复。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import os
|
||
import re
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional, Tuple
|
||
from urllib.parse import urlparse
|
||
|
||
# marker 语法：【cite:src_xxx】/【file:AGENTS.md】/同前缀多来源逗号并列（中英文逗号均可）
|
||
CITATION_MARK_RE = re.compile(r"【(cite|file):([^】]+)】")
|
||
|
||
# locator 后缀：#L120 / #L120-148 / #L120-L148（GitHub 风格）/ #p7
|
||
_LOCATOR_RE = re.compile(r"^(?P<path>.*?)#(?:L(?P<line_start>\d+)(?:-L?(?P<line_end>\d+))?|p(?P<page>\d+))$", re.IGNORECASE)
|
||
|
||
_SRC_ID_RE = re.compile(r"^src_[0-9a-f]{10}x*$")
|
||
|
||
SNIPPET_MAX_CHARS = 300
|
||
|
||
# 文件 snippet 读取上限（前 64KB 足够覆盖 locator 行段与开头摘要）
|
||
_FILE_SNIPPET_READ_BYTES = 65536
|
||
|
||
# 图片扩展名：引用弹层渲染 <img> 而不是文本摘要
|
||
_IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".gif", ".webp", ".svg", ".bmp", ".ico", ".avif"}
|
||
|
||
|
||
def _read_file_snippet(target: Path, ref: Dict[str, Any]) -> Optional[str]:
|
||
"""读取文本文件内容片段作为引用预览：有 #L 定位时取对应行段，否则取开头。
|
||
|
||
二进制文件（含空字节）返回 None。空白压缩为单行，限长 SNIPPET_MAX_CHARS。
|
||
"""
|
||
try:
|
||
with open(target, "rb") as f:
|
||
raw = f.read(_FILE_SNIPPET_READ_BYTES)
|
||
except Exception:
|
||
return None
|
||
if b"\x00" in raw:
|
||
return None
|
||
text = raw.decode("utf-8", errors="replace")
|
||
line_start = ref.get("line_start")
|
||
if line_start:
|
||
lines = text.splitlines()
|
||
start = max(1, int(line_start))
|
||
end = min(len(lines), int(ref.get("line_end") or start))
|
||
snippet = "\n".join(lines[start - 1 : end]) if start <= len(lines) else ""
|
||
else:
|
||
snippet = text
|
||
# 保留换行（弹层 pre-line 渲染），仅压缩行内空白、去空行
|
||
snippet = "\n".join(ln.strip() for ln in snippet.splitlines() if ln.strip())
|
||
return snippet[:SNIPPET_MAX_CHARS] or None
|
||
|
||
|
||
def normalize_url(url: str) -> str:
|
||
"""URL 规范化：小写 scheme/host、去末尾斜杠。用于同来源去重。"""
|
||
url = (url or "").strip()
|
||
try:
|
||
p = urlparse(url)
|
||
scheme = (p.scheme or "https").lower()
|
||
netloc = p.netloc.lower()
|
||
path = p.path.rstrip("/")
|
||
query = f"?{p.query}" if p.query else ""
|
||
return f"{scheme}://{netloc}{path}{query}"
|
||
except Exception:
|
||
return url
|
||
|
||
|
||
def domain_of(url: str) -> str:
|
||
try:
|
||
return urlparse(url).netloc.lower().lstrip("www.")
|
||
except Exception:
|
||
return ""
|
||
|
||
|
||
def make_url_citation_id(url: str) -> str:
|
||
digest = hashlib.sha1(normalize_url(url).encode("utf-8")).hexdigest()[:10]
|
||
return f"src_{digest}"
|
||
|
||
|
||
class CitationRegistry:
|
||
"""会话级网页来源注册表（运行期内存态，不落盘）。"""
|
||
|
||
def __init__(self) -> None:
|
||
self._by_id: Dict[str, Dict[str, Any]] = {}
|
||
self._id_by_url: Dict[str, str] = {}
|
||
|
||
def register_url(
|
||
self,
|
||
*,
|
||
title: str = "",
|
||
url: str = "",
|
||
snippet: str = "",
|
||
published_date: Optional[str] = None,
|
||
source_tool: Optional[str] = None,
|
||
) -> Dict[str, Any]:
|
||
"""注册（或按规范化 URL 去重返回已有）一个网页来源 annotation。"""
|
||
if not url:
|
||
return {}
|
||
norm = normalize_url(url)
|
||
existing_id = self._id_by_url.get(norm)
|
||
if existing_id is not None:
|
||
return self._by_id[existing_id]
|
||
|
||
cid = make_url_citation_id(url)
|
||
# 哈希碰撞兜底（同一 URL 已命中则上面已返回；不同 URL 撞 hash 时加后缀）
|
||
while cid in self._by_id and normalize_url(self._by_id[cid].get("url", "")) != norm:
|
||
cid += "x"
|
||
|
||
ann: Dict[str, Any] = {
|
||
"id": cid,
|
||
"type": "url_citation",
|
||
"title": title or url,
|
||
"url": url,
|
||
"domain": domain_of(url),
|
||
"snippet": (snippet or "")[:SNIPPET_MAX_CHARS],
|
||
}
|
||
if published_date:
|
||
ann["published_date"] = published_date
|
||
if source_tool:
|
||
ann["source_tool"] = source_tool
|
||
self._by_id[cid] = ann
|
||
self._id_by_url[norm] = cid
|
||
return ann
|
||
|
||
def resolve(self, citation_id: str) -> Optional[Dict[str, Any]]:
|
||
return self._by_id.get(citation_id)
|
||
|
||
|
||
def get_registry(terminal: Any) -> CitationRegistry:
|
||
"""取 terminal 实例上的会话级 registry（懒创建）。"""
|
||
reg = getattr(terminal, "_citation_registry", None)
|
||
if reg is None:
|
||
reg = CitationRegistry()
|
||
terminal._citation_registry = reg
|
||
return reg
|
||
|
||
|
||
def parse_marker_content(content: str) -> List[str]:
|
||
"""拆分 marker 内容为引用 token 列表（支持中英文逗号分隔）。"""
|
||
parts = re.split(r"[,，]", content or "")
|
||
return [p.strip() for p in parts if p.strip()]
|
||
|
||
|
||
def extract_citation_refs(text: str) -> List[str]:
|
||
"""提取正文中全部完整 marker 的引用 token（保持出现顺序，去重）。"""
|
||
refs: List[str] = []
|
||
for m in CITATION_MARK_RE.finditer(text or ""):
|
||
for token in parse_marker_content(m.group(2)):
|
||
if token not in refs:
|
||
refs.append(token)
|
||
return refs
|
||
|
||
|
||
def _parse_file_ref(token: str) -> Optional[Dict[str, Any]]:
|
||
"""解析文件引用 token（裸相对路径，可带 #L120-148 / #p7 locator）。非法返回 None。"""
|
||
body = token.strip()
|
||
if not body:
|
||
return None
|
||
m = _LOCATOR_RE.match(body)
|
||
if m:
|
||
ref: Dict[str, Any] = {"path": m.group("path")}
|
||
if m.group("page"):
|
||
ref["page"] = int(m.group("page"))
|
||
if m.group("line_start"):
|
||
ref["line_start"] = int(m.group("line_start"))
|
||
if m.group("line_end"):
|
||
ref["line_end"] = int(m.group("line_end"))
|
||
return ref
|
||
return {"path": body}
|
||
|
||
|
||
def _build_file_annotation(ref: Dict[str, Any], token: str, workspace_root: str) -> Optional[Dict[str, Any]]:
|
||
"""校验文件引用：路径必须落在工作区内且文件存在；富化 file_name/size。
|
||
|
||
annotation id 与 marker 内 token 完全一致，前端按 id 精确匹配。
|
||
"""
|
||
# 注意不能用 lstrip("./")：它是字符集语义，会把 .astrion 这类隐藏目录的前导点吃掉
|
||
rel_path = (ref.get("path") or "").strip()
|
||
if rel_path.startswith("./"):
|
||
rel_path = rel_path[2:]
|
||
if not rel_path or rel_path.startswith("/"):
|
||
return None
|
||
try:
|
||
root = Path(workspace_root).resolve()
|
||
target = (root / rel_path).resolve()
|
||
# 边界检查：必须在工作区内
|
||
if root != target and root not in target.parents:
|
||
return None
|
||
if not target.is_file():
|
||
return None
|
||
stat = target.stat()
|
||
except Exception:
|
||
return None
|
||
|
||
ann: Dict[str, Any] = {
|
||
"id": token,
|
||
"type": "file_citation",
|
||
"file_path": rel_path,
|
||
"file_name": target.name,
|
||
"size": stat.st_size,
|
||
}
|
||
snippet: Optional[str] = None
|
||
# 图片文件的 snippet 无意义（且 svg 这类文本型图片会把标记源码当摘要），
|
||
# 前端弹层按扩展名直接渲染 <img>
|
||
if target.suffix.lower() not in _IMAGE_EXTS:
|
||
snippet = _read_file_snippet(target, ref)
|
||
if snippet:
|
||
ann["snippet"] = snippet
|
||
if ref.get("page") is not None:
|
||
ann["page"] = ref["page"]
|
||
if ref.get("line_start") is not None:
|
||
ann["line_start"] = ref["line_start"]
|
||
if ref.get("line_end") is not None:
|
||
ann["line_end"] = ref["line_end"]
|
||
return ann
|
||
|
||
|
||
def finalize_message_citations(
|
||
text: str,
|
||
registry: Optional[CitationRegistry],
|
||
workspace_root: str,
|
||
) -> Tuple[str, List[Dict[str, Any]]]:
|
||
"""assistant 消息落库前处理：剥离无效 marker，返回 (clean_text, used_annotations)。
|
||
|
||
- src_xxx：registry 查得到才保留，否则剥离 marker；
|
||
- file: 引用：路径在工作区内且文件存在才保留；
|
||
- 一个 marker 里部分有效时，只保留有效 token 重建 marker；全部无效则整体移除。
|
||
"""
|
||
if not text or ("【cite:" not in text and "【file:" not in text):
|
||
return text, []
|
||
|
||
used: List[Dict[str, Any]] = []
|
||
seen_ids = set()
|
||
|
||
def _keep(ann: Optional[Dict[str, Any]]) -> Optional[str]:
|
||
if not ann:
|
||
return None
|
||
ann_id = ann.get("id")
|
||
if ann_id and ann_id not in seen_ids:
|
||
seen_ids.add(ann_id)
|
||
used.append(ann)
|
||
return ann_id
|
||
|
||
def _replace(m: re.Match) -> str:
|
||
kind = m.group(1)
|
||
tokens = parse_marker_content(m.group(2))
|
||
kept: List[str] = []
|
||
for token in tokens:
|
||
if kind == "cite":
|
||
# 网页来源：src_xxx 且 registry 查得到才保留
|
||
if not _SRC_ID_RE.match(token):
|
||
continue
|
||
ann = registry.resolve(token) if registry else None
|
||
if _keep(ann):
|
||
kept.append(token)
|
||
else:
|
||
# 文件来源：路径在工作区内且存在才保留
|
||
ref = _parse_file_ref(token)
|
||
if not ref:
|
||
continue
|
||
ann = _build_file_annotation(ref, token, workspace_root)
|
||
if _keep(ann):
|
||
kept.append(token)
|
||
if not kept:
|
||
return ""
|
||
return "【" + kind + ":" + ",".join(kept) + "】"
|
||
|
||
clean = CITATION_MARK_RE.sub(_replace, text)
|
||
return clean, used
|