""" 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) }