init: RAG 知识库服务初始提交
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
This commit is contained in:
lacerate551
237
core/agentic_citation.py
Normal file
237
core/agentic_citation.py
Normal file
@@ -0,0 +1,237 @@
|
||||
"""
|
||||
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)
|
||||
}
|
||||
Reference in New Issue
Block a user