Compare commits
23 Commits
server-bas
...
c6a17ad1e4
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c6a17ad1e4 | ||
|
|
0f339a0998 | ||
|
|
21470ed28c | ||
|
|
258be54df7 | ||
|
|
d6d27b70f3 | ||
|
|
184511f4ec | ||
|
|
b5bcb3c941 | ||
|
|
5bc86a7c1f | ||
|
|
8f72ac5da4 | ||
|
|
3a3627f73c | ||
|
|
2c84cb8a5f | ||
|
|
8d60616124 | ||
|
|
8f75c51c59 | ||
|
|
87f3fae1aa | ||
|
|
fc11b11dda | ||
|
|
99f5cf519e | ||
|
|
15c0aec9a6 | ||
|
|
8c7a6eb3fa | ||
|
|
fda1b2f049 | ||
|
|
148559ee3c | ||
|
|
431af0217a | ||
|
|
aa04bb94a6 | ||
|
|
84a8be0ce8 |
23
.dockerignore
Normal file
23
.dockerignore
Normal file
@@ -0,0 +1,23 @@
|
|||||||
|
# 大体积数据目录(通过 volume 挂载,不需要打进镜像)
|
||||||
|
models/
|
||||||
|
knowledge/vector_store/
|
||||||
|
documents/
|
||||||
|
.data/
|
||||||
|
data/
|
||||||
|
|
||||||
|
# Python 虚拟环境
|
||||||
|
venv/
|
||||||
|
__pycache__/
|
||||||
|
*.pyc
|
||||||
|
|
||||||
|
# Git
|
||||||
|
.git/
|
||||||
|
|
||||||
|
# IDE 和编辑器
|
||||||
|
.vscode/
|
||||||
|
.idea/
|
||||||
|
*.swp
|
||||||
|
|
||||||
|
# 其他
|
||||||
|
*.log
|
||||||
|
.env*
|
||||||
1041
api/chat_routes.py
1041
api/chat_routes.py
File diff suppressed because it is too large
Load Diff
@@ -775,7 +775,22 @@ def delete_document(doc_path: str) -> Tuple[Any, int]:
|
|||||||
if kb_manager:
|
if kb_manager:
|
||||||
kb_manager.delete_document(collection, filename)
|
kb_manager.delete_document(collection, filename)
|
||||||
|
|
||||||
# 2. 删除文件
|
# 2. 缓存失效
|
||||||
|
try:
|
||||||
|
from core.cache import get_cache_manager
|
||||||
|
_cm = get_cache_manager()
|
||||||
|
_cm.increment_kb_version(collection)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
try:
|
||||||
|
from core.semantic_cache import get_semantic_cache
|
||||||
|
_sc = get_semantic_cache()
|
||||||
|
if _sc:
|
||||||
|
_sc.clear()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
# 3. 删除文件
|
||||||
os.remove(filepath)
|
os.remove(filepath)
|
||||||
|
|
||||||
return jsonify({
|
return jsonify({
|
||||||
|
|||||||
@@ -260,6 +260,21 @@ def delete_collection(kb_name: str) -> Tuple[Any, int]:
|
|||||||
success, message = kb_manager.delete_collection(kb_name, delete_documents)
|
success, message = kb_manager.delete_collection(kb_name, delete_documents)
|
||||||
|
|
||||||
if success:
|
if success:
|
||||||
|
# 缓存失效
|
||||||
|
try:
|
||||||
|
from core.cache import get_cache_manager
|
||||||
|
_cm = get_cache_manager()
|
||||||
|
_cm.increment_kb_version(kb_name)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
try:
|
||||||
|
from core.semantic_cache import get_semantic_cache
|
||||||
|
_sc = get_semantic_cache()
|
||||||
|
if _sc:
|
||||||
|
_sc.clear()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
return jsonify({
|
return jsonify({
|
||||||
"success": True,
|
"success": True,
|
||||||
"message": message,
|
"message": message,
|
||||||
|
|||||||
207
cleanup_orphans.py
Normal file
207
cleanup_orphans.py
Normal file
@@ -0,0 +1,207 @@
|
|||||||
|
"""
|
||||||
|
孤儿文件清理工具
|
||||||
|
|
||||||
|
清理 .data/images/ 和 .data/cache/vlm/ 中不再被任何 ChromaDB 切片引用的文件。
|
||||||
|
|
||||||
|
用法:
|
||||||
|
# 预览模式(不删除,只显示孤儿文件)
|
||||||
|
python cleanup_orphans.py
|
||||||
|
|
||||||
|
# 实际删除
|
||||||
|
python cleanup_orphans.py --force
|
||||||
|
|
||||||
|
# 仅清理特定知识库
|
||||||
|
python cleanup_orphans.py --collections public_kb test_kb
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import hashlib
|
||||||
|
import logging
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
IMAGES_DIR = Path(".data/images")
|
||||||
|
VLM_CACHE_DIR = Path(".data/cache/vlm")
|
||||||
|
|
||||||
|
|
||||||
|
def compute_file_hash(file_path: str) -> str:
|
||||||
|
"""计算文件 MD5"""
|
||||||
|
with open(file_path, 'rb') as f:
|
||||||
|
return hashlib.md5(f.read()).hexdigest()
|
||||||
|
|
||||||
|
|
||||||
|
def collect_referenced_images(manager, collections=None) -> dict:
|
||||||
|
"""
|
||||||
|
从 ChromaDB 收集所有被引用的图片路径。
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
{image_filename: set of chunk_ids referencing it}
|
||||||
|
"""
|
||||||
|
referenced = {}
|
||||||
|
|
||||||
|
if collections:
|
||||||
|
kb_names = collections
|
||||||
|
else:
|
||||||
|
kb_names = [c.name if hasattr(c, 'name') else str(c)
|
||||||
|
for c in manager.list_collections()]
|
||||||
|
|
||||||
|
for kb_name in kb_names:
|
||||||
|
try:
|
||||||
|
col = manager.get_collection(kb_name)
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"无法获取 {kb_name}: {e}")
|
||||||
|
continue
|
||||||
|
|
||||||
|
# 查找所有带 image_path 的切片
|
||||||
|
results = col.get(include=['metadatas'])
|
||||||
|
if not results['ids']:
|
||||||
|
continue
|
||||||
|
|
||||||
|
for chunk_id, meta in zip(results['ids'], results['metadatas']):
|
||||||
|
image_path = meta.get('image_path', '')
|
||||||
|
if not image_path:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# image_path 可能是: "185a7a75d246.png" 或相对路径
|
||||||
|
filename = os.path.basename(image_path)
|
||||||
|
if filename not in referenced:
|
||||||
|
referenced[filename] = set()
|
||||||
|
referenced[filename].add(f"{kb_name}/{chunk_id}")
|
||||||
|
|
||||||
|
return referenced
|
||||||
|
|
||||||
|
|
||||||
|
def find_orphan_images(referenced: dict) -> list:
|
||||||
|
"""
|
||||||
|
查找 .data/images/ 中不再被引用的图片文件。
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
[(filepath, filename, size_bytes)]
|
||||||
|
"""
|
||||||
|
orphans = []
|
||||||
|
if not IMAGES_DIR.exists():
|
||||||
|
return orphans
|
||||||
|
|
||||||
|
for f in IMAGES_DIR.iterdir():
|
||||||
|
if not f.is_file():
|
||||||
|
continue
|
||||||
|
if f.name not in referenced:
|
||||||
|
orphans.append((str(f), f.name, f.stat().st_size))
|
||||||
|
|
||||||
|
return orphans
|
||||||
|
|
||||||
|
|
||||||
|
def find_orphan_vlm_caches(referenced: dict) -> list:
|
||||||
|
"""
|
||||||
|
查找 .data/cache/vlm/ 中对应的图片已不存在的缓存文件。
|
||||||
|
|
||||||
|
缓存文件以图片 MD5 命名,如果图片被删了,缓存也应该是孤儿。
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
[(filepath, filename, size_bytes)]
|
||||||
|
"""
|
||||||
|
orphans = []
|
||||||
|
if not VLM_CACHE_DIR.exists():
|
||||||
|
return orphans
|
||||||
|
|
||||||
|
# 构建 referenced 中所有图片的 MD5 集合
|
||||||
|
referenced_hashes = set()
|
||||||
|
for filename in referenced:
|
||||||
|
full_path = IMAGES_DIR / filename
|
||||||
|
if full_path.exists():
|
||||||
|
try:
|
||||||
|
img_hash = compute_file_hash(str(full_path))
|
||||||
|
referenced_hashes.add(img_hash)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
for f in VLM_CACHE_DIR.iterdir():
|
||||||
|
if not f.is_file() or not f.suffix == '.txt':
|
||||||
|
continue
|
||||||
|
cache_hash = f.stem # 文件名就是 MD5
|
||||||
|
if cache_hash not in referenced_hashes:
|
||||||
|
orphans.append((str(f), f.name, f.stat().st_size))
|
||||||
|
|
||||||
|
return orphans
|
||||||
|
|
||||||
|
|
||||||
|
def delete_files(file_list: list, dry_run: bool = True) -> int:
|
||||||
|
"""删除文件列表,返回删除数量"""
|
||||||
|
deleted = 0
|
||||||
|
for filepath, filename, size in file_list:
|
||||||
|
if dry_run:
|
||||||
|
logger.info(f" [DRY-RUN] 将删除: {filename} ({size/1024:.1f} KB)")
|
||||||
|
else:
|
||||||
|
try:
|
||||||
|
os.remove(filepath)
|
||||||
|
logger.info(f" 已删除: {filename} ({size/1024:.1f} KB)")
|
||||||
|
deleted += 1
|
||||||
|
except OSError as e:
|
||||||
|
logger.warning(f" 删除失败: {filename} - {e}")
|
||||||
|
return deleted
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(description="孤儿文件清理工具")
|
||||||
|
parser.add_argument("--force", action="store_true",
|
||||||
|
help="实际删除文件(默认仅预览)")
|
||||||
|
parser.add_argument("--collections", nargs="+",
|
||||||
|
help="仅检查指定知识库(默认检查全部)")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
os.chdir(os.path.dirname(os.path.abspath(__file__)))
|
||||||
|
|
||||||
|
mode = "删除" if args.force else "预览(不删除)"
|
||||||
|
print("=" * 60)
|
||||||
|
print(f"孤儿文件清理 - {mode}")
|
||||||
|
print("=" * 60)
|
||||||
|
|
||||||
|
# 1. 收集引用
|
||||||
|
print("\n[1/4] 扫描 ChromaDB 引用...")
|
||||||
|
sys.path.insert(0, '.')
|
||||||
|
from knowledge.manager import get_kb_manager
|
||||||
|
manager = get_kb_manager()
|
||||||
|
referenced = collect_referenced_images(manager, args.collections)
|
||||||
|
print(f" 被引用的图片: {len(referenced)} 个")
|
||||||
|
|
||||||
|
# 2. 查找孤儿图片
|
||||||
|
print("\n[2/4] 查找孤儿图片...")
|
||||||
|
orphan_images = find_orphan_images(referenced)
|
||||||
|
total_img_size = sum(s for _, _, s in orphan_images)
|
||||||
|
print(f" 孤儿图片: {len(orphan_images)} 个 ({total_img_size/1024:.1f} KB)")
|
||||||
|
|
||||||
|
# 3. 查找孤儿 VLM 缓存
|
||||||
|
print("\n[3/4] 查找孤儿 VLM 缓存...")
|
||||||
|
orphan_caches = find_orphan_vlm_caches(referenced)
|
||||||
|
total_cache_size = sum(s for _, _, s in orphan_caches)
|
||||||
|
print(f" 孤儿缓存: {len(orphan_caches)} 个 ({total_cache_size/1024:.1f} KB)")
|
||||||
|
|
||||||
|
# 4. 清理
|
||||||
|
print("\n[4/4] 清理...")
|
||||||
|
if not orphan_images and not orphan_caches:
|
||||||
|
print(" 没有需要清理的文件")
|
||||||
|
else:
|
||||||
|
if not args.force:
|
||||||
|
print(" 预览模式,以下文件将被删除:")
|
||||||
|
|
||||||
|
img_deleted = delete_files(orphan_images, dry_run=not args.force)
|
||||||
|
cache_deleted = delete_files(orphan_caches, dry_run=not args.force)
|
||||||
|
|
||||||
|
if args.force:
|
||||||
|
print(f"\n 删除了 {img_deleted} 个图片 + {cache_deleted} 个缓存")
|
||||||
|
freed = total_img_size + total_cache_size
|
||||||
|
print(f" 释放空间: {freed/1024:.1f} KB")
|
||||||
|
else:
|
||||||
|
print(f"\n 共 {len(orphan_images) + len(orphan_caches)} 个文件待清理")
|
||||||
|
print(f" 释放空间: {(total_img_size + total_cache_size)/1024:.1f} KB")
|
||||||
|
print(" 使用 --force 参数执行实际删除")
|
||||||
|
|
||||||
|
print("=" * 60)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
main()
|
||||||
@@ -109,6 +109,11 @@ USE_RERANK = True
|
|||||||
RERANK_CANDIDATES = 20 # 送入重排序的候选数
|
RERANK_CANDIDATES = 20 # 送入重排序的候选数
|
||||||
RERANK_TOP_K = 15 # 重排序后保留数
|
RERANK_TOP_K = 15 # 重排序后保留数
|
||||||
RERANK_USE_ONNX = os.getenv("RERANK_USE_ONNX", "true").lower() == "true"
|
RERANK_USE_ONNX = os.getenv("RERANK_USE_ONNX", "true").lower() == "true"
|
||||||
|
RERANK_BACKEND = os.getenv("RERANK_BACKEND", "local") # local / cloud / fallback
|
||||||
|
RERANK_CLOUD_MODEL = os.getenv("RERANK_CLOUD_MODEL", "qwen3-rerank")
|
||||||
|
RERANK_CLOUD_API_KEY = os.getenv("RERANK_CLOUD_API_KEY", "")
|
||||||
|
RERANK_CLOUD_BASE_URL = os.getenv("RERANK_CLOUD_BASE_URL", "https://dashscope.aliyuncs.com/compatible-api/v1/reranks")
|
||||||
|
RERANK_CLOUD_TIMEOUT = int(os.getenv("RERANK_CLOUD_TIMEOUT", "15"))
|
||||||
RERANK_CONTEXT_MIN_SCORE = 0.05 # Rerank 分数低于此值的切片不送入 LLM
|
RERANK_CONTEXT_MIN_SCORE = 0.05 # Rerank 分数低于此值的切片不送入 LLM
|
||||||
|
|
||||||
# ----- RRF 融合 -----
|
# ----- RRF 融合 -----
|
||||||
|
|||||||
@@ -31,17 +31,11 @@ from core.llm_utils import call_llm, quick_yes_no, parse_json_from_response
|
|||||||
from .agentic_base import (
|
from .agentic_base import (
|
||||||
API_KEY, BASE_URL, MODEL,
|
API_KEY, BASE_URL, MODEL,
|
||||||
HAS_SERPER,
|
HAS_SERPER,
|
||||||
HAS_BUDGET, SEMANTIC_CACHE_ENABLED,
|
HAS_BUDGET,
|
||||||
MAX_CONTEXT_TOKENS, MAX_CONTEXT_COUNT, RERANK_THRESHOLD,
|
MAX_CONTEXT_TOKENS, MAX_CONTEXT_COUNT, RERANK_THRESHOLD,
|
||||||
SOURCE_KB, SOURCE_WEB,
|
SOURCE_KB, SOURCE_WEB,
|
||||||
)
|
)
|
||||||
|
|
||||||
# 尝试导入语义缓存
|
|
||||||
try:
|
|
||||||
from core.semantic_cache import SemanticCache
|
|
||||||
except ImportError:
|
|
||||||
SemanticCache = None
|
|
||||||
|
|
||||||
# 导入 Mixin 类
|
# 导入 Mixin 类
|
||||||
from .agentic_query import QueryRewriteMixin
|
from .agentic_query import QueryRewriteMixin
|
||||||
from .agentic_search import SearchMixin
|
from .agentic_search import SearchMixin
|
||||||
@@ -120,30 +114,6 @@ class AgenticRAG(
|
|||||||
except ImportError:
|
except ImportError:
|
||||||
self.loop_guard = None
|
self.loop_guard = None
|
||||||
|
|
||||||
# 初始化语义缓存
|
|
||||||
self.semantic_cache = None
|
|
||||||
self.embedding_model = None
|
|
||||||
if SEMANTIC_CACHE_ENABLED and SemanticCache:
|
|
||||||
try:
|
|
||||||
engine = get_engine()
|
|
||||||
if engine and hasattr(engine, 'embedding_model'):
|
|
||||||
self.embedding_model = engine.embedding_model
|
|
||||||
emb_dim = 768
|
|
||||||
# 优先使用新 API,兼容旧版本
|
|
||||||
if hasattr(self.embedding_model, 'get_embedding_dimension'):
|
|
||||||
emb_dim = self.embedding_model.get_embedding_dimension()
|
|
||||||
elif hasattr(self.embedding_model, 'get_sentence_embedding_dimension'):
|
|
||||||
emb_dim = self.embedding_model.get_sentence_embedding_dimension()
|
|
||||||
self.semantic_cache = SemanticCache(
|
|
||||||
dim=emb_dim,
|
|
||||||
threshold=0.92,
|
|
||||||
max_size=5000
|
|
||||||
)
|
|
||||||
logger.info(f"语义缓存已启用,维度={emb_dim}")
|
|
||||||
except Exception as e:
|
|
||||||
logger.warning(f"语义缓存初始化失败: {e}")
|
|
||||||
self.semantic_cache = None
|
|
||||||
|
|
||||||
# Context Compression 配置
|
# Context Compression 配置
|
||||||
self.MAX_CONTEXT_TOKENS = MAX_CONTEXT_TOKENS
|
self.MAX_CONTEXT_TOKENS = MAX_CONTEXT_TOKENS
|
||||||
self.MAX_CONTEXT_COUNT = MAX_CONTEXT_COUNT
|
self.MAX_CONTEXT_COUNT = MAX_CONTEXT_COUNT
|
||||||
|
|||||||
@@ -25,22 +25,6 @@ except ImportError:
|
|||||||
HAS_BUDGET = False
|
HAS_BUDGET = False
|
||||||
CallType = None
|
CallType = None
|
||||||
|
|
||||||
# 语义缓存
|
|
||||||
try:
|
|
||||||
from config import SEMANTIC_CACHE_ENABLED, SEMANTIC_CACHE_THRESHOLD
|
|
||||||
HAS_SEMANTIC_CACHE_CONFIG = True
|
|
||||||
except ImportError:
|
|
||||||
SEMANTIC_CACHE_ENABLED = False
|
|
||||||
SEMANTIC_CACHE_THRESHOLD = 0.92
|
|
||||||
HAS_SEMANTIC_CACHE_CONFIG = False
|
|
||||||
|
|
||||||
try:
|
|
||||||
from core.semantic_cache import SemanticCache, get_semantic_cache
|
|
||||||
HAS_SEMANTIC_CACHE = True
|
|
||||||
except ImportError:
|
|
||||||
HAS_SEMANTIC_CACHE = False
|
|
||||||
SemanticCache = None
|
|
||||||
|
|
||||||
# LLM 配置
|
# LLM 配置
|
||||||
try:
|
try:
|
||||||
from config import API_KEY, BASE_URL, MODEL
|
from config import API_KEY, BASE_URL, MODEL
|
||||||
|
|||||||
@@ -189,6 +189,8 @@ class RAGCacheManager:
|
|||||||
# 失效旧版本缓存
|
# 失效旧版本缓存
|
||||||
self.query_cache.invalidate_by_version(old_version)
|
self.query_cache.invalidate_by_version(old_version)
|
||||||
self.embedding_cache.invalidate_by_version(old_version)
|
self.embedding_cache.invalidate_by_version(old_version)
|
||||||
|
# Rerank cache 无 kb_version 字段,文档变更时全量清空以防过时分数
|
||||||
|
self.rerank_cache.clear()
|
||||||
|
|
||||||
logger.info(f"知识库 {kb_name} 版本更新: {old_version} -> {new_version}")
|
logger.info(f"知识库 {kb_name} 版本更新: {old_version} -> {new_version}")
|
||||||
return new_version
|
return new_version
|
||||||
@@ -196,90 +198,41 @@ class RAGCacheManager:
|
|||||||
# ==================== Query Cache 方法 ====================
|
# ==================== Query Cache 方法 ====================
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _make_query_cache_key(query: str, kb_name: str, kb_version: int, doc_hash: str = "") -> str:
|
def _make_query_cache_key(query: str, kb_name: str, kb_version: int) -> str:
|
||||||
"""
|
"""
|
||||||
生成查询缓存 key
|
生成查询缓存 key(基于 kb_version 的粗粒度失效)
|
||||||
|
|
||||||
Args:
|
|
||||||
query: 查询文本
|
|
||||||
kb_name: 知识库名称
|
|
||||||
kb_version: 知识库版本号
|
|
||||||
doc_hash: 相关文档版本哈希(细粒度失效)
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
缓存 key
|
|
||||||
"""
|
"""
|
||||||
if doc_hash:
|
|
||||||
# 细粒度:只失效相关文档的缓存
|
|
||||||
return hashlib.md5(
|
|
||||||
f"query:{query}:{kb_name}:{doc_hash}".encode()
|
|
||||||
).hexdigest()
|
|
||||||
else:
|
|
||||||
# 粗粒度:整个知识库版本变化时失效
|
|
||||||
return hashlib.md5(
|
return hashlib.md5(
|
||||||
f"query:{query}:{kb_name}:{kb_version}".encode()
|
f"query:{query}:{kb_name}:{kb_version}".encode()
|
||||||
).hexdigest()
|
).hexdigest()
|
||||||
|
|
||||||
def get_query_result(self, query: str, kb_name: str, doc_ids: List[str] = None) -> Optional[Dict]:
|
def get_query_result(self, query: str, kb_name: str) -> Optional[Dict]:
|
||||||
"""
|
"""
|
||||||
获取查询缓存结果
|
获取查询缓存结果
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
query: 查询文本
|
query: 查询文本
|
||||||
kb_name: 知识库名称
|
kb_name: 知识库名称
|
||||||
doc_ids: 相关文档 ID 列表(用于细粒度缓存 key)
|
|
||||||
"""
|
"""
|
||||||
kb_version = self.get_kb_version(kb_name)
|
kb_version = self.get_kb_version(kb_name)
|
||||||
|
key = self._make_query_cache_key(query, kb_name, kb_version)
|
||||||
# 计算文档哈希(如果提供了 doc_ids)
|
|
||||||
doc_hash = ""
|
|
||||||
if doc_ids:
|
|
||||||
doc_hash = self._compute_doc_hash(kb_name, doc_ids)
|
|
||||||
|
|
||||||
key = self._make_query_cache_key(query, kb_name, kb_version, doc_hash)
|
|
||||||
return self.query_cache.get(key)
|
return self.query_cache.get(key)
|
||||||
|
|
||||||
def set_query_result(self, query: str, kb_name: str, result: Dict, doc_ids: List[str] = None) -> None:
|
def set_query_result(self, query: str, kb_name: str, result: Dict) -> None:
|
||||||
"""
|
"""
|
||||||
设置查询缓存结果
|
设置查询缓存结果
|
||||||
|
|
||||||
|
kb_version 在文档变更时自增,触发整个知识库的缓存失效。
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
query: 查询文本
|
query: 查询文本
|
||||||
kb_name: 知识库名称
|
kb_name: 知识库名称
|
||||||
result: 缓存结果
|
result: 缓存结果
|
||||||
doc_ids: 相关文档 ID 列表(用于细粒度失效)
|
|
||||||
"""
|
"""
|
||||||
kb_version = self.get_kb_version(kb_name)
|
kb_version = self.get_kb_version(kb_name)
|
||||||
|
key = self._make_query_cache_key(query, kb_name, kb_version)
|
||||||
# 计算相关文档的版本哈希(细粒度失效)
|
|
||||||
doc_hash = ""
|
|
||||||
if doc_ids:
|
|
||||||
doc_hash = self._compute_doc_hash(kb_name, doc_ids)
|
|
||||||
|
|
||||||
key = self._make_query_cache_key(query, kb_name, kb_version, doc_hash)
|
|
||||||
self.query_cache.set(key, result, kb_version=kb_version)
|
self.query_cache.set(key, result, kb_version=kb_version)
|
||||||
|
|
||||||
def _compute_doc_hash(self, kb_name: str, doc_ids: List[str]) -> str:
|
|
||||||
"""
|
|
||||||
计算文档版本哈希
|
|
||||||
|
|
||||||
用于细粒度缓存失效:只失效相关文档变化时的缓存
|
|
||||||
"""
|
|
||||||
if not doc_ids:
|
|
||||||
return ""
|
|
||||||
|
|
||||||
# 从文档 ID 中提取 source(文件名)
|
|
||||||
sources = set()
|
|
||||||
for doc_id in doc_ids:
|
|
||||||
# doc_id 格式通常为 "filename_text_0" 或类似
|
|
||||||
parts = doc_id.split('_')
|
|
||||||
if parts:
|
|
||||||
sources.add(parts[0])
|
|
||||||
|
|
||||||
# 生成哈希
|
|
||||||
sources_str = ','.join(sorted(sources))
|
|
||||||
return hashlib.md5(f"docs:{sources_str}".encode()).hexdigest()
|
|
||||||
|
|
||||||
# ==================== Embedding Cache 方法 ====================
|
# ==================== Embedding Cache 方法 ====================
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
|
|||||||
558
core/engine.py
558
core/engine.py
@@ -29,6 +29,7 @@ RAG 核心引擎
|
|||||||
|
|
||||||
import os
|
import os
|
||||||
import gc
|
import gc
|
||||||
|
import re
|
||||||
import time
|
import time
|
||||||
import logging
|
import logging
|
||||||
import threading
|
import threading
|
||||||
@@ -76,6 +77,10 @@ try:
|
|||||||
CONTEXT_EXPANSION_MAX_CHUNKS, ENUM_QUERY_DISABLE_TOPK_SHRINK, ENUM_QUERY_MMR_LAMBDA,
|
CONTEXT_EXPANSION_MAX_CHUNKS, ENUM_QUERY_DISABLE_TOPK_SHRINK, ENUM_QUERY_MMR_LAMBDA,
|
||||||
# Phase 3 扩展精细化
|
# Phase 3 扩展精细化
|
||||||
EXPANSION_SCORE_THRESHOLD, MAX_EXPANDED_NEIGHBORS,
|
EXPANSION_SCORE_THRESHOLD, MAX_EXPANDED_NEIGHBORS,
|
||||||
|
# 章节聚类救援
|
||||||
|
SECTION_CLUSTER_BOOST_ENABLED, CLUSTER_MIN_MEMBERS, CLUSTER_MIN_TYPES,
|
||||||
|
CLUSTER_SEED_FLOOR, CLUSTER_MAX_BOOST_PER_SECTION, CLUSTER_MAX_SECTIONS,
|
||||||
|
CLUSTER_SECTION_PREFIX_LEVELS,
|
||||||
# 上下文与生成
|
# 上下文与生成
|
||||||
LLM_TEMPERATURE, LLM_MAX_TOKENS, RECALL_MULTIPLIER,
|
LLM_TEMPERATURE, LLM_MAX_TOKENS, RECALL_MULTIPLIER,
|
||||||
# FAQ 与黑名单
|
# FAQ 与黑名单
|
||||||
@@ -101,20 +106,28 @@ except ImportError:
|
|||||||
MMR_TOP_K = 30
|
MMR_TOP_K = 30
|
||||||
CONTEXT_EXPANSION_ENABLED = True
|
CONTEXT_EXPANSION_ENABLED = True
|
||||||
CONTEXT_EXPANSION_BEFORE = 1
|
CONTEXT_EXPANSION_BEFORE = 1
|
||||||
CONTEXT_EXPANSION_AFTER = 5
|
CONTEXT_EXPANSION_AFTER = 8
|
||||||
CONTEXT_EXPANSION_MAX_CHUNKS = 24
|
CONTEXT_EXPANSION_MAX_CHUNKS = 24
|
||||||
EXPANSION_SCORE_THRESHOLD = 0.3
|
EXPANSION_SCORE_THRESHOLD = 0.3
|
||||||
MAX_EXPANDED_NEIGHBORS = 4
|
MAX_EXPANDED_NEIGHBORS = 8
|
||||||
|
# 章节聚类救援默认值
|
||||||
|
SECTION_CLUSTER_BOOST_ENABLED = True
|
||||||
|
CLUSTER_MIN_MEMBERS = 3
|
||||||
|
CLUSTER_MIN_TYPES = 2
|
||||||
|
CLUSTER_SEED_FLOOR = 0.35
|
||||||
|
CLUSTER_MAX_BOOST_PER_SECTION = 8
|
||||||
|
CLUSTER_MAX_SECTIONS = 3
|
||||||
|
CLUSTER_SECTION_PREFIX_LEVELS = 1
|
||||||
ENUM_QUERY_DISABLE_TOPK_SHRINK = True
|
ENUM_QUERY_DISABLE_TOPK_SHRINK = True
|
||||||
ENUM_QUERY_MMR_LAMBDA = 0.85
|
ENUM_QUERY_MMR_LAMBDA = 0.85
|
||||||
DYNAMIC_RRF_ENABLED = True
|
DYNAMIC_RRF_ENABLED = True
|
||||||
EMBEDDING_DEVICE = "auto"
|
EMBEDDING_DEVICE = "auto"
|
||||||
RERANK_DEVICE = "auto"
|
RERANK_DEVICE = "auto"
|
||||||
RERANK_USE_ONNX = False
|
RERANK_USE_ONNX = False
|
||||||
RERANK_BACKEND = "local"
|
RERANK_BACKEND = "cloud"
|
||||||
RERANK_CLOUD_MODEL = "qwen3-rerank"
|
RERANK_CLOUD_MODEL = "xop3qwen8breranker"
|
||||||
RERANK_CLOUD_API_KEY = ""
|
RERANK_CLOUD_API_KEY = ""
|
||||||
RERANK_CLOUD_BASE_URL = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks"
|
RERANK_CLOUD_BASE_URL = "https://maas-api.cn-huabei-1.xf-yun.com/v2/rerank"
|
||||||
RERANK_CLOUD_TIMEOUT = 15
|
RERANK_CLOUD_TIMEOUT = 15
|
||||||
|
|
||||||
|
|
||||||
@@ -272,8 +285,7 @@ class CloudReranker:
|
|||||||
body = {
|
body = {
|
||||||
"model": self.model,
|
"model": self.model,
|
||||||
"query": query,
|
"query": query,
|
||||||
"documents": documents,
|
"documents": documents
|
||||||
"top_n": len(documents)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
session = self._get_session()
|
session = self._get_session()
|
||||||
@@ -547,8 +559,7 @@ class RAGEngine:
|
|||||||
top_dist = result['distances'][0][0] if result.get('distances') and result['distances'][0] else 1.0
|
top_dist = result['distances'][0][0] if result.get('distances') and result['distances'][0] else 1.0
|
||||||
top_score = 1.0 - top_dist
|
top_score = 1.0 - top_dist
|
||||||
if top_score >= CACHE_MIN_SCORE:
|
if top_score >= CACHE_MIN_SCORE:
|
||||||
doc_ids = result.get('ids', [[]])[0] if result.get('ids') else []
|
cache.set_query_result(query, kb_name, result)
|
||||||
cache.set_query_result(query, kb_name, result, doc_ids=doc_ids)
|
|
||||||
result['_debug'] = _debug
|
result['_debug'] = _debug
|
||||||
return result
|
return result
|
||||||
|
|
||||||
@@ -569,8 +580,7 @@ class RAGEngine:
|
|||||||
top_dist = result['distances'][0][0] if result.get('distances') and result['distances'][0] else 1.0
|
top_dist = result['distances'][0][0] if result.get('distances') and result['distances'][0] else 1.0
|
||||||
top_score = 1.0 - top_dist
|
top_score = 1.0 - top_dist
|
||||||
if top_score >= CACHE_MIN_SCORE:
|
if top_score >= CACHE_MIN_SCORE:
|
||||||
doc_ids = result.get('ids', [[]])[0] if result.get('ids') else []
|
cache.set_query_result(query, kb_name, result)
|
||||||
cache.set_query_result(query, kb_name, result, doc_ids=doc_ids)
|
|
||||||
result['_debug'] = _debug
|
result['_debug'] = _debug
|
||||||
return result
|
return result
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -605,9 +615,7 @@ class RAGEngine:
|
|||||||
top_dist = result['distances'][0][0] if result.get('distances') and result['distances'][0] else 1.0
|
top_dist = result['distances'][0][0] if result.get('distances') and result['distances'][0] else 1.0
|
||||||
top_score = 1.0 - top_dist
|
top_score = 1.0 - top_dist
|
||||||
if top_score >= CACHE_MIN_SCORE:
|
if top_score >= CACHE_MIN_SCORE:
|
||||||
# 传递 doc_ids 实现细粒度缓存失效
|
cache.set_query_result(query, kb_name, result)
|
||||||
doc_ids = result.get('ids', [[]])[0] if result.get('ids') else []
|
|
||||||
cache.set_query_result(query, kb_name, result, doc_ids=doc_ids)
|
|
||||||
_debug['timing']['total_ms'] = int((time.time() - _overall_start) * 1000)
|
_debug['timing']['total_ms'] = int((time.time() - _overall_start) * 1000)
|
||||||
result['_debug'] = _debug
|
result['_debug'] = _debug
|
||||||
return result
|
return result
|
||||||
@@ -625,7 +633,7 @@ class RAGEngine:
|
|||||||
elif len(conditions) > 1:
|
elif len(conditions) > 1:
|
||||||
where_filter = {"$and": conditions}
|
where_filter = {"$and": conditions}
|
||||||
|
|
||||||
query_vector = self.embedding_model.encode(query).tolist()
|
query_vector = self._encode_cached(query).tolist()
|
||||||
recall_k = RERANK_CANDIDATES if (USE_RERANK or USE_HYBRID_SEARCH) else top_k
|
recall_k = RERANK_CANDIDATES if (USE_RERANK or USE_HYBRID_SEARCH) else top_k
|
||||||
recall_k = max(recall_k, top_k * RECALL_MULTIPLIER)
|
recall_k = max(recall_k, top_k * RECALL_MULTIPLIER)
|
||||||
|
|
||||||
@@ -657,6 +665,7 @@ class RAGEngine:
|
|||||||
|
|
||||||
results_list = [vector_results]
|
results_list = [vector_results]
|
||||||
weights = [VECTOR_WEIGHT]
|
weights = [VECTOR_WEIGHT]
|
||||||
|
bm25_results = None # 初始化,防止 NameError
|
||||||
|
|
||||||
if USE_HYBRID_SEARCH and self.bm25_index.bm25:
|
if USE_HYBRID_SEARCH and self.bm25_index.bm25:
|
||||||
bm25_results = self.bm25_index.search(query, top_k=recall_k)
|
bm25_results = self.bm25_index.search(query, top_k=recall_k)
|
||||||
@@ -671,6 +680,22 @@ class RAGEngine:
|
|||||||
vector_w, bm25_w = self._get_dynamic_rrf_weights(query)
|
vector_w, bm25_w = self._get_dynamic_rrf_weights(query)
|
||||||
weights = [vector_w, bm25_w]
|
weights = [vector_w, bm25_w]
|
||||||
|
|
||||||
|
# ========== 保留 BM25 原始 top-3 完整信息,用于下游分歧检测救援 ==========
|
||||||
|
_bm25_raw_top3 = []
|
||||||
|
if USE_HYBRID_SEARCH and bm25_results and bm25_results.get('ids') and bm25_results['ids'][0]:
|
||||||
|
_bm25_ids = bm25_results['ids'][0][:3]
|
||||||
|
_bm25_docs = bm25_results['documents'][0][:3]
|
||||||
|
_bm25_metas = bm25_results['metadatas'][0][:3]
|
||||||
|
_bm25_dists = (bm25_results.get('distances', [[]])[0] or [0]*3)[:3]
|
||||||
|
for i in range(len(_bm25_ids)):
|
||||||
|
_bm25_raw_top3.append({
|
||||||
|
'id': _bm25_ids[i],
|
||||||
|
'doc': _bm25_docs[i],
|
||||||
|
'meta': _bm25_metas[i],
|
||||||
|
'bm25_score': _bm25_dists[i],
|
||||||
|
'rank': i + 1
|
||||||
|
})
|
||||||
|
|
||||||
if len(results_list) > 1:
|
if len(results_list) > 1:
|
||||||
fused_results = self.reciprocal_rank_fusion(results_list, weights)
|
fused_results = self.reciprocal_rank_fusion(results_list, weights)
|
||||||
_debug['steps'].append({'name': 'rrf_fusion', 'count': len(fused_results['ids'][0]) if fused_results.get('ids') else 0, 'weights': [round(w, 2) for w in weights]})
|
_debug['steps'].append({'name': 'rrf_fusion', 'count': len(fused_results['ids'][0]) if fused_results.get('ids') else 0, 'weights': [round(w, 2) for w in weights]})
|
||||||
@@ -686,6 +711,8 @@ class RAGEngine:
|
|||||||
|
|
||||||
is_enum_query = self._is_enumeration_query(query)
|
is_enum_query = self._is_enumeration_query(query)
|
||||||
fused_results['_enum_query'] = is_enum_query
|
fused_results['_enum_query'] = is_enum_query
|
||||||
|
# 传递 BM25 原始 top-3 到路由层,用于分歧检测救援
|
||||||
|
fused_results['_bm25_top3'] = _bm25_raw_top3
|
||||||
|
|
||||||
# 章节过滤(如果查询中提到了章节)
|
# 章节过滤(如果查询中提到了章节)
|
||||||
fused_results = self._filter_by_section(fused_results, query)
|
fused_results = self._filter_by_section(fused_results, query)
|
||||||
@@ -732,11 +759,19 @@ class RAGEngine:
|
|||||||
# 时间衰减(Time Decay)
|
# 时间衰减(Time Decay)
|
||||||
fused_results = self._apply_time_decay(fused_results)
|
fused_results = self._apply_time_decay(fused_results)
|
||||||
|
|
||||||
|
# 提前附加 _debug,使聚类提升能写入调试步骤
|
||||||
|
fused_results['_debug'] = _debug
|
||||||
|
|
||||||
|
# ========== 章节聚类提升:在扩展前将低分但聚类的切片提升至种子阈值 ==========
|
||||||
|
if SECTION_CLUSTER_BOOST_ENABLED:
|
||||||
|
fused_results = self._section_cluster_boost(fused_results, query)
|
||||||
|
|
||||||
# ========== 上下文扩展:补充强命中切片周围的连续文本(rerank 之后,防止被截断)==========
|
# ========== 上下文扩展:补充强命中切片周围的连续文本(rerank 之后,防止被截断)==========
|
||||||
# Phase 3:仅对高分种子扩展邻居
|
# Phase 3:仅对高分种子扩展邻居
|
||||||
before_exp = len(fused_results['ids'][0]) if fused_results.get('ids') else 0
|
before_exp = len(fused_results['ids'][0]) if fused_results.get('ids') else 0
|
||||||
fused_results = self._expand_contiguous_chunks(fused_results, top_k=top_k,
|
fused_results = self._expand_contiguous_chunks(fused_results, top_k=top_k,
|
||||||
min_score=EXPANSION_SCORE_THRESHOLD)
|
min_score=EXPANSION_SCORE_THRESHOLD,
|
||||||
|
query=query)
|
||||||
after_exp = len(fused_results['ids'][0]) if fused_results.get('ids') else 0
|
after_exp = len(fused_results['ids'][0]) if fused_results.get('ids') else 0
|
||||||
_debug['steps'].append({'name': 'context_expansion', 'before': before_exp, 'after': after_exp})
|
_debug['steps'].append({'name': 'context_expansion', 'before': before_exp, 'after': after_exp})
|
||||||
|
|
||||||
@@ -748,7 +783,15 @@ class RAGEngine:
|
|||||||
and not (is_enum_query and ENUM_QUERY_DISABLE_TOPK_SHRINK)
|
and not (is_enum_query and ENUM_QUERY_DISABLE_TOPK_SHRINK)
|
||||||
and fused_results.get('_score_source') != 'rrf'
|
and fused_results.get('_score_source') != 'rrf'
|
||||||
):
|
):
|
||||||
top_score = 1.0 - fused_results['distances'][0][0] # 距离转相似度
|
# 根据分数来源计算相似度分数(越高越好)
|
||||||
|
score_source = fused_results.get('_score_source')
|
||||||
|
top_dist = fused_results['distances'][0][0]
|
||||||
|
if score_source == 'rerank':
|
||||||
|
# Rerank 后 distances 是相关性分数,越大越好,直接使用
|
||||||
|
top_score = top_dist
|
||||||
|
else:
|
||||||
|
# 向量距离,越小越好,转为相似度
|
||||||
|
top_score = 1.0 - top_dist
|
||||||
adjusted_k, should_retrieve, reason = self._adaptive_topk.adjust(top_score, top_k)
|
adjusted_k, should_retrieve, reason = self._adaptive_topk.adjust(top_score, top_k)
|
||||||
if "high_confidence" in reason:
|
if "high_confidence" in reason:
|
||||||
# 高置信度时截断结果
|
# 高置信度时截断结果
|
||||||
@@ -763,9 +806,7 @@ class RAGEngine:
|
|||||||
top_dist = fused_results['distances'][0][0] if fused_results.get('distances') and fused_results['distances'][0] else 1.0
|
top_dist = fused_results['distances'][0][0] if fused_results.get('distances') and fused_results['distances'][0] else 1.0
|
||||||
top_score = 1.0 - top_dist # 距离转相似度
|
top_score = 1.0 - top_dist # 距离转相似度
|
||||||
if top_score >= CACHE_MIN_SCORE: # 置信度阈值
|
if top_score >= CACHE_MIN_SCORE: # 置信度阈值
|
||||||
# 传递 doc_ids 实现细粒度缓存失效
|
cache.set_query_result(query, kb_name, fused_results)
|
||||||
doc_ids = fused_results.get('ids', [[]])[0] if fused_results.get('ids') else []
|
|
||||||
cache.set_query_result(query, kb_name, fused_results, doc_ids=doc_ids)
|
|
||||||
|
|
||||||
fused_results['_debug'] = _debug
|
fused_results['_debug'] = _debug
|
||||||
_debug['timing']['total_ms'] = int((time.time() - _overall_start) * 1000)
|
_debug['timing']['total_ms'] = int((time.time() - _overall_start) * 1000)
|
||||||
@@ -811,6 +852,76 @@ class RAGEngine:
|
|||||||
logger.warning(f"FAQ 集合查询失败: {e}")
|
logger.warning(f"FAQ 集合查询失败: {e}")
|
||||||
return get_empty_result()
|
return get_empty_result()
|
||||||
|
|
||||||
|
def _encode_cached(self, text):
|
||||||
|
"""
|
||||||
|
带缓存的 embedding 编码
|
||||||
|
|
||||||
|
优先从 Embedding Cache(LRU)读取,未命中再调用模型编码并写入缓存。
|
||||||
|
支持单文本和批量文本输入。
|
||||||
|
|
||||||
|
Args:
|
||||||
|
text: 单个文本字符串 或 文本列表
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
numpy 数组(单文本为一维,批量为二维)
|
||||||
|
"""
|
||||||
|
import numpy as _np
|
||||||
|
|
||||||
|
# 检查 embedding 缓存是否启用(缓存配置查询结果,避免每次重复导入)
|
||||||
|
if not hasattr(self, '_emb_cache_enabled'):
|
||||||
|
self._emb_cache_enabled = True # 默认启用
|
||||||
|
if CACHE_AVAILABLE:
|
||||||
|
try:
|
||||||
|
from config import EMBEDDING_CACHE_ENABLED
|
||||||
|
self._emb_cache_enabled = EMBEDDING_CACHE_ENABLED
|
||||||
|
except ImportError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
if not self._emb_cache_enabled:
|
||||||
|
return self.embedding_model.encode(text)
|
||||||
|
|
||||||
|
try:
|
||||||
|
_cache = get_cache_manager()
|
||||||
|
except Exception:
|
||||||
|
return self.embedding_model.encode(text)
|
||||||
|
|
||||||
|
# 批量输入
|
||||||
|
if isinstance(text, list):
|
||||||
|
try:
|
||||||
|
cached_embs, missed_indices = _cache.get_embeddings_batch(text)
|
||||||
|
if missed_indices:
|
||||||
|
missed_texts = [text[i] for i in missed_indices]
|
||||||
|
# encode(list) 始终返回 2D ndarray,直接按行索引即可
|
||||||
|
new_embs = self.embedding_model.encode(missed_texts)
|
||||||
|
if len(missed_indices) == 1:
|
||||||
|
# 单条时 encode 可能返回 1D,需统一处理
|
||||||
|
if new_embs.ndim == 1:
|
||||||
|
new_embs = new_embs.reshape(1, -1)
|
||||||
|
for idx, mi in enumerate(missed_indices):
|
||||||
|
emb_list = new_embs[idx].tolist()
|
||||||
|
cached_embs[mi] = emb_list
|
||||||
|
try:
|
||||||
|
_cache.set_embedding(text[mi], emb_list)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
return _np.array(cached_embs)
|
||||||
|
except Exception:
|
||||||
|
# 缓存故障时优雅降级为直接编码
|
||||||
|
return self.embedding_model.encode(text)
|
||||||
|
|
||||||
|
# 单文本输入
|
||||||
|
cached = _cache.get_embedding(text)
|
||||||
|
if cached is not None:
|
||||||
|
return _np.array(cached)
|
||||||
|
|
||||||
|
embedding = self.embedding_model.encode(text)
|
||||||
|
try:
|
||||||
|
emb_list = embedding.tolist() if hasattr(embedding, 'tolist') else list(embedding)
|
||||||
|
_cache.set_embedding(text, emb_list)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
return embedding
|
||||||
|
|
||||||
def _search_image_chunks(self, query_vector: list, top_k: int = 5, where_filter: dict = None) -> dict:
|
def _search_image_chunks(self, query_vector: list, top_k: int = 5, where_filter: dict = None) -> dict:
|
||||||
"""
|
"""
|
||||||
独立检索图片切片(P0:图片独立召回通道)
|
独立检索图片切片(P0:图片独立召回通道)
|
||||||
@@ -1017,7 +1128,7 @@ class RAGEngine:
|
|||||||
'metadatas': [f_metas],
|
'metadatas': [f_metas],
|
||||||
'distances': [f_scores]
|
'distances': [f_scores]
|
||||||
}
|
}
|
||||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'):
|
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context', '_bm25_top3'):
|
||||||
if key in results:
|
if key in results:
|
||||||
filtered[key] = results[key]
|
filtered[key] = results[key]
|
||||||
return filtered
|
return filtered
|
||||||
@@ -1040,7 +1151,7 @@ class RAGEngine:
|
|||||||
'metadatas': [f_metas],
|
'metadatas': [f_metas],
|
||||||
'distances': [f_scores]
|
'distances': [f_scores]
|
||||||
}
|
}
|
||||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'):
|
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context', '_bm25_top3'):
|
||||||
if key in results:
|
if key in results:
|
||||||
filtered[key] = results[key]
|
filtered[key] = results[key]
|
||||||
return filtered
|
return filtered
|
||||||
@@ -1055,7 +1166,7 @@ class RAGEngine:
|
|||||||
'metadatas': [results['metadatas'][0][:top_k]],
|
'metadatas': [results['metadatas'][0][:top_k]],
|
||||||
'distances': [results['distances'][0][:top_k]]
|
'distances': [results['distances'][0][:top_k]]
|
||||||
}
|
}
|
||||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'):
|
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context', '_bm25_top3'):
|
||||||
if key in results:
|
if key in results:
|
||||||
truncated[key] = results[key]
|
truncated[key] = results[key]
|
||||||
return truncated
|
return truncated
|
||||||
@@ -1090,14 +1201,139 @@ class RAGEngine:
|
|||||||
return None
|
return None
|
||||||
return self.collection
|
return self.collection
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _normalize_section_path(section_path: str, levels: int = None) -> str:
|
||||||
|
"""归一化 section_path:取前 N 级路径,容忍 MinerU 标题检测误差。
|
||||||
|
|
||||||
|
例如: "第三章 吸烟场所的功能设置 > 第三条 文明吸烟..." → "第三章 吸烟场所的功能设置"
|
||||||
|
"""
|
||||||
|
if not section_path:
|
||||||
|
return ''
|
||||||
|
if levels is None:
|
||||||
|
levels = CLUSTER_SECTION_PREFIX_LEVELS
|
||||||
|
parts = [p.strip() for p in section_path.split('>')]
|
||||||
|
return ' > '.join(parts[:levels])
|
||||||
|
|
||||||
|
def _section_cluster_boost(self, results: dict, query: str = '') -> dict:
|
||||||
|
"""章节聚类提升:当同一 section 下多个切片(text+table)同时出现在候选集中,
|
||||||
|
即使单个切片 CrossEncoder 分数很低,也将整组提升到种子阈值。
|
||||||
|
|
||||||
|
核心洞察:单个低分切片不可信,但同一 section 多个切片同时出现是强信号。
|
||||||
|
提升后的切片可以作为 _expand_contiguous_chunks 的种子,触发邻居扩展。
|
||||||
|
|
||||||
|
Args:
|
||||||
|
results: rerank 后的检索结果
|
||||||
|
query: 用户查询(用于后续扩展)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
修改后的 results(distances 被调整,meta 中标记 _cluster_boosted)
|
||||||
|
"""
|
||||||
|
if not results.get('ids') or not results['ids'][0]:
|
||||||
|
return results
|
||||||
|
|
||||||
|
ids = results['ids'][0]
|
||||||
|
metas = results.get('metadatas', [[]])[0]
|
||||||
|
distances = results.get('distances', [[]])[0] if results.get('distances') else None
|
||||||
|
|
||||||
|
if not distances:
|
||||||
|
return results
|
||||||
|
|
||||||
|
# 1. 按 (source, normalized_section) 分组
|
||||||
|
from collections import defaultdict
|
||||||
|
section_groups = defaultdict(list) # key → [(index, meta, dist)]
|
||||||
|
|
||||||
|
for i, (meta, dist) in enumerate(zip(metas, distances)):
|
||||||
|
source = meta.get('source', '')
|
||||||
|
section_path = meta.get('section', '') or meta.get('section_path', '')
|
||||||
|
norm_section = self._normalize_section_path(section_path)
|
||||||
|
if not source or not norm_section:
|
||||||
|
continue
|
||||||
|
key = (source, norm_section)
|
||||||
|
section_groups[key].append((i, meta, dist))
|
||||||
|
|
||||||
|
# 2. 检测聚类信号并提升
|
||||||
|
boost_target_dist = 1.0 - CLUSTER_SEED_FLOOR # score=0.35 → dist=0.65
|
||||||
|
boosted_sections = []
|
||||||
|
total_boosted = 0
|
||||||
|
|
||||||
|
# 按组成员数降序排列,优先处理最大聚类
|
||||||
|
sorted_groups = sorted(section_groups.items(), key=lambda x: len(x[1]), reverse=True)
|
||||||
|
|
||||||
|
for (source, norm_section), members in sorted_groups:
|
||||||
|
if len(boosted_sections) >= CLUSTER_MAX_SECTIONS:
|
||||||
|
break
|
||||||
|
|
||||||
|
# 聚类信号检测:成员数 >= 阈值 且 类型多样性 >= 阈值
|
||||||
|
chunk_types = set(m[1].get('chunk_type', 'text') for m in members)
|
||||||
|
if len(members) < CLUSTER_MIN_MEMBERS or len(chunk_types) < CLUSTER_MIN_TYPES:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# 提升组内切片分数(仅提升低于阈值的)
|
||||||
|
boost_count = 0
|
||||||
|
for idx, meta, dist in members:
|
||||||
|
if boost_count >= CLUSTER_MAX_BOOST_PER_SECTION:
|
||||||
|
break
|
||||||
|
# 只提升分数低于种子阈值的切片(高分切片不需要)
|
||||||
|
if dist > boost_target_dist:
|
||||||
|
distances[idx] = boost_target_dist
|
||||||
|
meta['_cluster_boosted'] = True
|
||||||
|
boost_count += 1
|
||||||
|
total_boosted += 1
|
||||||
|
|
||||||
|
if boost_count > 0:
|
||||||
|
boosted_sections.append({
|
||||||
|
'source': source,
|
||||||
|
'section': norm_section,
|
||||||
|
'members': len(members),
|
||||||
|
'types': list(chunk_types),
|
||||||
|
'boosted': boost_count
|
||||||
|
})
|
||||||
|
|
||||||
|
# 3. 写 debug 信息
|
||||||
|
if boosted_sections:
|
||||||
|
debug_info = results.get('_debug', {})
|
||||||
|
if 'steps' not in debug_info:
|
||||||
|
debug_info['steps'] = []
|
||||||
|
debug_info['steps'].append({
|
||||||
|
'name': 'section_cluster_boost',
|
||||||
|
'sections': boosted_sections,
|
||||||
|
'total_boosted': total_boosted
|
||||||
|
})
|
||||||
|
results['_debug'] = debug_info
|
||||||
|
logger.info(f"[章节聚类提升] 提升 {total_boosted} 个切片,"
|
||||||
|
f"涉及 {len(boosted_sections)} 个 section: "
|
||||||
|
f"{[s['section'][:30] for s in boosted_sections]}")
|
||||||
|
|
||||||
|
return results
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _chunk_lexical_score(chunk_text: str, query: str) -> float:
|
||||||
|
"""计算切片文本与查询的词法重叠度(bigram 命中率),用于辅助种子资格判定。"""
|
||||||
|
if not chunk_text or not query:
|
||||||
|
return 0.0
|
||||||
|
import re
|
||||||
|
clean_q = re.sub(r'[??!!。,,、;;::"""\'\s*#`]+', ' ', query).strip()
|
||||||
|
if len(clean_q) < 2:
|
||||||
|
return 0.0
|
||||||
|
bigrams = set()
|
||||||
|
for i in range(len(clean_q) - 1):
|
||||||
|
w = clean_q[i:i+2].strip()
|
||||||
|
if len(w) == 2:
|
||||||
|
bigrams.add(w)
|
||||||
|
if not bigrams:
|
||||||
|
return 0.0
|
||||||
|
matched = sum(1 for w in bigrams if w in chunk_text)
|
||||||
|
return matched / len(bigrams)
|
||||||
|
|
||||||
def _expand_contiguous_chunks(self, results: dict, top_k: int = None,
|
def _expand_contiguous_chunks(self, results: dict, top_k: int = None,
|
||||||
min_score: float = 0.0) -> dict:
|
min_score: float = 0.0, query: str = '') -> dict:
|
||||||
"""Add same-source same-section neighbor text chunks around strong hits.
|
"""Add same-source same-section neighbor text chunks around strong hits.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
results: 检索结果
|
results: 检索结果
|
||||||
top_k: 最大切片数
|
top_k: 最大切片数
|
||||||
min_score: Phase 3 最低分数阈值,仅对 Rerank 分数高于此值的种子扩展
|
min_score: Phase 3 最低分数阈值,仅对 Rerank 分数高于此值的种子扩展
|
||||||
|
query: 查询文本,用于词法匹配辅助种子资格判定
|
||||||
"""
|
"""
|
||||||
if not CONTEXT_EXPANSION_ENABLED:
|
if not CONTEXT_EXPANSION_ENABLED:
|
||||||
return results
|
return results
|
||||||
@@ -1127,7 +1363,7 @@ class RAGEngine:
|
|||||||
seeds = [
|
seeds = [
|
||||||
(doc_id, doc, meta, dist)
|
(doc_id, doc, meta, dist)
|
||||||
for doc_id, doc, meta, dist in items[:base_limit]
|
for doc_id, doc, meta, dist in items[:base_limit]
|
||||||
if meta.get('chunk_type', 'text') == 'text'
|
if (meta.get('chunk_type', 'text') == 'text' or meta.get('_cluster_boosted'))
|
||||||
and meta.get('source')
|
and meta.get('source')
|
||||||
and self._to_int(meta.get('chunk_index')) is not None
|
and self._to_int(meta.get('chunk_index')) is not None
|
||||||
]
|
]
|
||||||
@@ -1138,7 +1374,11 @@ class RAGEngine:
|
|||||||
break
|
break
|
||||||
|
|
||||||
# Phase 3:跳过分数低于阈值的种子(仅当 min_score > 0 时生效)
|
# Phase 3:跳过分数低于阈值的种子(仅当 min_score > 0 时生效)
|
||||||
|
# 词法匹配豁免:CrossEncoder 低分但关键词重叠度高时仍允许作为种子
|
||||||
if min_score > 0 and seed_dist < min_score:
|
if min_score > 0 and seed_dist < min_score:
|
||||||
|
if query and self._chunk_lexical_score(_seed_doc, query) > 0.3:
|
||||||
|
pass # 词法匹配度高,允许作为种子
|
||||||
|
else:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
source = seed_meta.get('source')
|
source = seed_meta.get('source')
|
||||||
@@ -1161,6 +1401,7 @@ class RAGEngine:
|
|||||||
logger.warning(f"扩展连续切片失败: {e}")
|
logger.warning(f"扩展连续切片失败: {e}")
|
||||||
return {'ids': [], 'documents': [], 'metadatas': []}
|
return {'ids': [], 'documents': [], 'metadatas': []}
|
||||||
|
|
||||||
|
# 扩展同 section 的 text 邻居
|
||||||
where_filter = {"$and": [{"source": source}, {"chunk_type": "text"}]}
|
where_filter = {"$and": [{"source": source}, {"chunk_type": "text"}]}
|
||||||
if section:
|
if section:
|
||||||
where_filter["$and"].append({"section": section})
|
where_filter["$and"].append({"section": section})
|
||||||
@@ -1171,6 +1412,20 @@ class RAGEngine:
|
|||||||
if not neighbors.get('ids') or len(neighbors.get('ids', [])) <= 1:
|
if not neighbors.get('ids') or len(neighbors.get('ids', [])) <= 1:
|
||||||
neighbors = _get_neighbors({"$and": [{"source": source}, {"chunk_type": "text"}]})
|
neighbors = _get_neighbors({"$and": [{"source": source}, {"chunk_type": "text"}]})
|
||||||
|
|
||||||
|
# 同时扩展同 section 的 table 邻居(table 切片的 rerank 分数往往偏低,
|
||||||
|
# 但与同 section 的 text 切片属于同一语义单元,不应割裂)
|
||||||
|
table_where = {"$and": [{"source": source}, {"chunk_type": "table"}]}
|
||||||
|
if section:
|
||||||
|
table_where["$and"].append({"section": section})
|
||||||
|
table_neighbors = _get_neighbors(table_where)
|
||||||
|
|
||||||
|
# 当 section 为空时,table 查询只有 source 条件,可能拉入大量无关表格,
|
||||||
|
# 缩小 chunk_index 窗口至 ±1 以降低噪音;有 section 时使用正常窗口
|
||||||
|
if section:
|
||||||
|
_t_before, _t_after = CONTEXT_EXPANSION_BEFORE, CONTEXT_EXPANSION_AFTER
|
||||||
|
else:
|
||||||
|
_t_before, _t_after = 1, 1
|
||||||
|
|
||||||
neighbor_rows = []
|
neighbor_rows = []
|
||||||
for n_id, n_doc, n_meta in zip(
|
for n_id, n_doc, n_meta in zip(
|
||||||
neighbors.get('ids', []),
|
neighbors.get('ids', []),
|
||||||
@@ -1183,6 +1438,18 @@ class RAGEngine:
|
|||||||
if seed_index - CONTEXT_EXPANSION_BEFORE <= n_index <= seed_index + CONTEXT_EXPANSION_AFTER:
|
if seed_index - CONTEXT_EXPANSION_BEFORE <= n_index <= seed_index + CONTEXT_EXPANSION_AFTER:
|
||||||
neighbor_rows.append((n_index, n_id, n_doc, n_meta))
|
neighbor_rows.append((n_index, n_id, n_doc, n_meta))
|
||||||
|
|
||||||
|
# 同 section 的 table 邻居也加入扩展范围
|
||||||
|
for n_id, n_doc, n_meta in zip(
|
||||||
|
table_neighbors.get('ids', []),
|
||||||
|
table_neighbors.get('documents', []),
|
||||||
|
table_neighbors.get('metadatas', [])
|
||||||
|
):
|
||||||
|
n_index = self._to_int(n_meta.get('chunk_index'))
|
||||||
|
if n_index is None:
|
||||||
|
continue
|
||||||
|
if seed_index - _t_before <= n_index <= seed_index + _t_after:
|
||||||
|
neighbor_rows.append((n_index, n_id, n_doc, n_meta))
|
||||||
|
|
||||||
seed_neighbors_added = 0
|
seed_neighbors_added = 0
|
||||||
for n_index, n_id, n_doc, n_meta in sorted(neighbor_rows, key=lambda row: row[0]):
|
for n_index, n_id, n_doc, n_meta in sorted(neighbor_rows, key=lambda row: row[0]):
|
||||||
if len(items) >= max_chunks:
|
if len(items) >= max_chunks:
|
||||||
@@ -1214,7 +1481,7 @@ class RAGEngine:
|
|||||||
'distances': [[item[3] for item in items]],
|
'distances': [[item[3] for item in items]],
|
||||||
'_expanded_context': {'added': added}
|
'_expanded_context': {'added': added}
|
||||||
}
|
}
|
||||||
for key in ('_debug', '_score_source', '_enum_query'):
|
for key in ('_debug', '_score_source', '_enum_query', '_bm25_top3'):
|
||||||
if key in results:
|
if key in results:
|
||||||
expanded[key] = results[key]
|
expanded[key] = results[key]
|
||||||
return expanded
|
return expanded
|
||||||
@@ -1241,6 +1508,7 @@ class RAGEngine:
|
|||||||
sub_top_k = max(top_k, 5)
|
sub_top_k = max(top_k, 5)
|
||||||
|
|
||||||
all_results = []
|
all_results = []
|
||||||
|
_all_bm25_top3 = [] # 收集各子查询的 BM25 top3
|
||||||
for sub_q in sub_queries:
|
for sub_q in sub_queries:
|
||||||
try:
|
try:
|
||||||
sub_result = self.search_knowledge(
|
sub_result = self.search_knowledge(
|
||||||
@@ -1251,6 +1519,9 @@ class RAGEngine:
|
|||||||
)
|
)
|
||||||
if sub_result and sub_result.get('ids') and sub_result['ids'][0]:
|
if sub_result and sub_result.get('ids') and sub_result['ids'][0]:
|
||||||
all_results.append(sub_result)
|
all_results.append(sub_result)
|
||||||
|
# 收集子查询的 BM25 top3
|
||||||
|
if sub_result.get('_bm25_top3'):
|
||||||
|
_all_bm25_top3.extend(sub_result['_bm25_top3'])
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"子查询检索失败: '{sub_q}' - {e}")
|
logger.warning(f"子查询检索失败: '{sub_q}' - {e}")
|
||||||
|
|
||||||
@@ -1259,9 +1530,19 @@ class RAGEngine:
|
|||||||
|
|
||||||
# 合并去重
|
# 合并去重
|
||||||
if len(all_results) == 1:
|
if len(all_results) == 1:
|
||||||
return all_results[0]
|
merged = all_results[0]
|
||||||
|
else:
|
||||||
|
merged = self._merge_and_deduplicate(all_results, top_k)
|
||||||
|
|
||||||
return self._merge_and_deduplicate(all_results, top_k)
|
# 将收集的 BM25 top3 传递到合并结果中
|
||||||
|
if _all_bm25_top3:
|
||||||
|
_all_bm25_top3.sort(key=lambda x: x.get('bm25_score', 0), reverse=True)
|
||||||
|
_all_bm25_top3 = _all_bm25_top3[:3]
|
||||||
|
for rank, item in enumerate(_all_bm25_top3):
|
||||||
|
item['rank'] = rank + 1
|
||||||
|
merged['_bm25_top3'] = _all_bm25_top3
|
||||||
|
|
||||||
|
return merged
|
||||||
|
|
||||||
def _search_with_decomposition(
|
def _search_with_decomposition(
|
||||||
self, query, decomposer, top_k=5, allowed_levels=None,
|
self, query, decomposer, top_k=5, allowed_levels=None,
|
||||||
@@ -1293,6 +1574,7 @@ class RAGEngine:
|
|||||||
|
|
||||||
# 并行检索各子查询
|
# 并行检索各子查询
|
||||||
all_results = []
|
all_results = []
|
||||||
|
_all_bm25_top3 = [] # 收集各子查询的 BM25 top3
|
||||||
for sub_q in sub_queries:
|
for sub_q in sub_queries:
|
||||||
try:
|
try:
|
||||||
sub_result = self.search_knowledge(
|
sub_result = self.search_knowledge(
|
||||||
@@ -1303,6 +1585,9 @@ class RAGEngine:
|
|||||||
)
|
)
|
||||||
if sub_result and sub_result.get('ids') and sub_result['ids'][0]:
|
if sub_result and sub_result.get('ids') and sub_result['ids'][0]:
|
||||||
all_results.append(sub_result)
|
all_results.append(sub_result)
|
||||||
|
# 收集子查询的 BM25 top3
|
||||||
|
if sub_result.get('_bm25_top3'):
|
||||||
|
_all_bm25_top3.extend(sub_result['_bm25_top3'])
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"子查询检索失败: '{sub_q}' - {e}")
|
logger.warning(f"子查询检索失败: '{sub_q}' - {e}")
|
||||||
|
|
||||||
@@ -1315,6 +1600,14 @@ class RAGEngine:
|
|||||||
else:
|
else:
|
||||||
merged = self._merge_and_deduplicate(all_results, top_k)
|
merged = self._merge_and_deduplicate(all_results, top_k)
|
||||||
|
|
||||||
|
# 将收集的 BM25 top3 传递到合并结果中
|
||||||
|
if _all_bm25_top3:
|
||||||
|
_all_bm25_top3.sort(key=lambda x: x.get('bm25_score', 0), reverse=True)
|
||||||
|
_all_bm25_top3 = _all_bm25_top3[:3]
|
||||||
|
for rank, item in enumerate(_all_bm25_top3):
|
||||||
|
item['rank'] = rank + 1
|
||||||
|
merged['_bm25_top3'] = _all_bm25_top3
|
||||||
|
|
||||||
return merged
|
return merged
|
||||||
|
|
||||||
def _merge_and_deduplicate(self, results_list, top_k):
|
def _merge_and_deduplicate(self, results_list, top_k):
|
||||||
@@ -1387,7 +1680,7 @@ class RAGEngine:
|
|||||||
if not target_collections:
|
if not target_collections:
|
||||||
return get_empty_result()
|
return get_empty_result()
|
||||||
|
|
||||||
query_vector = self.embedding_model.encode(query).tolist()
|
query_vector = self._encode_cached(query).tolist()
|
||||||
# 扩大召回数量,以便过滤废止切片后仍有足够结果
|
# 扩大召回数量,以便过滤废止切片后仍有足够结果
|
||||||
recall_k = RERANK_CANDIDATES if (USE_RERANK or USE_HYBRID_SEARCH) else top_k
|
recall_k = RERANK_CANDIDATES if (USE_RERANK or USE_HYBRID_SEARCH) else top_k
|
||||||
recall_k = max(recall_k, top_k * RECALL_MULTIPLIER)
|
recall_k = max(recall_k, top_k * RECALL_MULTIPLIER)
|
||||||
@@ -1399,12 +1692,13 @@ class RAGEngine:
|
|||||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||||
|
|
||||||
def _query_single_collection(coll_name):
|
def _query_single_collection(coll_name):
|
||||||
"""查询单个向量库(向量 + BM25)"""
|
"""查询单个向量库(向量 + BM25),返回 (coll_results, bm25_raw_items)"""
|
||||||
coll_results = []
|
coll_results = []
|
||||||
|
bm25_raw_items = [] # 该 collection 的 BM25 原始结果
|
||||||
try:
|
try:
|
||||||
coll = self.kb_manager.get_collection(coll_name)
|
coll = self.kb_manager.get_collection(coll_name)
|
||||||
if not coll:
|
if not coll:
|
||||||
return coll_results
|
return coll_results, bm25_raw_items
|
||||||
|
|
||||||
query_kwargs = {
|
query_kwargs = {
|
||||||
"query_embeddings": [query_vector],
|
"query_embeddings": [query_vector],
|
||||||
@@ -1422,25 +1716,61 @@ class RAGEngine:
|
|||||||
if USE_HYBRID_SEARCH:
|
if USE_HYBRID_SEARCH:
|
||||||
try:
|
try:
|
||||||
bm25 = self.kb_manager.get_bm25_index(coll_name)
|
bm25 = self.kb_manager.get_bm25_index(coll_name)
|
||||||
if bm25.bm25:
|
if bm25 and bm25.bm25:
|
||||||
bm25_res = bm25.search(query, top_k=recall_k)
|
bm25_res = bm25.search(query, top_k=recall_k)
|
||||||
if source_filter and bm25_res['metadatas'] and bm25_res['metadatas'][0]:
|
# 兼容两种 BM25Index:core.bm25_index 返回 dict,knowledge.base 返回 tuple
|
||||||
|
if isinstance(bm25_res, tuple):
|
||||||
|
_ids, _docs, _metas, _dists = bm25_res
|
||||||
|
bm25_res = {
|
||||||
|
'ids': [_ids],
|
||||||
|
'documents': [_docs],
|
||||||
|
'metadatas': [_metas],
|
||||||
|
'distances': [_dists]
|
||||||
|
}
|
||||||
|
if source_filter and bm25_res['metadatas'][0]:
|
||||||
bm25_res = self._filter_results(bm25_res, lambda meta: meta.get('source') == source_filter)
|
bm25_res = self._filter_results(bm25_res, lambda meta: meta.get('source') == source_filter)
|
||||||
if bm25_res['metadatas'] and bm25_res['metadatas'][0]:
|
if bm25_res['metadatas'] and bm25_res['metadatas'][0]:
|
||||||
for meta in bm25_res['metadatas'][0]:
|
for meta in bm25_res['metadatas'][0]:
|
||||||
meta['_collection'] = coll_name
|
meta['_collection'] = coll_name
|
||||||
coll_results.append(bm25_res)
|
coll_results.append(bm25_res)
|
||||||
|
# 提取 BM25 原始 top-3(在此处直接捕获,避免与向量结果混淆)
|
||||||
|
_bm25_ids = bm25_res['ids'][0][:3]
|
||||||
|
_bm25_docs = bm25_res['documents'][0][:3]
|
||||||
|
_bm25_metas = bm25_res['metadatas'][0][:3]
|
||||||
|
_bm25_dists = (bm25_res.get('distances', [[]])[0] or [0]*3)[:3]
|
||||||
|
for i in range(len(_bm25_ids)):
|
||||||
|
# 确保 meta 包含 _collection(用于路由层注入时下游处理)
|
||||||
|
bm25_meta = _bm25_metas[i]
|
||||||
|
if '_collection' not in bm25_meta:
|
||||||
|
bm25_meta = {**bm25_meta, '_collection': coll_name}
|
||||||
|
bm25_raw_items.append({
|
||||||
|
'id': _bm25_ids[i],
|
||||||
|
'doc': _bm25_docs[i],
|
||||||
|
'meta': bm25_meta,
|
||||||
|
'bm25_score': _bm25_dists[i],
|
||||||
|
})
|
||||||
|
logger.debug(f"[BM25] {coll_name}: captured {len(bm25_raw_items)} raw items")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.debug(f"向量库 {coll_name} 检索失败: {e}")
|
logger.debug(f"向量库 {coll_name} BM25检索失败: {e}")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.debug(f"多向量库检索失败: {e}")
|
logger.debug(f"多向量库检索失败: {e}")
|
||||||
return coll_results
|
return coll_results, bm25_raw_items
|
||||||
|
|
||||||
all_results = []
|
all_results = []
|
||||||
|
_bm25_raw_top3 = []
|
||||||
with ThreadPoolExecutor(max_workers=len(target_collections)) as executor:
|
with ThreadPoolExecutor(max_workers=len(target_collections)) as executor:
|
||||||
futures = {executor.submit(_query_single_collection, name): name for name in target_collections}
|
futures = {executor.submit(_query_single_collection, name): name for name in target_collections}
|
||||||
for future in as_completed(futures):
|
for future in as_completed(futures):
|
||||||
all_results.extend(future.result())
|
coll_results, bm25_raw_items = future.result()
|
||||||
|
all_results.extend(coll_results)
|
||||||
|
_bm25_raw_top3.extend(bm25_raw_items)
|
||||||
|
|
||||||
|
# 按 bm25_score 降序取全局 top-3
|
||||||
|
if _bm25_raw_top3:
|
||||||
|
_bm25_raw_top3.sort(key=lambda x: x['bm25_score'], reverse=True)
|
||||||
|
_bm25_raw_top3 = _bm25_raw_top3[:3]
|
||||||
|
for rank, item in enumerate(_bm25_raw_top3):
|
||||||
|
item['rank'] = rank + 1
|
||||||
|
|
||||||
# ========== FAQ 检索 ==========
|
# ========== FAQ 检索 ==========
|
||||||
faq_results = self._search_faq_collection(query_vector, top_k=FAQ_RECALL_TOP_K)
|
faq_results = self._search_faq_collection(query_vector, top_k=FAQ_RECALL_TOP_K)
|
||||||
@@ -1488,6 +1818,8 @@ class RAGEngine:
|
|||||||
|
|
||||||
is_enum_query = self._is_enumeration_query(query)
|
is_enum_query = self._is_enumeration_query(query)
|
||||||
fused_results['_enum_query'] = is_enum_query
|
fused_results['_enum_query'] = is_enum_query
|
||||||
|
# 传递 BM25 原始 top-3 到路由层,用于分歧检测救援
|
||||||
|
fused_results['_bm25_top3'] = _bm25_raw_top3
|
||||||
|
|
||||||
# 章节过滤(如果查询中提到了章节)
|
# 章节过滤(如果查询中提到了章节)
|
||||||
fused_results = self._filter_by_section(fused_results, query)
|
fused_results = self._filter_by_section(fused_results, query)
|
||||||
@@ -1530,11 +1862,19 @@ class RAGEngine:
|
|||||||
# 时间衰减
|
# 时间衰减
|
||||||
fused_results = self._apply_time_decay(fused_results)
|
fused_results = self._apply_time_decay(fused_results)
|
||||||
|
|
||||||
|
# 提前附加 _debug,使聚类提升能写入调试步骤
|
||||||
|
fused_results['_debug'] = _debug
|
||||||
|
|
||||||
|
# ========== 章节聚类提升:在扩展前将低分但聚类的切片提升至种子阈值 ==========
|
||||||
|
if SECTION_CLUSTER_BOOST_ENABLED:
|
||||||
|
fused_results = self._section_cluster_boost(fused_results, query)
|
||||||
|
|
||||||
# ========== 上下文扩展:补充强命中切片周围的连续文本(rerank 之后,防止被截断)==========
|
# ========== 上下文扩展:补充强命中切片周围的连续文本(rerank 之后,防止被截断)==========
|
||||||
# Phase 3:仅对高分种子扩展邻居
|
# Phase 3:仅对高分种子扩展邻居
|
||||||
before_exp = len(fused_results['ids'][0]) if fused_results.get('ids') else 0
|
before_exp = len(fused_results['ids'][0]) if fused_results.get('ids') else 0
|
||||||
fused_results = self._expand_contiguous_chunks(fused_results, top_k=top_k,
|
fused_results = self._expand_contiguous_chunks(fused_results, top_k=top_k,
|
||||||
min_score=EXPANSION_SCORE_THRESHOLD)
|
min_score=EXPANSION_SCORE_THRESHOLD,
|
||||||
|
query=query)
|
||||||
after_exp = len(fused_results['ids'][0]) if fused_results.get('ids') else 0
|
after_exp = len(fused_results['ids'][0]) if fused_results.get('ids') else 0
|
||||||
if _debug is not None:
|
if _debug is not None:
|
||||||
_debug['steps'].append({'name': 'context_expansion', 'before': before_exp, 'after': after_exp})
|
_debug['steps'].append({'name': 'context_expansion', 'before': before_exp, 'after': after_exp})
|
||||||
@@ -1547,7 +1887,15 @@ class RAGEngine:
|
|||||||
and not (is_enum_query and ENUM_QUERY_DISABLE_TOPK_SHRINK)
|
and not (is_enum_query and ENUM_QUERY_DISABLE_TOPK_SHRINK)
|
||||||
and fused_results.get('_score_source') != 'rrf'
|
and fused_results.get('_score_source') != 'rrf'
|
||||||
):
|
):
|
||||||
top_score = 1.0 - fused_results['distances'][0][0] # 距离转相似度
|
# 根据分数来源计算相似度分数(越高越好)
|
||||||
|
score_source = fused_results.get('_score_source')
|
||||||
|
top_dist = fused_results['distances'][0][0]
|
||||||
|
if score_source == 'rerank':
|
||||||
|
# Rerank 后 distances 是相关性分数,越大越好,直接使用
|
||||||
|
top_score = top_dist
|
||||||
|
else:
|
||||||
|
# 向量距离,越小越好,转为相似度
|
||||||
|
top_score = 1.0 - top_dist
|
||||||
adjusted_k, should_retrieve, reason = self._adaptive_topk.adjust(top_score, top_k)
|
adjusted_k, should_retrieve, reason = self._adaptive_topk.adjust(top_score, top_k)
|
||||||
if "high_confidence" in reason:
|
if "high_confidence" in reason:
|
||||||
# 高置信度时截断结果
|
# 高置信度时截断结果
|
||||||
@@ -1597,7 +1945,7 @@ class RAGEngine:
|
|||||||
'metadatas': [filtered_metas],
|
'metadatas': [filtered_metas],
|
||||||
'distances': [filtered_distances]
|
'distances': [filtered_distances]
|
||||||
}
|
}
|
||||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'):
|
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context', '_bm25_top3'):
|
||||||
if key in results:
|
if key in results:
|
||||||
filtered[key] = results[key]
|
filtered[key] = results[key]
|
||||||
return filtered
|
return filtered
|
||||||
@@ -1670,7 +2018,7 @@ class RAGEngine:
|
|||||||
'metadatas': [filtered_metas],
|
'metadatas': [filtered_metas],
|
||||||
'distances': [filtered_distances]
|
'distances': [filtered_distances]
|
||||||
}
|
}
|
||||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'):
|
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context', '_bm25_top3'):
|
||||||
if key in results:
|
if key in results:
|
||||||
filtered[key] = results[key]
|
filtered[key] = results[key]
|
||||||
return filtered
|
return filtered
|
||||||
@@ -1739,12 +2087,12 @@ class RAGEngine:
|
|||||||
# === 高精度版:基于语义向量 ===
|
# === 高精度版:基于语义向量 ===
|
||||||
from core.mmr import mmr_rerank
|
from core.mmr import mmr_rerank
|
||||||
|
|
||||||
# 获取查询向量
|
# 获取查询向量(使用 embedding 缓存)
|
||||||
query_emb = np.array(self.embedding_model.encode(query))
|
query_emb = np.array(self._encode_cached(query))
|
||||||
|
|
||||||
# 批量编码所有文档
|
# 批量编码所有文档(使用 embedding 缓存)
|
||||||
docs_list = results['documents'][0]
|
docs_list = results['documents'][0]
|
||||||
all_embeddings = self.embedding_model.encode(docs_list)
|
all_embeddings = self._encode_cached(docs_list)
|
||||||
|
|
||||||
# 构建候选列表
|
# 构建候选列表
|
||||||
candidates = []
|
candidates = []
|
||||||
@@ -1786,7 +2134,7 @@ class RAGEngine:
|
|||||||
'metadatas': [[c['metadata'] for c in selected]],
|
'metadatas': [[c['metadata'] for c in selected]],
|
||||||
'distances': [[id_to_dist.get(doc_id, 0) for doc_id in selected_ids]]
|
'distances': [[id_to_dist.get(doc_id, 0) for doc_id in selected_ids]]
|
||||||
}
|
}
|
||||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'):
|
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context', '_bm25_top3'):
|
||||||
if key in results:
|
if key in results:
|
||||||
filtered[key] = results[key]
|
filtered[key] = results[key]
|
||||||
return filtered
|
return filtered
|
||||||
@@ -1826,7 +2174,7 @@ class RAGEngine:
|
|||||||
'metadatas': [[c['metadata'] for c in selected]],
|
'metadatas': [[c['metadata'] for c in selected]],
|
||||||
'distances': [[id_to_dist.get(c['id'], 0) for c in selected]]
|
'distances': [[id_to_dist.get(c['id'], 0) for c in selected]]
|
||||||
}
|
}
|
||||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'):
|
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context', '_bm25_top3'):
|
||||||
if key in results:
|
if key in results:
|
||||||
filtered[key] = results[key]
|
filtered[key] = results[key]
|
||||||
return filtered
|
return filtered
|
||||||
@@ -1935,109 +2283,15 @@ class RAGEngine:
|
|||||||
'_rerank_cached': cache_hit
|
'_rerank_cached': cache_hit
|
||||||
}
|
}
|
||||||
# 保留原有标记字段
|
# 保留原有标记字段
|
||||||
for key in ('_debug', '_score_source', '_enum_query', '_expanded_context'):
|
for key in ('_debug', '_enum_query', '_expanded_context', '_bm25_top3'):
|
||||||
if key in results:
|
if key in results:
|
||||||
reranked[key] = results[key]
|
reranked[key] = results[key]
|
||||||
|
# Rerank 后 distances 语义变为 CrossEncoder 分数,更新 _score_source
|
||||||
|
# 使自适应 TopK 能正确应用(之前 _score_source='rrf' 会导致自适应 TopK 被跳过)
|
||||||
|
reranked['_score_source'] = 'rerank'
|
||||||
return reranked
|
return reranked
|
||||||
|
|
||||||
# ---------------- 安全与工具 ----------------
|
# ---------------- 流式生成 ----------------
|
||||||
|
|
||||||
def check_restricted_documents(self, query, allowed_levels, top_k=3, role=None, department=None):
|
|
||||||
if not self._initialized:
|
|
||||||
self.initialize()
|
|
||||||
|
|
||||||
if USE_MULTI_KB and self.kb_manager and role and department:
|
|
||||||
from auth.gateway import get_accessible_collections
|
|
||||||
all_colls = [c.name for c in self.kb_manager.list_collections()]
|
|
||||||
accessible = set(get_accessible_collections(role, department, 'read'))
|
|
||||||
restricted = set(all_colls) - accessible
|
|
||||||
|
|
||||||
if not restricted:
|
|
||||||
return {"has_restricted": False, "restricted_levels": [], "restricted_sources": []}
|
|
||||||
|
|
||||||
query_vector = self.embedding_model.encode(query).tolist()
|
|
||||||
found_sources = set()
|
|
||||||
top_score = 0.0
|
|
||||||
|
|
||||||
for coll_name in restricted:
|
|
||||||
try:
|
|
||||||
coll = self.kb_manager.get_collection(coll_name)
|
|
||||||
if not coll: continue
|
|
||||||
res = coll.query(query_embeddings=[query_vector], n_results=top_k)
|
|
||||||
if res['metadatas'] and res['metadatas'][0]:
|
|
||||||
for meta in res['metadatas'][0]:
|
|
||||||
found_sources.add(meta.get('source', '未知'))
|
|
||||||
for dist in (res.get('distances', [[]])[0] or []):
|
|
||||||
if dist > top_score: top_score = dist
|
|
||||||
except Exception as e:
|
|
||||||
logger.debug(f"权限检查遍历失败: {e}")
|
|
||||||
|
|
||||||
return {
|
|
||||||
"has_restricted": len(found_sources) > 0,
|
|
||||||
"restricted_levels": [c.replace('dept_', '') for c in restricted if True][:3],
|
|
||||||
"restricted_sources": list(found_sources)[:3],
|
|
||||||
"top_restricted_score": top_score
|
|
||||||
}
|
|
||||||
|
|
||||||
if not allowed_levels:
|
|
||||||
return {"has_restricted": False, "restricted_levels": [], "restricted_sources": [], "top_restricted_score": 0.0}
|
|
||||||
|
|
||||||
restricted_levels = {"public", "internal", "confidential", "secret"} - set(allowed_levels)
|
|
||||||
if not restricted_levels:
|
|
||||||
return {"has_restricted": False, "restricted_levels": [], "restricted_sources": [], "top_restricted_score": 0.0}
|
|
||||||
|
|
||||||
query_vector = self.embedding_model.encode(query).tolist()
|
|
||||||
try:
|
|
||||||
res = self.collection.query(
|
|
||||||
query_embeddings=[query_vector],
|
|
||||||
n_results=top_k,
|
|
||||||
where={"security_level": {"$in": list(restricted_levels)}}
|
|
||||||
)
|
|
||||||
docs = res.get('documents', [[]])[0]
|
|
||||||
if not docs:
|
|
||||||
return {"has_restricted": False, "restricted_levels": [], "restricted_sources": [], "top_restricted_score": 0.0}
|
|
||||||
|
|
||||||
metas = res.get('metadatas', [[]])[0]
|
|
||||||
dists = res.get('distances', [[]])[0]
|
|
||||||
found_levels, found_sources, top_score = set(), set(), 0.0
|
|
||||||
|
|
||||||
for meta, dist in zip(metas, dists):
|
|
||||||
found_levels.add(meta.get('security_level', 'public'))
|
|
||||||
found_sources.add(meta.get('source', '未知'))
|
|
||||||
if dist > top_score: top_score = dist
|
|
||||||
|
|
||||||
return {
|
|
||||||
"has_restricted": True,
|
|
||||||
"restricted_levels": list(found_levels),
|
|
||||||
"restricted_sources": list(found_sources)[:3],
|
|
||||||
"top_restricted_score": top_score
|
|
||||||
}
|
|
||||||
except Exception as e:
|
|
||||||
logger.warning(f"受限内容检查失败: {e}")
|
|
||||||
return {"has_restricted": False, "restricted_levels": [], "restricted_sources": [], "top_restricted_score": 0.0}
|
|
||||||
|
|
||||||
def generate_answer(self, query, context):
|
|
||||||
"""底层生成答复能力"""
|
|
||||||
prompt = f"""你是一个严谨的智能助手,请根据以下参考资料回答用户的问题。
|
|
||||||
...
|
|
||||||
参考资料:
|
|
||||||
{context}
|
|
||||||
|
|
||||||
用户问题:{query}
|
|
||||||
|
|
||||||
请回答:"""
|
|
||||||
try:
|
|
||||||
from core.llm_utils import call_llm
|
|
||||||
result = call_llm(
|
|
||||||
self.llm_client,
|
|
||||||
prompt,
|
|
||||||
MODEL,
|
|
||||||
temperature=LLM_TEMPERATURE,
|
|
||||||
max_tokens=LLM_MAX_TOKENS
|
|
||||||
)
|
|
||||||
return result or f"调用大模型失败: 返回结果为空"
|
|
||||||
except Exception as e:
|
|
||||||
return f"调用大模型失败: {str(e)}"
|
|
||||||
|
|
||||||
def generate_answer_stream(self, query, context, history=None):
|
def generate_answer_stream(self, query, context, history=None):
|
||||||
"""
|
"""
|
||||||
@@ -2068,21 +2322,33 @@ class RAGEngine:
|
|||||||
"content": (
|
"content": (
|
||||||
"你是一个严谨的知识库问答助手。"
|
"你是一个严谨的知识库问答助手。"
|
||||||
"你必须且只能根据用户提供的【参考资料】回答问题。"
|
"你必须且只能根据用户提供的【参考资料】回答问题。"
|
||||||
|
"参考资料中每段内容前标有章节路径(━格式),请注意区分不同章节的内容,"
|
||||||
|
"特别当不同章节标题相似或包含相同关键词时,务必根据章节路径准确定位,不要混淆。"
|
||||||
"如果参考资料中有答案,必须引用对应内容回答,并在回答末尾标注引用编号(如[1]、[2])。"
|
"如果参考资料中有答案,必须引用对应内容回答,并在回答末尾标注引用编号(如[1]、[2])。"
|
||||||
"如果参考资料中确实没有相关信息,简短说明即可,不要编造或补充资料外的内容。"
|
"如果参考资料中确实没有相关信息,简短说明即可,不要编造或补充资料外的内容。"
|
||||||
"禁止使用参考资料以外的知识进行补充或推测。"
|
"禁止使用参考资料以外的知识进行补充或推测。"
|
||||||
|
"【重要-表格处理规则】当用户询问表格、要求展示表格内容时,你必须将参考资料中的 Markdown 表格原样输出(保留 | 分隔符和表格结构),"
|
||||||
|
"不要仅用文字描述表格存在或仅列出章节名称。如果参考资料中多个章节都有表格,"
|
||||||
|
"优先展示与用户问题最相关的表格完整内容。"
|
||||||
)
|
)
|
||||||
})
|
})
|
||||||
|
|
||||||
# 添加当前问题(带上下文)- 强化指令
|
# 添加当前问题(带上下文)- 强化指令
|
||||||
if context:
|
if context:
|
||||||
|
# 检测用户问题是否涉及表格,加入针对性指令
|
||||||
|
_table_hint = ""
|
||||||
|
# 检测上下文中是否包含 Markdown 表格(数据驱动,无需硬编码关键词)
|
||||||
|
_has_table_in_context = bool(re.search(r'\|.+\|', context)) if context else False
|
||||||
|
if _has_table_in_context:
|
||||||
|
_table_hint = "\n注意:参考资料中包含 Markdown 格式的表格数据,请务必将相关表格以原始 Markdown 表格格式完整展示在回答中,不要仅用文字描述。"
|
||||||
|
|
||||||
user_message = f"""【参考资料】
|
user_message = f"""【参考资料】
|
||||||
{context}
|
{context}
|
||||||
|
|
||||||
【用户问题】
|
【用户问题】
|
||||||
{query}
|
{query}
|
||||||
|
|
||||||
请仔细阅读以上全部参考资料后回答。如果参考资料中包含相关内容,必须引用回答并标注编号。如果资料中没有相关信息,请明确说明。"""
|
请仔细阅读以上全部参考资料后回答。注意参考资料中标有章节路径,请根据章节路径准确定位相关内容。如果参考资料中包含相关内容,必须引用回答并标注编号。如果资料中没有相关信息,请明确说明。{_table_hint}"""
|
||||||
else:
|
else:
|
||||||
user_message = query
|
user_message = query
|
||||||
|
|
||||||
|
|||||||
@@ -83,9 +83,16 @@ class IntentAnalyzer:
|
|||||||
根据对话历史和当前用户消息,输出一个 JSON 对象,包含以下字段:
|
根据对话历史和当前用户消息,输出一个 JSON 对象,包含以下字段:
|
||||||
|
|
||||||
1. **rewritten_query**: 改写后的完整问题
|
1. **rewritten_query**: 改写后的完整问题
|
||||||
- 如果问题包含指代(如"这两张图片"、"继续说"),将其改写为完整、独立的问题
|
- **指代消解**:如果问题包含指代(如"这两张图片"、"继续说"),将其改写为完整、独立的问题
|
||||||
- 例如:"分析一下这两张图片" → "分析一下对话历史中提到的图片"
|
- 例如:"分析一下这两张图片" → "分析一下对话历史中提到的图片"
|
||||||
- 如果问题本身已经完整,直接返回原文
|
- **追问补全**:如果问题是省略式追问(省略了上一轮讨论的主题实体),必须补全为完整问题
|
||||||
|
- 判断方法:当前问题缺少主语/宾语,且对话历史中可以推断出省略的实体
|
||||||
|
- 补全方法:从上一轮用户问题中提取主题实体,与追问组合成完整问题
|
||||||
|
- 例如:
|
||||||
|
- 上一轮问"吸烟点C1类是什么区?",追问"有完整表格吗?" → "吸烟点C1类有完整表格吗?"
|
||||||
|
- 上一轮问"三峡工程的投资情况",追问"建设地点在哪?" → "三峡工程的建设地点在哪?"
|
||||||
|
- 上一轮问"货源投放有哪些原则?",追问"具体内容是什么?" → "货源投放原则的具体内容是什么?"
|
||||||
|
- 如果问题本身已经完整且独立,直接返回原文
|
||||||
|
|
||||||
2. **use_context**: 布尔值
|
2. **use_context**: 布尔值
|
||||||
- true: 问题依赖历史对话中的信息,答案已经在历史回答中
|
- true: 问题依赖历史对话中的信息,答案已经在历史回答中
|
||||||
@@ -105,7 +112,9 @@ class IntentAnalyzer:
|
|||||||
- 推理类(intent="reasoning"):生成最多2个子查询
|
- 推理类(intent="reasoning"):生成最多2个子查询
|
||||||
* 原问题的检索查询
|
* 原问题的检索查询
|
||||||
* 一个补充角度的检索查询(如原因、背景、影响等),帮助获取更全面的上下文
|
* 一个补充角度的检索查询(如原因、背景、影响等),帮助获取更全面的上下文
|
||||||
- 其他类(factual/instruction/other):严格只生成1个子查询(原问题)
|
- 其他类(factual/instruction/other):严格只生成1个子查询
|
||||||
|
* 子查询应基于 rewritten_query(改写后的完整问题),而非用户原始输入
|
||||||
|
* 例如:追问"有完整表格吗?"改写为"吸烟点C1类有完整表格吗?"后,子查询应为"吸烟点C1类的完整表格内容"
|
||||||
- 不要为同一实体生成语义重叠的查询
|
- 不要为同一实体生成语义重叠的查询
|
||||||
- 子查询应保持原问题的关键词,长度20-60字符为宜
|
- 子查询应保持原问题的关键词,长度20-60字符为宜
|
||||||
|
|
||||||
@@ -224,10 +233,10 @@ class IntentAnalyzer:
|
|||||||
self._exact_cache_max = 500
|
self._exact_cache_max = 500
|
||||||
|
|
||||||
def _get_client(self):
|
def _get_client(self):
|
||||||
"""获取 LLM 客户端"""
|
"""获取 LLM 客户端(百炼快速模型)"""
|
||||||
if self._client is None:
|
if self._client is None:
|
||||||
from config import get_llm_client
|
from config import get_intent_client
|
||||||
self._client = get_llm_client()
|
self._client = get_intent_client()
|
||||||
return self._client
|
return self._client
|
||||||
|
|
||||||
def _get_cache(self):
|
def _get_cache(self):
|
||||||
@@ -299,17 +308,30 @@ class IntentAnalyzer:
|
|||||||
return self._exact_cache[exact_key]
|
return self._exact_cache[exact_key]
|
||||||
|
|
||||||
# 2. 尝试从语义缓存获取
|
# 2. 尝试从语义缓存获取
|
||||||
|
# 关键:语义缓存只用原始 query 做 embedding(不含历史),
|
||||||
|
# 避免同会话中不同问题因历史上下文污染导致误命中
|
||||||
cache = self._get_cache()
|
cache = self._get_cache()
|
||||||
if cache:
|
if cache:
|
||||||
# 使用 query + 历史关键信息作为缓存键
|
query_emb = self._get_embedding(query)
|
||||||
cache_key = self._build_cache_key(query, history)
|
|
||||||
cache_emb = self._get_embedding(cache_key)
|
|
||||||
|
|
||||||
if cache_emb is not None:
|
if query_emb is not None:
|
||||||
cached = cache.get(cache_emb)
|
cached = cache.get(query_emb)
|
||||||
if cached:
|
# 确保缓存条目是意图分析结果(包含式校验,避免新增缓存类型时误命中)
|
||||||
|
if cached and cached.get("cache_type") == "intent_analysis":
|
||||||
|
# 二次验证:检查原始 query 文本相似度
|
||||||
|
cached_query = cached.get("_raw_query", "")
|
||||||
|
if cached_query and self._query_text_similar(query, cached_query):
|
||||||
logger.info(f"意图分析缓存命中: {cached.get('reason', '')[:50]}")
|
logger.info(f"意图分析缓存命中: {cached.get('reason', '')[:50]}")
|
||||||
return IntentAnalysis.from_dict(cached)
|
return IntentAnalysis.from_dict(cached)
|
||||||
|
elif cached_query:
|
||||||
|
logger.info(
|
||||||
|
f"意图分析缓存二次验证拒绝: "
|
||||||
|
f"query='{query[:30]}' vs cached='{cached_query[:30]}'"
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# 旧缓存无 _raw_query 字段,兼容放行
|
||||||
|
logger.info(f"意图分析缓存命中(无验证): {cached.get('reason', '')[:50]}")
|
||||||
|
return IntentAnalysis.from_dict(cached)
|
||||||
else:
|
else:
|
||||||
logger.debug(f"意图分析缓存未命中,缓存状态: {cache.get_stats()}")
|
logger.debug(f"意图分析缓存未命中,缓存状态: {cache.get_stats()}")
|
||||||
else:
|
else:
|
||||||
@@ -377,9 +399,13 @@ class IntentAnalyzer:
|
|||||||
intent=intent_type
|
intent=intent_type
|
||||||
)
|
)
|
||||||
|
|
||||||
# 存入语义缓存
|
# 存入语义缓存(标记类型,避免与 RAG 回答缓存混淆)
|
||||||
if cache and cache_emb is not None:
|
# 使用仅含 query 的 embedding,不含历史,防止同会话误命中
|
||||||
cache.set(cache_emb, analysis.to_dict())
|
if cache and query_emb is not None:
|
||||||
|
cache_data = analysis.to_dict()
|
||||||
|
cache_data["cache_type"] = "intent_analysis"
|
||||||
|
cache_data["_raw_query"] = query # 供二次验证使用
|
||||||
|
cache.set(query_emb, cache_data)
|
||||||
|
|
||||||
# 存入精确匹配缓存
|
# 存入精确匹配缓存
|
||||||
if len(self._exact_cache) < self._exact_cache_max:
|
if len(self._exact_cache) < self._exact_cache_max:
|
||||||
@@ -419,6 +445,32 @@ class IntentAnalyzer:
|
|||||||
|
|
||||||
return " | ".join(parts)
|
return " | ".join(parts)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _query_text_similar(query: str, cached_query: str, threshold: float = 0.5) -> bool:
|
||||||
|
"""
|
||||||
|
判断两个 query 文本是否足够相似(字符级 Jaccard)。
|
||||||
|
用于语义缓存命中后的二次验证,防止语义相近但实际意图不同的问题误命中。
|
||||||
|
|
||||||
|
Args:
|
||||||
|
query: 当前查询
|
||||||
|
cached_query: 缓存中的原始查询
|
||||||
|
threshold: 相似度阈值,默认 0.5
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
True 表示足够相似,可以命中缓存
|
||||||
|
"""
|
||||||
|
# 精确匹配快速路径
|
||||||
|
if query.strip() == cached_query.strip():
|
||||||
|
return True
|
||||||
|
# 字符级 Jaccard 相似度
|
||||||
|
set_a = set(query)
|
||||||
|
set_b = set(cached_query)
|
||||||
|
if not set_a or not set_b:
|
||||||
|
return False
|
||||||
|
intersection = len(set_a & set_b)
|
||||||
|
union = len(set_a | set_b)
|
||||||
|
return (intersection / union) >= threshold
|
||||||
|
|
||||||
def _build_history_summary(
|
def _build_history_summary(
|
||||||
self,
|
self,
|
||||||
history: List[dict],
|
history: List[dict],
|
||||||
@@ -437,6 +489,7 @@ class IntentAnalyzer:
|
|||||||
if not history:
|
if not history:
|
||||||
return "(无历史对话)"
|
return "(无历史对话)"
|
||||||
|
|
||||||
|
import re
|
||||||
parts = []
|
parts = []
|
||||||
|
|
||||||
# 提取最近 3 轮对话
|
# 提取最近 3 轮对话
|
||||||
@@ -445,6 +498,7 @@ class IntentAnalyzer:
|
|||||||
for msg in recent_history:
|
for msg in recent_history:
|
||||||
role = "用户" if msg.get("role") == "user" else "助手"
|
role = "用户" if msg.get("role") == "user" else "助手"
|
||||||
content = msg.get("content", "")
|
content = msg.get("content", "")
|
||||||
|
original_content = content # 保留原始内容用于结构化提取
|
||||||
|
|
||||||
# 截断过长的内容
|
# 截断过长的内容
|
||||||
if len(content) > 500:
|
if len(content) > 500:
|
||||||
@@ -452,9 +506,10 @@ class IntentAnalyzer:
|
|||||||
|
|
||||||
parts.append(f"【{role}】{content}")
|
parts.append(f"【{role}】{content}")
|
||||||
|
|
||||||
# 提取图片信息
|
# 提取结构化信息(从 assistant 消息中提取章节、表格、来源等)
|
||||||
metadata = msg.get("metadata", {})
|
metadata = msg.get("metadata", {})
|
||||||
if isinstance(metadata, dict):
|
if isinstance(metadata, dict):
|
||||||
|
# 已有的图片提取
|
||||||
images = metadata.get("images", [])
|
images = metadata.get("images", [])
|
||||||
if images:
|
if images:
|
||||||
for img in images[:3]:
|
for img in images[:3]:
|
||||||
@@ -463,6 +518,43 @@ class IntentAnalyzer:
|
|||||||
img_type = img.get("type", "图片")
|
img_type = img.get("type", "图片")
|
||||||
parts.append(f" └─ {img_type}: {desc}")
|
parts.append(f" └─ {img_type}: {desc}")
|
||||||
|
|
||||||
|
# 来源文件提取
|
||||||
|
sources = metadata.get("sources", [])
|
||||||
|
if sources:
|
||||||
|
source_names = []
|
||||||
|
for s in sources[:3]:
|
||||||
|
if isinstance(s, dict):
|
||||||
|
name = s.get("source", "") or s.get("name", "")
|
||||||
|
if name:
|
||||||
|
source_names.append(name)
|
||||||
|
elif isinstance(s, str):
|
||||||
|
source_names.append(s)
|
||||||
|
if source_names:
|
||||||
|
parts.append(f" └─ 来源文件: {', '.join(source_names)}")
|
||||||
|
|
||||||
|
# collections(检索知识库)提取
|
||||||
|
colls = metadata.get("collections", [])
|
||||||
|
if colls:
|
||||||
|
parts.append(f" └─ 检索知识库: {', '.join(colls)}")
|
||||||
|
|
||||||
|
# 从 assistant 原始内容中提取章节路径和表格结构
|
||||||
|
if role == "助手" and original_content:
|
||||||
|
# 提取章节路径:━ xxx ━ 格式
|
||||||
|
sections = re.findall(r'━\s*(.+?)\s*━', original_content)
|
||||||
|
if sections:
|
||||||
|
unique_sections = list(dict.fromkeys(sections)) # 去重保序
|
||||||
|
parts.append(f" └─ 涉及章节: {'; '.join(unique_sections[:3])}")
|
||||||
|
|
||||||
|
# 提取表格列名:| A | B | C | 格式的表头行
|
||||||
|
table_headers = re.findall(r'^\|\s*(.+?)\s*\|', original_content, re.MULTILINE)
|
||||||
|
if table_headers:
|
||||||
|
# 取第一个表格的列名
|
||||||
|
first_header = table_headers[0]
|
||||||
|
cols = [c.strip() for c in first_header.split('|') if c.strip()]
|
||||||
|
# 排除分隔符行(--- 格式)
|
||||||
|
if cols and not all(re.match(r'^[-:]+$', c) for c in cols):
|
||||||
|
parts.append(f" └─ 含表格,列名: {', '.join(cols[:6])}")
|
||||||
|
|
||||||
# 添加图片上下文
|
# 添加图片上下文
|
||||||
if context_images:
|
if context_images:
|
||||||
parts.append("\n【上下文中的图片】")
|
parts.append("\n【上下文中的图片】")
|
||||||
|
|||||||
@@ -72,16 +72,35 @@ def call_llm(
|
|||||||
|
|
||||||
content = response.choices[0].message.content
|
content = response.choices[0].message.content
|
||||||
|
|
||||||
# 推理模型兼容:content 为空时尝试从 reasoning_content 提取
|
# 推理模型兼容(mimo-v2.5 等):
|
||||||
|
# 推理模型思考链消耗大量 token(~1000),max_tokens 不足时 content 为空,
|
||||||
|
# 全部输出进入 reasoning_content。此处从思考链中提取有效内容。
|
||||||
if not content or not content.strip():
|
if not content or not content.strip():
|
||||||
reasoning = getattr(response.choices[0].message, 'reasoning_content', None)
|
reasoning = getattr(response.choices[0].message, 'reasoning_content', None)
|
||||||
if reasoning and reasoning.strip():
|
if reasoning and reasoning.strip():
|
||||||
# 从思维链中提取 JSON 块作为内容
|
# 先去掉 <think>...</think> 标签
|
||||||
json_match = re.search(r'\{[\s\S]*\}', reasoning)
|
cleaned = re.sub(r'', '', reasoning, flags=re.DOTALL).strip()
|
||||||
|
if cleaned:
|
||||||
|
logger.info("LLM: content为空,从reasoning_content提取内容")
|
||||||
|
# 尝试提取 JSON 对象(兼容结构化响应场景)
|
||||||
|
json_match = re.search(r'\{[\s\S]*\}', cleaned)
|
||||||
if json_match:
|
if json_match:
|
||||||
logger.info("LLM: content为空,从reasoning_content提取JSON")
|
try:
|
||||||
|
json.loads(json_match.group())
|
||||||
return json_match.group().strip()
|
return json_match.group().strip()
|
||||||
logger.warning("LLM 返回空 content(可能需要增大 max_tokens)")
|
except (json.JSONDecodeError, ValueError):
|
||||||
|
pass
|
||||||
|
# 尝试提取 JSON 数组
|
||||||
|
bracket_match = re.search(r'\[[\s\S]*\]', cleaned)
|
||||||
|
if bracket_match:
|
||||||
|
try:
|
||||||
|
json.loads(bracket_match.group())
|
||||||
|
return bracket_match.group().strip()
|
||||||
|
except (json.JSONDecodeError, ValueError):
|
||||||
|
pass
|
||||||
|
# 纯文本响应:直接返回清理后的内容
|
||||||
|
return cleaned
|
||||||
|
logger.warning("LLM 返回空 content 且 reasoning_content 也无法提取(可能需要增大 max_tokens)")
|
||||||
return None
|
return None
|
||||||
|
|
||||||
return content.strip()
|
return content.strip()
|
||||||
@@ -95,7 +114,7 @@ def call_llm_stream(
|
|||||||
prompt: str,
|
prompt: str,
|
||||||
model: str,
|
model: str,
|
||||||
temperature: float = 0.3,
|
temperature: float = 0.3,
|
||||||
max_tokens: int = 1000,
|
max_tokens: int = 3000,
|
||||||
messages: List[dict] = None,
|
messages: List[dict] = None,
|
||||||
error_prefix: str = "[错误]",
|
error_prefix: str = "[错误]",
|
||||||
**kwargs
|
**kwargs
|
||||||
@@ -104,13 +123,14 @@ def call_llm_stream(
|
|||||||
流式 LLM 调用(生成器封装)
|
流式 LLM 调用(生成器封装)
|
||||||
|
|
||||||
自动处理流式响应,逐块 yield 文本内容。
|
自动处理流式响应,逐块 yield 文本内容。
|
||||||
|
兼容推理模型(mimo-v2.5 等):当 content 为空时回退到 reasoning_content。
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
client: OpenAI 客户端实例
|
client: OpenAI 客户端实例
|
||||||
prompt: 用户提示
|
prompt: 用户提示
|
||||||
model: 模型名称
|
model: 模型名称
|
||||||
temperature: 温度参数
|
temperature: 温度参数
|
||||||
max_tokens: 最大 token 数
|
max_tokens: 最大 token 数(推理模型需留足思考链预算)
|
||||||
messages: 完整消息列表
|
messages: 完整消息列表
|
||||||
error_prefix: 错误时的前缀
|
error_prefix: 错误时的前缀
|
||||||
**kwargs: 其他参数
|
**kwargs: 其他参数
|
||||||
@@ -135,9 +155,33 @@ def call_llm_stream(
|
|||||||
**kwargs
|
**kwargs
|
||||||
)
|
)
|
||||||
|
|
||||||
|
content_yielded = False
|
||||||
|
reasoning_buffer = []
|
||||||
|
|
||||||
for chunk in stream:
|
for chunk in stream:
|
||||||
if chunk.choices and chunk.choices[0].delta.content:
|
if not chunk.choices:
|
||||||
yield chunk.choices[0].delta.content
|
continue
|
||||||
|
delta = chunk.choices[0].delta
|
||||||
|
|
||||||
|
# 正常 content 输出
|
||||||
|
if hasattr(delta, 'content') and delta.content:
|
||||||
|
content_yielded = True
|
||||||
|
yield delta.content
|
||||||
|
continue
|
||||||
|
|
||||||
|
# 推理模型:reasoning_content(思考链)
|
||||||
|
rc = getattr(delta, 'reasoning_content', None)
|
||||||
|
if rc:
|
||||||
|
reasoning_buffer.append(rc)
|
||||||
|
|
||||||
|
# 回退:content 为空但 reasoning_content 有内容(推理模型 token 不足时)
|
||||||
|
if not content_yielded and reasoning_buffer:
|
||||||
|
reasoning_text = ''.join(reasoning_buffer)
|
||||||
|
# 去掉 <think>...</think> 标签
|
||||||
|
cleaned = re.sub(r'', '', reasoning_text, flags=re.DOTALL).strip()
|
||||||
|
if cleaned:
|
||||||
|
logger.info("流式 LLM: content为空,从reasoning_content提取内容")
|
||||||
|
yield cleaned
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"LLM 流式调用失败: {e}")
|
logger.error(f"LLM 流式调用失败: {e}")
|
||||||
@@ -352,7 +396,7 @@ def quick_yes_no(
|
|||||||
if keywords is None:
|
if keywords is None:
|
||||||
keywords = ["是", "需要", "yes", "true"]
|
keywords = ["是", "需要", "yes", "true"]
|
||||||
|
|
||||||
result = call_llm(client, prompt, model, temperature=0, max_tokens=10)
|
result = call_llm(client, prompt, model, temperature=0, max_tokens=128)
|
||||||
if result is None:
|
if result is None:
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
|||||||
73
core/mmr.py
73
core/mmr.py
@@ -108,22 +108,44 @@ def mmr_rerank(
|
|||||||
return selected
|
return selected
|
||||||
|
|
||||||
|
|
||||||
|
def _tokenize_words(text: str) -> set:
|
||||||
|
"""
|
||||||
|
使用 jieba 分词并过滤噪声,返回有意义的词集合。
|
||||||
|
|
||||||
|
过滤规则:
|
||||||
|
- 去除单字符词(如 "的", "了", "在")—— 这些是停用词,对区分文档无意义
|
||||||
|
- 去除纯数字 / 纯标点
|
||||||
|
- 保留 2 字及以上的实词
|
||||||
|
"""
|
||||||
|
import jieba
|
||||||
|
words = set()
|
||||||
|
for w in jieba.cut(text):
|
||||||
|
w = w.strip()
|
||||||
|
if len(w) >= 2 and not w.isdigit():
|
||||||
|
words.add(w)
|
||||||
|
return words
|
||||||
|
|
||||||
|
|
||||||
def mmr_filter_by_content(
|
def mmr_filter_by_content(
|
||||||
candidates: List[Dict],
|
candidates: List[Dict],
|
||||||
top_k: int = 30,
|
top_k: int = 30,
|
||||||
similarity_threshold: float = 0.9
|
similarity_threshold: float = 0.85
|
||||||
) -> List[Dict]:
|
) -> List[Dict]:
|
||||||
"""
|
"""
|
||||||
基于内容相似度的去重(简化版,不需要 embedding)
|
基于 jieba 词级 Jaccard 相似度的去重(不需要 embedding)
|
||||||
|
|
||||||
|
与旧版字符级 set(text) 的区别:
|
||||||
|
- 旧版:set("安全生产管理制度") → {'安','全','生','产',...},中文文档间字符集合高度重叠
|
||||||
|
- 新版:jieba 分词 → {"安全生产", "管理制度", ...},词级集合区分度高
|
||||||
|
|
||||||
适用于:
|
适用于:
|
||||||
- 没有 embedding 的情况
|
- MMR_USE_EMBEDDING=False 时的快速去重
|
||||||
- 快速去重场景
|
- 避免 CPU 编码 100+ 文档的 50 秒开销
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
candidates: 候选文档列表
|
candidates: 候选文档列表
|
||||||
top_k: 返回数量
|
top_k: 返回数量
|
||||||
similarity_threshold: 相似度阈值,超过则视为重复
|
similarity_threshold: 相似度阈值,超过则视为重复(默认 0.85)
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
去重后的候选文档列表
|
去重后的候选文档列表
|
||||||
@@ -134,25 +156,32 @@ def mmr_filter_by_content(
|
|||||||
if len(candidates) <= top_k:
|
if len(candidates) <= top_k:
|
||||||
return candidates
|
return candidates
|
||||||
|
|
||||||
selected = []
|
# 预分词:对所有候选文档一次性分词,避免重复调用 jieba.cut
|
||||||
remaining = candidates.copy()
|
word_sets = []
|
||||||
|
for c in candidates:
|
||||||
|
content = c.get('content', c.get('document', ''))[:500]
|
||||||
|
word_sets.append(_tokenize_words(content))
|
||||||
|
|
||||||
while len(selected) < top_k and remaining:
|
selected_indices = []
|
||||||
current = remaining.pop(0)
|
|
||||||
|
for i in range(len(candidates)):
|
||||||
|
if len(selected_indices) >= top_k:
|
||||||
|
break
|
||||||
|
|
||||||
|
current_words = word_sets[i]
|
||||||
|
if not current_words:
|
||||||
|
# 空内容直接保留
|
||||||
|
selected_indices.append(i)
|
||||||
|
continue
|
||||||
|
|
||||||
# 检查是否与已选内容重复
|
|
||||||
is_duplicate = False
|
is_duplicate = False
|
||||||
current_content = current.get('content', current.get('document', ''))[:200]
|
for j in selected_indices:
|
||||||
|
selected_words = word_sets[j]
|
||||||
|
if not selected_words:
|
||||||
|
continue
|
||||||
|
|
||||||
for s in selected:
|
intersection = len(current_words & selected_words)
|
||||||
s_content = s.get('content', s.get('document', ''))[:200]
|
union = len(current_words | selected_words)
|
||||||
|
|
||||||
# 简单的 Jaccard 相似度
|
|
||||||
words1 = set(current_content)
|
|
||||||
words2 = set(s_content)
|
|
||||||
if words1 and words2:
|
|
||||||
intersection = len(words1 & words2)
|
|
||||||
union = len(words1 | words2)
|
|
||||||
similarity = intersection / union if union > 0 else 0
|
similarity = intersection / union if union > 0 else 0
|
||||||
|
|
||||||
if similarity > similarity_threshold:
|
if similarity > similarity_threshold:
|
||||||
@@ -160,9 +189,9 @@ def mmr_filter_by_content(
|
|||||||
break
|
break
|
||||||
|
|
||||||
if not is_duplicate:
|
if not is_duplicate:
|
||||||
selected.append(current)
|
selected_indices.append(i)
|
||||||
|
|
||||||
return selected
|
return [candidates[i] for i in selected_indices]
|
||||||
|
|
||||||
|
|
||||||
# ==================== 测试 ====================
|
# ==================== 测试 ====================
|
||||||
|
|||||||
@@ -382,18 +382,34 @@ class BM25Index:
|
|||||||
|
|
||||||
def add_documents(self, ids: List[str], documents: List[str], metadatas: List[dict]) -> None:
|
def add_documents(self, ids: List[str], documents: List[str], metadatas: List[dict]) -> None:
|
||||||
"""
|
"""
|
||||||
添加文档到索引(会覆盖原有索引)
|
添加文档到索引(追加模式,自动去重)
|
||||||
|
|
||||||
|
如果 ID 已存在则更新对应文档,否则追加新文档。
|
||||||
|
添加后自动重建 BM25 索引。
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
ids: 文档 ID 列表
|
ids: 文档 ID 列表
|
||||||
documents: 文档内容列表
|
documents: 文档内容列表
|
||||||
metadatas: 文档元数据列表
|
metadatas: 文档元数据列表
|
||||||
"""
|
"""
|
||||||
self.ids = ids
|
# 建立已有 ID -> 索引位置 的映射,用于去重
|
||||||
self.documents = documents
|
existing_map = {doc_id: idx for idx, doc_id in enumerate(self.ids)}
|
||||||
self.metadatas = metadatas
|
|
||||||
if documents:
|
for i, doc_id in enumerate(ids):
|
||||||
tokenized = [self.tokenize(doc) for doc in documents]
|
if doc_id in existing_map:
|
||||||
|
# 更新已有文档
|
||||||
|
pos = existing_map[doc_id]
|
||||||
|
self.documents[pos] = documents[i]
|
||||||
|
self.metadatas[pos] = metadatas[i]
|
||||||
|
else:
|
||||||
|
# 追加新文档
|
||||||
|
existing_map[doc_id] = len(self.ids)
|
||||||
|
self.ids.append(doc_id)
|
||||||
|
self.documents.append(documents[i])
|
||||||
|
self.metadatas.append(metadatas[i])
|
||||||
|
|
||||||
|
if self.documents:
|
||||||
|
tokenized = [self.tokenize(doc) for doc in self.documents]
|
||||||
self.bm25 = BM25Okapi(tokenized)
|
self.bm25 = BM25Okapi(tokenized)
|
||||||
|
|
||||||
def search(self, query: str, top_k: int = 10) -> Tuple[List[str], List[str], List[dict], List[float]]:
|
def search(self, query: str, top_k: int = 10) -> Tuple[List[str], List[str], List[dict], List[float]]:
|
||||||
|
|||||||
@@ -134,6 +134,9 @@ class CollectionMixin:
|
|||||||
"""
|
"""
|
||||||
from .base import BM25Index
|
from .base import BM25Index
|
||||||
|
|
||||||
|
# 从磁盘重新加载元数据,确保多 worker 进程间状态一致
|
||||||
|
self._metadata = self._load_metadata()
|
||||||
|
|
||||||
if not kb_name or not kb_name.replace('_', '').isalnum():
|
if not kb_name or not kb_name.replace('_', '').isalnum():
|
||||||
return False, "向量库名称只能包含字母、数字和下划线"
|
return False, "向量库名称只能包含字母、数字和下划线"
|
||||||
|
|
||||||
@@ -190,6 +193,9 @@ class CollectionMixin:
|
|||||||
Returns:
|
Returns:
|
||||||
更新成功返回 True,向量库不存在返回 False
|
更新成功返回 True,向量库不存在返回 False
|
||||||
"""
|
"""
|
||||||
|
# 从磁盘重新加载元数据,确保多 worker 进程间状态一致
|
||||||
|
self._metadata = self._load_metadata()
|
||||||
|
|
||||||
collections = self._metadata.get("collections", {})
|
collections = self._metadata.get("collections", {})
|
||||||
if kb_name not in collections:
|
if kb_name not in collections:
|
||||||
return False
|
return False
|
||||||
@@ -228,6 +234,9 @@ class CollectionMixin:
|
|||||||
"""
|
"""
|
||||||
import shutil
|
import shutil
|
||||||
|
|
||||||
|
# 从磁盘重新加载元数据,确保多 worker 进程间状态一致
|
||||||
|
self._metadata = self._load_metadata()
|
||||||
|
|
||||||
if kb_name == PUBLIC_KB_NAME:
|
if kb_name == PUBLIC_KB_NAME:
|
||||||
return False, "公开知识库不能删除"
|
return False, "公开知识库不能删除"
|
||||||
|
|
||||||
@@ -322,6 +331,21 @@ class CollectionMixin:
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"清理版本记录失败: {e}")
|
logger.warning(f"清理版本记录失败: {e}")
|
||||||
|
|
||||||
|
# 清理不再被引用的图片和 VLM 缓存文件
|
||||||
|
# 注意:此时 ChromaDB collection 已删除,cleanup_image_orphans 会扫描
|
||||||
|
# 所有剩余 collection,仅该 collection 引用的图片会被识别为孤儿
|
||||||
|
try:
|
||||||
|
from knowledge.image_cleanup import cleanup_image_orphans
|
||||||
|
cleanup_result = cleanup_image_orphans(self)
|
||||||
|
if cleanup_result['deleted_images'] or cleanup_result['deleted_caches']:
|
||||||
|
logger.info(
|
||||||
|
f"清理孤儿文件: {cleanup_result['deleted_images']} 图片 + "
|
||||||
|
f"{cleanup_result['deleted_caches']} VLM缓存, "
|
||||||
|
f"释放 {cleanup_result['freed_bytes']/1024:.1f} KB"
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"清理孤儿文件失败: {e}")
|
||||||
|
|
||||||
if kb_name in self._metadata.get("collections", {}):
|
if kb_name in self._metadata.get("collections", {}):
|
||||||
del self._metadata["collections"][kb_name]
|
del self._metadata["collections"][kb_name]
|
||||||
self._save_metadata()
|
self._save_metadata()
|
||||||
@@ -348,6 +372,8 @@ class CollectionMixin:
|
|||||||
- department: 所属部门
|
- department: 所属部门
|
||||||
- description: 描述
|
- description: 描述
|
||||||
"""
|
"""
|
||||||
|
# 从磁盘重新加载元数据,确保多 worker 进程间状态一致
|
||||||
|
self._metadata = self._load_metadata()
|
||||||
result = []
|
result = []
|
||||||
|
|
||||||
# 扫描 base_path 下的所有子目录作为向量库
|
# 扫描 base_path 下的所有子目录作为向量库
|
||||||
@@ -382,7 +408,10 @@ class CollectionMixin:
|
|||||||
|
|
||||||
self._save_metadata()
|
self._save_metadata()
|
||||||
|
|
||||||
|
stale_collections = []
|
||||||
|
|
||||||
for name, info in self._metadata.get("collections", {}).items():
|
for name, info in self._metadata.get("collections", {}).items():
|
||||||
|
try:
|
||||||
collection = self.get_collection(name)
|
collection = self.get_collection(name)
|
||||||
result.append(CollectionInfo(
|
result.append(CollectionInfo(
|
||||||
name=name,
|
name=name,
|
||||||
@@ -392,6 +421,18 @@ class CollectionMixin:
|
|||||||
department=info.get("department", ""),
|
department=info.get("department", ""),
|
||||||
description=info.get("description", "")
|
description=info.get("description", "")
|
||||||
))
|
))
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(
|
||||||
|
f"跳过异常向量库 '{name}': {e},可能是 ChromaDB 数据目录已丢失"
|
||||||
|
)
|
||||||
|
stale_collections.append(name)
|
||||||
|
|
||||||
|
# 清理元数据中指向已失效集合的条目
|
||||||
|
if stale_collections:
|
||||||
|
for name in stale_collections:
|
||||||
|
self._metadata.get("collections", {}).pop(name, None)
|
||||||
|
logger.info(f"清理失效向量库元数据: {name}")
|
||||||
|
self._save_metadata()
|
||||||
|
|
||||||
return result
|
return result
|
||||||
|
|
||||||
@@ -405,4 +446,6 @@ class CollectionMixin:
|
|||||||
Returns:
|
Returns:
|
||||||
存在返回 True,不存在返回 False
|
存在返回 True,不存在返回 False
|
||||||
"""
|
"""
|
||||||
|
# 从磁盘重新加载元数据,确保多 worker 进程间状态一致
|
||||||
|
self._metadata = self._load_metadata()
|
||||||
return kb_name in self._metadata.get("collections", {})
|
return kb_name in self._metadata.get("collections", {})
|
||||||
|
|||||||
@@ -72,6 +72,18 @@ class DocumentMixin:
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"清理版本记录失败: {e}")
|
logger.warning(f"清理版本记录失败: {e}")
|
||||||
|
|
||||||
|
# 清理不再被引用的图片和 VLM 缓存文件
|
||||||
|
try:
|
||||||
|
from knowledge.image_cleanup import cleanup_image_orphans
|
||||||
|
cleanup_result = cleanup_image_orphans(self, collections=[kb_name])
|
||||||
|
if cleanup_result['deleted_images'] or cleanup_result['deleted_caches']:
|
||||||
|
logger.info(
|
||||||
|
f"清理孤儿文件: {cleanup_result['deleted_images']} 图片 + "
|
||||||
|
f"{cleanup_result['deleted_caches']} VLM缓存"
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"清理孤儿文件失败: {e}")
|
||||||
|
|
||||||
logger.info(f"从 {kb_name} 删除文档: {filename}, 片段数: {deleted}")
|
logger.info(f"从 {kb_name} 删除文档: {filename}, 片段数: {deleted}")
|
||||||
return deleted
|
return deleted
|
||||||
|
|
||||||
|
|||||||
140
knowledge/image_cleanup.py
Normal file
140
knowledge/image_cleanup.py
Normal file
@@ -0,0 +1,140 @@
|
|||||||
|
"""
|
||||||
|
图片/VLM缓存孤儿文件清理模块
|
||||||
|
|
||||||
|
提供可被 document.py / collection.py 调用的清理函数,
|
||||||
|
也可被 cleanup_orphans.py 独立脚本使用。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import hashlib
|
||||||
|
import logging
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
IMAGES_DIR = Path(".data/images")
|
||||||
|
VLM_CACHE_DIR = Path(".data/cache/vlm")
|
||||||
|
|
||||||
|
|
||||||
|
def compute_file_hash(file_path: str) -> str:
|
||||||
|
"""计算文件 MD5"""
|
||||||
|
with open(file_path, 'rb') as f:
|
||||||
|
return hashlib.md5(f.read()).hexdigest()
|
||||||
|
|
||||||
|
|
||||||
|
def collect_referenced_images(manager, collections=None) -> set:
|
||||||
|
"""
|
||||||
|
从 ChromaDB 收集所有被引用的图片文件名。
|
||||||
|
|
||||||
|
Args:
|
||||||
|
manager: KnowledgeBaseManager 实例
|
||||||
|
collections: 限定知识库列表,None 表示全部
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
set of image filenames (e.g., {"185a7a75d246.png", ...})
|
||||||
|
"""
|
||||||
|
referenced = set()
|
||||||
|
|
||||||
|
if collections:
|
||||||
|
kb_names = collections
|
||||||
|
else:
|
||||||
|
kb_names = [c.name if hasattr(c, 'name') else str(c)
|
||||||
|
for c in manager.list_collections()]
|
||||||
|
|
||||||
|
for kb_name in kb_names:
|
||||||
|
try:
|
||||||
|
col = manager.get_collection(kb_name)
|
||||||
|
except Exception:
|
||||||
|
continue
|
||||||
|
|
||||||
|
results = col.get(include=['metadatas'])
|
||||||
|
if not results['ids']:
|
||||||
|
continue
|
||||||
|
|
||||||
|
for meta in results['metadatas']:
|
||||||
|
image_path = meta.get('image_path', '')
|
||||||
|
if image_path:
|
||||||
|
referenced.add(os.path.basename(image_path))
|
||||||
|
|
||||||
|
return referenced
|
||||||
|
|
||||||
|
|
||||||
|
def cleanup_image_orphans(manager, collections=None, dry_run=False) -> dict:
|
||||||
|
"""
|
||||||
|
清理不再被任何 ChromaDB 切片引用的图片和 VLM 缓存文件。
|
||||||
|
|
||||||
|
Args:
|
||||||
|
manager: KnowledgeBaseManager 实例
|
||||||
|
collections: 限定知识库列表,None 表示全部
|
||||||
|
dry_run: True 时只返回孤儿列表不实际删除
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
{
|
||||||
|
'orphan_images': [(filepath, filename, size_bytes)],
|
||||||
|
'orphan_caches': [(filepath, filename, size_bytes)],
|
||||||
|
'deleted_images': int,
|
||||||
|
'deleted_caches': int,
|
||||||
|
'freed_bytes': int
|
||||||
|
}
|
||||||
|
"""
|
||||||
|
result = {
|
||||||
|
'orphan_images': [],
|
||||||
|
'orphan_caches': [],
|
||||||
|
'deleted_images': 0,
|
||||||
|
'deleted_caches': 0,
|
||||||
|
'freed_bytes': 0
|
||||||
|
}
|
||||||
|
|
||||||
|
# 1. 收集引用
|
||||||
|
referenced = collect_referenced_images(manager, collections)
|
||||||
|
|
||||||
|
# 2. 查找孤儿图片
|
||||||
|
if IMAGES_DIR.exists():
|
||||||
|
for f in IMAGES_DIR.iterdir():
|
||||||
|
if f.is_file() and f.name not in referenced:
|
||||||
|
result['orphan_images'].append((str(f), f.name, f.stat().st_size))
|
||||||
|
|
||||||
|
# 3. 查找孤儿 VLM 缓存(图片已删除则缓存也应是孤儿)
|
||||||
|
if VLM_CACHE_DIR.exists():
|
||||||
|
referenced_hashes = set()
|
||||||
|
for filename in referenced:
|
||||||
|
full_path = IMAGES_DIR / filename
|
||||||
|
if full_path.exists():
|
||||||
|
try:
|
||||||
|
img_hash = compute_file_hash(str(full_path))
|
||||||
|
referenced_hashes.add(img_hash)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
for f in VLM_CACHE_DIR.iterdir():
|
||||||
|
if f.is_file() and f.suffix == '.txt':
|
||||||
|
cache_hash = f.stem
|
||||||
|
if cache_hash not in referenced_hashes:
|
||||||
|
result['orphan_caches'].append((str(f), f.name, f.stat().st_size))
|
||||||
|
|
||||||
|
# 4. 删除
|
||||||
|
if not dry_run:
|
||||||
|
for filepath, filename, size in result['orphan_images']:
|
||||||
|
try:
|
||||||
|
os.remove(filepath)
|
||||||
|
result['deleted_images'] += 1
|
||||||
|
result['freed_bytes'] += size
|
||||||
|
except OSError as e:
|
||||||
|
logger.warning(f"删除孤儿图片失败: {filename} - {e}")
|
||||||
|
|
||||||
|
for filepath, filename, size in result['orphan_caches']:
|
||||||
|
try:
|
||||||
|
os.remove(filepath)
|
||||||
|
result['deleted_caches'] += 1
|
||||||
|
result['freed_bytes'] += size
|
||||||
|
except OSError as e:
|
||||||
|
logger.warning(f"删除孤儿缓存失败: {filename} - {e}")
|
||||||
|
|
||||||
|
if result['deleted_images'] or result['deleted_caches']:
|
||||||
|
logger.info(
|
||||||
|
f"清理孤儿: {result['deleted_images']} 图片 + "
|
||||||
|
f"{result['deleted_caches']} 缓存, "
|
||||||
|
f"释放 {result['freed_bytes']/1024:.1f} KB"
|
||||||
|
)
|
||||||
|
|
||||||
|
return result
|
||||||
@@ -25,7 +25,20 @@ def compute_file_hash(file_path: str) -> str:
|
|||||||
return hashlib.md5(file_path.encode()).hexdigest()
|
return hashlib.md5(file_path.encode()).hexdigest()
|
||||||
|
|
||||||
|
|
||||||
async def lazy_vlm_description(chunk_id: str, image_path: str, kb_name: str, metadata: dict = None) -> str:
|
def _get_embedding_model():
|
||||||
|
"""从 RAGEngine 获取 embedding 模型(KnowledgeBaseManager 上没有此属性)"""
|
||||||
|
try:
|
||||||
|
from core.engine import get_engine
|
||||||
|
engine = get_engine()
|
||||||
|
if not engine._initialized:
|
||||||
|
engine.initialize()
|
||||||
|
return engine.embedding_model
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"获取 embedding 模型失败: {e}")
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
async def lazy_vlm_description(chunk_id: str, image_path: str, kb_name: str, metadata: dict = None, defer_chromadb: bool = False) -> str:
|
||||||
"""
|
"""
|
||||||
懒加载 VLM 描述
|
懒加载 VLM 描述
|
||||||
|
|
||||||
@@ -36,6 +49,7 @@ async def lazy_vlm_description(chunk_id: str, image_path: str, kb_name: str, met
|
|||||||
image_path: 图片路径(相对路径或绝对路径)
|
image_path: 图片路径(相对路径或绝对路径)
|
||||||
kb_name: 知识库名称
|
kb_name: 知识库名称
|
||||||
metadata: 图片元数据(包含 section、page、caption、上下文等)
|
metadata: 图片元数据(包含 section、page、caption、上下文等)
|
||||||
|
defer_chromadb: 为 True 时跳过 ChromaDB 更新(仅写文件缓存),避免后台线程写锁竞争
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
VLM 生成的图片描述
|
VLM 生成的图片描述
|
||||||
@@ -49,23 +63,45 @@ async def lazy_vlm_description(chunk_id: str, image_path: str, kb_name: str, met
|
|||||||
else:
|
else:
|
||||||
full_image_path = image_path
|
full_image_path = image_path
|
||||||
|
|
||||||
# 1. 检查缓存
|
# 1. 检查缓存(空缓存视为无效,需重新生成)
|
||||||
img_hash = compute_file_hash(full_image_path)
|
img_hash = compute_file_hash(full_image_path)
|
||||||
cache_file = VLM_CACHE_DIR / f"{img_hash}.txt"
|
cache_file = VLM_CACHE_DIR / f"{img_hash}.txt"
|
||||||
if cache_file.exists():
|
if cache_file.exists():
|
||||||
|
cached = cache_file.read_text(encoding='utf-8')
|
||||||
|
if len(cached.strip()) >= 5:
|
||||||
logger.info(f"VLM 缓存命中: {image_path}")
|
logger.info(f"VLM 缓存命中: {image_path}")
|
||||||
return cache_file.read_text(encoding='utf-8')
|
return cached
|
||||||
|
else:
|
||||||
|
logger.warning(f"VLM 缓存内容过短({len(cached.strip())}字符),删除并重新生成: {image_path}")
|
||||||
|
try:
|
||||||
|
cache_file.unlink()
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
|
|
||||||
# 2. 调用 VLM(传入元数据)
|
# 2. 调用 VLM(传入元数据)
|
||||||
logger.info(f"VLM 懒加载: {image_path}")
|
logger.info(f"VLM 懒加载: {image_path}")
|
||||||
kb_manager = get_kb_manager()
|
kb_manager = get_kb_manager()
|
||||||
description = kb_manager._generate_image_description(full_image_path, metadata=metadata)
|
description = kb_manager._generate_image_description(full_image_path, metadata=metadata)
|
||||||
|
|
||||||
# 3. 写入缓存
|
# 3. 空描述保护:VLM 返回内容过短时不写入缓存和向量库
|
||||||
|
if not description or len(description.strip()) < 5:
|
||||||
|
logger.warning(f"VLM 返回描述过短({len(description.strip()) if description else 0}字符),跳过缓存和向量库更新: {image_path}")
|
||||||
|
return description or ''
|
||||||
|
|
||||||
|
# 4. 写入缓存
|
||||||
VLM_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
VLM_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||||
cache_file.write_text(description, encoding='utf-8')
|
cache_file.write_text(description, encoding='utf-8')
|
||||||
|
|
||||||
# 4. 更新向量库(metadata + embedding)
|
# 5. 更新向量库(metadata + embedding),需校验 chunk_id 非空
|
||||||
|
# defer_chromadb=True 时跳过(后台线程只写缓存,避免 SQLite 写锁竞争)
|
||||||
|
if defer_chromadb:
|
||||||
|
logger.info(f"延迟 ChromaDB 更新(仅写缓存): {chunk_id}")
|
||||||
|
return description
|
||||||
|
|
||||||
|
if not chunk_id:
|
||||||
|
logger.warning("chunk_id 为空,跳过向量库更新")
|
||||||
|
return description
|
||||||
|
|
||||||
try:
|
try:
|
||||||
collection = kb_manager.get_collection(kb_name)
|
collection = kb_manager.get_collection(kb_name)
|
||||||
result = collection.get(ids=[chunk_id], include=['metadatas'])
|
result = collection.get(ids=[chunk_id], include=['metadatas'])
|
||||||
@@ -79,7 +115,7 @@ async def lazy_vlm_description(chunk_id: str, image_path: str, kb_name: str, met
|
|||||||
|
|
||||||
# 更新 embedding(使用 VLM 描述重新计算向量)
|
# 更新 embedding(使用 VLM 描述重新计算向量)
|
||||||
# 这样 VLM 描述中的关键词(如"发电量")才能参与相似度检索
|
# 这样 VLM 描述中的关键词(如"发电量")才能参与相似度检索
|
||||||
embedding_model = kb_manager.embedding_model
|
embedding_model = _get_embedding_model()
|
||||||
if embedding_model:
|
if embedding_model:
|
||||||
new_vector = embedding_model.encode(description).tolist()
|
new_vector = embedding_model.encode(description).tolist()
|
||||||
if isinstance(new_vector[0], list):
|
if isinstance(new_vector[0], list):
|
||||||
@@ -91,20 +127,21 @@ async def lazy_vlm_description(chunk_id: str, image_path: str, kb_name: str, met
|
|||||||
embeddings=[new_vector],
|
embeddings=[new_vector],
|
||||||
documents=[description] # 同时更新 document 字段
|
documents=[description] # 同时更新 document 字段
|
||||||
)
|
)
|
||||||
logger.info(f"已更新向量库 embedding: {chunk_id}")
|
logger.info(f"已更新向量库(embedding+metadata): {chunk_id}")
|
||||||
else:
|
else:
|
||||||
# 无 embedding 模型时只更新 metadata
|
# 无 embedding 模型时只更新 metadata
|
||||||
collection.update(
|
collection.update(
|
||||||
ids=[chunk_id],
|
ids=[chunk_id],
|
||||||
metadatas=[new_metadata]
|
metadatas=[new_metadata]
|
||||||
)
|
)
|
||||||
|
logger.info(f"已更新向量库(仅metadata,无embedding模型): {chunk_id}")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"更新向量库失败: {e}")
|
logger.warning(f"更新向量库失败: {e}")
|
||||||
|
|
||||||
return description
|
return description
|
||||||
|
|
||||||
|
|
||||||
async def lazy_table_summary(chunk_id: str, table_md: str, kb_name: str) -> str:
|
async def lazy_table_summary(chunk_id: str, table_md: str, kb_name: str, defer_chromadb: bool = False) -> str:
|
||||||
"""
|
"""
|
||||||
懒加载表格摘要
|
懒加载表格摘要
|
||||||
|
|
||||||
@@ -114,35 +151,58 @@ async def lazy_table_summary(chunk_id: str, table_md: str, kb_name: str) -> str:
|
|||||||
chunk_id: 切片 ID
|
chunk_id: 切片 ID
|
||||||
table_md: 表格 Markdown 内容
|
table_md: 表格 Markdown 内容
|
||||||
kb_name: 知识库名称
|
kb_name: 知识库名称
|
||||||
|
defer_chromadb: 为 True 时跳过 ChromaDB 更新(仅写文件缓存),避免后台线程写锁竞争
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
LLM 生成的表格摘要
|
LLM 生成的表格摘要
|
||||||
"""
|
"""
|
||||||
from knowledge.manager import get_kb_manager
|
from knowledge.manager import get_kb_manager
|
||||||
|
|
||||||
# 1. 检查缓存
|
# 1. 检查缓存(空缓存视为无效)
|
||||||
table_hash = hashlib.md5(table_md.encode()).hexdigest()
|
table_hash = hashlib.md5(table_md.encode()).hexdigest()
|
||||||
cache_file = LLM_CACHE_DIR / f"{table_hash}.txt"
|
cache_file = LLM_CACHE_DIR / f"{table_hash}.txt"
|
||||||
if cache_file.exists():
|
if cache_file.exists():
|
||||||
|
cached = cache_file.read_text(encoding='utf-8')
|
||||||
|
if len(cached.strip()) >= 5:
|
||||||
logger.info(f"LLM 缓存命中: {chunk_id}")
|
logger.info(f"LLM 缓存命中: {chunk_id}")
|
||||||
return cache_file.read_text(encoding='utf-8')
|
return cached
|
||||||
|
else:
|
||||||
|
logger.warning(f"LLM 缓存内容过短({len(cached.strip())}字符),删除并重新生成: {chunk_id}")
|
||||||
|
try:
|
||||||
|
cache_file.unlink()
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
|
|
||||||
# 2. 调用 LLM
|
# 2. 调用 LLM
|
||||||
logger.info(f"LLM 懒加载: {chunk_id}")
|
logger.info(f"LLM 懒加载: {chunk_id}")
|
||||||
kb_manager = get_kb_manager()
|
kb_manager = get_kb_manager()
|
||||||
summary = kb_manager._generate_table_summary(table_md, None)
|
summary = kb_manager._generate_table_summary(table_md, None)
|
||||||
|
|
||||||
|
# 空摘要保护
|
||||||
|
if not summary or len(summary.strip()) < 5:
|
||||||
|
logger.warning(f"LLM 返回摘要过短,跳过缓存和向量库更新: {chunk_id}")
|
||||||
|
return summary or ''
|
||||||
|
|
||||||
# 3. 写入缓存
|
# 3. 写入缓存
|
||||||
LLM_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
LLM_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||||
cache_file.write_text(summary, encoding='utf-8')
|
cache_file.write_text(summary, encoding='utf-8')
|
||||||
|
|
||||||
# 4. 更新向量库(可选)
|
# 4. 更新向量库,需校验 chunk_id 非空
|
||||||
|
# defer_chromadb=True 时跳过(后台线程只写缓存,避免 SQLite 写锁竞争)
|
||||||
|
if defer_chromadb:
|
||||||
|
logger.info(f"延迟 ChromaDB 更新(仅写缓存): {chunk_id}")
|
||||||
|
return summary
|
||||||
|
|
||||||
|
if not chunk_id:
|
||||||
|
logger.warning("chunk_id 为空,跳过表格向量库更新")
|
||||||
|
return summary
|
||||||
try:
|
try:
|
||||||
collection = kb_manager.get_collection(kb_name)
|
collection = kb_manager.get_collection(kb_name)
|
||||||
result = collection.get(ids=[chunk_id], include=['metadatas'])
|
result = collection.get(ids=[chunk_id], include=['metadatas'])
|
||||||
if result['metadatas']:
|
if result['metadatas']:
|
||||||
# 新增摘要切片
|
# 新增摘要切片(需要 embedding 模型)
|
||||||
embedding_model = kb_manager.embedding_model
|
embedding_model = _get_embedding_model()
|
||||||
|
if embedding_model:
|
||||||
vector = embedding_model.encode(summary).tolist()
|
vector = embedding_model.encode(summary).tolist()
|
||||||
if isinstance(vector[0], list):
|
if isinstance(vector[0], list):
|
||||||
vector = vector[0]
|
vector = vector[0]
|
||||||
@@ -157,7 +217,10 @@ async def lazy_table_summary(chunk_id: str, table_md: str, kb_name: str) -> str:
|
|||||||
'original_doc_id': chunk_id
|
'original_doc_id': chunk_id
|
||||||
}]
|
}]
|
||||||
)
|
)
|
||||||
# 更新原切片标记
|
logger.info(f"已新增摘要切片(embedding): {chunk_id}_summary")
|
||||||
|
else:
|
||||||
|
logger.info(f"跳过摘要切片(无embedding模型): {chunk_id}")
|
||||||
|
# 更新原切片标记(不依赖 embedding 模型)
|
||||||
collection.update(
|
collection.update(
|
||||||
ids=[chunk_id],
|
ids=[chunk_id],
|
||||||
metadatas=[{**result['metadatas'][0], 'has_summary': True}]
|
metadatas=[{**result['metadatas'][0], 'has_summary': True}]
|
||||||
@@ -168,7 +231,7 @@ async def lazy_table_summary(chunk_id: str, table_md: str, kb_name: str) -> str:
|
|||||||
return summary
|
return summary
|
||||||
|
|
||||||
|
|
||||||
async def enhance_retrieved_chunks(contexts: list, query: str, kb_name: str):
|
async def enhance_retrieved_chunks(contexts: list, query: str, kb_name: str, defer_chromadb: bool = False):
|
||||||
"""
|
"""
|
||||||
检索后增强:按需调用 LLM/VLM
|
检索后增强:按需调用 LLM/VLM
|
||||||
|
|
||||||
@@ -176,8 +239,12 @@ async def enhance_retrieved_chunks(contexts: list, query: str, kb_name: str):
|
|||||||
contexts: 检索上下文列表
|
contexts: 检索上下文列表
|
||||||
query: 用户查询
|
query: 用户查询
|
||||||
kb_name: 知识库名称
|
kb_name: 知识库名称
|
||||||
|
defer_chromadb: 为 True 时后台线程只写文件缓存,不更新 ChromaDB(避免写锁竞争)
|
||||||
"""
|
"""
|
||||||
|
import re
|
||||||
|
|
||||||
for ctx in contexts:
|
for ctx in contexts:
|
||||||
|
try:
|
||||||
meta = ctx.get('meta', {})
|
meta = ctx.get('meta', {})
|
||||||
chunk_type = meta.get('chunk_type', 'text')
|
chunk_type = meta.get('chunk_type', 'text')
|
||||||
image_path = meta.get('image_path', '')
|
image_path = meta.get('image_path', '')
|
||||||
@@ -185,14 +252,11 @@ async def enhance_retrieved_chunks(contexts: list, query: str, kb_name: str):
|
|||||||
# 图片切片:懒加载 VLM 描述
|
# 图片切片:懒加载 VLM 描述
|
||||||
if chunk_type in ('image', 'chart') and not meta.get('has_vlm_desc'):
|
if chunk_type in ('image', 'chart') and not meta.get('has_vlm_desc'):
|
||||||
if image_path:
|
if image_path:
|
||||||
try:
|
|
||||||
# 从 doc 字段中提取图号(上下文可能包含"见图2.5"等)
|
# 从 doc 字段中提取图号(上下文可能包含"见图2.5"等)
|
||||||
doc_text = ctx.get('doc', '')
|
doc_text = ctx.get('doc', '')
|
||||||
import re
|
|
||||||
|
|
||||||
# 提取图号(从前文/后文中)
|
# 提取图号(从前文/后文中)
|
||||||
figure_number = ""
|
figure_number = ""
|
||||||
# 匹配 "见图2.5"、"图2.5"、"见图 2.5" 等
|
|
||||||
fig_match = re.search(r'[见如]?图\s*(\d+\.?\d*)', doc_text)
|
fig_match = re.search(r'[见如]?图\s*(\d+\.?\d*)', doc_text)
|
||||||
if fig_match:
|
if fig_match:
|
||||||
figure_number = fig_match.group(1)
|
figure_number = fig_match.group(1)
|
||||||
@@ -210,19 +274,19 @@ async def enhance_retrieved_chunks(contexts: list, query: str, kb_name: str):
|
|||||||
'page': meta.get('page'),
|
'page': meta.get('page'),
|
||||||
'caption': meta.get('caption', ''),
|
'caption': meta.get('caption', ''),
|
||||||
'source': meta.get('source', ''),
|
'source': meta.get('source', ''),
|
||||||
'figure_number': figure_number, # 添加提取的图号
|
'figure_number': figure_number,
|
||||||
'doc_text': doc_text # 添加完整文档文本
|
'doc_text': doc_text
|
||||||
}
|
}
|
||||||
vlm_desc = await lazy_vlm_description(
|
vlm_desc = await lazy_vlm_description(
|
||||||
meta.get('id', ''),
|
meta.get('chunk_id', ''),
|
||||||
image_path,
|
image_path,
|
||||||
kb_name,
|
kb_name,
|
||||||
metadata=image_metadata
|
metadata=image_metadata,
|
||||||
|
defer_chromadb=defer_chromadb
|
||||||
)
|
)
|
||||||
|
if vlm_desc:
|
||||||
ctx['doc'] = vlm_desc
|
ctx['doc'] = vlm_desc
|
||||||
ctx['vlm_enhanced'] = True
|
ctx['vlm_enhanced'] = True
|
||||||
except Exception as e:
|
|
||||||
logger.warning(f"VLM 懒加载失败: {e}")
|
|
||||||
|
|
||||||
# 表格切片:同时处理摘要和关联图片的 VLM 描述
|
# 表格切片:同时处理摘要和关联图片的 VLM 描述
|
||||||
elif chunk_type == 'table':
|
elif chunk_type == 'table':
|
||||||
@@ -230,58 +294,53 @@ async def enhance_retrieved_chunks(contexts: list, query: str, kb_name: str):
|
|||||||
|
|
||||||
# 1. 懒加载表格摘要(高分切片)
|
# 1. 懒加载表格摘要(高分切片)
|
||||||
if not meta.get('has_summary'):
|
if not meta.get('has_summary'):
|
||||||
score = meta.get('score', 0)
|
score = ctx.get('score', 0)
|
||||||
if score > 0.7: # 只对高相关表格生成摘要
|
if score > 0.7:
|
||||||
try:
|
|
||||||
summary = await lazy_table_summary(
|
summary = await lazy_table_summary(
|
||||||
meta.get('id', ''),
|
meta.get('chunk_id', ''),
|
||||||
doc_text,
|
doc_text,
|
||||||
kb_name
|
kb_name,
|
||||||
|
defer_chromadb=defer_chromadb
|
||||||
)
|
)
|
||||||
# 摘要作为补充信息
|
if summary:
|
||||||
ctx['summary'] = summary
|
ctx['summary'] = summary
|
||||||
ctx['llm_enhanced'] = True
|
ctx['llm_enhanced'] = True
|
||||||
except Exception as e:
|
|
||||||
logger.warning(f"表格摘要懒加载失败: {e}")
|
|
||||||
|
|
||||||
# 2. 表格有关联图片时,懒加载 VLM 描述
|
# 2. 表格有关联图片时,懒加载 VLM 描述
|
||||||
if image_path and not meta.get('has_vlm_desc'):
|
if image_path and not meta.get('has_vlm_desc'):
|
||||||
try:
|
|
||||||
import re
|
|
||||||
|
|
||||||
# 提取表号(如 "表2.2"、"见表2.1")
|
# 提取表号(如 "表2.2"、"见表2.1")
|
||||||
table_number = ""
|
table_number = ""
|
||||||
# 匹配 "表2.2"、"见表2.2"、"见表 2.2" 等
|
|
||||||
table_match = re.search(r'[见如]?表\s*(\d+\.?\d*)', doc_text)
|
table_match = re.search(r'[见如]?表\s*(\d+\.?\d*)', doc_text)
|
||||||
if table_match:
|
if table_match:
|
||||||
table_number = table_match.group(1)
|
table_number = table_match.group(1)
|
||||||
|
|
||||||
# 如果 doc 中没有,尝试从 section 中提取
|
|
||||||
section = meta.get('section') or meta.get('section_path', '')
|
section = meta.get('section') or meta.get('section_path', '')
|
||||||
if not table_number and section:
|
if not table_number and section:
|
||||||
table_match = re.search(r'[见如]?表\s*(\d+\.?\d*)', section)
|
table_match = re.search(r'[见如]?表\s*(\d+\.?\d*)', section)
|
||||||
if table_match:
|
if table_match:
|
||||||
table_number = table_match.group(1)
|
table_number = table_match.group(1)
|
||||||
|
|
||||||
# 构建表格图片元数据
|
|
||||||
table_image_metadata = {
|
table_image_metadata = {
|
||||||
'section': section,
|
'section': section,
|
||||||
'page': meta.get('page'),
|
'page': meta.get('page'),
|
||||||
'caption': meta.get('caption', ''),
|
'caption': meta.get('caption', ''),
|
||||||
'source': meta.get('source', ''),
|
'source': meta.get('source', ''),
|
||||||
'table_number': table_number, # 表号
|
'table_number': table_number,
|
||||||
'figure_number': table_number, # 兼容字段
|
'figure_number': table_number,
|
||||||
'doc_text': doc_text,
|
'doc_text': doc_text,
|
||||||
'is_table': True # 标记为表格图片
|
'is_table': True
|
||||||
}
|
}
|
||||||
vlm_desc = await lazy_vlm_description(
|
vlm_desc = await lazy_vlm_description(
|
||||||
meta.get('id', ''),
|
meta.get('chunk_id', ''),
|
||||||
image_path,
|
image_path,
|
||||||
kb_name,
|
kb_name,
|
||||||
metadata=table_image_metadata
|
metadata=table_image_metadata,
|
||||||
|
defer_chromadb=defer_chromadb
|
||||||
)
|
)
|
||||||
# 表格图片描述作为补充信息
|
if vlm_desc:
|
||||||
ctx['image_description'] = vlm_desc
|
ctx['image_description'] = vlm_desc
|
||||||
ctx['vlm_enhanced'] = True
|
ctx['vlm_enhanced'] = True
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"表格图片 VLM 懒加载失败: {e}")
|
chunk_id = ctx.get('meta', {}).get('chunk_id', '?')
|
||||||
|
logger.warning(f"增强切片失败(chunk_id={chunk_id}): {e}")
|
||||||
|
|||||||
@@ -29,11 +29,17 @@
|
|||||||
import os
|
import os
|
||||||
import json
|
import json
|
||||||
import threading
|
import threading
|
||||||
|
try:
|
||||||
|
import fcntl
|
||||||
|
_HAS_FCNTL = True
|
||||||
|
except ImportError:
|
||||||
|
_HAS_FCNTL = False # Windows 环境无 fcntl
|
||||||
from typing import List, Dict, Optional, Tuple
|
from typing import List, Dict, Optional, Tuple
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
import logging
|
import logging
|
||||||
|
|
||||||
import chromadb
|
import chromadb
|
||||||
|
from bs4 import BeautifulSoup
|
||||||
|
|
||||||
# 从 base.py 导入基础类和常量
|
# 从 base.py 导入基础类和常量
|
||||||
from .base import (
|
from .base import (
|
||||||
@@ -134,22 +140,36 @@ class KnowledgeBaseManager(
|
|||||||
logger.info(f"知识库管理器初始化完成,路径: {self.base_path},发现 {len(existing_kbs)} 个向量库: {existing_kbs}")
|
logger.info(f"知识库管理器初始化完成,路径: {self.base_path},发现 {len(existing_kbs)} 个向量库: {existing_kbs}")
|
||||||
|
|
||||||
def _load_metadata(self) -> dict:
|
def _load_metadata(self) -> dict:
|
||||||
"""加载元数据"""
|
"""加载元数据(带文件锁,确保多 worker 进程间一致)"""
|
||||||
metadata_path = os.path.join(self.base_path, KB_METADATA_FILE)
|
metadata_path = os.path.join(self.base_path, KB_METADATA_FILE)
|
||||||
if os.path.exists(metadata_path):
|
if os.path.exists(metadata_path):
|
||||||
try:
|
try:
|
||||||
with open(metadata_path, 'r', encoding='utf-8') as f:
|
with open(metadata_path, 'r', encoding='utf-8') as f:
|
||||||
|
if _HAS_FCNTL:
|
||||||
|
fcntl.flock(f.fileno(), fcntl.LOCK_SH)
|
||||||
|
try:
|
||||||
return json.load(f)
|
return json.load(f)
|
||||||
|
finally:
|
||||||
|
if _HAS_FCNTL:
|
||||||
|
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"加载元数据失败: {e}")
|
logger.error(f"加载元数据失败: {e}")
|
||||||
return {"collections": {}}
|
return {"collections": {}}
|
||||||
|
|
||||||
def _save_metadata(self):
|
def _save_metadata(self):
|
||||||
"""保存元数据"""
|
"""保存元数据(带文件锁,防止并发写入数据覆盖)"""
|
||||||
metadata_path = os.path.join(self.base_path, KB_METADATA_FILE)
|
metadata_path = os.path.join(self.base_path, KB_METADATA_FILE)
|
||||||
try:
|
try:
|
||||||
with open(metadata_path, 'w', encoding='utf-8') as f:
|
with open(metadata_path, 'w', encoding='utf-8') as f:
|
||||||
|
if _HAS_FCNTL:
|
||||||
|
fcntl.flock(f.fileno(), fcntl.LOCK_EX)
|
||||||
|
try:
|
||||||
json.dump(self._metadata, f, ensure_ascii=False, indent=2)
|
json.dump(self._metadata, f, ensure_ascii=False, indent=2)
|
||||||
|
f.flush()
|
||||||
|
os.fsync(f.fileno())
|
||||||
|
finally:
|
||||||
|
if _HAS_FCNTL:
|
||||||
|
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"保存元数据失败: {e}")
|
logger.error(f"保存元数据失败: {e}")
|
||||||
|
|
||||||
@@ -316,6 +336,7 @@ class KnowledgeBaseManager(
|
|||||||
"section": section_path,
|
"section": section_path,
|
||||||
"status": "active",
|
"status": "active",
|
||||||
"version": "v1",
|
"version": "v1",
|
||||||
|
"text_level": getattr(chunk, 'text_level', 0), # 标题级别(0=正文,1=h1,2=h2,3=h3),供检索管线层次感知
|
||||||
}
|
}
|
||||||
|
|
||||||
if extra_metadata:
|
if extra_metadata:
|
||||||
@@ -328,6 +349,25 @@ class KnowledgeBaseManager(
|
|||||||
if hasattr(chunk, 'image_path') and chunk.image_path:
|
if hasattr(chunk, 'image_path') and chunk.image_path:
|
||||||
metadata['image_path'] = chunk.image_path
|
metadata['image_path'] = chunk.image_path
|
||||||
|
|
||||||
|
# bbox 坐标(PDF 有,DOCX 无,需 None 保护)
|
||||||
|
# _build_citation() 从 metadata 读取 bbox 做引用定位
|
||||||
|
chunk_bbox = getattr(chunk, 'bbox', None)
|
||||||
|
if chunk_bbox:
|
||||||
|
metadata['bbox'] = json.dumps(chunk_bbox)
|
||||||
|
|
||||||
|
# MinerU 结构化元数据(表格类型、嵌套层级、图片子类型)
|
||||||
|
chunk_table_type = getattr(chunk, 'table_type', '')
|
||||||
|
if chunk_table_type:
|
||||||
|
metadata['table_type'] = chunk_table_type
|
||||||
|
|
||||||
|
chunk_nest_level = getattr(chunk, 'table_nest_level', '')
|
||||||
|
if chunk_nest_level:
|
||||||
|
metadata['table_nest_level'] = str(chunk_nest_level)
|
||||||
|
|
||||||
|
chunk_sub_type = getattr(chunk, 'sub_type', '')
|
||||||
|
if chunk_sub_type:
|
||||||
|
metadata['sub_type'] = chunk_sub_type
|
||||||
|
|
||||||
# 生成向量
|
# 生成向量
|
||||||
try:
|
try:
|
||||||
embedding = embedding_model.encode(semantic_content).tolist()
|
embedding = embedding_model.encode(semantic_content).tolist()
|
||||||
@@ -393,6 +433,11 @@ class KnowledgeBaseManager(
|
|||||||
if len(chunks) < 2:
|
if len(chunks) < 2:
|
||||||
return chunks
|
return chunks
|
||||||
|
|
||||||
|
# 检测页码是否可靠:若所有 chunk 的 page_start 相同(如 Word 文档 page_idx 全为 0),
|
||||||
|
# 则页码信息不可用,需要启用降级合并规则
|
||||||
|
page_values = set(getattr(c, 'page_start', 0) for c in chunks)
|
||||||
|
pages_unavailable = len(page_values) <= 1
|
||||||
|
|
||||||
merged_chunks = []
|
merged_chunks = []
|
||||||
i = 0
|
i = 0
|
||||||
merge_count = 0
|
merge_count = 0
|
||||||
@@ -405,6 +450,7 @@ class KnowledgeBaseManager(
|
|||||||
# 查找下一个表格(跳过中间的"续表"文本)
|
# 查找下一个表格(跳过中间的"续表"文本)
|
||||||
next_table_idx = None
|
next_table_idx = None
|
||||||
next_chunk = None
|
next_chunk = None
|
||||||
|
intermediate_texts = [] # 收集中间文本用于降级判断
|
||||||
|
|
||||||
for j in range(i + 1, min(i + 4, len(chunks))): # 最多向前看3个切片
|
for j in range(i + 1, min(i + 4, len(chunks))): # 最多向前看3个切片
|
||||||
candidate = chunks[j]
|
candidate = chunks[j]
|
||||||
@@ -419,7 +465,11 @@ class KnowledgeBaseManager(
|
|||||||
elif candidate_type == 'text' and ('续表' in candidate_title or '续表' in candidate_content):
|
elif candidate_type == 'text' and ('续表' in candidate_title or '续表' in candidate_content):
|
||||||
# 遇到"续表"文本,继续查找下一个表格
|
# 遇到"续表"文本,继续查找下一个表格
|
||||||
continue
|
continue
|
||||||
elif candidate_type not in ('text',):
|
elif candidate_type == 'text':
|
||||||
|
# 非"续表"文本,收集后停止查找
|
||||||
|
intermediate_texts.append(candidate)
|
||||||
|
break
|
||||||
|
else:
|
||||||
# 遇到非文本类型,停止查找
|
# 遇到非文本类型,停止查找
|
||||||
break
|
break
|
||||||
|
|
||||||
@@ -439,12 +489,22 @@ class KnowledgeBaseManager(
|
|||||||
# 获取内容(用于检测"续表")
|
# 获取内容(用于检测"续表")
|
||||||
next_content = getattr(next_chunk, 'content', '')
|
next_content = getattr(next_chunk, 'content', '')
|
||||||
|
|
||||||
|
# 通用/无意义标题集合,这些标题不能用于"标题相似"判定
|
||||||
|
_GENERIC_TITLES = {'表格', 'table', '表格', ''}
|
||||||
|
|
||||||
# 判断是否为跨页表格
|
# 判断是否为跨页表格
|
||||||
is_cross_page = False
|
is_cross_page = False
|
||||||
|
|
||||||
# 规则1: 页码连续(如果页码有效)
|
# 规则1: 页码连续(如果页码有效)
|
||||||
page_valid = curr_page_end > 0 and next_page_start > 0
|
page_valid = curr_page_end > 0 and next_page_start > 0
|
||||||
if page_valid and curr_page_end + 1 == next_page_start:
|
if page_valid and curr_page_end + 1 == next_page_start:
|
||||||
|
# 页码连续时,还需标题匹配或为通用标题才合并
|
||||||
|
# 避免把不同页面上不相关的表格错误合并
|
||||||
|
if curr_title == next_title or curr_title in _GENERIC_TITLES and next_title in _GENERIC_TITLES:
|
||||||
|
is_cross_page = True
|
||||||
|
elif curr_title and next_title:
|
||||||
|
clean_next_r1 = next_title.replace('续表', '').strip()
|
||||||
|
if curr_title in clean_next_r1 or clean_next_r1 in curr_title:
|
||||||
is_cross_page = True
|
is_cross_page = True
|
||||||
|
|
||||||
# 规则2: 第二个表格标题或内容包含"续表"
|
# 规则2: 第二个表格标题或内容包含"续表"
|
||||||
@@ -452,11 +512,36 @@ class KnowledgeBaseManager(
|
|||||||
is_cross_page = True
|
is_cross_page = True
|
||||||
|
|
||||||
# 规则3: 标题相似(去掉"续表"后比较)
|
# 规则3: 标题相似(去掉"续表"后比较)
|
||||||
elif curr_title and next_title:
|
# 排除通用标题(如"表格"),防止把所有标题为"表格"的相邻表格都误合并
|
||||||
|
elif (curr_title and next_title
|
||||||
|
and curr_title not in _GENERIC_TITLES
|
||||||
|
and next_title not in _GENERIC_TITLES):
|
||||||
clean_next = next_title.replace('续表', '').strip()
|
clean_next = next_title.replace('续表', '').strip()
|
||||||
if curr_title in clean_next or clean_next in curr_title:
|
if clean_next and (curr_title in clean_next or clean_next in curr_title):
|
||||||
is_cross_page = True
|
is_cross_page = True
|
||||||
|
|
||||||
|
# 规则4(降级): 页码不可用(如 Word 文档 page_idx 全为 0)
|
||||||
|
# 仅当页码信息缺失时才启用此规则,避免 PDF 正常页码时被误合并
|
||||||
|
if (not is_cross_page
|
||||||
|
and pages_unavailable
|
||||||
|
and curr_title in _GENERIC_TITLES
|
||||||
|
and next_title in _GENERIC_TITLES):
|
||||||
|
# 检查中间文本是否暗示跨页延续(空、短文本、续表标记等)
|
||||||
|
has_separating_content = False
|
||||||
|
for text_chunk in intermediate_texts:
|
||||||
|
tc = (getattr(text_chunk, 'content', '') or '').strip()
|
||||||
|
tt = (getattr(text_chunk, 'title', '') or '').strip()
|
||||||
|
if not tc:
|
||||||
|
continue # 空文本不算分隔
|
||||||
|
if '续表' in tc or '续表' in tt:
|
||||||
|
continue # 续表标记,说明是跨页
|
||||||
|
# 有实质性中间内容(如分类标题"A3类:xxx"),不合并
|
||||||
|
has_separating_content = True
|
||||||
|
break
|
||||||
|
if not has_separating_content:
|
||||||
|
is_cross_page = True
|
||||||
|
logger.debug(f"降级合并(页码不可用): '{curr_title}' + '{next_title}'")
|
||||||
|
|
||||||
if is_cross_page:
|
if is_cross_page:
|
||||||
# 执行合并
|
# 执行合并
|
||||||
merge_count += 1
|
merge_count += 1
|
||||||
@@ -466,17 +551,61 @@ class KnowledgeBaseManager(
|
|||||||
curr_html = getattr(current, 'table_html', '') or ''
|
curr_html = getattr(current, 'table_html', '') or ''
|
||||||
next_html = getattr(next_chunk, 'table_html', '') or ''
|
next_html = getattr(next_chunk, 'table_html', '') or ''
|
||||||
if curr_html and next_html:
|
if curr_html and next_html:
|
||||||
# 合并两个表格的 HTML
|
# 正确合并两个表格的 HTML:
|
||||||
|
# 将第二个表格的 <tr> 行追加到第一个表格中
|
||||||
|
# (而非简单拼接两个 <table>,否则 html_table_to_markdown
|
||||||
|
# 的 soup.find('table') 只能找到第一个表格)
|
||||||
|
try:
|
||||||
|
soup1 = BeautifulSoup(curr_html, 'html.parser')
|
||||||
|
soup2 = BeautifulSoup(next_html, 'html.parser')
|
||||||
|
table1 = soup1.find('table')
|
||||||
|
table2 = soup2.find('table')
|
||||||
|
if table1 and table2:
|
||||||
|
# 从第二个表格提取数据行
|
||||||
|
next_rows = table2.find_all('tr')
|
||||||
|
# 跳过与第一个表格表头重复的行
|
||||||
|
# 对比第一行而非所有 th(find_all('th') 会匹配
|
||||||
|
# 整个表格的 th,无法与单行做列表比较)
|
||||||
|
first_row_t1 = table1.find('tr')
|
||||||
|
if first_row_t1 and next_rows:
|
||||||
|
row1_texts = [c.get_text(strip=True) for c in first_row_t1.find_all(['th', 'td'])]
|
||||||
|
row2_texts = [c.get_text(strip=True) for c in next_rows[0].find_all(['th', 'td'])]
|
||||||
|
if row1_texts and row2_texts and row1_texts == row2_texts:
|
||||||
|
next_rows = next_rows[1:]
|
||||||
|
logger.debug("跨页表格合并: 跳过了重复的表头行")
|
||||||
|
for row in next_rows:
|
||||||
|
table1.append(row)
|
||||||
|
current.table_html = str(soup1)
|
||||||
|
logger.debug(f"跨页表格 HTML 合并成功: 追加了 {len(next_rows)} 行")
|
||||||
|
else:
|
||||||
current.table_html = curr_html + '\n' + next_html
|
current.table_html = curr_html + '\n' + next_html
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"跨页表格 HTML 合并异常: {e},回退到简单拼接")
|
||||||
|
current.table_html = curr_html + '\n' + next_html
|
||||||
|
elif not curr_html and next_html:
|
||||||
|
current.table_html = next_html
|
||||||
|
|
||||||
# 合并 image_path 到 images
|
# 合并 image_path 和嵌入图片到 images
|
||||||
curr_img = getattr(current, 'image_path', None)
|
curr_img = getattr(current, 'image_path', None)
|
||||||
next_img = getattr(next_chunk, 'image_path', None)
|
next_img = getattr(next_chunk, 'image_path', None)
|
||||||
merged_images = []
|
curr_images = getattr(current, 'images', None) or []
|
||||||
if curr_img:
|
next_images = getattr(next_chunk, 'images', None) or []
|
||||||
|
|
||||||
|
# 合并两个表格的所有图片(image_path + 嵌入图片)
|
||||||
|
merged_images = list(curr_images) # 保留当前表格的嵌入图片
|
||||||
|
# 添加 image_path 图片(如果不在列表中)
|
||||||
|
existing_ids = {img.get('id', '') for img in merged_images if isinstance(img, dict)}
|
||||||
|
if curr_img and curr_img not in existing_ids:
|
||||||
merged_images.append({'id': curr_img, 'page': curr_page_end})
|
merged_images.append({'id': curr_img, 'page': curr_page_end})
|
||||||
if next_img:
|
existing_ids.add(curr_img)
|
||||||
|
for img in next_images: # 添加下一个表格的嵌入图片
|
||||||
|
img_id = img.get('id', '') if isinstance(img, dict) else ''
|
||||||
|
if img_id and img_id not in existing_ids:
|
||||||
|
merged_images.append(img)
|
||||||
|
existing_ids.add(img_id)
|
||||||
|
if next_img and next_img not in existing_ids:
|
||||||
merged_images.append({'id': next_img, 'page': next_page_start})
|
merged_images.append({'id': next_img, 'page': next_page_start})
|
||||||
|
|
||||||
if merged_images:
|
if merged_images:
|
||||||
current.images = merged_images
|
current.images = merged_images
|
||||||
# 保留第一个图片作为主 image_path
|
# 保留第一个图片作为主 image_path
|
||||||
@@ -525,7 +654,7 @@ class KnowledgeBaseManager(
|
|||||||
try:
|
try:
|
||||||
from config import get_llm_client, DASHSCOPE_MODEL
|
from config import get_llm_client, DASHSCOPE_MODEL
|
||||||
client = get_llm_client()
|
client = get_llm_client()
|
||||||
summary = call_llm(client, prompt, DASHSCOPE_MODEL, max_tokens=100)
|
summary = call_llm(client, prompt, DASHSCOPE_MODEL, max_tokens=2048)
|
||||||
return summary.strip() if summary else ""
|
return summary.strip() if summary else ""
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"生成表格摘要失败: {e}")
|
logger.warning(f"生成表格摘要失败: {e}")
|
||||||
@@ -584,13 +713,31 @@ class KnowledgeBaseManager(
|
|||||||
]
|
]
|
||||||
}
|
}
|
||||||
],
|
],
|
||||||
max_tokens=200
|
max_tokens=2048 # mimo-v2.5 推理模型思考链消耗 ~1000 token,需留足输出空间
|
||||||
)
|
)
|
||||||
|
|
||||||
description = response.choices[0].message.content
|
description = response.choices[0].message.content
|
||||||
|
|
||||||
|
# 推理模型兼容:content 为空时从 reasoning_content 提取
|
||||||
|
if not description or not description.strip():
|
||||||
|
reasoning = getattr(response.choices[0].message, 'reasoning_content', None)
|
||||||
|
if reasoning and reasoning.strip():
|
||||||
|
import re
|
||||||
|
# 尝试从思考链中提取有用文本(去掉 <think> 标签后的内容)
|
||||||
|
cleaned = re.sub(r'', '', reasoning, flags=re.DOTALL).strip()
|
||||||
|
if cleaned:
|
||||||
|
logger.info(f"VLM content为空,从reasoning_content提取描述: {image_path}")
|
||||||
|
description = cleaned
|
||||||
|
else:
|
||||||
|
description = reasoning.strip()
|
||||||
|
|
||||||
|
if not description:
|
||||||
|
logger.warning(f"VLM 返回空描述: {image_path}")
|
||||||
|
return ""
|
||||||
|
|
||||||
# 缓存结果
|
# 缓存结果
|
||||||
import hashlib
|
import hashlib
|
||||||
|
import re as _re
|
||||||
img_hash = hashlib.md5(img_path.read_bytes()).hexdigest()
|
img_hash = hashlib.md5(img_path.read_bytes()).hexdigest()
|
||||||
cache_dir = Path('.data/cache/vlm')
|
cache_dir = Path('.data/cache/vlm')
|
||||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|||||||
@@ -283,7 +283,7 @@ class KnowledgeBaseRouter:
|
|||||||
content = call_llm(
|
content = call_llm(
|
||||||
self.llm_client, prompt, MODEL,
|
self.llm_client, prompt, MODEL,
|
||||||
temperature=0.1,
|
temperature=0.1,
|
||||||
max_tokens=100
|
max_tokens=512
|
||||||
)
|
)
|
||||||
|
|
||||||
if content is None:
|
if content is None:
|
||||||
|
|||||||
@@ -664,6 +664,17 @@ class KnowledgeSyncService:
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"递增缓存版本号失败: {e}")
|
logger.warning(f"递增缓存版本号失败: {e}")
|
||||||
|
|
||||||
|
# 语义缓存无版本号机制,文档变更后必须清空,
|
||||||
|
# 否则可能返回过时的 images/sources/citations(如已删除的图片 404)
|
||||||
|
try:
|
||||||
|
from core.semantic_cache import get_semantic_cache
|
||||||
|
_sc = get_semantic_cache()
|
||||||
|
if _sc:
|
||||||
|
_sc.clear()
|
||||||
|
logger.debug(f"已清空语义缓存(文档变更触发): {kb_name}")
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"清空语义缓存失败: {e}")
|
||||||
|
|
||||||
return True
|
return True
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
374
parsers/heading_rules.py
Normal file
374
parsers/heading_rules.py
Normal file
@@ -0,0 +1,374 @@
|
|||||||
|
# -*- coding: utf-8 -*-
|
||||||
|
"""
|
||||||
|
标题识别规则引擎
|
||||||
|
|
||||||
|
将 _detect_heading_level 的硬编码正则提取为可配置的规则列表。
|
||||||
|
规则按优先级从高到低排序,第一个匹配即返回。
|
||||||
|
|
||||||
|
MinerU 解析 DOCX 等 Office 格式时通常不提供 text_level(全部为 0),
|
||||||
|
此时需要启发式识别标题层级。本模块提供可配置的规则引擎替代原来的
|
||||||
|
硬编码 if-elif 链。
|
||||||
|
|
||||||
|
设计要点:
|
||||||
|
- HeadingRule 数据类支持正向匹配(pattern)和反向排除(exclude_pattern)
|
||||||
|
- 长度约束(min_length / max_length)可精确控制匹配范围
|
||||||
|
- 规则可单独禁用(enabled=False),便于调试
|
||||||
|
- 全局单例通过 config.py 覆盖默认值
|
||||||
|
|
||||||
|
MinerU v2 格式备注:
|
||||||
|
content_list_v2.json 中的 paragraph_content 包含 style=["bold"] 信息,
|
||||||
|
layout.json 中的 spans 也有 style 信息。这些信息比正则匹配 **加粗** 更可靠,
|
||||||
|
但当前代码使用 v1 格式(content_list.json),暂不利用 v2 的 style。
|
||||||
|
HeadingRuleEngine.detect() 签名预留了 style 参数,未来切换到 v2 格式后
|
||||||
|
可直接利用 style 信息辅助判断。
|
||||||
|
"""
|
||||||
|
|
||||||
|
import re
|
||||||
|
import logging
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import Optional, List, Tuple
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class HeadingRule:
|
||||||
|
"""
|
||||||
|
标题识别规则
|
||||||
|
|
||||||
|
每条规则定义一个文本模式到标题级别的映射。
|
||||||
|
规则引擎按列表顺序逐条匹配,第一个命中即返回。
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
pattern: 编译后的正则(match 语义,从文本开头匹配)
|
||||||
|
level: 匹配时返回的标题级别 (1=h1, 2=h2, 3=h3)
|
||||||
|
name: 规则名称(用于日志和配置覆盖)
|
||||||
|
max_length: 文本最大长度,0=不限
|
||||||
|
min_length: 文本最小长度,0=不限
|
||||||
|
enabled: 是否启用
|
||||||
|
exclude_pattern: 匹配此模式则排除(反向过滤)
|
||||||
|
|
||||||
|
Example:
|
||||||
|
>>> rule = HeadingRule(
|
||||||
|
... pattern=re.compile(r'^第[一二三四五六七八九十百千万]+[章节篇部]'),
|
||||||
|
... level=1,
|
||||||
|
... name="chinese_chapter",
|
||||||
|
... )
|
||||||
|
>>> rule.match("第一章 总则")
|
||||||
|
1
|
||||||
|
>>> rule.match("这是正文")
|
||||||
|
0
|
||||||
|
"""
|
||||||
|
|
||||||
|
pattern: re.Pattern
|
||||||
|
level: int
|
||||||
|
name: str
|
||||||
|
max_length: int = 0
|
||||||
|
min_length: int = 0
|
||||||
|
enabled: bool = True
|
||||||
|
exclude_pattern: Optional[re.Pattern] = None
|
||||||
|
|
||||||
|
def match(self, text: str) -> int:
|
||||||
|
"""
|
||||||
|
检查文本是否匹配此规则
|
||||||
|
|
||||||
|
Args:
|
||||||
|
text: 待检测文本(调用前应已 strip)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
标题级别,0 表示不匹配
|
||||||
|
"""
|
||||||
|
if not self.enabled:
|
||||||
|
return 0
|
||||||
|
if self.min_length > 0 and len(text) < self.min_length:
|
||||||
|
return 0
|
||||||
|
if self.max_length > 0 and len(text) > self.max_length:
|
||||||
|
return 0
|
||||||
|
if self.exclude_pattern and self.exclude_pattern.search(text):
|
||||||
|
return 0
|
||||||
|
if self.pattern.match(text):
|
||||||
|
return self.level
|
||||||
|
return 0
|
||||||
|
|
||||||
|
|
||||||
|
# 默认规则列表(按优先级从高到低)
|
||||||
|
#
|
||||||
|
# 注意事项:
|
||||||
|
# - 数字三级标题 (1.1.1) 必须在二级 (1.1) 之前,因为 1.1.1 也匹配 ^\d+\.\d+
|
||||||
|
# - short_chinese_heading 是最宽泛的规则,放在最后作为兜底
|
||||||
|
# - 第 9 条规则相比原版增加了 exclude_pattern,排除以句末标点结尾的短文本
|
||||||
|
DEFAULT_HEADING_RULES: List[HeadingRule] = [
|
||||||
|
# 1. 中文章节标题 -> h1
|
||||||
|
# 匹配:第一章、第二章、第十节、第三篇 等
|
||||||
|
HeadingRule(
|
||||||
|
pattern=re.compile(r'^第[一二三四五六七八九十百千万]+[章节篇部]'),
|
||||||
|
level=1,
|
||||||
|
name="chinese_chapter",
|
||||||
|
),
|
||||||
|
# 2. 中文条款编号 -> h2
|
||||||
|
# 匹配:第一条、第三款 等(仅短标题,长正文段落不算标题)
|
||||||
|
HeadingRule(
|
||||||
|
pattern=re.compile(r'^第[一二三四五六七八九十百千万]+[条款]'),
|
||||||
|
level=2,
|
||||||
|
name="chinese_article",
|
||||||
|
max_length=30,
|
||||||
|
),
|
||||||
|
# 3. 数字三级标题 -> h3(必须在二级之前匹配)
|
||||||
|
# 匹配:1.1.1 背景、2.3.4 方案 等
|
||||||
|
HeadingRule(
|
||||||
|
pattern=re.compile(r'^\d+\.\d+\.\d+[\.、\s]'),
|
||||||
|
level=3,
|
||||||
|
name="numeric_level3",
|
||||||
|
max_length=100,
|
||||||
|
),
|
||||||
|
# 4. 数字二级标题 -> h2(必须在一级之前匹配)
|
||||||
|
# 匹配:1.1 背景、2.3 方案、2.1运行调度(无空格) 等
|
||||||
|
# 使用负向前瞻排除三级标题(由 numeric_level3 处理)
|
||||||
|
HeadingRule(
|
||||||
|
pattern=re.compile(r'^\d+\.\d+(?!\.\d)'),
|
||||||
|
level=2,
|
||||||
|
name="numeric_level2",
|
||||||
|
max_length=80,
|
||||||
|
),
|
||||||
|
# 5. 数字一级标题 -> h1
|
||||||
|
# 匹配:1. 概述、2、背景 等
|
||||||
|
# 排除:以 ;;。,、: 结尾的文本(这些是编号列表项/子条目,不是独立标题)
|
||||||
|
HeadingRule(
|
||||||
|
pattern=re.compile(r'^\d+[\.、\s]'),
|
||||||
|
level=1,
|
||||||
|
name="numeric_level1",
|
||||||
|
max_length=50,
|
||||||
|
exclude_pattern=re.compile(r'[;;。,、::]$'),
|
||||||
|
),
|
||||||
|
# 6. 英文章节标题 -> h1
|
||||||
|
# 匹配:Chapter 1、Section 2、Part 3 等
|
||||||
|
HeadingRule(
|
||||||
|
pattern=re.compile(r'^(Chapter|Section|Part|Chapter\s+\d+|Section\s+\d+)', re.IGNORECASE),
|
||||||
|
level=1,
|
||||||
|
name="english_chapter",
|
||||||
|
),
|
||||||
|
# 7. 分类标题 -> h3(必须在 bold_short_text 之前,否则 **A2类:** 会被加粗规则抢先匹配)
|
||||||
|
# 匹配:A1类:公园、**A2类**:各类卫生医疗机构、**B1类:** 道路 等
|
||||||
|
HeadingRule(
|
||||||
|
pattern=re.compile(r'^\*{0,2}[A-Z]\d+[类類]\*{0,2}[::]'),
|
||||||
|
level=3,
|
||||||
|
name="category_heading",
|
||||||
|
),
|
||||||
|
# 8. 加粗短文本 -> h2
|
||||||
|
# 匹配:**重要通知**、**概述** 等(Markdown 加粗标记)
|
||||||
|
# 注意:**A2类:** 已被分类标题规则优先匹配,不会误判为 h2
|
||||||
|
HeadingRule(
|
||||||
|
pattern=re.compile(r'^\*\*.+\*\*$'),
|
||||||
|
level=2,
|
||||||
|
name="bold_short_text",
|
||||||
|
max_length=50,
|
||||||
|
),
|
||||||
|
# 9. 短中文文本 -> h2(替代原"任何 <20 字符含中文"规则)
|
||||||
|
# 关键改进:排除以句末标点结尾的文本
|
||||||
|
# 原规则将 "这是一段正文。" 也识别为 h2,导致大量误判
|
||||||
|
# 新规则:包含中文 + 长度 2-20 + 不以句末标点结尾 → h2
|
||||||
|
HeadingRule(
|
||||||
|
pattern=re.compile(r'[一-鿿]'),
|
||||||
|
level=2,
|
||||||
|
name="short_chinese_heading",
|
||||||
|
max_length=20,
|
||||||
|
min_length=2,
|
||||||
|
exclude_pattern=re.compile(r'[。!?;…]$'),
|
||||||
|
enabled=True,
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
class HeadingRuleEngine:
|
||||||
|
"""
|
||||||
|
标题识别规则引擎
|
||||||
|
|
||||||
|
按规则列表顺序逐条匹配,第一个命中即返回标题级别。
|
||||||
|
支持从 config.py 加载自定义规则或覆盖默认规则参数。
|
||||||
|
|
||||||
|
Example:
|
||||||
|
>>> engine = HeadingRuleEngine()
|
||||||
|
>>> engine.detect("第一章 总则")
|
||||||
|
(1, 'chinese_chapter')
|
||||||
|
>>> engine.detect("这是普通正文。")
|
||||||
|
(0, None)
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, rules: Optional[List[HeadingRule]] = None) -> None:
|
||||||
|
"""
|
||||||
|
Args:
|
||||||
|
rules: 规则列表,None 则使用默认规则的深拷贝
|
||||||
|
"""
|
||||||
|
if rules is not None:
|
||||||
|
self.rules: List[HeadingRule] = rules
|
||||||
|
else:
|
||||||
|
import copy
|
||||||
|
self.rules = copy.deepcopy(DEFAULT_HEADING_RULES)
|
||||||
|
|
||||||
|
def _validate_level(self, level: int, text: str, rule_name=None):
|
||||||
|
"""各级别标题长度防护:超长文本不应作为标题,降为正文。
|
||||||
|
H1 > 40字, H2 > 60字, H3 > 50字 → 降为正文。
|
||||||
|
统一覆盖 v1 常规匹配、v2 style 匹配、bold_short_text 兜底所有返回路径。"""
|
||||||
|
text_len = len(text)
|
||||||
|
if level == 1 and text_len > 40:
|
||||||
|
logger.debug(f"标题识别: '{text[:30]}...' H1 但超长({text_len}字),降为正文")
|
||||||
|
return 0, None
|
||||||
|
if level == 2 and text_len > 60:
|
||||||
|
logger.debug(f"标题识别: '{text[:30]}...' H2 但超长({text_len}字),降为正文")
|
||||||
|
return 0, None
|
||||||
|
if level == 3 and text_len > 50:
|
||||||
|
logger.debug(f"标题识别: '{text[:30]}...' H3 但超长({text_len}字),降为正文")
|
||||||
|
return 0, None
|
||||||
|
return level, rule_name
|
||||||
|
|
||||||
|
def detect(self, text: str, style: Optional[List[str]] = None) -> Tuple[int, Optional[str]]:
|
||||||
|
"""
|
||||||
|
检测文本的标题级别
|
||||||
|
|
||||||
|
Args:
|
||||||
|
text: 待检测文本
|
||||||
|
style: MinerU v2 格式中的 style 信息(如 ["bold"]),
|
||||||
|
当文本标记为 bold 且较短时,可直接判定为标题,
|
||||||
|
无需依赖 Markdown **...** 标记。
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
(level, rule_name): 标题级别和匹配的规则名
|
||||||
|
level=0 表示不是标题
|
||||||
|
"""
|
||||||
|
text = text.strip()
|
||||||
|
if not text:
|
||||||
|
return 0, None
|
||||||
|
|
||||||
|
# v2 style 信息:如果文本标记为 bold 且较短,优先尝试加粗规则
|
||||||
|
if style and 'bold' in style and 2 <= len(text) <= 50:
|
||||||
|
# 先检查是否匹配更高优先级的分类标题规则
|
||||||
|
for rule in self.rules:
|
||||||
|
if rule.name == 'category_heading' and rule.enabled:
|
||||||
|
level = rule.match(text)
|
||||||
|
if level > 0:
|
||||||
|
logger.debug(f"标题识别(v2 style): '{text[:30]}' -> h{level} (规则: {rule.name})")
|
||||||
|
return self._validate_level(level, text, rule.name)
|
||||||
|
|
||||||
|
# 再检查是否匹配中文章节/条款等高优先级规则
|
||||||
|
for rule in self.rules:
|
||||||
|
if rule.name in ('chinese_chapter', 'chinese_article', 'numeric_level3',
|
||||||
|
'numeric_level2', 'numeric_level1', 'english_chapter') and rule.enabled:
|
||||||
|
level = rule.match(text)
|
||||||
|
if level > 0:
|
||||||
|
logger.debug(f"标题识别(v2 style): '{text[:30]}' -> h{level} (规则: {rule.name})")
|
||||||
|
return self._validate_level(level, text, rule.name)
|
||||||
|
|
||||||
|
# 最后兜底:加粗短文本 → h2
|
||||||
|
# 但需先检查所有规则的 exclude_pattern,防止编号列表项被误判为标题
|
||||||
|
# 例如 "3.完全满足品规。指..." 虽有 bold 样式,但属于列表项而非标题
|
||||||
|
for rule in self.rules:
|
||||||
|
if rule.enabled and rule.exclude_pattern and rule.exclude_pattern.search(text):
|
||||||
|
logger.debug(f"标题识别(v2 style): '{text[:30]}' 被 {rule.name} 的 exclude_pattern 排除")
|
||||||
|
return 0, None
|
||||||
|
|
||||||
|
logger.debug(f"标题识别(v2 style): '{text[:30]}' -> h2 (规则: bold_short_text_via_style)")
|
||||||
|
return 2, 'bold_short_text'
|
||||||
|
|
||||||
|
# 常规规则匹配(v1 格式或无 style 信息时)
|
||||||
|
for rule in self.rules:
|
||||||
|
level = rule.match(text)
|
||||||
|
if level > 0:
|
||||||
|
logger.debug(f"标题识别: '{text[:30]}' -> h{level} (规则: {rule.name})")
|
||||||
|
return self._validate_level(level, text, rule.name)
|
||||||
|
|
||||||
|
return 0, None
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== 全局单例 ====================
|
||||||
|
|
||||||
|
_engine: Optional[HeadingRuleEngine] = None
|
||||||
|
|
||||||
|
|
||||||
|
def get_heading_engine() -> HeadingRuleEngine:
|
||||||
|
"""获取全局标题识别引擎(延迟初始化,线程安全)"""
|
||||||
|
global _engine
|
||||||
|
if _engine is None:
|
||||||
|
_engine = _create_engine_from_config()
|
||||||
|
return _engine
|
||||||
|
|
||||||
|
|
||||||
|
def _create_engine_from_config() -> HeadingRuleEngine:
|
||||||
|
"""
|
||||||
|
从 config 创建引擎(支持配置覆盖)
|
||||||
|
|
||||||
|
优先级:
|
||||||
|
1. config.HEADING_RULES_CONFIG 不为 None → 使用自定义规则
|
||||||
|
2. config 细粒度参数覆盖默认规则(如 HEADING_SHORT_TEXT_ENABLED)
|
||||||
|
3. 使用默认规则
|
||||||
|
"""
|
||||||
|
# 尝试加载完整自定义规则
|
||||||
|
try:
|
||||||
|
from config import HEADING_RULES_CONFIG
|
||||||
|
if HEADING_RULES_CONFIG is not None:
|
||||||
|
rules = _build_rules_from_config(HEADING_RULES_CONFIG)
|
||||||
|
logger.info(f"使用自定义标题规则: {len(rules)} 条")
|
||||||
|
return HeadingRuleEngine(rules)
|
||||||
|
except ImportError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
# 使用默认规则,应用细粒度配置覆盖
|
||||||
|
rules = list(DEFAULT_HEADING_RULES)
|
||||||
|
try:
|
||||||
|
from config import HEADING_SHORT_TEXT_ENABLED
|
||||||
|
for rule in rules:
|
||||||
|
if rule.name == "short_chinese_heading":
|
||||||
|
rule.enabled = HEADING_SHORT_TEXT_ENABLED
|
||||||
|
logger.debug(f"配置覆盖: short_chinese_heading.enabled={HEADING_SHORT_TEXT_ENABLED}")
|
||||||
|
except ImportError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
try:
|
||||||
|
from config import HEADING_SHORT_TEXT_MAX_LENGTH
|
||||||
|
for rule in rules:
|
||||||
|
if rule.name == "short_chinese_heading":
|
||||||
|
rule.max_length = HEADING_SHORT_TEXT_MAX_LENGTH
|
||||||
|
logger.debug(f"配置覆盖: short_chinese_heading.max_length={HEADING_SHORT_TEXT_MAX_LENGTH}")
|
||||||
|
except ImportError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
return HeadingRuleEngine(rules)
|
||||||
|
|
||||||
|
|
||||||
|
def _build_rules_from_config(config: list) -> List[HeadingRule]:
|
||||||
|
"""
|
||||||
|
从配置字典列表构建规则列表
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config: 规则配置列表,每项为 dict,包含:
|
||||||
|
- pattern (str): 正则表达式字符串
|
||||||
|
- level (int): 标题级别
|
||||||
|
- name (str): 规则名称
|
||||||
|
- max_length (int, 可选): 文本最大长度
|
||||||
|
- min_length (int, 可选): 文本最小长度
|
||||||
|
- enabled (bool, 可选): 是否启用
|
||||||
|
- exclude_pattern (str, 可选): 排除正则
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
规则列表
|
||||||
|
"""
|
||||||
|
rules = []
|
||||||
|
for item in config:
|
||||||
|
exclude = None
|
||||||
|
if 'exclude_pattern' in item:
|
||||||
|
exclude = re.compile(item['exclude_pattern'])
|
||||||
|
rules.append(HeadingRule(
|
||||||
|
pattern=re.compile(item['pattern']),
|
||||||
|
level=item['level'],
|
||||||
|
name=item['name'],
|
||||||
|
max_length=item.get('max_length', 0),
|
||||||
|
min_length=item.get('min_length', 0),
|
||||||
|
enabled=item.get('enabled', True),
|
||||||
|
exclude_pattern=exclude,
|
||||||
|
))
|
||||||
|
return rules
|
||||||
|
|
||||||
|
|
||||||
|
def reset_heading_engine() -> None:
|
||||||
|
"""重置引擎(用于测试)"""
|
||||||
|
global _engine
|
||||||
|
_engine = None
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -616,7 +616,7 @@ class FeedbackService:
|
|||||||
prompt=prompt,
|
prompt=prompt,
|
||||||
model=self.model,
|
model=self.model,
|
||||||
temperature=0.7,
|
temperature=0.7,
|
||||||
max_tokens=200
|
max_tokens=512
|
||||||
)
|
)
|
||||||
|
|
||||||
if not response:
|
if not response:
|
||||||
|
|||||||
@@ -159,7 +159,7 @@ class SessionManager:
|
|||||||
SELECT id, role, content, metadata, created_at
|
SELECT id, role, content, metadata, created_at
|
||||||
FROM messages
|
FROM messages
|
||||||
WHERE session_id = ?
|
WHERE session_id = ?
|
||||||
ORDER BY created_at DESC
|
ORDER BY created_at DESC, id DESC
|
||||||
LIMIT ?
|
LIMIT ?
|
||||||
''', (session_id, limit))
|
''', (session_id, limit))
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user