367 lines
13 KiB
Python
367 lines
13 KiB
Python
"""Build stable per-turn citation batches from LangGraph messages."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
from typing import Any
|
||
|
||
|
||
def build_reference_batches(messages: list[Any]) -> list[dict[str, Any]]:
|
||
batches: list[dict[str, Any]] = []
|
||
for index, message in enumerate(messages):
|
||
if _message_type(message) != "human":
|
||
continue
|
||
end = len(messages)
|
||
for cursor in range(index + 1, len(messages)):
|
||
if _message_type(messages[cursor]) == "human":
|
||
end = cursor
|
||
break
|
||
assistant = _find_last_assistant_answer(messages, index + 1, end)
|
||
if assistant is None:
|
||
continue
|
||
sources = _collect_sources(messages[index + 1 : end])
|
||
if sources:
|
||
message_id = _message_id(assistant) or str(index)
|
||
batches.append(
|
||
{
|
||
"id": message_id,
|
||
"assistant_message_id": message_id,
|
||
"sources": sources,
|
||
}
|
||
)
|
||
return batches
|
||
|
||
|
||
def build_latest_reference_prompt(messages: list[Any]) -> tuple[str, dict[str, Any] | None]:
|
||
human_index = _find_latest_human_index(messages)
|
||
if human_index is None:
|
||
return "", None
|
||
sources = _collect_sources(messages[human_index + 1 :])
|
||
if not sources:
|
||
return "", None
|
||
assistant = _find_last_assistant_answer(messages, human_index + 1, len(messages))
|
||
message_id = _message_id(assistant) or _message_id(messages[human_index]) or str(human_index)
|
||
batch = {
|
||
"id": message_id,
|
||
"assistant_message_id": message_id,
|
||
"sources": sources,
|
||
}
|
||
lines = [
|
||
"本轮检索结果已由系统统一合并并重新编号。最终回答必须使用下面这套编号。",
|
||
"如果你仍需要补充检索，请先继续调用工具；最终回答只能引用最新合并后的编号。",
|
||
"",
|
||
]
|
||
for source in sources:
|
||
title = source.get("title") or f"参考文献 {source.get('index')}"
|
||
meta = " · ".join(
|
||
str(item)
|
||
for item in (source.get("source"), source.get("time"))
|
||
if item
|
||
)
|
||
content = str(source.get("content") or source.get("snippet") or "").strip()
|
||
if len(content) > 520:
|
||
content = content[:517] + "..."
|
||
head = f"[{source.get('index')}] {title}"
|
||
if meta:
|
||
head += f"（{meta}）"
|
||
lines.append(head)
|
||
if content:
|
||
lines.append(content)
|
||
lines.append("")
|
||
lines.extend(
|
||
[
|
||
"引用规则：",
|
||
"1. 回答正文中的关键结论必须标注对应编号，例如 [1]、[2]。",
|
||
"2. 多个来源可写为 [1][3][5]。",
|
||
"3. 只能引用上方存在的编号，不要编造编号。",
|
||
"4. 不要使用工具原始编号，也不要重新排序编号。",
|
||
"5. 不要在正文末尾堆完整参考文献列表，前端会展示右侧参考文献。",
|
||
]
|
||
)
|
||
return "\n".join(lines).strip(), batch
|
||
|
||
|
||
def _find_latest_human_index(messages: list[Any]) -> int | None:
|
||
for index in range(len(messages) - 1, -1, -1):
|
||
if _message_type(messages[index]) == "human":
|
||
return index
|
||
return None
|
||
|
||
|
||
def build_reference_debug_summary(messages: list[Any], batches: list[dict[str, Any]] | None = None) -> dict[str, Any]:
|
||
"""Return compact counts for citation extraction logs."""
|
||
turn_summaries: list[dict[str, Any]] = []
|
||
total_tool_messages = 0
|
||
total_raw_items = 0
|
||
total_normalized_items = 0
|
||
total_kept_items = 0
|
||
for index, message in enumerate(messages):
|
||
if _message_type(message) != "human":
|
||
continue
|
||
end = len(messages)
|
||
for cursor in range(index + 1, len(messages)):
|
||
if _message_type(messages[cursor]) == "human":
|
||
end = cursor
|
||
break
|
||
tool_messages = 0
|
||
raw_items = 0
|
||
normalized_items = 0
|
||
kept_items = 0
|
||
for tool_message in messages[index + 1 : end]:
|
||
if _message_type(tool_message) != "tool":
|
||
continue
|
||
tool_messages += 1
|
||
items = _first_list(_extract_tool_payload(tool_message))
|
||
raw_items += len(items)
|
||
for raw_item in items:
|
||
normalized = _normalize_reference_item(raw_item, normalized_items + 1)
|
||
if not normalized:
|
||
continue
|
||
normalized_items += 1
|
||
kept_items += 1
|
||
total_tool_messages += tool_messages
|
||
total_raw_items += raw_items
|
||
total_normalized_items += normalized_items
|
||
total_kept_items += kept_items
|
||
turn_summaries.append(
|
||
{
|
||
"human_index": index,
|
||
"tool_messages": tool_messages,
|
||
"raw_items": raw_items,
|
||
"normalized_items": normalized_items,
|
||
"kept_items": kept_items,
|
||
}
|
||
)
|
||
batch_summaries = [
|
||
{
|
||
"id": batch.get("id"),
|
||
"assistant_message_id": batch.get("assistant_message_id"),
|
||
"source_count": len(batch.get("sources") or []),
|
||
}
|
||
for batch in (batches or build_reference_batches(messages))
|
||
]
|
||
return {
|
||
"messages": len(messages),
|
||
"turns": len(turn_summaries),
|
||
"tool_messages": total_tool_messages,
|
||
"raw_items": total_raw_items,
|
||
"normalized_items": total_normalized_items,
|
||
"kept_items": total_kept_items,
|
||
"turns_detail": turn_summaries,
|
||
"batches": batch_summaries,
|
||
}
|
||
|
||
|
||
def _collect_sources(messages: list[Any]) -> list[dict[str, Any]]:
|
||
sources: list[dict[str, Any]] = []
|
||
for message in messages:
|
||
if _message_type(message) != "tool":
|
||
continue
|
||
payload = _extract_tool_payload(message)
|
||
for raw_item in _first_list(payload):
|
||
normalized = _normalize_reference_item(raw_item, len(sources) + 1)
|
||
if not normalized:
|
||
continue
|
||
normalized["index"] = len(sources) + 1
|
||
sources.append(normalized)
|
||
return sources
|
||
|
||
|
||
def _find_last_assistant_answer(messages: list[Any], start: int, end: int) -> Any | None:
|
||
for index in range(end - 1, start - 1, -1):
|
||
candidate = messages[index]
|
||
if (
|
||
_message_type(candidate) == "ai"
|
||
and _message_id(candidate)
|
||
and _has_answer_content(candidate)
|
||
and not _has_tool_calls(candidate)
|
||
):
|
||
return candidate
|
||
return None
|
||
|
||
|
||
def _message_type(message: Any) -> str:
|
||
if isinstance(message, dict):
|
||
return str(message.get("type") or message.get("role") or "")
|
||
return str(getattr(message, "type", "") or "")
|
||
|
||
|
||
def _message_id(message: Any) -> str | None:
|
||
value = message.get("id") if isinstance(message, dict) else getattr(message, "id", None)
|
||
return str(value) if value else None
|
||
|
||
|
||
def _message_content(message: Any) -> Any:
|
||
return message.get("content") if isinstance(message, dict) else getattr(message, "content", None)
|
||
|
||
|
||
def _has_tool_calls(message: Any) -> bool:
|
||
if isinstance(message, dict):
|
||
direct = message.get("tool_calls")
|
||
chunks = message.get("tool_call_chunks")
|
||
extra = (message.get("additional_kwargs") or {}).get("tool_calls")
|
||
else:
|
||
direct = getattr(message, "tool_calls", None)
|
||
chunks = getattr(message, "tool_call_chunks", None)
|
||
extra = (getattr(message, "additional_kwargs", None) or {}).get("tool_calls")
|
||
return bool(direct or chunks or extra)
|
||
|
||
|
||
def _has_answer_content(message: Any) -> bool:
|
||
content = _message_content(message)
|
||
if isinstance(content, str):
|
||
return bool(content.strip())
|
||
if isinstance(content, list):
|
||
return any(
|
||
isinstance(part, dict)
|
||
and part.get("type") == "text"
|
||
and str(part.get("text") or "").strip()
|
||
for part in content
|
||
)
|
||
return False
|
||
|
||
|
||
def _extract_tool_payload(message: Any) -> Any:
|
||
content = _message_content(message)
|
||
if isinstance(content, dict):
|
||
return _unwrap_tool_payload(content)
|
||
if isinstance(content, list):
|
||
texts: list[str] = []
|
||
for part in content:
|
||
if not isinstance(part, dict):
|
||
continue
|
||
if isinstance(part.get("json"), dict):
|
||
return _unwrap_tool_payload(part["json"])
|
||
if part.get("type") == "json" and isinstance(part.get("json"), dict):
|
||
return _unwrap_tool_payload(part["json"])
|
||
if part.get("type") == "text" and isinstance(part.get("text"), str):
|
||
texts.append(part["text"])
|
||
unwrapped = _unwrap_tool_payload(part)
|
||
if unwrapped is not None:
|
||
return unwrapped
|
||
content = "\n".join(texts)
|
||
if isinstance(content, str):
|
||
return _parse_json_like(content)
|
||
return None
|
||
|
||
|
||
def _unwrap_tool_payload(value: dict[str, Any]) -> Any:
|
||
list_keys = ("results", "data", "items", "hits", "documents", "records", "sources")
|
||
if any(isinstance(value.get(key), list) for key in list_keys):
|
||
return value
|
||
for key in ("result", "output", "stdout", "content", "text", "payload", "response"):
|
||
nested = value.get(key)
|
||
if isinstance(nested, str):
|
||
parsed = _parse_json_like(nested)
|
||
if parsed is not None:
|
||
return parsed
|
||
if isinstance(nested, dict):
|
||
unwrapped = _unwrap_tool_payload(nested)
|
||
if unwrapped is not None:
|
||
return unwrapped
|
||
return value
|
||
|
||
|
||
def _parse_json_like(text: str) -> Any:
|
||
trimmed = text.strip()
|
||
if not trimmed:
|
||
return None
|
||
try:
|
||
return json.loads(trimmed)
|
||
except Exception:
|
||
pass
|
||
first = trimmed.find("{")
|
||
last = trimmed.rfind("}")
|
||
if 0 <= first < last:
|
||
try:
|
||
return json.loads(trimmed[first : last + 1])
|
||
except Exception:
|
||
return None
|
||
return None
|
||
|
||
|
||
def _first_list(payload: Any) -> list[Any]:
|
||
if isinstance(payload, list):
|
||
return payload
|
||
if not isinstance(payload, dict):
|
||
return []
|
||
for key in ("results", "data", "items", "hits", "documents", "records", "entries", "matches", "chunks", "sources"):
|
||
value = payload.get(key)
|
||
if isinstance(value, list):
|
||
return value
|
||
for key in ("result", "payload", "response", "output", "data"):
|
||
found = _first_list(payload.get(key))
|
||
if found:
|
||
return found
|
||
return []
|
||
|
||
|
||
def _nested_records(raw: Any) -> list[dict[str, Any]]:
|
||
if not isinstance(raw, dict):
|
||
return []
|
||
layers = [raw]
|
||
for key in ("metadata", "meta", "attrs", "attributes", "properties", "extra"):
|
||
nested = raw.get(key)
|
||
if isinstance(nested, dict):
|
||
layers.append(nested)
|
||
return layers
|
||
|
||
|
||
def _pick(layers: list[dict[str, Any]], keys: tuple[str, ...]) -> str | None:
|
||
for layer in layers:
|
||
for key in keys:
|
||
value = layer.get(key)
|
||
if isinstance(value, (str, int, float)) and str(value).strip():
|
||
return str(value).strip()
|
||
return None
|
||
|
||
|
||
def _normalize_reference_item(raw: Any, fallback_index: int) -> dict[str, Any] | None:
|
||
if not isinstance(raw, dict):
|
||
return None
|
||
layers = _nested_records(raw)
|
||
title = _pick(layers, ("title", "m_title", "name", "question", "keyword", "query", "text"))
|
||
content = _pick(layers, ("content", "page_content", "summary", "snippet", "description", "content_preview", "abstract", "body"))
|
||
url = _pick(layers, ("url", "link", "href", "source_url", "web_url"))
|
||
source = _pick(layers, ("source", "source1", "site", "provider", "from"))
|
||
time = _pick(layers, ("time", "date", "published_date", "publishedDate", "m_publish", "publish_time", "created_at"))
|
||
item_id = _pick(layers, ("recUuid", "rec_uuid", "uuid", "docId", "doc_id", "id", "recordId", "taskId"))
|
||
score_raw = _pick(layers, ("score", "relevance", "similarity"))
|
||
score = None
|
||
if score_raw is not None:
|
||
try:
|
||
score = float(score_raw)
|
||
except Exception:
|
||
score = None
|
||
if not title:
|
||
title = source or f"结果 {fallback_index}"
|
||
if not content:
|
||
content = ""
|
||
if not title and not content and not url:
|
||
return None
|
||
return {
|
||
"id": item_id or url,
|
||
"skillName": "backend-reference",
|
||
"mode": "citation",
|
||
"title": title,
|
||
"type": "result",
|
||
"snippet": content,
|
||
"content": content,
|
||
"source": source,
|
||
"time": time,
|
||
"url": url,
|
||
"score": score,
|
||
"raw": raw,
|
||
}
|
||
|
||
|
||
def _reference_key(source: dict[str, Any]) -> str:
|
||
raw = source.get("raw") if isinstance(source.get("raw"), dict) else {}
|
||
item_id = source.get("id") or source.get("url") or raw.get("recUuid") or raw.get("uuid") or raw.get("docId") or raw.get("id")
|
||
if item_id:
|
||
return str(item_id).casefold().strip()
|
||
return "|".join(
|
||
str(source.get(key) or "").casefold().strip()
|
||
for key in ("title", "source", "time", "content", "snippet")
|
||
)
|