- 将全部路由文件(12个)的 jsonify 响应迁移至 success_response/error_response 统一格式 - 修复 sync_routes.py error_response 参数错误(P0) - 新增异步任务系统:task_registry + task_routes - 新增状态码:TASK_NOT_FOUND(4014)、TASK_CONFLICT(4015)、REINDEX_ERROR(5040) - 修正 task_routes/exam_pkg 中语义不匹配的状态码 - 更新 curl 测试手册、后端对接规范文档 - 添加缓存性能报告和 Redis 迁移计划
2139 lines
90 KiB
Python
2139 lines
90 KiB
Python
"""
|
||
核心聊天与检索 API
|
||
|
||
本模块提供 RAG 系统的核心问答接口,包括:
|
||
- 普通聊天模式(/chat)
|
||
- 知识库问答模式(/rag,支持 SSE 流式返回)
|
||
- 混合检索接口(/search,供外部系统调用)
|
||
|
||
路由列表:
|
||
POST /chat : 普通聊天模式(JSON 响应)
|
||
POST /rag : 知识库问答模式(SSE 流式返回)
|
||
POST /search : 混合检索接口(供 Dify 调用)
|
||
|
||
架构说明:
|
||
- 会话管理由后端服务负责,RAG 服务不存储对话历史
|
||
- 权限验证由后端网关完成(通过 request.current_user 获取用户信息)
|
||
- /rag 接口已升级为 SSE 流式返回,支持实时输出
|
||
|
||
Example:
|
||
curl -X POST http://localhost:5001/rag \\
|
||
-H "Content-Type: application/json" \\
|
||
-H "Authorization: Bearer mock-token-admin" \\
|
||
-d '{"query": "公司报销制度是什么?"}'
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import queue
|
||
import threading
|
||
import time as _time
|
||
from typing import List, Dict, Any, Optional, Tuple
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
from flask import Blueprint, request, jsonify, Response, current_app
|
||
import logging
|
||
|
||
logger = logging.getLogger(__name__)
|
||
from auth.gateway import require_gateway_auth
|
||
from core.status_codes import SUCCESS, BAD_REQUEST, NOT_FOUND, INTERNAL_ERROR
|
||
from api.response_utils import success_response, error_response
|
||
|
||
from auth.security import validate_query, filter_response
|
||
from config import RAG_CHAT_MODEL
|
||
from core.llm_utils import call_llm, call_llm_stream
|
||
from core.prompt_guard import detect_injection, sanitize_user_input
|
||
|
||
|
||
def _get_vlm_cache(image_path: str) -> Optional[str]:
|
||
"""
|
||
从缓存获取图片的 VLM 描述
|
||
|
||
Args:
|
||
image_path: 图片文件名(不含路径),如 '0cd5e156ff0a.png'
|
||
|
||
Returns:
|
||
VLM 描述文本,无缓存返回 None
|
||
"""
|
||
import hashlib
|
||
|
||
try:
|
||
images_dir = Path(".data/images")
|
||
full_path = images_dir / image_path
|
||
|
||
if not full_path.exists():
|
||
return None
|
||
|
||
img_hash = hashlib.md5(full_path.read_bytes()).hexdigest()
|
||
cache_file = Path(f".data/cache/vlm/{img_hash}.txt")
|
||
|
||
if cache_file.exists():
|
||
return cache_file.read_text(encoding='utf-8')
|
||
except Exception as e:
|
||
logger.debug(f"读取 VLM 缓存失败: {e}")
|
||
|
||
return None
|
||
|
||
|
||
def _check_vlm_relevance(query: str, vlm_desc: str) -> float:
|
||
"""
|
||
使用 LLM 评估图片描述与查询的相关性
|
||
|
||
Args:
|
||
query: 用户查询
|
||
vlm_desc: 图片的 VLM 描述
|
||
|
||
Returns:
|
||
相关性分数 (0-1),0=不相关,1=高度相关
|
||
"""
|
||
if not vlm_desc or len(vlm_desc) < 20:
|
||
return 0.5 # 无有效描述,中性评分
|
||
|
||
# 使用 jieba 分词提取关键词
|
||
import jieba
|
||
stop_words = {'图', '表', '图片', '图表', '如图', '所示', '的', '是', '在', '有', '了', '和', '与', '或', '及', '等', '中', '对'}
|
||
query_keywords = [w for w in jieba.lcut(query) if len(w) >= 2 and w not in stop_words]
|
||
if not query_keywords:
|
||
return 0.5
|
||
|
||
# 检查关键词是否在 VLM 描述中出现
|
||
matches = sum(1 for kw in query_keywords if kw in vlm_desc)
|
||
if matches == 0:
|
||
return 0.0 # 完全不相关
|
||
|
||
# 按匹配比例评分
|
||
ratio = matches / len(query_keywords)
|
||
return min(ratio, 1.0)
|
||
|
||
chat_bp = Blueprint('chat', __name__)
|
||
|
||
|
||
def _safe_int(value: Any, default: int = 10**9) -> int:
|
||
"""
|
||
安全的整数转换
|
||
|
||
Args:
|
||
value: 待转换的值
|
||
default: 转换失败时的默认值
|
||
|
||
Returns:
|
||
转换后的整数值,失败返回默认值
|
||
"""
|
||
try:
|
||
if value is None or value == "":
|
||
return default
|
||
return int(value)
|
||
except (TypeError, ValueError):
|
||
return default
|
||
|
||
|
||
def _is_enum_query(query: str) -> bool:
|
||
"""
|
||
判断是否为枚举型查询
|
||
|
||
枚举型查询通常包含"哪些"、"列出"等关键词,
|
||
需要特殊的上下文排序策略。
|
||
|
||
Args:
|
||
query: 用户查询文本
|
||
|
||
Returns:
|
||
是枚举型查询返回 True,否则返回 False
|
||
"""
|
||
try:
|
||
from core.query_classifier import is_enumeration_query
|
||
return is_enumeration_query(query)
|
||
except Exception as e:
|
||
logger.debug(f"枚举查询检测失败: {e}")
|
||
markers = ("哪些", "有哪些", "列出", "严禁", "禁止", "不得", "包括", "要求", "情形", "场景")
|
||
return bool(query) and any(marker in query for marker in markers)
|
||
|
||
|
||
def _get_full_table_from_docstore(chunk_id: str) -> Optional[str]:
|
||
"""
|
||
从 DocStore 获取表格的完整 Markdown 内容
|
||
|
||
Args:
|
||
chunk_id: 切片 ID,如 '组织架构.xlsx_0' 或 'test_report.pdf_3'
|
||
|
||
Returns:
|
||
完整表格 Markdown,未找到返回 None
|
||
"""
|
||
docstore_dir = Path(".data/docstore")
|
||
if not docstore_dir.exists():
|
||
return None
|
||
|
||
# chunk_id 格式: {filename}_{index}
|
||
# DocStore 格式: {filename}_table_{index}.json 或 {filename}_{index}.json
|
||
possible_paths = [
|
||
docstore_dir / f"{chunk_id}.json", # 直接匹配
|
||
docstore_dir / f"{chunk_id.replace('_', '_table_', 1)}.json", # xxx_0 -> xxx_table_0
|
||
]
|
||
|
||
# 如果 chunk_id 是 {filename}_{num} 格式,尝试 {filename}_table_{num}
|
||
parts = chunk_id.rsplit('_', 1)
|
||
if len(parts) == 2:
|
||
filename, idx = parts
|
||
possible_paths.append(docstore_dir / f"{filename}_table_{idx}.json")
|
||
|
||
for doc_path in possible_paths:
|
||
if doc_path.exists():
|
||
try:
|
||
with open(doc_path, 'r', encoding='utf-8') as f:
|
||
record = json.load(f)
|
||
return record.get('markdown', '')
|
||
except Exception as e:
|
||
logger.debug(f"读取 DocStore 失败: {doc_path}, {e}")
|
||
return None
|
||
|
||
|
||
def _order_text_contexts_for_prompt(contexts: List[Dict], query: str, max_chunks: int,
|
||
min_score: float = 0.0) -> List[Dict]:
|
||
"""
|
||
对文本上下文排序,优化提示词构建
|
||
|
||
对于列举型查询,保持同一文档章节的切片连续排列,
|
||
便于 LLM 理解完整的语义上下文。
|
||
|
||
Args:
|
||
contexts: 检索到的上下文列表
|
||
query: 用户查询
|
||
max_chunks: 最大切片数量
|
||
|
||
Returns:
|
||
排序后的上下文列表
|
||
|
||
Example:
|
||
>>> ordered = _order_text_contexts_for_prompt(contexts, "有哪些禁止情形?", 10)
|
||
"""
|
||
# 文本切片 + 表格切片(从 DocStore 获取完整内容)
|
||
text_contexts = []
|
||
for ctx in contexts:
|
||
chunk_type = ctx.get('meta', {}).get('chunk_type', '')
|
||
if chunk_type in ('image', 'chart'):
|
||
continue # 图片/图表单独处理
|
||
|
||
if chunk_type == 'table':
|
||
# 表格:从 DocStore 获取完整 Markdown
|
||
chunk_id = ctx.get('meta', {}).get('chunk_id', '')
|
||
full_table = _get_full_table_from_docstore(chunk_id) if chunk_id else None
|
||
if full_table:
|
||
# 替换为完整内容,保留 score
|
||
text_contexts.append({
|
||
'doc': full_table,
|
||
'meta': {**ctx.get('meta', {}), '_from_docstore': True},
|
||
'score': ctx.get('score', 0)
|
||
})
|
||
else:
|
||
# DocStore 中没有,使用原始摘要
|
||
text_contexts.append(ctx)
|
||
else:
|
||
# 普通文本切片
|
||
text_contexts.append(ctx)
|
||
|
||
# 图片/图表切片:提取其 doc 内容(包含前文/后文上下文)
|
||
# 这些切片虽然 chunk_type 是 image/chart,但其 doc 字段包含有价值的上下文信息
|
||
chart_contexts = []
|
||
for ctx in contexts:
|
||
chunk_type = ctx.get('meta', {}).get('chunk_type', '')
|
||
if chunk_type in ('image', 'chart'):
|
||
doc = ctx.get('doc', '')
|
||
if doc and len(doc) > 20: # 有实际内容
|
||
# 标记来源,避免重复
|
||
chart_contexts.append({
|
||
'doc': doc,
|
||
'meta': {**ctx.get('meta', {}), '_is_chart_context': True},
|
||
'score': ctx.get('score', 0)
|
||
})
|
||
|
||
# Phase 1:按 Rerank 分数过滤低分切片
|
||
if min_score > 0:
|
||
text_contexts = [c for c in text_contexts if c.get('score', 0) >= min_score]
|
||
chart_contexts = [c for c in chart_contexts if c.get('score', 0) >= min_score]
|
||
|
||
# 合并:文本切片优先,图表切片补充
|
||
# 限制图表切片数量,避免过多
|
||
max_chart_contexts = 3
|
||
combined_contexts = text_contexts + chart_contexts[:max_chart_contexts]
|
||
|
||
if not _is_enum_query(query):
|
||
return combined_contexts[:max_chunks]
|
||
|
||
def sort_key(ctx):
|
||
meta = ctx.get('meta', {})
|
||
source = meta.get('source', '')
|
||
section = meta.get('section', '') or meta.get('section_path', '')
|
||
rank = _safe_int(meta.get('_retrieval_rank'), 10**6)
|
||
return (
|
||
source != primary_source,
|
||
section != primary_section,
|
||
source,
|
||
section,
|
||
_safe_int(meta.get('chunk_index')),
|
||
_safe_int(meta.get('section_chunk_id')),
|
||
rank,
|
||
)
|
||
|
||
primary = text_contexts[0].get('meta', {}) if text_contexts else {}
|
||
primary_source = primary.get('source', '')
|
||
primary_section = primary.get('section', '') or primary.get('section_path', '')
|
||
primary_index = _safe_int(primary.get('chunk_index'), None)
|
||
|
||
try:
|
||
from config import CONTEXT_EXPANSION_BEFORE, CONTEXT_EXPANSION_AFTER
|
||
except Exception as e:
|
||
logger.debug(f"读取上下文扩展配置失败: {e}")
|
||
CONTEXT_EXPANSION_BEFORE, CONTEXT_EXPANSION_AFTER = 1, 5
|
||
|
||
if primary_source and primary_index is not None:
|
||
window_start = primary_index - CONTEXT_EXPANSION_BEFORE
|
||
window_end = primary_index + CONTEXT_EXPANSION_AFTER + 1
|
||
window = []
|
||
for ctx in text_contexts:
|
||
meta = ctx.get('meta', {})
|
||
idx = _safe_int(meta.get('chunk_index'), None)
|
||
if meta.get('source') == primary_source and idx is not None and window_start <= idx <= window_end:
|
||
window.append(ctx)
|
||
window_ids = {
|
||
ctx.get('meta', {}).get('chunk_id') or (ctx.get('meta', {}).get('source'), ctx.get('meta', {}).get('chunk_index'))
|
||
for ctx in window
|
||
}
|
||
window = sorted(window, key=lambda ctx: _safe_int(ctx.get('meta', {}).get('chunk_index')))
|
||
remainder = [
|
||
ctx for ctx in text_contexts
|
||
if (ctx.get('meta', {}).get('chunk_id') or (ctx.get('meta', {}).get('source'), ctx.get('meta', {}).get('chunk_index'))) not in window_ids
|
||
]
|
||
# 枚举查询不截断 max_chunks:需要跨 section 信息,
|
||
# window 优先保证主命中文档连续,remainder 按相关性补充
|
||
return window + sorted(remainder, key=sort_key)
|
||
|
||
ordered = sorted(text_contexts, key=sort_key)
|
||
return ordered
|
||
|
||
|
||
def _build_context_with_budget(contexts: List[Dict], max_chars: int, soft_limit: int) -> str:
|
||
"""
|
||
Phase 2:按字符预算构建上下文文本。
|
||
|
||
1. 按 (source, section) 分组,组内按 chunk_index 保持连续
|
||
2. 组间按组内最高 Rerank 分数降序排列
|
||
3. 逐组加入直到达到 max_chars 预算
|
||
4. 超过 soft_limit 后仅接受组内最高分 >= soft_min_score 的组
|
||
|
||
对于列举类查询(_is_enum_query),保持原始顺序直接拼接,
|
||
因为 _order_text_contexts_for_prompt 已经优化了排序。
|
||
|
||
Args:
|
||
contexts: 已排序的上下文列表
|
||
max_chars: 硬性字符上限
|
||
soft_limit: 软限制,超过后收紧准入
|
||
|
||
Returns:
|
||
拼接好的上下文文本
|
||
"""
|
||
if not contexts:
|
||
return ""
|
||
|
||
# 按 (source, section) 分组
|
||
groups = {}
|
||
group_order = []
|
||
for ctx in contexts:
|
||
meta = ctx.get('meta', {})
|
||
key = (meta.get('source', ''), meta.get('section', '') or meta.get('section_path', ''))
|
||
if key not in groups:
|
||
groups[key] = []
|
||
group_order.append(key)
|
||
groups[key].append(ctx)
|
||
|
||
# 组内按 chunk_index 排序
|
||
for key in groups:
|
||
groups[key].sort(key=lambda c: _safe_int(c.get('meta', {}).get('chunk_index')))
|
||
|
||
# 组间按组内最高 score 降序
|
||
def group_max_score(key):
|
||
return max((c.get('score', 0) for c in groups[key]), default=0)
|
||
|
||
group_order.sort(key=group_max_score, reverse=True)
|
||
|
||
# 贪心加入直到预算满
|
||
parts = []
|
||
total_chars = 0
|
||
for key in group_order:
|
||
group = groups[key]
|
||
group_text = "\n\n".join(ctx.get('doc', '') for ctx in group)
|
||
|
||
# 超过软限制后,只接受高分组
|
||
if total_chars > soft_limit and group_max_score(key) < 0.1:
|
||
continue
|
||
|
||
if total_chars + len(group_text) > max_chars:
|
||
# 尝试逐条加入该组,直到预算满
|
||
for ctx in group:
|
||
doc = ctx.get('doc', '')
|
||
if total_chars + len(doc) + 2 > max_chars: # +2 for "\n\n"
|
||
break
|
||
parts.append(doc)
|
||
total_chars += len(doc) + 2
|
||
break
|
||
|
||
parts.append(group_text)
|
||
total_chars += len(group_text) + 2 # +2 for "\n\n"
|
||
|
||
return "\n\n".join(parts)
|
||
|
||
|
||
def _attach_citations(answer: str, contexts: List[Dict]) -> Dict[str, Any]:
|
||
"""
|
||
自动为回答添加引用标记(按段落级别匹配,jieba 分词精准匹配)
|
||
|
||
流程:
|
||
1. 将 answer 按自然段落分割(\n\n)
|
||
2. 对每个段落用 jieba 分词后计算词级重叠度
|
||
3. Phase 6:动态阈值(短段落 0.55 / 长段落 0.45),每段最多 2 个引用
|
||
4. 前端负责对引用进行重新编号
|
||
|
||
Args:
|
||
answer: LLM 生成的回答
|
||
contexts: 检索到的上下文列表
|
||
|
||
Returns:
|
||
{
|
||
"answer_with_refs": "回答文本(含 [ref:chunk_id] 标记)",
|
||
"citations": [引用列表]
|
||
}
|
||
"""
|
||
import re
|
||
|
||
if not contexts:
|
||
return {"answer_with_refs": answer, "citations": []}
|
||
|
||
# 按 (collection, chunk_id) 复合键组织 contexts,防止跨库同名文件覆盖
|
||
ctx_by_chunk = {}
|
||
for ctx in contexts:
|
||
meta = ctx.get('meta', {})
|
||
chunk_id = meta.get('chunk_id') or f"{meta.get('source')}_{meta.get('chunk_index', 0)}"
|
||
coll = meta.get('_collection') or meta.get('collection') or ''
|
||
composite_key = f"{coll}/{chunk_id}" if coll else chunk_id
|
||
# 保存原始 chunk_id,用于对外输出(ref tag / citation)
|
||
ctx['_raw_chunk_id'] = chunk_id
|
||
ctx_by_chunk[composite_key] = ctx
|
||
|
||
# jieba 分词函数(fallback 到字符级)
|
||
try:
|
||
import jieba
|
||
def _tokenize(text: str) -> set:
|
||
return set(w for w in jieba.lcut(text) if len(w) >= 2)
|
||
except ImportError:
|
||
def _tokenize(text: str) -> set:
|
||
return set(text)
|
||
|
||
# 按自然段落分割(\n\n 为分隔符,保留分隔符用于重组)
|
||
parts = re.split(r'(\n\n+)', answer)
|
||
|
||
cited_chunks_ordered: List[str] = [] # 保序去重的 chunk_id 列表
|
||
cited_set: set = set()
|
||
result_parts = []
|
||
|
||
for i in range(0, len(parts), 2):
|
||
para = parts[i]
|
||
sep = parts[i + 1] if i + 1 < len(parts) else ''
|
||
|
||
stripped = para.strip()
|
||
# 跳过过短段落或 markdown 表格/标题行
|
||
is_table = stripped.startswith('|') or '|---' in stripped
|
||
is_short = len(stripped) < 15
|
||
|
||
if is_table or is_short:
|
||
result_parts.append(para + sep)
|
||
continue
|
||
|
||
# 词级重叠匹配:找最相关的 chunk
|
||
para_words = _tokenize(stripped[:300])
|
||
|
||
# Phase 6:动态阈值 — 短段落用更高阈值避免误匹配
|
||
overlap_threshold = 0.55 if len(stripped) < 50 else 0.45
|
||
|
||
# 收集所有超过阈值的候选 chunk,按分数降序
|
||
candidates = []
|
||
for chunk_id, ctx in ctx_by_chunk.items():
|
||
ctx_doc = ctx.get('doc', '')
|
||
if not ctx_doc:
|
||
continue
|
||
ctx_words = _tokenize(ctx_doc[:400])
|
||
if not para_words:
|
||
continue
|
||
overlap = len(para_words & ctx_words)
|
||
score = overlap / len(para_words)
|
||
if score >= overlap_threshold:
|
||
candidates.append((chunk_id, score))
|
||
|
||
candidates.sort(key=lambda x: x[1], reverse=True)
|
||
|
||
# Phase 6:允许最多 2 个引用(分数差距 < 0.1 时附加第二引用)
|
||
selected_ids = []
|
||
if candidates:
|
||
selected_ids.append(candidates[0][0])
|
||
if len(candidates) > 1 and (candidates[0][1] - candidates[1][1]) < 0.1:
|
||
selected_ids.append(candidates[1][0])
|
||
|
||
if selected_ids:
|
||
for cid in selected_ids:
|
||
if cid not in cited_set:
|
||
cited_set.add(cid)
|
||
cited_chunks_ordered.append(cid)
|
||
# 在段落末尾插入引用标记(使用原始 chunk_id,不暴露复合键)
|
||
ref_tags = "".join(
|
||
f"[ref:{ctx_by_chunk[cid].get('_raw_chunk_id', cid)}]"
|
||
for cid in selected_ids
|
||
)
|
||
result_parts.append(f"{para}{ref_tags}{sep}")
|
||
else:
|
||
result_parts.append(para + sep)
|
||
|
||
# 构建引用列表(按出现顺序),使用原始 chunk_id 构建 citation
|
||
citations = []
|
||
for composite_key in cited_chunks_ordered:
|
||
ctx = ctx_by_chunk.get(composite_key)
|
||
if ctx:
|
||
meta = ctx.get('meta', {})
|
||
full_content = ctx.get('doc', '')
|
||
citation = _build_citation(meta, full_content)
|
||
# 确保 citation 中的 chunk_id 使用原始值(不含 collection 前缀)
|
||
raw_id = ctx.get('_raw_chunk_id')
|
||
if raw_id:
|
||
citation['chunk_id'] = raw_id
|
||
citations.append(citation)
|
||
|
||
return {
|
||
"answer_with_refs": "".join(result_parts),
|
||
"citations": citations
|
||
}
|
||
|
||
|
||
def _clean_section(raw: str) -> str:
|
||
"""清洗 section 字段:过滤掉像正文内容而非章节路径的值"""
|
||
if not raw:
|
||
return ''
|
||
import re as _re
|
||
|
||
def _is_heading(part: str) -> bool:
|
||
"""判断一个字符串是否像章节标题(而非正文内容)"""
|
||
p = part.strip()
|
||
if not p:
|
||
return False
|
||
# 以句号/问号/感叹号结尾 → 正文句子
|
||
if p.endswith(('。', '!', '?', '.', '!', '?')):
|
||
return False
|
||
# 含冒号且偏长 → 更像是内容摘要而非标题
|
||
if ':' in p and len(p) > 25:
|
||
return False
|
||
# 长度超过 30 字 → 不像是标题
|
||
if len(p) > 30:
|
||
return False
|
||
return True
|
||
|
||
# 去掉 Markdown 加粗标记
|
||
cleaned = raw.replace('**', '').strip()
|
||
|
||
# 如果含 ' > ' 分隔符,验证每一级都是合法标题
|
||
if ' > ' in cleaned:
|
||
parts = [p.strip() for p in cleaned.split('>') if p.strip()]
|
||
valid = [p for p in parts if _is_heading(p)]
|
||
if valid:
|
||
return ' > '.join(valid[:3])
|
||
# 所有部分都不像标题 → 清空
|
||
return ''
|
||
|
||
# 匹配 【xxx】 或 第X篇/章/节 格式
|
||
if _re.match(r'^(【[^】]+】|第[一二三四五六七八九十\d]+[篇章节部])', cleaned):
|
||
return cleaned[:40]
|
||
# 短字符串(< 30字)且不像句子(无句号/逗号),保留
|
||
if len(cleaned) < 30 and not any(c in cleaned for c in '。,;:'):
|
||
return cleaned
|
||
# 其余视为正文内容,清空
|
||
return ''
|
||
|
||
|
||
def _build_citation(meta: Dict, full_content: str = '') -> Dict[str, Any]:
|
||
"""
|
||
根据文档类型构建定位信息
|
||
|
||
不同文档类型使用不同的定位策略:
|
||
- PDF: 坐标定位(page + bbox)
|
||
- Word: 语义定位(section + section_chunk_id + preview)
|
||
- Excel: 表格定位(sheet + preview)
|
||
|
||
Args:
|
||
meta: 切片元数据
|
||
full_content: 完整切片内容(可选)
|
||
|
||
Returns:
|
||
引用信息字典,包含定位信息
|
||
|
||
Example:
|
||
>>> citation = _build_citation(meta, "完整内容...")
|
||
>>> print(citation["page"]) # PDF: 页码
|
||
>>> print(citation["section"]) # Word: 章节
|
||
"""
|
||
# 从 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') or meta.get('collection', ''), # 所属向量库,用于前端文档预览跳转
|
||
"doc_type": meta.get('doc_type', 'other'),
|
||
"section": _clean_section(meta.get('section', '')),
|
||
"preview": meta.get('preview', ''),
|
||
"content": (full_content or meta.get('preview', ''))[:300], # 截断至 300 字避免冒返大量数据
|
||
"chunk_type": meta.get('chunk_type', 'text'),
|
||
}
|
||
|
||
doc_type = meta.get('doc_type', 'other')
|
||
|
||
if doc_type == 'pdf':
|
||
# 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':
|
||
# Word: 语义定位
|
||
citation.update({
|
||
"section_chunk_id": meta.get('section_chunk_id'), # 章节内段落序号
|
||
})
|
||
elif doc_type == 'excel':
|
||
# 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 score_image_relevance(query: str, meta: Dict, doc: str = '') -> float:
|
||
"""
|
||
图片相关性打分(语义增强版)
|
||
|
||
通过多维度特征评估图片与查询的相关性:
|
||
1. 图片编号精确匹配(如 "图2.1")
|
||
2. 关键词匹配(年份、数值+单位、中文词组)
|
||
3. 整体文本相似度
|
||
4. 章节匹配
|
||
5. 图片类型加分
|
||
|
||
Args:
|
||
query: 用户查询
|
||
meta: 切片元数据
|
||
doc: document 字段(包含图片描述和上下文)
|
||
|
||
Returns:
|
||
相关性分数(>= 3.0 推荐展示)
|
||
|
||
Example:
|
||
>>> score = score_image_relevance("图2.1是什么?", meta, doc)
|
||
>>> if score >= 3.0:
|
||
... # 推荐展示该图片
|
||
"""
|
||
import re
|
||
score = 0.0
|
||
|
||
# 优先使用 doc 字段(包含完整描述和上下文)
|
||
search_text = doc or meta.get('caption', '')
|
||
section = meta.get('section', '') or meta.get('section_path', '')
|
||
source = meta.get('source', '')
|
||
|
||
# 1. 图片编号精确匹配(最高优先级)
|
||
figure_pattern = r'图\s*(\d+\.?\d*)'
|
||
figure_matches = re.findall(figure_pattern, query)
|
||
|
||
if figure_matches:
|
||
for fig_num in figure_matches:
|
||
# 在所有文本中查找图号
|
||
all_text = f"{search_text} {section} {source}"
|
||
if f"图{fig_num}" in all_text or f"图 {fig_num}" in all_text or f"见图{fig_num}" in all_text:
|
||
score += 10.0 # 精确匹配,直接返回
|
||
return score
|
||
|
||
# 1.5. 表格编号精确匹配(新增:支持表格图片)
|
||
table_pattern = r'表\s*(\d+\.?\d*)'
|
||
table_matches = re.findall(table_pattern, query)
|
||
|
||
if table_matches:
|
||
for table_num in table_matches:
|
||
# 在所有文本中查找表号
|
||
all_text = f"{search_text} {section} {source}"
|
||
if f"表{table_num}" in all_text or f"表 {table_num}" in all_text or f"见表{table_num}" in all_text:
|
||
score += 10.0 # 精确匹配,直接返回
|
||
return score
|
||
|
||
# 2. 查询词匹配(通用方式,不硬编码关键词)
|
||
# 从查询中提取有意义的词:中文词组、数字+单位、年份等
|
||
# 使用 jieba 分词(如果可用)或简单的正则提取
|
||
query_keywords = []
|
||
|
||
# 提取年份(如 "2003年")
|
||
year_matches = re.findall(r'(\d{4})\s*年', query)
|
||
query_keywords.extend(year_matches)
|
||
|
||
# 提取数值+单位(如 "100亿"、"50万千瓦时")
|
||
num_unit_matches = re.findall(r'(\d+\.?\d*\s*[亿万万千百吨米秒])', query)
|
||
query_keywords.extend(num_unit_matches)
|
||
|
||
# 使用 jieba 分词提取中文关键词
|
||
import jieba
|
||
jieba_words = [w for w in jieba.lcut(query) if len(w) >= 2]
|
||
query_keywords.extend(jieba_words)
|
||
|
||
# 过滤掉泛词(图、表、图片等)
|
||
stop_words = {'图', '表', '图片', '图表', '如图', '所示', '如下', '如下表', '如下图', '的', '是', '在', '有', '了', '和', '与', '或', '及', '等', '中', '对'}
|
||
query_keywords = [kw for kw in query_keywords if kw not in stop_words]
|
||
|
||
# 在图片描述中匹配关键词
|
||
keyword_match_score = 0.0
|
||
for kw in query_keywords:
|
||
if kw in search_text or kw in section:
|
||
keyword_match_score += 2.0
|
||
|
||
score += min(keyword_match_score, 8.0) # 最多加 8 分
|
||
|
||
# 3. 整体文本相似度(字符级别)
|
||
if search_text:
|
||
# 检查查询的核心词是否在描述中
|
||
query_core = re.sub(r'[图表图片如图所示]', '', query) # 去掉泛词
|
||
if query_core:
|
||
overlap = len(set(query_core) & set(search_text))
|
||
score += min(overlap * 0.2, 3.0)
|
||
|
||
# 4. 章节匹配
|
||
if section:
|
||
# 从查询中提取章节关键词(复用 jieba 分词)
|
||
section_keywords = query_keywords
|
||
for kw in section_keywords:
|
||
if kw in section:
|
||
score += 1.5
|
||
|
||
# 5. 图片类型加分
|
||
if meta.get('chunk_type') == 'chart':
|
||
score += 2.0
|
||
elif meta.get('chunk_type') == 'image':
|
||
score += 1.0
|
||
|
||
# 6. 检索相似度(如果有)
|
||
retrieval_score = meta.get('score', 0)
|
||
if retrieval_score > 0:
|
||
score += min(retrieval_score * 2, 2.0)
|
||
|
||
return score
|
||
|
||
|
||
def select_images(contexts: List[Dict], query: str) -> List[Dict]:
|
||
"""
|
||
选择要展示的图片(打分排序 + 预算控制)
|
||
|
||
根据查询意图动态调整图片数量上限:
|
||
- 精确查图(指定图号): 最多 2 张
|
||
- 强图片意图(示意图、流程图等): 最多 3 张
|
||
- 列举型查询: 最多 5 张
|
||
- 普通查询: 最多 2 张
|
||
|
||
核心逻辑:
|
||
1. 检测查询中的图号引用
|
||
2. 从检索文本中提取图表引用
|
||
3. 对图片打分并过滤低分图片
|
||
4. 通过章节关联排除不相关图片
|
||
|
||
Args:
|
||
contexts: 检索上下文列表
|
||
query: 用户查询
|
||
|
||
Returns:
|
||
精选图片列表(每项含 score, id, url, type, source 等字段)
|
||
|
||
Example:
|
||
>>> images = select_images(contexts, "图2.1展示了什么?")
|
||
>>> print(len(images)) # <= 2
|
||
"""
|
||
import re
|
||
|
||
# 动态预算:列举型查询允许更多图片
|
||
list_keywords = ["哪些", "有什么", "包含", "列出", "所有", "全部"]
|
||
is_list_query = any(kw in query for kw in list_keywords)
|
||
|
||
# 检测查询中的图片编号(如 "图2.1")
|
||
figure_pattern = r'图\s*(\d+\.?\d*)'
|
||
figure_matches = re.findall(figure_pattern, query)
|
||
has_figure_query = bool(figure_matches)
|
||
|
||
# 新增:从检索文本中提取图表引用(见表2.2、见图2.5 等)
|
||
# 同时记录引用所在的文件来源
|
||
# 重要:只从语义相关的 top 5 文本块提取,避免不相关引用干扰
|
||
referenced_figures = {} # {图号: set(文件来源)}
|
||
referenced_tables = {} # {表号: set(文件来源)}
|
||
|
||
for ctx in contexts[:5]: # 只从前5个最相关的文本块提取引用
|
||
doc_text = ctx.get('doc', '')
|
||
source = ctx.get('meta', {}).get('source', '')
|
||
|
||
# 提取 "见图X.X"、"如图X.X" 或单独的 "图X.X"
|
||
fig_refs = re.findall(r'(?:[见如])?图\s*(\d+\.?\d*)', doc_text)
|
||
for fig_num in fig_refs:
|
||
if fig_num not in referenced_figures:
|
||
referenced_figures[fig_num] = set()
|
||
if source:
|
||
referenced_figures[fig_num].add(source)
|
||
|
||
# 提取 "见表X.X"、"如表X.X" 或单独的 "表X.X"
|
||
table_refs = re.findall(r'(?:[见如])?表\s*(\d+\.?\d*)', doc_text)
|
||
for table_num in table_refs:
|
||
if table_num not in referenced_tables:
|
||
referenced_tables[table_num] = set()
|
||
if source:
|
||
referenced_tables[table_num].add(source)
|
||
|
||
has_referenced_figures = bool(referenced_figures or referenced_tables)
|
||
|
||
# 获取检索结果中涉及的主要文件来源
|
||
primary_sources = set()
|
||
for ctx in contexts[:5]: # 只看前5个最相关的
|
||
source = ctx.get('meta', {}).get('source', '')
|
||
if source:
|
||
primary_sources.add(source)
|
||
|
||
# 图片意图检测:区分精确查图和泛指
|
||
strong_image_keywords = ["示意图", "流程图", "结构图", "过程线", "曲线图", "分布图", "图示", "看图", "显示图"]
|
||
weak_image_keywords = ["图片", "图表", "如图", "图", "统计"]
|
||
|
||
has_strong_image_intent = any(kw in query for kw in strong_image_keywords)
|
||
has_weak_image_intent = any(kw in query for kw in weak_image_keywords)
|
||
|
||
# 精确查图:用户指定了具体图号(如 "图2.3"),只返回最匹配的 1-2 张
|
||
if has_figure_query:
|
||
MAX_IMAGES = 2
|
||
MIN_SCORE = 5.0 # 精确匹配应该高分
|
||
# 强图片意图:明确要看某种图
|
||
elif has_strong_image_intent:
|
||
MAX_IMAGES = 3
|
||
MIN_SCORE = 5.0 # 提高阈值,避免不相关图片通过
|
||
# 列举型查询
|
||
elif is_list_query:
|
||
MAX_IMAGES = 5
|
||
MIN_SCORE = 3.0
|
||
# 有图表引用:检索文本中提到了图表
|
||
elif has_referenced_figures:
|
||
MAX_IMAGES = 3
|
||
MIN_SCORE = 2.0 # 降低阈值,让引用的图表能通过
|
||
# 弱图片意图:只是提到"图"字,可能是泛指(如 "发电量图")
|
||
elif has_weak_image_intent:
|
||
MAX_IMAGES = 1 # 只返回最相关的一张
|
||
MIN_SCORE = 2.0 # 降低阈值,让语义相关的图片能通过
|
||
# 普通查询
|
||
else:
|
||
MAX_IMAGES = 2
|
||
MIN_SCORE = 3.0
|
||
|
||
# 获取检索结果中涉及的主要章节(只看前 3 个最相关的文本块)
|
||
primary_sections = set()
|
||
for ctx in contexts[:3]:
|
||
section = ctx.get('meta', {}).get('section', '') or ctx.get('meta', {}).get('section_path', '')
|
||
if section:
|
||
# 提取章节编号
|
||
# 优先匹配 X.X 格式(如 "2.3发电"),再匹配 第X章 格式
|
||
section_num = re.search(r'(\d+\.\d+)', section)
|
||
if not section_num:
|
||
# 尝试匹配 "第X章" 格式
|
||
chapter_match = re.search(r'第\s*(\d+)\s*章', section)
|
||
if chapter_match:
|
||
section_num = chapter_match
|
||
if section_num:
|
||
primary_sections.add(section_num.group(1))
|
||
|
||
scored_images = []
|
||
for ctx in contexts:
|
||
meta = ctx.get('meta', {})
|
||
chunk_type = meta.get('chunk_type', 'text')
|
||
|
||
# 处理图片类型和有关联图片的表格类型
|
||
if meta.get('image_path') and chunk_type in ('image', 'chart', 'table'):
|
||
# 优先使用 VLM 详细描述
|
||
# 1. lazy_enhance 对 image/chart 更新 ctx['doc']
|
||
# 2. lazy_enhance 对 table 更新 ctx['image_description'](而非 doc)
|
||
doc = ctx.get('image_description', '') or meta.get('vlm_desc', '') or ctx.get('doc', '')
|
||
s = score_image_relevance(query, meta, doc)
|
||
|
||
# ========== VLM 相关性筛选(方案 C)==========
|
||
# 用 VLM 描述判断图片内容是否与查询相关
|
||
# 优先级:meta.vlm_desc(已同步) > .data/cache/vlm/(懒加载缓存)
|
||
image_path = meta.get('image_path', '')
|
||
vlm_desc = meta.get('vlm_desc', '') or _get_vlm_cache(image_path)
|
||
if vlm_desc:
|
||
vlm_relevance = _check_vlm_relevance(query, vlm_desc)
|
||
if vlm_relevance < 0.3:
|
||
# VLM 描述与查询不相关,适度降分
|
||
s -= 3.0
|
||
logger.debug(f"图片 {image_path} VLM 不相关,降分: {vlm_relevance:.2f}")
|
||
elif vlm_relevance >= 0.5:
|
||
# 相关,小幅加分
|
||
s += 2.0
|
||
|
||
# 图片来源
|
||
img_source = meta.get('source', '')
|
||
|
||
# 图片章节
|
||
img_section = meta.get('section', '') or meta.get('section_path', '')
|
||
# 只匹配 X.X 格式的章节号,避免匹配年份
|
||
img_section_num = re.search(r'(\d+\.\d+)', img_section)
|
||
img_section_id = img_section_num.group(1) if img_section_num else None
|
||
|
||
# ========== 核心修复:图片必须与主要文本切片章节关联 ==========
|
||
# 如果有主要章节,且图片章节不在其中,大幅降低分数
|
||
section_penalty = 0.0
|
||
if primary_sections and img_section_id and img_section_id not in primary_sections:
|
||
# 图片章节与主要检索结果不匹配,惩罚
|
||
section_penalty = -5.0 # 大幅降低分数
|
||
# 除非图片被文本切片明确引用
|
||
is_referenced = False
|
||
for fig_num in referenced_figures:
|
||
if f"图{fig_num}" in doc or f"图 {fig_num}" in doc:
|
||
is_referenced = True
|
||
break
|
||
if not is_referenced:
|
||
for table_num in referenced_tables:
|
||
if f"表{table_num}" in doc or f"表 {table_num}" in doc:
|
||
is_referenced = True
|
||
break
|
||
if is_referenced:
|
||
section_penalty = 0.0 # 被引用则不惩罚
|
||
|
||
s += section_penalty
|
||
|
||
# 新增:如果图片描述中包含检索文本引用的图号,大幅加分
|
||
# 前提:图片章节与主要检索结果的章节相关
|
||
if referenced_figures:
|
||
for fig_num, sources in referenced_figures.items():
|
||
# 只检查 doc 字段,不检查 meta(避免 section 中的误匹配)
|
||
if f"图{fig_num}" in doc or f"图 {fig_num}" in doc:
|
||
# 检查图片章节是否与主要章节匹配
|
||
section_match = img_section_id and img_section_id in primary_sections
|
||
|
||
# P2:只有章节匹配才加分,移除"s >= 5.0"漏洞
|
||
if section_match:
|
||
# 图号匹配加分
|
||
s += 8.0
|
||
# 如果图片来源与引用来源一致,额外加分
|
||
if img_source in sources:
|
||
s += 5.0 # 文件匹配额外加分
|
||
break
|
||
|
||
# 新增:如果表格描述中包含检索文本引用的表号,大幅加分
|
||
# 同样要求章节相关性
|
||
if referenced_tables:
|
||
for table_num, sources in referenced_tables.items():
|
||
# 只检查 doc 字段
|
||
if f"表{table_num}" in doc or f"表 {table_num}" in doc:
|
||
# 检查图片章节是否与主要章节匹配
|
||
section_match = img_section_id and img_section_id in primary_sections
|
||
|
||
# P2:只有章节匹配才加分,移除"s >= 5.0"漏洞
|
||
if section_match:
|
||
# 表号匹配加分
|
||
s += 8.0
|
||
# 如果图片来源与引用来源一致,额外加分
|
||
if img_source in sources:
|
||
s += 5.0 # 文件匹配额外加分
|
||
break
|
||
|
||
# 新增:如果图片来源在主要检索结果中,加分
|
||
if img_source in primary_sources:
|
||
s += 2.0
|
||
|
||
if s >= MIN_SCORE:
|
||
scored_images.append({
|
||
'score': s,
|
||
'id': os.path.basename(meta['image_path']),
|
||
'url': f"/images/{os.path.basename(meta['image_path'])}",
|
||
'type': meta['chunk_type'],
|
||
'source': meta.get('source'),
|
||
'page': meta.get('page'),
|
||
'description': doc[:100], # 短描述用于 UI 展示
|
||
'full_description': doc # Bug 6b 修复:完整描述用于 LLM 上下文
|
||
})
|
||
|
||
# ========== P1.5:处理表格切片的 images_json(跨页表格多图)==========
|
||
# 当表格切片有 images_json 字段时,添加所有关联图片
|
||
if chunk_type == 'table' and meta.get('images_json'):
|
||
try:
|
||
images_list = json.loads(meta['images_json'])
|
||
for img_info in images_list:
|
||
if isinstance(img_info, dict):
|
||
img_id = img_info.get('id') or img_info.get('path', '')
|
||
img_page = img_info.get('page', meta.get('page'))
|
||
else:
|
||
img_id = str(img_info)
|
||
img_page = meta.get('page')
|
||
|
||
# 跳过已添加的图片(避免重复)
|
||
existing_ids = {img['id'] for img in scored_images}
|
||
if img_id and img_id not in existing_ids:
|
||
# 为关联图片计算分数(继承主表格分数,略低)
|
||
assoc_score = s - 1.0 if s >= MIN_SCORE else MIN_SCORE - 1.0
|
||
if assoc_score >= MIN_SCORE:
|
||
scored_images.append({
|
||
'score': assoc_score,
|
||
'id': img_id,
|
||
'url': f"/images/{img_id}",
|
||
'type': 'table_image',
|
||
'source': meta.get('source'),
|
||
'page': img_page,
|
||
'description': doc[:100],
|
||
'full_description': doc
|
||
})
|
||
except (json.JSONDecodeError, TypeError):
|
||
pass
|
||
|
||
# ========== P2:通过文本切片的 referenced_images 补充图片 ==========
|
||
# 检查 top 5 文本切片的 referenced_images,补充未选中的关联图片
|
||
existing_image_ids = {img['id'] for img in scored_images}
|
||
|
||
for ctx in contexts[:5]:
|
||
meta = ctx.get('meta', {})
|
||
if meta.get('chunk_type') != 'text':
|
||
continue
|
||
|
||
referenced = meta.get('referenced_images', [])
|
||
if not referenced:
|
||
continue
|
||
|
||
# 查找对应的图片切片
|
||
for fig_num in referenced:
|
||
# 在所有 contexts 中查找匹配的图片
|
||
for img_ctx in contexts:
|
||
img_meta = img_ctx.get('meta', {})
|
||
if img_meta.get('chunk_type') not in ('image', 'chart', 'table'):
|
||
continue
|
||
|
||
img_path = img_meta.get('image_path', '')
|
||
img_id = os.path.basename(img_path)
|
||
|
||
# 检查是否已存在
|
||
if img_id in existing_image_ids:
|
||
continue
|
||
|
||
# 检查图号/表号是否匹配
|
||
img_doc = img_ctx.get('doc', '')
|
||
if f"图{fig_num}" in img_doc or f"表{fig_num}" in img_doc:
|
||
# 添加到结果中
|
||
scored_images.append({
|
||
'score': 8.0, # 基础分
|
||
'id': img_id,
|
||
'url': f"/images/{img_id}",
|
||
'type': img_meta.get('chunk_type'),
|
||
'source': img_meta.get('source'),
|
||
'page': img_meta.get('page'),
|
||
'description': img_doc[:100], # 短描述用于 UI 展示
|
||
'full_description': img_doc # Bug 6b 修复:完整描述用于 LLM 上下文
|
||
})
|
||
existing_image_ids.add(img_id)
|
||
break
|
||
|
||
scored_images.sort(key=lambda x: x['score'], reverse=True)
|
||
return scored_images[:MAX_IMAGES]
|
||
|
||
def chat_with_llm(message: str, history: List[Dict] = None, enable_web_search: bool = True) -> Dict[str, Any]:
|
||
"""
|
||
普通聊天 - 使用 LLM 直接回复
|
||
|
||
当查询不需要知识库检索时,直接调用 LLM 进行回复。
|
||
可选启用网络搜索增强。
|
||
|
||
Args:
|
||
message: 用户消息
|
||
history: 对话历史(由后端传入)
|
||
enable_web_search: 是否启用网络搜索
|
||
|
||
Returns:
|
||
{
|
||
"answer": str,
|
||
"sources": list,
|
||
"web_searched": bool
|
||
}
|
||
|
||
Example:
|
||
>>> result = chat_with_llm("你好", enable_web_search=False)
|
||
>>> print(result["answer"])
|
||
"""
|
||
from config import get_llm_client, LLM_MAX_TOKENS
|
||
|
||
client = get_llm_client()
|
||
|
||
# 构建消息
|
||
messages = []
|
||
|
||
# 添加历史
|
||
if history:
|
||
for h in history[-MAX_HISTORY_ROUNDS:]:
|
||
messages.append({"role": h["role"], "content": h["content"]})
|
||
|
||
messages.append({"role": "user", "content": message})
|
||
|
||
# 调用 LLM
|
||
answer = call_llm(client, prompt="", model=RAG_CHAT_MODEL, messages=messages, max_tokens=LLM_MAX_TOKENS)
|
||
|
||
return {
|
||
"answer": answer or "",
|
||
"sources": [],
|
||
"web_searched": False
|
||
}
|
||
|
||
|
||
def reciprocal_rank_fusion(results_list, weights=None, k=60):
|
||
"""
|
||
倒数排名融合算法
|
||
|
||
Args:
|
||
results_list: 多个检索结果列表
|
||
weights: 各结果权重
|
||
k: RRF 参数
|
||
|
||
Returns:
|
||
融合后的排序结果
|
||
"""
|
||
if weights is None:
|
||
weights = [1.0] * len(results_list)
|
||
|
||
fused_scores = {}
|
||
doc_data = {}
|
||
|
||
for results, weight in zip(results_list, weights):
|
||
if not results or not results.get('ids'):
|
||
continue
|
||
|
||
ids = results['ids'][0]
|
||
docs = results['documents'][0] if results.get('documents') else [''] * len(ids)
|
||
metas = results['metadatas'][0] if results.get('metadatas') else [{}] * len(ids)
|
||
distances = results['distances'][0] if results.get('distances') else [0] * len(ids)
|
||
|
||
for rank, (doc_id, doc, meta, dist) in enumerate(zip(ids, docs, metas, distances)):
|
||
if doc_id not in fused_scores:
|
||
fused_scores[doc_id] = 0
|
||
doc_data[doc_id] = {'doc': doc, 'meta': meta, 'dist': dist}
|
||
|
||
# RRF 分数
|
||
fused_scores[doc_id] += weight / (rank + k)
|
||
|
||
# 按分数排序
|
||
sorted_ids = sorted(fused_scores.keys(), key=lambda x: fused_scores[x], reverse=True)
|
||
|
||
return {
|
||
'ids': sorted_ids,
|
||
'documents': [doc_data[i]['doc'] for i in sorted_ids],
|
||
'metadatas': [doc_data[i]['meta'] for i in sorted_ids],
|
||
'scores': [fused_scores[i] for i in sorted_ids],
|
||
'distances': [doc_data[i]['dist'] for i in sorted_ids]
|
||
}
|
||
|
||
|
||
def search_hybrid(query: str, top_k: int = 5, candidates: int = 15,
|
||
allowed_levels: list = None, allowed_collections: list = None,
|
||
sub_queries: list = None):
|
||
"""
|
||
混合检索:直接调用生产环境引擎,确保测试效果与生产一致
|
||
|
||
Args:
|
||
query: 查询文本
|
||
top_k: 返回数量
|
||
candidates: 候选数量(用于 RERANK_CANDIDATES,由 config 控制)
|
||
allowed_levels: 允许的安全级别
|
||
allowed_collections: 允许的向量库列表
|
||
sub_queries: 意图分析器生成的子查询列表(对比类查询用)
|
||
|
||
Returns:
|
||
融合后的检索结果
|
||
"""
|
||
from core.engine import get_engine
|
||
|
||
engine = get_engine()
|
||
|
||
# 直接调用生产环境的检索方法
|
||
result = engine.search_knowledge(
|
||
query=query,
|
||
top_k=top_k,
|
||
allowed_levels=allowed_levels,
|
||
collections=allowed_collections,
|
||
sub_queries=sub_queries
|
||
)
|
||
|
||
# 添加 scores 字段(用于前端显示和上下文过滤)
|
||
if result and result.get('ids') and result['ids'][0]:
|
||
distances = result.get('distances', [[]])[0]
|
||
if result.get('_reranked'):
|
||
# Rerank 后 distances 中存储的是相关性分数(越高越好),直接使用
|
||
scores = [float(d) for d in distances]
|
||
else:
|
||
# 未 Rerank:将向量距离转换为相似度分数(距离越小,分数越高)
|
||
scores = [1.0 - d if d <= 1.0 else 1.0 / (1.0 + d) for d in distances]
|
||
result['scores'] = [scores]
|
||
else:
|
||
result = {
|
||
'ids': [[]],
|
||
'documents': [[]],
|
||
'metadatas': [[]],
|
||
'distances': [[]],
|
||
'scores': [[]]
|
||
}
|
||
|
||
return result
|
||
|
||
|
||
# ==================== 路由 ====================
|
||
|
||
@chat_bp.route('/chat', methods=['POST'])
|
||
@require_gateway_auth
|
||
def chat():
|
||
"""
|
||
普通聊天模式 - 直接使用LLM回复
|
||
|
||
请求体:
|
||
{
|
||
"message": "消息内容",
|
||
"history": [{"role": "user/assistant", "content": "..."}] // 可选
|
||
}
|
||
"""
|
||
data = request.json or {}
|
||
message = data.get('message')
|
||
history = data.get('history', [])
|
||
|
||
if not message:
|
||
return error_response("MISSING_PARAMS", BAD_REQUEST, "缺少 message", http_status=400)
|
||
|
||
# 输入安全验证
|
||
is_valid, reason = validate_query(message)
|
||
if not is_valid:
|
||
return error_response("INVALID_QUERY", BAD_REQUEST, reason, http_status=400)
|
||
|
||
# 智能聊天
|
||
result = chat_with_llm(message, history)
|
||
|
||
# 过滤敏感信息
|
||
answer = filter_response(result["answer"])
|
||
|
||
return success_response(data={
|
||
"answer": answer,
|
||
"mode": "chat",
|
||
"sources": result.get("sources", []),
|
||
"web_searched": result.get("web_searched", False)
|
||
})
|
||
|
||
|
||
@chat_bp.route('/rag', methods=['POST'])
|
||
@require_gateway_auth
|
||
def rag():
|
||
"""
|
||
知识库问答模式 - SSE 流式返回
|
||
|
||
请求体:
|
||
{
|
||
"message": "消息内容",
|
||
"history": [{"role": "user/assistant", "content": "..."}], // 可选(开发环境)
|
||
"chat_history": [{"role": "user/assistant", "content": "..."}], // 可选(生产环境)
|
||
"collections": ["public_kb"], // 可选,知识库列表
|
||
"session_id": "xxx" // 可选,会话ID
|
||
}
|
||
|
||
SSE 事件序列:
|
||
1. start: 开始处理
|
||
2. sources: 检索到的来源
|
||
3. chunk: 每个 token
|
||
4. finish: 完成响应(包含完整 answer 和 sources)
|
||
5. error: 错误事件
|
||
"""
|
||
import re
|
||
from config import (
|
||
IS_PROD, IS_DEV, ENABLE_SESSION,
|
||
RAG_SEARCH_TOP_K, RAG_SEARCH_CANDIDATES,
|
||
MAX_CONTEXT_CHUNKS, MAX_SOURCES_RETURNED,
|
||
LLM_TEMPERATURE, LLM_MAX_TOKENS,
|
||
MAX_HISTORY_ROUNDS, IMAGE_CONTEXT_HISTORY,
|
||
DIRECT_CONTEXT_MAX_CHARS,
|
||
RERANK_CONTEXT_MIN_SCORE,
|
||
CONTEXT_MAX_CHARS, CONTEXT_SOFT_LIMIT,
|
||
CONFIDENCE_WARN_THRESHOLD, CONFIDENCE_CAUTION_THRESHOLD
|
||
)
|
||
|
||
data = request.json or {}
|
||
|
||
message = data.get('message')
|
||
# 兼容两种参数名:history(旧)和 chat_history(新)
|
||
# 用 is not None 判断,避免 [] 被 or 吞掉
|
||
if 'chat_history' in data:
|
||
history = data['chat_history']
|
||
elif 'history' in data:
|
||
history = data['history']
|
||
else:
|
||
history = None
|
||
collections = data.get('collections')
|
||
session_id = data.get('session_id')
|
||
|
||
if not message:
|
||
return error_response("MISSING_PARAMS", BAD_REQUEST, "缺少 message", http_status=400)
|
||
|
||
# 生产环境强制校验 chat_history
|
||
if IS_PROD and history is None:
|
||
return error_response("MISSING_HISTORY", BAD_REQUEST, "chat_history is required in production", http_status=400)
|
||
|
||
# 输入安全验证
|
||
is_valid, reason = validate_query(message)
|
||
if not is_valid:
|
||
return error_response("INVALID_QUERY", BAD_REQUEST, reason, http_status=400)
|
||
|
||
# 如果没有指定 collections,使用默认的公开库
|
||
if not collections:
|
||
collections = ['public_kb']
|
||
|
||
# ==================== 会话历史加载 ====================
|
||
# 优先使用传入的 history,否则从 session_repo 加载
|
||
user_id = request.current_user.get("user_id")
|
||
|
||
if history is not None:
|
||
# 使用传入的历史(生产环境必须传入)
|
||
pass
|
||
elif session_id and ENABLE_SESSION:
|
||
# 开发环境:从本地数据库加载
|
||
try:
|
||
session_repo = current_app.session_repo
|
||
history = session_repo.get_history(session_id)
|
||
# 限制历史长度
|
||
history = history[-MAX_HISTORY_ROUNDS:] if len(history) > MAX_HISTORY_ROUNDS else history
|
||
except Exception as e:
|
||
logger.debug(f"解析历史记录失败: {e}")
|
||
history = []
|
||
else:
|
||
history = []
|
||
|
||
# 如果没有 session_id,创建新会话(仅开发环境)
|
||
if not session_id and ENABLE_SESSION:
|
||
try:
|
||
session_repo = current_app.session_repo
|
||
session_id = session_repo.create_session(user_id)
|
||
except Exception as e:
|
||
logger.debug(f"创建会话失败: {e}")
|
||
|
||
# 提前获取 session_repo 引用,避免在生成器内部访问 current_app
|
||
# (生成器执行时应用上下文可能已结束)
|
||
session_repo_ref = None
|
||
if ENABLE_SESSION:
|
||
try:
|
||
session_repo_ref = current_app.session_repo
|
||
except Exception as e:
|
||
logger.debug(f"获取会话仓库失败: {e}")
|
||
|
||
def generate():
|
||
"""生成 SSE 流"""
|
||
import re
|
||
start_time = _time.time()
|
||
full_answer = []
|
||
|
||
try:
|
||
# 0. 意图分析(改写 + 双层判断)
|
||
context_images = []
|
||
if history:
|
||
# 从历史中提取图片信息
|
||
for msg in reversed(history[-IMAGE_CONTEXT_HISTORY:]):
|
||
metadata = msg.get("metadata", {})
|
||
if isinstance(metadata, dict):
|
||
images = metadata.get("images", [])
|
||
if images:
|
||
context_images.extend(images[:3])
|
||
|
||
intent = None
|
||
try:
|
||
from core.intent_analyzer import analyze_intent
|
||
intent = analyze_intent(message, history or [], context_images)
|
||
|
||
logger.info(f"[意图分析] use_context={intent.use_context}, need_retrieval={intent.need_retrieval}, intent={intent.intent}, sub_queries={intent.sub_queries}")
|
||
|
||
# 调试事件:意图分析结果
|
||
if IS_DEV:
|
||
yield f"data: {json.dumps({'type': 'intent_result', 'data': {'intent': intent.intent, 'confidence': round(intent.confidence, 2), 'rewritten_query': intent.rewritten_query, 'sub_queries': intent.sub_queries, 'need_retrieval': intent.need_retrieval, 'reason': intent.reason}}, ensure_ascii=False)}\n\n"
|
||
|
||
# 如果不需要检索,直接使用上下文回答
|
||
if not intent.need_retrieval and intent.use_context:
|
||
yield f"data: {json.dumps({'type': 'start', 'message': '正在分析...'}, ensure_ascii=False)}\n\n"
|
||
|
||
# 构建上下文
|
||
context_text = ""
|
||
if history:
|
||
# 提取最近的助手回答
|
||
for msg in reversed(history):
|
||
if msg.get("role") == "assistant":
|
||
context_text = msg.get("content", "")
|
||
break
|
||
|
||
# 构建图片上下文
|
||
image_context = ""
|
||
if context_images:
|
||
image_context = "\n\n【上下文中的图片】\n"
|
||
for img in context_images[:5]:
|
||
if isinstance(img, dict):
|
||
desc = img.get("description", "")
|
||
img_type = img.get("type", "图片")
|
||
image_context += f"- {img_type}: {desc}\n"
|
||
|
||
# 直接调用 LLM
|
||
from config import get_llm_client, DASHSCOPE_MODEL
|
||
client = get_llm_client()
|
||
|
||
system_prompt = f"""你是一个专业的知识库问答助手。请根据对话历史和上下文回答用户问题。
|
||
|
||
如果用户问题是关于图片的,请根据上下文中的图片描述进行分析。
|
||
|
||
{image_context}"""
|
||
|
||
user_prompt = f"""对话历史:
|
||
{context_text[:DIRECT_CONTEXT_MAX_CHARS] if context_text else '(无历史上下文)'}
|
||
|
||
用户问题:{intent.rewritten_query}
|
||
|
||
请直接回答用户问题。"""
|
||
|
||
# 流式生成回答
|
||
for content in call_llm_stream(
|
||
client,
|
||
prompt=user_prompt,
|
||
model=DASHSCOPE_MODEL,
|
||
messages=[
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": user_prompt}
|
||
],
|
||
temperature=LLM_TEMPERATURE
|
||
):
|
||
full_answer.append(content)
|
||
yield f"data: {json.dumps({'type': 'chunk', 'content': content}, ensure_ascii=False)}\n\n"
|
||
|
||
# 发送完成事件
|
||
yield f"data: {json.dumps({'type': 'finish', 'answer': ''.join(full_answer), 'sources': []}, ensure_ascii=False)}\n\n"
|
||
return # 直接返回,不执行后续检索
|
||
|
||
except Exception as e:
|
||
logger.warning(f"意图分析失败: {e},继续执行检索流程")
|
||
|
||
# 1.5 语义缓存检查(跳过检索+生成全流程)
|
||
_semantic_cache_emb = None
|
||
try:
|
||
from core.semantic_cache import get_semantic_cache
|
||
from core.engine import get_engine as _get_eng
|
||
_sc = get_semantic_cache()
|
||
_eng = _get_eng()
|
||
if _sc and _eng and hasattr(_eng, 'embedding_model'):
|
||
_semantic_cache_emb = _eng.embedding_model.encode(message)
|
||
cached = _sc.get(_semantic_cache_emb)
|
||
if cached is not None:
|
||
logger.info(f"[语义缓存] 命中: {message[:50]}...")
|
||
cached_answer = cached.get("answer", "")
|
||
# 流式返回缓存的答案
|
||
yield f"data: {json.dumps({'type': 'start', 'message': '正在检索知识库...'}, ensure_ascii=False)}\n\n"
|
||
# 分块发送缓存答案
|
||
chunk_size = 20
|
||
for i in range(0, len(cached_answer), chunk_size):
|
||
chunk = cached_answer[i:i+chunk_size]
|
||
full_answer.append(chunk)
|
||
yield f"data: {json.dumps({'type': 'chunk', 'content': chunk}, ensure_ascii=False)}\n\n"
|
||
finish_event = {
|
||
"type": "finish",
|
||
"answer": cached_answer,
|
||
"mode": "rag",
|
||
"session_id": session_id,
|
||
"sources": cached.get("sources", []),
|
||
"citations": cached.get("citations", []),
|
||
"images": cached.get("images", []),
|
||
"tables": cached.get("tables", []),
|
||
"sections": [],
|
||
"duration_ms": int((_time.time() - start_time) * 1000),
|
||
"confidence_score": 1.0,
|
||
"semantic_cache_hit": True
|
||
}
|
||
yield f"data: {json.dumps(finish_event, ensure_ascii=False)}\n\n"
|
||
return
|
||
except Exception as e:
|
||
logger.warning(f"[语义缓存] 检查失败: {e}")
|
||
_semantic_cache_emb = None
|
||
|
||
# 1. 发送开始事件
|
||
yield f"data: {json.dumps({'type': 'start', 'message': '正在检索知识库...'}, ensure_ascii=False)}\n\n"
|
||
|
||
# 2. 执行混合检索(扩大召回数量,确保图片切片有机会被召回)
|
||
# 如果意图分析生成了子查询(对比类),传给搜索引擎并行检索
|
||
sub_queries = None
|
||
if intent and intent.sub_queries and len(intent.sub_queries) > 1:
|
||
sub_queries = intent.sub_queries
|
||
|
||
search_result = search_hybrid(
|
||
message,
|
||
top_k=RAG_SEARCH_TOP_K,
|
||
candidates=RAG_SEARCH_CANDIDATES,
|
||
allowed_collections=collections,
|
||
sub_queries=sub_queries
|
||
)
|
||
|
||
# 调试事件:检索管线详情
|
||
if IS_DEV:
|
||
debug_info = search_result.get('_debug', {})
|
||
yield f"data: {json.dumps({'type': 'retrieval_debug', 'data': {'steps': debug_info.get('steps', []), 'collections_searched': collections, 'total_candidates': len(search_result.get('ids', [[]])[0])}}, ensure_ascii=False)}\n\n"
|
||
|
||
# 提取上下文
|
||
contexts = []
|
||
sources = []
|
||
|
||
if search_result.get('documents') and search_result['documents'][0]:
|
||
docs = search_result['documents'][0]
|
||
metas = search_result.get('metadatas', [[]])[0]
|
||
scores = search_result.get('scores', [[]])[0]
|
||
|
||
# 图片相关性提升:检测图片编号或强意图
|
||
import re
|
||
figure_pattern = r'图\s*(\d+\.?\d*)'
|
||
figure_matches = re.findall(figure_pattern, message)
|
||
has_figure_query = bool(figure_matches)
|
||
|
||
# Bug 4 修复:缩小 strong_image_keywords,移除歧义词
|
||
# "过程线" 是水文术语不是图片意图,"图片"/"图表"/"如图" 太泛
|
||
strong_image_keywords = ["示意图", "流程图", "结构图", "曲线图", "分布图", "图示", "看图", "显示图", "给我看"]
|
||
has_image_intent = has_figure_query or any(kw in message for kw in strong_image_keywords)
|
||
|
||
# Bug 2 修复:不重排 contexts,只在 meta 里打标记给后续的 select_images 用
|
||
if has_image_intent:
|
||
for i, (doc, meta, score) in enumerate(zip(docs, metas, scores)):
|
||
if meta.get('chunk_type') in ('image', 'chart'):
|
||
# 检查 caption 是否与查询相关
|
||
caption = meta.get('caption', '') or ''
|
||
should_boost = False
|
||
boost_factor = 1.0
|
||
|
||
# 如果查询包含图片编号,检查 caption 是否匹配
|
||
if has_figure_query:
|
||
for fig_num in figure_matches:
|
||
if f"图{fig_num}" in caption or f"图 {fig_num}" in caption:
|
||
should_boost = True
|
||
boost_factor = 2.0 # 强提升
|
||
break
|
||
|
||
# 或者 caption 与查询有足够重叠
|
||
if not should_boost:
|
||
overlap = len(set(message) & set(caption))
|
||
if overlap >= 3 or any(word in caption for word in message if len(word) >= 3):
|
||
should_boost = True
|
||
boost_factor = 1.5 # 提升 50%
|
||
|
||
if should_boost:
|
||
meta['_image_boost'] = boost_factor # 打标记,不重排
|
||
# 不做 sort!保持检索引擎的原始排序
|
||
|
||
# 按 source 去重,保留最高分
|
||
seen_sources = {}
|
||
for rank, (doc, meta, score) in enumerate(zip(docs, metas, scores)):
|
||
meta['_retrieval_rank'] = rank
|
||
# 确保 _collection 字段存在(单知识库路径下 ChromaDB 原生不返回此字段)
|
||
if not meta.get('_collection'):
|
||
# 优先使用入库时写入的 collection 字段
|
||
meta['_collection'] = meta.get('collection') or (collections[0] if collections else 'public_kb')
|
||
source_name = meta.get('source', '未知')
|
||
if source_name not in seen_sources or score > seen_sources[source_name]['score']:
|
||
doc_type = meta.get('doc_type', 'other')
|
||
page = meta.get('page', 0)
|
||
page_end = meta.get('page_end')
|
||
# 仅 PDF 的页码是真实可靠的;Word 等的页码是 MinerU 合成的,无意义,不外露
|
||
if doc_type == 'pdf' and page:
|
||
page_range = f"{page}-{page_end}" if (page_end and page_end > page) else str(page)
|
||
else:
|
||
page = None
|
||
page_end = None
|
||
page_range = ''
|
||
|
||
seen_sources[source_name] = {
|
||
'source': source_name,
|
||
'page': page,
|
||
'page_end': page_end,
|
||
'page_range': page_range,
|
||
'section': meta.get('section', '') or meta.get('section_path', ''),
|
||
'chunk_type': meta.get('chunk_type', 'text'),
|
||
'doc_type': doc_type, # 文档类型
|
||
'section_chunk_id': meta.get('section_chunk_id'), # 章节内序号
|
||
'score': round(score, 3) if isinstance(score, float) else score
|
||
}
|
||
# ========== P1:图片使用 full_description ==========
|
||
# 图片切片使用完整描述(用于 LLM 上下文),而非短摘要
|
||
display_doc = doc
|
||
if meta.get('chunk_type') in ('image', 'chart', 'table'):
|
||
full_desc = meta.get('full_description', '')
|
||
if full_desc:
|
||
display_doc = full_desc
|
||
|
||
# contexts 仍然保留所有结果用于生成答案
|
||
contexts.append({'doc': display_doc, 'meta': meta, 'score': score})
|
||
|
||
sources = list(seen_sources.values())
|
||
|
||
# 调试事件:召回切片详情
|
||
if IS_DEV:
|
||
chunks_debug = []
|
||
for i, ctx in enumerate(contexts[:20]):
|
||
m = ctx.get('meta', {})
|
||
chunks_debug.append({
|
||
'rank': i + 1,
|
||
'source': m.get('source', ''),
|
||
'page': m.get('page', 0),
|
||
'chunk_type': m.get('chunk_type', 'text'),
|
||
'section': m.get('section', ''),
|
||
'score': round(ctx.get('score', 0), 4) if ctx.get('score') else None,
|
||
'content': (ctx.get('doc', '') or '')[:300]
|
||
})
|
||
yield f"data: {json.dumps({'type': 'chunks_retrieved', 'data': {'count': len(contexts), 'chunks': chunks_debug}}, ensure_ascii=False)}\n\n"
|
||
|
||
# 补充检索:从文本切片中提取图号/表号引用,补充检索对应的图片
|
||
# 重要:只从最相关的 top 5 文本切片提取引用,避免不相关引用干扰
|
||
import re
|
||
referenced_figures = set()
|
||
referenced_tables = set()
|
||
|
||
# 只检查 top 5 文本切片(与 select_images 逻辑一致)
|
||
text_contexts = [ctx for ctx in contexts if ctx.get('meta', {}).get('chunk_type') == 'text'][:5]
|
||
for ctx in text_contexts:
|
||
doc_text = ctx.get('doc', '')
|
||
fig_refs = re.findall(r'(?:[见如及和与])?图\s*(\d+\.?\d*)', doc_text)
|
||
referenced_figures.update(fig_refs)
|
||
table_refs = re.findall(r'(?:[见如及和与])?表\s*(\d+\.?\d*)', doc_text)
|
||
referenced_tables.update(table_refs)
|
||
|
||
# 检查哪些图号/表号对应的图片不在 contexts 中
|
||
existing_figure_images = set()
|
||
existing_table_images = set()
|
||
for ctx in contexts:
|
||
doc = ctx.get('doc', '')
|
||
meta = ctx.get('meta', {})
|
||
if meta.get('chunk_type') in ('image', 'chart'):
|
||
for fig_num in referenced_figures:
|
||
if f"图{fig_num}" in doc:
|
||
existing_figure_images.add(fig_num)
|
||
for table_num in referenced_tables:
|
||
if f"表{table_num}" in doc:
|
||
existing_table_images.add(table_num)
|
||
|
||
# 需要补充检索的图号/表号
|
||
missing_figures = referenced_figures - existing_figure_images
|
||
missing_tables = referenced_tables - existing_table_images
|
||
|
||
# 计算主要章节(用于补充检索过滤)
|
||
primary_sections_for_supplement = set()
|
||
for ctx in text_contexts[:3]:
|
||
section = ctx.get('meta', {}).get('section', '') or ctx.get('meta', {}).get('section_path', '')
|
||
if section:
|
||
section_num = re.search(r'(\d+\.\d+)', section)
|
||
if not section_num:
|
||
chapter_match = re.search(r'第\s*(\d+)\s*章', section)
|
||
if chapter_match:
|
||
section_num = chapter_match
|
||
if section_num:
|
||
primary_sections_for_supplement.add(section_num.group(1))
|
||
|
||
if missing_figures or missing_tables:
|
||
# 补充检索
|
||
from knowledge.manager import get_kb_manager
|
||
kb_manager = get_kb_manager()
|
||
kb_name = collections[0] if collections else 'public_kb'
|
||
collection = kb_manager.get_collection(kb_name)
|
||
|
||
if collection:
|
||
# 构建补充查询
|
||
supplement_queries = []
|
||
for fig_num in missing_figures:
|
||
supplement_queries.append(f"图{fig_num}")
|
||
for table_num in missing_tables:
|
||
supplement_queries.append(f"表{table_num}")
|
||
|
||
supplement_query = " ".join(supplement_queries)
|
||
|
||
# 使用 embedding 检索
|
||
# P4:复用 engine 的 embedding 模型,避免重复加载
|
||
try:
|
||
from core.engine import get_engine
|
||
engine = get_engine()
|
||
query_vector = engine.embedding_model.encode(supplement_query).tolist()
|
||
if isinstance(query_vector[0], list):
|
||
query_vector = query_vector[0]
|
||
|
||
supplement_result = collection.query(
|
||
query_embeddings=[query_vector],
|
||
n_results=10,
|
||
include=['documents', 'metadatas', 'distances']
|
||
)
|
||
|
||
# 添加匹配的图片切片
|
||
for supp_doc, supp_meta, supp_dist in zip(
|
||
supplement_result['documents'][0],
|
||
supplement_result['metadatas'][0],
|
||
supplement_result['distances'][0]
|
||
):
|
||
chunk_type = supp_meta.get('chunk_type', '')
|
||
if chunk_type in ('image', 'chart'):
|
||
# 检查是否匹配缺失的图号/表号
|
||
is_match = False
|
||
matched_fig = None
|
||
for fig_num in missing_figures:
|
||
if f"图{fig_num}" in supp_doc:
|
||
is_match = True
|
||
matched_fig = fig_num
|
||
break
|
||
for table_num in missing_tables:
|
||
if f"表{table_num}" in supp_doc:
|
||
is_match = True
|
||
break
|
||
|
||
if is_match:
|
||
# 额外检查:图片章节是否与主要章节匹配
|
||
# 避免补充检索到不相关的图片
|
||
supp_section = supp_meta.get('section', '') or supp_meta.get('section_path', '')
|
||
supp_section_num = re.search(r'(\d+\.\d+)', supp_section)
|
||
supp_section_id = supp_section_num.group(1) if supp_section_num else None
|
||
|
||
# 如果图片章节不在主要章节中,跳过
|
||
if primary_sections_for_supplement and supp_section_id and supp_section_id not in primary_sections_for_supplement:
|
||
# 不是主要章节的图片,跳过
|
||
continue
|
||
|
||
# Bug 6a 修复:补充检索的图片也要做 full_description 替换
|
||
# 与正常检索保持一致
|
||
display_doc = supp_doc
|
||
if supp_meta.get('chunk_type') in ('image', 'chart', 'table'):
|
||
full_desc = supp_meta.get('full_description', '')
|
||
if full_desc:
|
||
display_doc = full_desc
|
||
|
||
contexts.append({
|
||
'doc': display_doc,
|
||
'meta': {**supp_meta, '_collection': supp_meta.get('_collection') or (collections[0] if collections else 'public_kb')},
|
||
'score': 1.0 - supp_dist
|
||
})
|
||
logger.info(f"[补充检索] 添加图片: {supp_meta.get('image_path', '')}")
|
||
except Exception as e:
|
||
logger.warning(f"补充检索失败: {e}")
|
||
|
||
# 发送来源事件
|
||
yield f"data: {json.dumps({'type': 'sources', 'sources': sources[:MAX_SOURCES_RETURNED]}, ensure_ascii=False)}\n\n"
|
||
|
||
# 调试:检查 contexts 中是否有图片切片
|
||
image_count = sum(1 for ctx in contexts if ctx.get('meta', {}).get('chunk_type') in ('image', 'chart'))
|
||
if image_count > 0:
|
||
logger.info(f"[图片检索] contexts 中包含 {image_count} 个图片/图表切片")
|
||
for ctx in contexts:
|
||
meta = ctx.get('meta', {})
|
||
if meta.get('chunk_type') in ('image', 'chart'):
|
||
logging.info(f" - 图片: {meta.get('caption', '')[:50]}, path: {meta.get('image_path', '')}")
|
||
|
||
# 2.5. 懒加载增强(Phase 4)
|
||
# 暂时禁用:VLM 调用耗时过长,可能导致请求超时
|
||
# TODO: 后续可改为异步后台任务
|
||
try:
|
||
import asyncio
|
||
from knowledge.lazy_enhance import enhance_retrieved_chunks
|
||
kb_name = collections[0] if collections else 'public_kb'
|
||
asyncio.run(enhance_retrieved_chunks(contexts, message, kb_name))
|
||
except Exception as e:
|
||
logger.warning(f"懒加载增强失败: {e}")
|
||
|
||
# 3. 选择要展示的图片(Phase 5)
|
||
selected_images = select_images(contexts, message)
|
||
|
||
# 调试事件:图片选择详情
|
||
if IS_DEV:
|
||
yield f"data: {json.dumps({'type': 'images_selected', 'data': {'total_scored': len([c for c in contexts if c.get('meta',{}).get('chunk_type') in ('image','chart','table')]), 'selected_count': len(selected_images), 'images': [{'source': img.get('source',''), 'page': img.get('page',0), 'score': round(img.get('_image_boost', 1.0), 2)} for img in selected_images]}}, ensure_ascii=False)}\n\n"
|
||
|
||
# 4. 构建 prompt(Phase 6:LLM 图片感知)
|
||
# Bug 1 修复:文本切片用于 top 5 名额竞争,图片描述不参与竞争
|
||
text_contexts = _order_text_contexts_for_prompt(contexts, message, MAX_CONTEXT_CHUNKS,
|
||
min_score=RERANK_CONTEXT_MIN_SCORE)
|
||
# Phase 2:按字符预算构建上下文
|
||
_is_comparison = intent and intent.intent == "comparison"
|
||
if _is_enum_query(message) or _is_comparison:
|
||
# 列举类 / 对比类查询:保持原始顺序,不做预算截断
|
||
context_text = "\n\n".join([ctx.get('doc', '') for ctx in text_contexts])
|
||
else:
|
||
context_text = _build_context_with_budget(text_contexts, CONTEXT_MAX_CHARS, CONTEXT_SOFT_LIMIT)
|
||
|
||
# Phase 4:计算置信度分数(top-3 平均 Rerank 分数)
|
||
_top_scores = [ctx.get('score', 0) for ctx in text_contexts[:3]]
|
||
_confidence_score = round(sum(_top_scores) / len(_top_scores), 4) if _top_scores else 0.0
|
||
|
||
# Bug 6b 优化:直接使用 selected_images 中的 full_description
|
||
# 这样 LLM 既能看到文本切片,也能知道图片内容
|
||
if selected_images:
|
||
image_descriptions = []
|
||
for i, img in enumerate(selected_images, 1):
|
||
# 直接使用 select_images 时带上的 full_description
|
||
full_desc = img.get('full_description', '') or img.get('description', '')
|
||
if full_desc:
|
||
# 添加图片来源信息
|
||
img_source = img.get('source', '')
|
||
img_page = img.get('page', '')
|
||
source_info = f"(来源:{img_source} 第{img_page}页)" if img_source and img_page else ""
|
||
image_descriptions.append(f"【图片{i}】{full_desc}{source_info}")
|
||
if image_descriptions:
|
||
context_text += "\n\n【相关图片信息】\n" + "\n\n".join(image_descriptions)
|
||
# 添加指令让 LLM 介绍图片
|
||
context_text += "\n\n【回答要求】回答时请简要介绍每张图片的内容和用途。"
|
||
|
||
enhanced_context = context_text
|
||
|
||
# 对比类查询:添加结构化对比指令
|
||
if intent and intent.intent == "comparison" and intent.sub_queries:
|
||
comparison_instruction = (
|
||
"\n\n【回答要求】这是一个对比类问题。"
|
||
"请根据参考资料,从多个角度对比分析,"
|
||
"使用表格或分点形式清晰呈现差异和共同点。"
|
||
"如果参考资料中缺少某一方面的信息,请如实说明。"
|
||
)
|
||
enhanced_context = comparison_instruction + "\n\n" + enhanced_context
|
||
|
||
# 推理类查询:添加因果分析指令
|
||
elif intent and intent.intent == "reasoning":
|
||
reasoning_instruction = (
|
||
"\n\n【回答要求】这是一个需要分析原因或推理的问题。"
|
||
"请根据参考资料,先梳理相关事实和数据,再给出逻辑清晰的分析。"
|
||
"如果涉及因果关系,请明确标注原因和结果;如果资料不足以支撑推理,请如实说明。"
|
||
)
|
||
enhanced_context = reasoning_instruction + "\n\n" + enhanced_context
|
||
|
||
# 操作指导类查询:添加步骤化指令
|
||
elif intent and intent.intent == "instruction":
|
||
instruction_instruction = (
|
||
"\n\n【回答要求】这是一个操作指导类问题。"
|
||
"请根据参考资料,以清晰的步骤或流程形式组织回答。"
|
||
"如有前置条件或注意事项,请在步骤前说明。"
|
||
)
|
||
enhanced_context = instruction_instruction + "\n\n" + enhanced_context
|
||
|
||
if _is_enum_query(message):
|
||
enum_instruction = (
|
||
"\n\n【回答要求】如果参考资料中包含编号列表、禁止情形、要求或条款,"
|
||
"请按资料中的原始顺序完整列出;不要合并相邻条目,不要跳项,"
|
||
"资料不足时明确说明缺少哪部分依据。\n\n"
|
||
)
|
||
enhanced_context = enum_instruction + enhanced_context
|
||
|
||
# Phase 4:根据置信度注入谨慎回答指令
|
||
if _confidence_score < CONFIDENCE_WARN_THRESHOLD:
|
||
enhanced_context += (
|
||
"\n\n【重要提示】参考资料与问题的相关性较低。"
|
||
"请仅基于参考资料中明确包含的信息回答,"
|
||
'如果资料不足以回答问题,请直接说明"知识库中未找到直接相关的信息"。'
|
||
)
|
||
elif _confidence_score < CONFIDENCE_CAUTION_THRESHOLD:
|
||
enhanced_context += (
|
||
"\n\n【提示】参考资料的相关性一般,请优先引用资料中的原文,避免推测。"
|
||
)
|
||
|
||
# 调试事件:最终上下文
|
||
if IS_DEV:
|
||
text_used = text_contexts
|
||
_scores = [round(ctx.get('score', 0), 4) for ctx in text_used]
|
||
yield f"data: {json.dumps({'type': 'context_built', 'data': {'chunk_count': len(text_used), 'context_length': len(enhanced_context), 'budget_max_chars': CONTEXT_MAX_CHARS, 'min_score_filter': RERANK_CONTEXT_MIN_SCORE, 'confidence_top3': _confidence_score, 'score_stats': {'max': max(_scores) if _scores else 0, 'min': min(_scores) if _scores else 0, 'avg': round(sum(_scores)/len(_scores), 4) if _scores else 0}, 'context_preview': enhanced_context[:500], 'chunks_used': [{'source': ctx.get('meta',{}).get('source',''), 'page': ctx.get('meta',{}).get('page',0), 'score': ctx.get('score',0), 'preview': (ctx.get('doc','') or '')[:100]} for ctx in text_used]}}, ensure_ascii=False)}\n\n"
|
||
|
||
# 5. 流式生成回答
|
||
from core.engine import get_engine
|
||
engine = get_engine()
|
||
|
||
for token in engine.generate_answer_stream(message, enhanced_context, history):
|
||
full_answer.append(token)
|
||
yield f"data: {json.dumps({'type': 'chunk', 'content': token}, ensure_ascii=False)}\n\n"
|
||
|
||
# 6. P0:答案对齐过滤器
|
||
# 从 LLM 回答中提取图号引用,过滤图片选择结果
|
||
full_answer_text = "".join(full_answer)
|
||
|
||
# 提取回答中引用的图号/表号
|
||
mentioned = set()
|
||
# 中文图号:图2.1、图 2-1、见图2.1 等
|
||
mentioned.update(re.findall(r'(?:[见如])?图\s*(\d+[\.\-]?\d*)', full_answer_text))
|
||
# 中文表号:表2.1、表 2-1、见表2.1 等
|
||
mentioned.update(re.findall(r'(?:[见如])?表\s*(\d+[\.\-]?\d*)', full_answer_text))
|
||
# 英文图号:Figure 2.1、Fig.2.1 等
|
||
mentioned.update(re.findall(r'(?:Fig(?:ure)?\.?\s*)(\d+[\.\-]?\d*)', full_answer_text, re.I))
|
||
|
||
# 根据回答中的引用过滤图片
|
||
if mentioned:
|
||
aligned_images = []
|
||
for img in selected_images:
|
||
desc = img.get('description', '')
|
||
# 标准化图号格式(将连字符转为点)
|
||
for ref in mentioned:
|
||
ref_normalized = ref.replace('-', '.')
|
||
if (f"图{ref_normalized}" in desc or
|
||
f"表{ref_normalized}" in desc or
|
||
f"图 {ref_normalized}" in desc or
|
||
f"表 {ref_normalized}" in desc):
|
||
aligned_images.append(img)
|
||
break
|
||
# 如果有匹配的图片,使用对齐后的结果
|
||
if aligned_images:
|
||
selected_images = aligned_images
|
||
# else: 没有匹配到,保留原选择(不再截断到1张)
|
||
# else: LLM 没有提图号,保留原选择(不再截断到1张)
|
||
|
||
rich_media = {'images': selected_images, 'tables': [], 'sections': []}
|
||
|
||
# 7. 去掉 LLM 添加的数字引用标记,避免与后端引用重复
|
||
clean_answer = re.sub(r'\[\d+\]', '', full_answer_text)
|
||
|
||
# 8. 添加引用标注(自动插入 [ref:chunk_id])
|
||
citation_result = _attach_citations(clean_answer, contexts)
|
||
|
||
# 9. 过滤敏感信息(违禁词等)
|
||
filtered_answer = filter_response(citation_result.get("answer_with_refs", clean_answer))
|
||
|
||
# 9. 保存消息到会话(仅开发环境)
|
||
if session_id and session_repo_ref:
|
||
try:
|
||
# 保存用户消息
|
||
session_repo_ref.add_message(session_id, 'user', message)
|
||
# 保存 AI 回答(包含完整 metadata:图片、来源、引用等)
|
||
assistant_metadata = {
|
||
'is_rag': True,
|
||
'mode': 'rag',
|
||
}
|
||
if rich_media.get('images'):
|
||
assistant_metadata['images'] = rich_media['images']
|
||
if sources:
|
||
assistant_metadata['sources'] = sources
|
||
if citation_result.get('citations'):
|
||
assistant_metadata['citations'] = citation_result['citations']
|
||
session_repo_ref.add_message(session_id, 'assistant', filtered_answer, assistant_metadata)
|
||
# 更新会话最后活跃时间
|
||
if hasattr(session_repo_ref, 'update_last_active'):
|
||
session_repo_ref.update_last_active(session_id)
|
||
except Exception as e:
|
||
logger.warning(f"保存会话消息失败: {e}")
|
||
|
||
# 10. 发送完成事件
|
||
duration_ms = int((_time.time() - start_time) * 1000)
|
||
finish_event = {
|
||
"type": "finish",
|
||
"answer": filtered_answer,
|
||
"mode": "rag",
|
||
"session_id": session_id,
|
||
"sources": sources,
|
||
"citations": citation_result.get("citations", []), # 结构化引用列表
|
||
"images": rich_media["images"],
|
||
"tables": rich_media["tables"],
|
||
"sections": rich_media["sections"],
|
||
"duration_ms": duration_ms,
|
||
"confidence_score": _confidence_score # Phase 4:top-3 平均 Rerank 分数
|
||
}
|
||
|
||
# 添加分阶段耗时信息(仅开发环境)
|
||
if IS_DEV and hasattr(search_result, 'get'):
|
||
debug_info = search_result.get('_debug', {})
|
||
timing_info = debug_info.get('timing', {})
|
||
# 从 _debug steps 中提取 Rerank 耗时
|
||
rerank_time = 0
|
||
rerank_cached = False
|
||
for step in debug_info.get('steps', []):
|
||
if step.get('name') == 'rerank' and step.get('applied'):
|
||
rerank_time = step.get('time_ms', 0)
|
||
rerank_cached = step.get('cached', False)
|
||
finish_event["timing"] = {
|
||
"total_search_ms": timing_info.get('total_ms', 0),
|
||
"rerank_ms": rerank_time,
|
||
"rerank_cached": rerank_cached,
|
||
"total_ms": duration_ms
|
||
}
|
||
# 10.5 写入语义缓存
|
||
try:
|
||
from core.semantic_cache import get_semantic_cache
|
||
_sc = get_semantic_cache()
|
||
if _sc and filtered_answer:
|
||
if _semantic_cache_emb is None:
|
||
from core.engine import get_engine as _get_eng
|
||
_eng = _get_eng()
|
||
if _eng and hasattr(_eng, 'embedding_model'):
|
||
_semantic_cache_emb = _eng.embedding_model.encode(message)
|
||
if _semantic_cache_emb is not None:
|
||
_sc.set(_semantic_cache_emb, {
|
||
"answer": filtered_answer,
|
||
"sources": sources,
|
||
"citations": citation_result.get("citations", []),
|
||
"images": rich_media.get("images", []),
|
||
"tables": rich_media.get("tables", []),
|
||
})
|
||
logger.debug(f"[语义缓存] 写入成功: {message[:50]}...")
|
||
except Exception as e:
|
||
logger.info(f"[语义缓存] 写入失败: {e}")
|
||
|
||
yield f"data: {json.dumps(finish_event, ensure_ascii=False)}\n\n"
|
||
|
||
except Exception as e:
|
||
import logging as _logging
|
||
_logging.getLogger(__name__).error(f"[SSE] RAG 流异常: {e}", exc_info=True)
|
||
error_event = {
|
||
"type": "error",
|
||
"message": "服务内部错误,请稍后重试"
|
||
}
|
||
yield f"data: {json.dumps(error_event, ensure_ascii=False)}\n\n"
|
||
|
||
return Response(
|
||
generate(),
|
||
mimetype='text/event-stream',
|
||
headers={
|
||
'Cache-Control': 'no-cache',
|
||
'X-Accel-Buffering': 'no'
|
||
}
|
||
)
|
||
|
||
|
||
@chat_bp.route('/search', methods=['POST'])
|
||
@require_gateway_auth
|
||
def search():
|
||
"""
|
||
混合检索接口 - 供 Dify 工作流调用
|
||
|
||
请求体:
|
||
{
|
||
"query": "查询文本",
|
||
"top_k": 5,
|
||
"collections": ["public_kb"] // 可选
|
||
}
|
||
"""
|
||
data = request.json or {}
|
||
query = data.get('query', '')
|
||
query = sanitize_user_input(query)
|
||
injection_matches = detect_injection(query)
|
||
if injection_matches:
|
||
logger.warning(f"[Chat] 检测到可疑注入: {injection_matches}")
|
||
top_k = data.get('top_k', 5)
|
||
collections = data.get('collections') # 后端传入的知识库列表
|
||
|
||
if not query:
|
||
return error_response("MISSING_PARAMS", BAD_REQUEST, "query is required", http_status=400)
|
||
|
||
# 输入安全校验(注入检测、违禁词、长度限制)
|
||
is_valid, reason = validate_query(query)
|
||
if not is_valid:
|
||
return error_response("INVALID_QUERY", BAD_REQUEST, reason, http_status=400)
|
||
|
||
# top_k 范围校验
|
||
try:
|
||
top_k = max(1, min(int(top_k), 50))
|
||
except (ValueError, TypeError):
|
||
top_k = 5
|
||
|
||
# 如果没有指定 collections,使用默认的公开库
|
||
if not collections:
|
||
collections = ['public_kb']
|
||
|
||
results = search_hybrid(query, top_k=top_k, allowed_collections=collections)
|
||
|
||
return success_response(data={
|
||
'contexts': results['documents'][0],
|
||
'metadatas': results['metadatas'][0],
|
||
'scores': results['scores'][0]
|
||
})
|
||
|
||
|
||
# ==================== 缓存调试接口(临时) ====================
|
||
|
||
@chat_bp.route('/cache/stats', methods=['GET'])
|
||
def cache_stats():
|
||
"""缓存统计接口(调试用)"""
|
||
stats = {}
|
||
try:
|
||
from core.cache import get_cache_manager
|
||
cache = get_cache_manager()
|
||
all_stats = cache.get_all_stats()
|
||
for name, s in all_stats.items():
|
||
stats[name] = {
|
||
'total_entries': s.total_entries,
|
||
'hits': s.hits,
|
||
'misses': s.misses,
|
||
'hit_rate': f"{s.hit_rate:.2%}",
|
||
'evictions': s.evictions,
|
||
}
|
||
except Exception as e:
|
||
stats['_error'] = str(e)
|
||
|
||
# 语义缓存统计(全局单例 + 意图分析器实例)
|
||
try:
|
||
from core.semantic_cache import get_semantic_cache
|
||
sc = get_semantic_cache()
|
||
if sc:
|
||
stats['semantic_cache'] = sc.get_stats()
|
||
except Exception:
|
||
pass
|
||
|
||
try:
|
||
from core.intent_analyzer import IntentAnalyzer
|
||
if hasattr(IntentAnalyzer, '_instance'):
|
||
ia = IntentAnalyzer._instance
|
||
if hasattr(ia, 'semantic_cache') and ia.semantic_cache:
|
||
stats['semantic_cache_intent'] = ia.semantic_cache.get_stats()
|
||
except Exception:
|
||
pass
|
||
|
||
return jsonify(stats)
|
||
|
||
|
||
@chat_bp.route('/cache/clear', methods=['POST'])
|
||
def cache_clear():
|
||
"""缓存清空接口(调试用)"""
|
||
result = {}
|
||
try:
|
||
from core.cache import get_cache_manager
|
||
cache = get_cache_manager()
|
||
cache.clear_all()
|
||
result['lru_cache'] = 'cleared'
|
||
except Exception as e:
|
||
result['lru_cache'] = f'error: {e}'
|
||
|
||
try:
|
||
from core.semantic_cache import get_semantic_cache
|
||
sc = get_semantic_cache()
|
||
sc.clear()
|
||
result['semantic_cache'] = 'cleared'
|
||
except Exception:
|
||
result['semantic_cache'] = 'not_available'
|
||
|
||
return jsonify({'status': 'ok', 'cleared': result})
|