vent_agent_2.0/runtime/source_tracker.py

"""来源标注：从工具结果里抽取、聚合、去重。

## 为什么换掉 1.5 的做法

1.5 让**模型在正文里自己写引用**，后端再用启发式规则去正文里刮。这条
路径有两个固有失败模式：模型漏写（内容有来源但没标注）和模型编造
（写出一个不存在的 url）。两者都无法在事后可靠地区分。

## 2.0 的做法

来源由**工具侧结构化产出**：工具返回 ``tools.registry.sourced(data, sources)``
构造的结构，本模块从 ToolMessage 的原始内容里把 ``sources`` 抽出来。
模型是否在正文里写引用，不影响 citations 的准确性——因为它读的是工具
的返回值，不是模型的叙述。

本轮没抽到任何来源时，发 ``model_only=True``，前端会显示「本回答未引用
外部数据源」而不是一个空卡片——如实说明「没有来源」比假装有来源重要。
"""

from __future__ import annotations

import json
from typing import Any, Iterable

from obs import get_logger
from tools.registry import get_label

logger = get_logger(__name__)

# 一条来源最多保留的摘要长度
_SNIPPET_LIMIT = 200
# 单轮最多输出多少条引用，避免工具刷屏时把前端塞爆
_MAX_CITATIONS = 50


def _classify(url: str, tool_name: str) -> tuple[str, str]:
    """给来源分类，返回 ``(type, type_cn)``。

    分类只依赖 url 形态与工具名，不看内容——这是一条纯函数式的判定，
    可以被测试覆盖，不会因为模型措辞不同而漂移。
    """
    lowered = (url or "").lower()
    if lowered.startswith(("http://", "https://")):
        if "kb" in lowered or "knowledge" in lowered:
            return "knowledge_base", "知识库"
        return "web", "网页"
    if tool_name.startswith(("search_", "query_")):
        return "knowledge_base", "知识库"
    if lowered.startswith("file://") or lowered.endswith((".pdf", ".docx", ".md")):
        return "file", "文件"
    return "tool", "工具"


def _coerce_sources(payload: Any) -> list[dict[str, Any]]:
    """从工具返回值里挖出 sources 列表。

    容忍几种常见形态：直接 dict、JSON 字符串、被包在 content 块列表里、
    或 ``{"data": {"sources": [...]}}`` 这种多套一层的情况。
    """
    if payload is None:
        return []

    if isinstance(payload, str):
        text = payload.strip()
        if not text or text[0] not in "{[":
            return []
        try:
            payload = json.loads(text)
        except (TypeError, ValueError):
            return []

    if isinstance(payload, list):
        collected: list[dict[str, Any]] = []
        for item in payload[:10]:
            collected.extend(_coerce_sources(item))
        return collected

    if not isinstance(payload, dict):
        return []

    sources = payload.get("sources")
    if isinstance(sources, list):
        return [s for s in sources if isinstance(s, dict)]

    # 多套一层：{"data": {...}} 或 {"result": {...}}
    for key in ("data", "result", "output"):
        nested = payload.get(key)
        if isinstance(nested, (dict, list, str)):
            found = _coerce_sources(nested)
            if found:
                return found
    return []


def _normalize(record: dict[str, Any], tool_name: str) -> dict[str, Any] | None:
    """把一条原始来源记录规整成引用条目；无有效标识则返回 None。"""
    title = str(record.get("title") or record.get("doc_name") or "").strip()
    url = str(record.get("url") or record.get("doc_path") or "").strip()
    snippet = str(record.get("snippet") or record.get("text") or "").strip()
    if not title and not url:
        # 既没标题也没链接的来源无法核实，丢弃比展示一条空壳有用
        return None

    kind, kind_cn = _classify(url, tool_name)
    item: dict[str, Any] = {
        "title": title or url,
        "url": url,
        "snippet": snippet[:_SNIPPET_LIMIT],
        "type": str(record.get("type") or kind),
        "type_cn": str(record.get("type_cn") or kind_cn),
        "tool": tool_name,
        "tool_cn": get_label(tool_name),
    }
    score = record.get("score")
    if isinstance(score, (int, float)):
        item["score"] = round(float(score), 4)
    return item


def _dedup_key(item: dict[str, Any]) -> tuple[str, str]:
    """去重键：优先 url，退而用标题 + 摘要前缀。"""
    url = item.get("url") or ""
    if url:
        return ("url", url)
    return ("title", f"{item.get('title', '')}|{item.get('snippet', '')[:40]}")


class SourceTracker:
    """一次请求的来源台账。"""

    def __init__(self) -> None:
        self._items: list[dict[str, Any]] = []
        self._seen: set[tuple[str, str]] = set()

    def note_tool_result(self, tool_name: str, payload: Any) -> list[dict[str, Any]]:
        """观测一次工具返回，抽取其中的来源。返回本次**新增**的条目。"""
        added: list[dict[str, Any]] = []
        for record in _coerce_sources(payload):
            item = _normalize(record, tool_name)
            if item is None:
                continue
            key = _dedup_key(item)
            if key in self._seen:
                continue
            self._seen.add(key)
            self._items.append(item)
            added.append(item)
            if len(self._items) >= _MAX_CITATIONS:
                logger.warning("引用条目达到上限 %d，后续来源不再收集", _MAX_CITATIONS)
                break
        return added

    def build_citations(self) -> list[dict[str, Any]]:
        """产出发给前端的引用列表。按相关度倒序（无分数的排在后面）。"""
        return sorted(
            self._items,
            key=lambda i: i.get("score", -1.0),
            reverse=True,
        )

    @property
    def count(self) -> int:
        return len(self._items)

    def __len__(self) -> int:
        return len(self._items)

    def is_empty(self) -> bool:
        return not self._items


def collect_from_tool_messages(
    messages: Iterable[Any], tracker: SourceTracker | None = None
) -> SourceTracker:
    """从一串消息里批量抽取来源（离线场景/补历史时用）。"""
    tracker = tracker or SourceTracker()
    for message in messages:
        name = getattr(message, "name", None) or ""
        content = getattr(message, "content", None)
        if name:
            tracker.note_tool_result(name, content)
    return tracker