100d1a06eb
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
238 lines
9.0 KiB
Python
238 lines
9.0 KiB
Python
"""
|
||
Agentic RAG - 引用处理 Mixin
|
||
|
||
包含来源提取、引用构建、引用附加等方法
|
||
"""
|
||
|
||
import json
|
||
import logging
|
||
|
||
from .agentic_base import logger, SOURCE_KB
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class CitationMixin:
|
||
"""引用处理方法"""
|
||
|
||
def _extract_sources(self, contexts: list) -> list:
|
||
"""提取来源列表，返回结构化定位信息"""
|
||
source_map = {}
|
||
|
||
for c in contexts:
|
||
meta = c.get('meta', {})
|
||
source_type = c.get('source_type', '未知')
|
||
|
||
if source_type == self.SOURCE_KB:
|
||
source_key = meta.get('source', '未知')
|
||
page = meta.get('page')
|
||
page_end = meta.get('page_end', page)
|
||
section = meta.get('section', '')
|
||
doc_type = meta.get('doc_type', 'other')
|
||
preview = meta.get('preview', '')
|
||
section_chunk_id = meta.get('section_chunk_id')
|
||
else:
|
||
source_key = meta.get('title', meta.get('source', '未知'))
|
||
page = None
|
||
page_end = None
|
||
section = ''
|
||
doc_type = 'other'
|
||
preview = ''
|
||
section_chunk_id = None
|
||
|
||
if source_key not in source_map:
|
||
source_map[source_key] = {
|
||
"source": source_key,
|
||
"type": source_type,
|
||
"doc_type": doc_type,
|
||
"count": 0,
|
||
"pages": [],
|
||
"sections": set(),
|
||
"previews": [],
|
||
"section_chunk_ids": set()
|
||
}
|
||
|
||
source_map[source_key]["count"] += 1
|
||
|
||
if page:
|
||
page_range = (page, page_end if page_end else page)
|
||
if page_range not in source_map[source_key]["pages"]:
|
||
source_map[source_key]["pages"].append(page_range)
|
||
|
||
if section:
|
||
source_map[source_key]["sections"].add(section)
|
||
|
||
if preview and len(source_map[source_key]["previews"]) < 3:
|
||
if preview not in source_map[source_key]["previews"]:
|
||
source_map[source_key]["previews"].append(preview)
|
||
|
||
if section_chunk_id:
|
||
source_map[source_key]["section_chunk_ids"].add(section_chunk_id)
|
||
|
||
sources = []
|
||
for key, info in source_map.items():
|
||
source_str = info["source"]
|
||
doc_type = info.get("doc_type", "other")
|
||
location_parts = []
|
||
|
||
if doc_type == 'pdf':
|
||
if info["pages"]:
|
||
valid_pages = [(s, e) for s, e in info["pages"] if s > 1 or e > 1]
|
||
if valid_pages or not info["sections"]:
|
||
page_strs = []
|
||
for start, end in sorted(info["pages"], key=lambda x: x[0]):
|
||
if start == end:
|
||
page_strs.append(f"第{start}页")
|
||
else:
|
||
page_strs.append(f"第{start}-{end}页")
|
||
location_parts.append(", ".join(page_strs))
|
||
|
||
if info["sections"]:
|
||
sections_list = sorted(info["sections"])[:3]
|
||
sections_str = "、".join(sections_list)
|
||
if len(info["sections"]) > 3:
|
||
sections_str += f"等{len(info['sections'])}个章节"
|
||
location_parts.append(sections_str)
|
||
|
||
elif doc_type == 'word':
|
||
if info["sections"]:
|
||
sections_list = sorted(info["sections"])[:3]
|
||
sections_str = "、".join(sections_list)
|
||
if len(info["sections"]) > 3:
|
||
sections_str += f"等{len(info['sections'])}个章节"
|
||
location_parts.append(sections_str)
|
||
|
||
if info.get("section_chunk_ids"):
|
||
chunk_ids = sorted(info["section_chunk_ids"])[:5]
|
||
if chunk_ids:
|
||
chunk_str = f"第{chunk_ids[0]}"
|
||
if len(chunk_ids) > 1:
|
||
chunk_str = f"第{chunk_ids[0]}-{chunk_ids[-1]}段"
|
||
location_parts.append(chunk_str)
|
||
|
||
elif doc_type == 'excel':
|
||
if info["sections"]:
|
||
sections_list = sorted(info["sections"])[:3]
|
||
sections_str = "、".join(sections_list)
|
||
location_parts.append(sections_str)
|
||
|
||
else:
|
||
if info["pages"]:
|
||
valid_pages = [(s, e) for s, e in info["pages"] if s > 1 or e > 1]
|
||
if valid_pages or not info["sections"]:
|
||
page_strs = []
|
||
for start, end in sorted(info["pages"], key=lambda x: x[0]):
|
||
if start == end:
|
||
page_strs.append(f"第{start}页")
|
||
else:
|
||
page_strs.append(f"第{start}-{end}页")
|
||
location_parts.append(", ".join(page_strs))
|
||
|
||
if info["sections"]:
|
||
sections_list = sorted(info["sections"])[:3]
|
||
sections_str = "、".join(sections_list)
|
||
if len(info["sections"]) > 3:
|
||
sections_str += f"等{len(info['sections'])}个章节"
|
||
location_parts.append(sections_str)
|
||
|
||
if location_parts:
|
||
source_str = f"{source_str} ({' | '.join(location_parts)})"
|
||
|
||
sources.append({
|
||
"source": source_str,
|
||
"type": info["type"],
|
||
"count": info["count"],
|
||
"doc_type": doc_type,
|
||
"previews": info.get("previews", []),
|
||
"section_chunk_ids": sorted(info.get("section_chunk_ids", []))[:5]
|
||
})
|
||
|
||
return sources
|
||
|
||
def _build_citation(self, meta: dict) -> dict:
|
||
"""根据文档类型构建定位信息"""
|
||
# 从 chunk_id 中提取全局切片序号（格式: "filename_N"）
|
||
chunk_id_raw = meta.get('chunk_id', '')
|
||
chunk_index = None
|
||
if chunk_id_raw and '_' in str(chunk_id_raw):
|
||
try:
|
||
chunk_index = int(str(chunk_id_raw).rsplit('_', 1)[-1])
|
||
except (ValueError, IndexError):
|
||
chunk_index = meta.get('chunk_index')
|
||
else:
|
||
chunk_index = meta.get('chunk_index')
|
||
|
||
citation = {
|
||
"chunk_id": chunk_id_raw,
|
||
"chunk_index": chunk_index, # 全局切片序号，用于前端文档预览跳转
|
||
"source": meta.get('source', ''),
|
||
"collection": meta.get('_collection', ''), # 所属向量库，用于前端文档预览跳转
|
||
"doc_type": meta.get('doc_type', 'other'),
|
||
"section": meta.get('section', ''),
|
||
"preview": meta.get('preview', ''),
|
||
"content": meta.get('preview', ''), # 初始用 preview，_attach_citations 中会用完整内容覆盖
|
||
"chunk_type": meta.get('chunk_type', 'text'),
|
||
}
|
||
|
||
doc_type = meta.get('doc_type', 'other')
|
||
|
||
if doc_type == 'pdf':
|
||
bbox_raw = meta.get('bbox')
|
||
bbox = None
|
||
if bbox_raw:
|
||
try:
|
||
bbox = json.loads(bbox_raw) if isinstance(bbox_raw, str) else bbox_raw
|
||
except (json.JSONDecodeError, TypeError):
|
||
bbox = bbox_raw
|
||
|
||
citation.update({
|
||
"page": meta.get('page'),
|
||
"page_end": meta.get('page_end'),
|
||
"bbox": bbox,
|
||
"bbox_mode": meta.get('bbox_mode'),
|
||
})
|
||
elif doc_type == 'word':
|
||
citation.update({
|
||
"section_chunk_id": meta.get('section_chunk_id'),
|
||
})
|
||
elif doc_type == 'excel':
|
||
citation.update({
|
||
"page": meta.get('page'),
|
||
})
|
||
else:
|
||
bbox_raw = meta.get('bbox')
|
||
bbox = None
|
||
if bbox_raw:
|
||
try:
|
||
bbox = json.loads(bbox_raw) if isinstance(bbox_raw, str) else bbox_raw
|
||
except (json.JSONDecodeError, TypeError):
|
||
bbox = bbox_raw
|
||
|
||
citation.update({
|
||
"page": meta.get('page'),
|
||
"page_end": meta.get('page_end'),
|
||
"bbox": bbox,
|
||
"bbox_mode": meta.get('bbox_mode'),
|
||
})
|
||
|
||
return citation
|
||
|
||
def _attach_citations(self, answer: str, contexts: list) -> dict:
|
||
"""将引用信息附加到答案"""
|
||
citations = []
|
||
|
||
for c in contexts:
|
||
meta = c.get('meta', {})
|
||
full_content = c.get('doc', '')
|
||
citation = self._build_citation(meta)
|
||
# 用上下文中的完整文档内容覆盖 content 字段
|
||
if full_content:
|
||
citation['content'] = full_content[:300]
|
||
citations.append(citation)
|
||
|
||
return {
|
||
"answer": answer,
|
||
"citations": citations,
|
||
"sources": self._extract_sources(contexts)
|
||
}
|