init: RAG 知识库服务初始提交
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
This commit is contained in:
lacerate551
717
knowledge/manager.py
Normal file
717
knowledge/manager.py
Normal file
@@ -0,0 +1,717 @@
|
||||
"""
|
||||
多向量库管理器 - 支持按部门/权限隔离的向量知识库
|
||||
|
||||
功能：
|
||||
1. 多向量库管理 - 创建、删除、列举向量库
|
||||
2. 多 BM25 索引管理 - 每个向量库独立的 BM25 索引
|
||||
3. 权限过滤 - 根据用户角色和部门返回可访问的向量库
|
||||
4. 并行检索 - 支持同时检索多个向量库
|
||||
|
||||
向量库命名规范：
|
||||
- public_kb: 公开知识库，所有人可访问
|
||||
- dept_{部门名}: 部门知识库，如 dept_finance, dept_hr, dept_tech
|
||||
|
||||
使用方式：
|
||||
from knowledge.manager import KnowledgeBaseManager
|
||||
|
||||
kb_manager = KnowledgeBaseManager()
|
||||
|
||||
# 获取向量库
|
||||
collection = kb_manager.get_collection("dept_finance")
|
||||
|
||||
# 列出所有向量库
|
||||
collections = kb_manager.list_collections()
|
||||
|
||||
# 获取用户可访问的向量库
|
||||
accessible = kb_manager.get_accessible_collections("manager", "finance")
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import threading
|
||||
from typing import List, Dict, Optional, Tuple
|
||||
from pathlib import Path
|
||||
import logging
|
||||
|
||||
import chromadb
|
||||
|
||||
# 从 base.py 导入基础类和常量
|
||||
from .base import (
|
||||
BM25Index,
|
||||
CollectionInfo,
|
||||
SearchResult,
|
||||
_get_doc_type,
|
||||
_extract_figure_number,
|
||||
_extract_section,
|
||||
_build_semantic_content_for_text,
|
||||
_build_semantic_content_for_table,
|
||||
CHROMA_DB_BASE_PATH,
|
||||
BM25_INDEX_BASE_PATH,
|
||||
KB_METADATA_FILE,
|
||||
PUBLIC_KB_NAME,
|
||||
DEFAULT_DEPARTMENTS,
|
||||
DEPARTMENT_NAME_MAP,
|
||||
normalize_department_name,
|
||||
)
|
||||
|
||||
# 导入 Mixin 类
|
||||
from .collection import CollectionMixin
|
||||
from .document import DocumentMixin
|
||||
from .chunk import ChunkMixin
|
||||
from .index import IndexMixin
|
||||
from .search import SearchMixin
|
||||
from .processing import ProcessingMixin
|
||||
from .permission import PermissionMixin
|
||||
|
||||
# 导入 LLM 工具函数
|
||||
from core.llm_utils import call_llm
|
||||
|
||||
# 设置日志
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ==================== 多向量库管理器 ====================
|
||||
|
||||
class KnowledgeBaseManager(
|
||||
CollectionMixin,
|
||||
DocumentMixin,
|
||||
ChunkMixin,
|
||||
IndexMixin,
|
||||
SearchMixin,
|
||||
ProcessingMixin,
|
||||
PermissionMixin
|
||||
):
|
||||
"""
|
||||
多向量库管理器
|
||||
|
||||
管理多个独立的 ChromaDB 集合，每个集合对应一个知识库。
|
||||
支持按部门隔离，每个部门有独立的向量库和 BM25 索引。
|
||||
|
||||
通过 Mixin 模式组合功能：
|
||||
- CollectionMixin: 向量库管理
|
||||
- DocumentMixin: 文档管理
|
||||
- ChunkMixin: 切片管理
|
||||
- IndexMixin: BM25 索引管理
|
||||
- SearchMixin: 检索功能
|
||||
- ProcessingMixin: 图片/表格处理
|
||||
- PermissionMixin: 权限控制
|
||||
"""
|
||||
|
||||
def __init__(self, base_path: str = None, bm25_base_path: str = None):
|
||||
"""
|
||||
初始化
|
||||
|
||||
Args:
|
||||
base_path: 向量库存储路径
|
||||
bm25_base_path: BM25 索引存储路径
|
||||
"""
|
||||
self.base_path = base_path or CHROMA_DB_BASE_PATH
|
||||
self.bm25_base_path = bm25_base_path or BM25_INDEX_BASE_PATH
|
||||
|
||||
# 缓存
|
||||
self._collections: Dict[str, chromadb.Collection] = {}
|
||||
self._bm25_indexes: Dict[str, BM25Index] = {}
|
||||
self._clients: Dict[str, chromadb.PersistentClient] = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
# 确保目录存在
|
||||
os.makedirs(self.base_path, exist_ok=True)
|
||||
os.makedirs(self.bm25_base_path, exist_ok=True)
|
||||
|
||||
# 加载元数据
|
||||
self._metadata = self._load_metadata()
|
||||
|
||||
# 初始化公开知识库
|
||||
self._ensure_public_kb()
|
||||
|
||||
# 扫描并列出所有已存在的向量库
|
||||
existing_kbs = []
|
||||
if os.path.exists(self.base_path):
|
||||
for item in os.listdir(self.base_path):
|
||||
if os.path.isdir(os.path.join(self.base_path, item)) and not item.startswith('.'):
|
||||
if os.path.exists(os.path.join(self.base_path, item, 'chroma.sqlite3')):
|
||||
existing_kbs.append(item)
|
||||
logger.info(f"知识库管理器初始化完成，路径: {self.base_path}，发现 {len(existing_kbs)} 个向量库: {existing_kbs}")
|
||||
|
||||
def _load_metadata(self) -> dict:
|
||||
"""加载元数据"""
|
||||
metadata_path = os.path.join(self.base_path, KB_METADATA_FILE)
|
||||
if os.path.exists(metadata_path):
|
||||
try:
|
||||
with open(metadata_path, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
except Exception as e:
|
||||
logger.error(f"加载元数据失败: {e}")
|
||||
return {"collections": {}}
|
||||
|
||||
def _save_metadata(self):
|
||||
"""保存元数据"""
|
||||
metadata_path = os.path.join(self.base_path, KB_METADATA_FILE)
|
||||
try:
|
||||
with open(metadata_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(self._metadata, f, ensure_ascii=False, indent=2)
|
||||
except Exception as e:
|
||||
logger.error(f"保存元数据失败: {e}")
|
||||
|
||||
def _ensure_public_kb(self):
|
||||
"""确保公开知识库存在"""
|
||||
if PUBLIC_KB_NAME not in self._metadata.get("collections", {}):
|
||||
self.create_collection(
|
||||
PUBLIC_KB_NAME,
|
||||
display_name="公开知识库",
|
||||
department="",
|
||||
description="所有人可访问的公开文档"
|
||||
)
|
||||
|
||||
# ==================== 文件添加 ====================
|
||||
|
||||
def add_file_to_kb(
|
||||
self,
|
||||
kb_name: str,
|
||||
filepath: str,
|
||||
embedding_model=None,
|
||||
extra_metadata: dict = None,
|
||||
enable_table_summary: bool = True,
|
||||
enable_image_description: bool = False,
|
||||
file_content: bytes = None
|
||||
) -> int:
|
||||
"""
|
||||
添加文件到指定向量库（v6 支持企业文件系统）
|
||||
|
||||
使用统一的 parse_document() 入口，支持：
|
||||
- PDF/DOCX/PPTX/图片 → MinerU 解析
|
||||
- XLSX/XLS → Pandas 解析
|
||||
- TXT → 文本解析
|
||||
|
||||
Args:
|
||||
kb_name: 向量库名称
|
||||
filepath: 文件路径（相对路径或绝对路径）
|
||||
embedding_model: 向量模型（可选，默认使用 engine 的）
|
||||
extra_metadata: 额外的元数据（如 status, version 等）
|
||||
enable_table_summary: 是否启用表格摘要管道（LLM 生成摘要）
|
||||
enable_image_description: 是否启用图片描述管道（VLM 生成描述）
|
||||
file_content: 文件二进制内容（可选，用于企业文件系统集成）
|
||||
|
||||
Returns:
|
||||
添加的片段数量
|
||||
"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
raise ValueError(f"向量库 '{kb_name}' 不存在")
|
||||
|
||||
# 解析文档
|
||||
from parsers import parse_document
|
||||
result = parse_document(filepath, file_content=file_content)
|
||||
|
||||
if not result:
|
||||
logger.warning(f"文档解析结果为空: {filepath}")
|
||||
return 0
|
||||
|
||||
chunks = result.get('chunks', [])
|
||||
if not chunks:
|
||||
logger.warning(f"文档解析后无有效切片: {filepath}")
|
||||
return 0
|
||||
|
||||
# 合并跨页表格
|
||||
chunks = self._merge_cross_page_tables(chunks)
|
||||
|
||||
# 准备向量模型
|
||||
if embedding_model is None:
|
||||
from core.engine import get_engine
|
||||
engine = get_engine()
|
||||
if not engine._initialized:
|
||||
engine.initialize()
|
||||
embedding_model = engine.embedding_model
|
||||
|
||||
# 处理切片
|
||||
ids = []
|
||||
documents = []
|
||||
metadatas = []
|
||||
embeddings = []
|
||||
|
||||
filename = Path(filepath).name
|
||||
doc_type = _get_doc_type(filename)
|
||||
|
||||
for i, chunk in enumerate(chunks):
|
||||
chunk_id = f"{filename}_{i}"
|
||||
|
||||
# 获取切片类型（优先从 chunk 直接获取，兼容 MinerUChunk 和其他格式）
|
||||
chunk_type = getattr(chunk, 'chunk_type', None)
|
||||
if not chunk_type:
|
||||
page_info = getattr(chunk, 'page_info', {}) or {}
|
||||
chunk_type = page_info.get('chunk_type', 'text')
|
||||
|
||||
# 获取页码信息（兼容 MinerUChunk 和 page_info 格式）
|
||||
page_start = getattr(chunk, 'page_start', None)
|
||||
page_end = getattr(chunk, 'page_end', None)
|
||||
if page_start is None:
|
||||
page_info = getattr(chunk, 'page_info', {}) or {}
|
||||
page_start = page_info.get('page', 0)
|
||||
page_end = page_info.get('page_end', page_start)
|
||||
page_start = page_start or 0
|
||||
page_end = page_end or page_start
|
||||
|
||||
# 获取章节信息
|
||||
section_path = getattr(chunk, 'section_path', '') or ''
|
||||
if not section_path:
|
||||
page_info = getattr(chunk, 'page_info', {}) or {}
|
||||
section_path = page_info.get('section_path', '') or page_info.get('section', '')
|
||||
|
||||
if chunk_type == 'table':
|
||||
# 表格内容优先使用 table_html（完整表格），fallback 到 content（标题）
|
||||
table_md = getattr(chunk, 'table_html', None) or chunk.content
|
||||
semantic_content = _build_semantic_content_for_table(
|
||||
table_md, {'page': page_start, 'page_end': page_end, 'section_path': section_path}, chunk, doc_type
|
||||
)
|
||||
elif chunk_type in ('image', 'chart'):
|
||||
semantic_content = self.generate_lightweight_image_description(
|
||||
chunk.content, chunk, {'page': page_start, 'page_end': page_end, 'section_path': section_path}
|
||||
)
|
||||
else:
|
||||
semantic_content = _build_semantic_content_for_text(
|
||||
chunk, {'page': page_start, 'page_end': page_end, 'section_path': section_path}, doc_type
|
||||
)
|
||||
|
||||
# 构建元数据
|
||||
metadata = {
|
||||
"chunk_id": chunk_id,
|
||||
"chunk_index": i,
|
||||
"chunk_type": chunk_type,
|
||||
"source": filename,
|
||||
"collection": kb_name,
|
||||
"doc_type": doc_type, # 文档类型(pdf/word/excel/ppt/other)，驱动前端差异化溯源展示
|
||||
"page": page_start,
|
||||
"page_end": page_end,
|
||||
"section": section_path,
|
||||
"status": "active",
|
||||
"version": "v1",
|
||||
}
|
||||
|
||||
if extra_metadata:
|
||||
metadata.update(extra_metadata)
|
||||
|
||||
# 序列化图片信息（修复图片召回断链）
|
||||
if hasattr(chunk, 'images') and chunk.images:
|
||||
metadata['images_json'] = json.dumps(chunk.images, ensure_ascii=False)
|
||||
|
||||
if hasattr(chunk, 'image_path') and chunk.image_path:
|
||||
metadata['image_path'] = chunk.image_path
|
||||
|
||||
# 生成向量
|
||||
try:
|
||||
embedding = embedding_model.encode(semantic_content).tolist()
|
||||
except Exception as e:
|
||||
logger.error(f"生成向量失败: {chunk_id}, 错误: {e}")
|
||||
continue
|
||||
|
||||
ids.append(chunk_id)
|
||||
documents.append(semantic_content)
|
||||
metadatas.append(metadata)
|
||||
embeddings.append(embedding)
|
||||
|
||||
# 处理表格摘要
|
||||
if chunk_type == 'table' and enable_table_summary:
|
||||
table_md = chunk.content
|
||||
summary = self._generate_table_summary(table_md, chunk)
|
||||
if summary:
|
||||
# 存储原始表格到 DocStore
|
||||
self._store_original_table(chunk_id, table_md, metadata)
|
||||
|
||||
# 处理图片描述
|
||||
if chunk_type in ('image', 'chart') and enable_image_description:
|
||||
image_path = chunk.content
|
||||
if self.should_process_image(image_path, "", page_info.get('caption', '')):
|
||||
desc = self._generate_image_description(image_path, metadata)
|
||||
if desc:
|
||||
self._store_image_reference(chunk_id, image_path, metadata)
|
||||
|
||||
if not ids:
|
||||
return 0
|
||||
|
||||
# 批量添加到向量库
|
||||
collection.add(
|
||||
ids=ids,
|
||||
documents=documents,
|
||||
metadatas=metadatas,
|
||||
embeddings=embeddings
|
||||
)
|
||||
|
||||
# 更新 BM25 索引
|
||||
bm25 = self.get_bm25_index(kb_name)
|
||||
if bm25:
|
||||
bm25.add_documents(ids, documents, metadatas)
|
||||
self.save_bm25_index(kb_name)
|
||||
|
||||
logger.info(f"添加文件: {filename} -> {kb_name}, 片段数: {len(ids)}")
|
||||
return len(ids)
|
||||
|
||||
def _merge_cross_page_tables(self, chunks: list) -> list:
|
||||
"""
|
||||
合并跨页表格
|
||||
|
||||
检测规则：
|
||||
1. 表格切片后面跟着"续表"文本或另一个表格
|
||||
2. 页码连续或无法判断时，通过标题匹配
|
||||
3. 第二个表格标题包含"续表"或标题相似
|
||||
|
||||
合并操作：
|
||||
1. 合并 table_html
|
||||
2. 将两个 image_path 存入 images 字段
|
||||
3. 更新 page_end
|
||||
"""
|
||||
if len(chunks) < 2:
|
||||
return chunks
|
||||
|
||||
merged_chunks = []
|
||||
i = 0
|
||||
merge_count = 0
|
||||
|
||||
while i < len(chunks):
|
||||
current = chunks[i]
|
||||
|
||||
# 检查是否为表格类型
|
||||
if getattr(current, 'chunk_type', '') == 'table':
|
||||
# 查找下一个表格（跳过中间的"续表"文本）
|
||||
next_table_idx = None
|
||||
next_chunk = None
|
||||
|
||||
for j in range(i + 1, min(i + 4, len(chunks))): # 最多向前看3个切片
|
||||
candidate = chunks[j]
|
||||
candidate_type = getattr(candidate, 'chunk_type', '')
|
||||
candidate_title = getattr(candidate, 'title', '') or ''
|
||||
candidate_content = getattr(candidate, 'content', '')
|
||||
|
||||
if candidate_type == 'table':
|
||||
next_table_idx = j
|
||||
next_chunk = candidate
|
||||
break
|
||||
elif candidate_type == 'text' and ('续表' in candidate_title or '续表' in candidate_content):
|
||||
# 遇到"续表"文本，继续查找下一个表格
|
||||
continue
|
||||
elif candidate_type not in ('text',):
|
||||
# 遇到非文本类型，停止查找
|
||||
break
|
||||
|
||||
if next_chunk is not None:
|
||||
# 获取页码信息
|
||||
curr_page_end = getattr(current, 'page_end', getattr(current, 'page_start', 0))
|
||||
next_page_start = getattr(next_chunk, 'page_start', 0)
|
||||
|
||||
# 获取标题
|
||||
curr_title = getattr(current, 'title', '') or ''
|
||||
next_title = getattr(next_chunk, 'title', '') or ''
|
||||
if isinstance(curr_title, list):
|
||||
curr_title = curr_title[0] if curr_title else ''
|
||||
if isinstance(next_title, list):
|
||||
next_title = next_title[0] if next_title else ''
|
||||
|
||||
# 获取内容（用于检测"续表"）
|
||||
next_content = getattr(next_chunk, 'content', '')
|
||||
|
||||
# 判断是否为跨页表格
|
||||
is_cross_page = False
|
||||
|
||||
# 规则1: 页码连续（如果页码有效）
|
||||
page_valid = curr_page_end > 0 and next_page_start > 0
|
||||
if page_valid and curr_page_end + 1 == next_page_start:
|
||||
is_cross_page = True
|
||||
|
||||
# 规则2: 第二个表格标题或内容包含"续表"
|
||||
elif '续表' in next_title or '续表' in next_content:
|
||||
is_cross_page = True
|
||||
|
||||
# 规则3: 标题相似（去掉"续表"后比较）
|
||||
elif curr_title and next_title:
|
||||
clean_next = next_title.replace('续表', '').strip()
|
||||
if curr_title in clean_next or clean_next in curr_title:
|
||||
is_cross_page = True
|
||||
|
||||
if is_cross_page:
|
||||
# 执行合并
|
||||
merge_count += 1
|
||||
logger.info(f"合并跨页表格: {curr_title} (页{curr_page_end}) + {next_title} (页{next_page_start})")
|
||||
|
||||
# 合并 table_html
|
||||
curr_html = getattr(current, 'table_html', '') or ''
|
||||
next_html = getattr(next_chunk, 'table_html', '') or ''
|
||||
if curr_html and next_html:
|
||||
# 合并两个表格的 HTML
|
||||
current.table_html = curr_html + '\n' + next_html
|
||||
|
||||
# 合并 image_path 到 images
|
||||
curr_img = getattr(current, 'image_path', None)
|
||||
next_img = getattr(next_chunk, 'image_path', None)
|
||||
merged_images = []
|
||||
if curr_img:
|
||||
merged_images.append({'id': curr_img, 'page': curr_page_end})
|
||||
if next_img:
|
||||
merged_images.append({'id': next_img, 'page': next_page_start})
|
||||
if merged_images:
|
||||
current.images = merged_images
|
||||
# 保留第一个图片作为主 image_path
|
||||
current.image_path = curr_img
|
||||
|
||||
# 更新页码范围
|
||||
current.page_end = getattr(next_chunk, 'page_end', next_page_start)
|
||||
|
||||
# 添加合并后的切片，跳过中间所有切片
|
||||
merged_chunks.append(current)
|
||||
i = next_table_idx + 1
|
||||
continue
|
||||
|
||||
# 不需要合并，直接添加
|
||||
merged_chunks.append(current)
|
||||
i += 1
|
||||
|
||||
if merge_count > 0:
|
||||
logger.info(f"跨页表格合并完成: {merge_count} 组表格被合并")
|
||||
|
||||
return merged_chunks
|
||||
|
||||
def _generate_table_summary(self, table_md: str, chunk) -> str:
|
||||
"""生成表格摘要（LLM）"""
|
||||
# 提取表头
|
||||
lines = table_md.split('\n')
|
||||
headers = []
|
||||
for line in lines:
|
||||
if line.startswith('|') and '---' not in line:
|
||||
headers = [h.strip() for h in line.split('|') if h.strip()]
|
||||
break
|
||||
|
||||
if not headers:
|
||||
return ""
|
||||
|
||||
# 构建提示词
|
||||
prompt = f"""请用一句话总结以下表格的主要内容，不超过50字。
|
||||
|
||||
表头：{', '.join(headers)}
|
||||
|
||||
表格内容（前5行）：
|
||||
{chr(10).join(lines[:6])}
|
||||
|
||||
摘要："""
|
||||
|
||||
try:
|
||||
from config import get_llm_client, DASHSCOPE_MODEL
|
||||
client = get_llm_client()
|
||||
summary = call_llm(client, prompt, DASHSCOPE_MODEL, max_tokens=100)
|
||||
return summary.strip() if summary else ""
|
||||
except Exception as e:
|
||||
logger.warning(f"生成表格摘要失败: {e}")
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def _extract_table_title(table_md: str) -> str:
|
||||
"""从表格 Markdown 中提取标题"""
|
||||
lines = table_md.split('\n')
|
||||
for line in lines[:3]:
|
||||
if line.startswith('#'):
|
||||
return line.lstrip('#').strip()
|
||||
return ""
|
||||
|
||||
def _generate_image_description(self, image_path: str, metadata: dict = None) -> str:
|
||||
"""生成图片描述（VLM）"""
|
||||
# 检查缓存
|
||||
cached = self._get_vlm_cache(image_path)
|
||||
if cached:
|
||||
return cached
|
||||
|
||||
# 调用 VLM
|
||||
try:
|
||||
import base64
|
||||
from pathlib import Path
|
||||
|
||||
# 读取图片
|
||||
img_path = Path(image_path)
|
||||
if not img_path.exists():
|
||||
return ""
|
||||
|
||||
img_data = base64.b64encode(img_path.read_bytes()).decode()
|
||||
|
||||
# 构建提示词
|
||||
prompt = """请描述这张图片的内容，包括：
|
||||
1. 图片类型（如流程图、架构图、数据图表等）
|
||||
2. 主要内容和关键信息
|
||||
3. 如果是图表，描述数据趋势或关键数值
|
||||
|
||||
描述应简洁，不超过100字。"""
|
||||
|
||||
# 调用 VLM（需要支持视觉的模型）
|
||||
from config import DASHSCOPE_API_KEY, DASHSCOPE_BASE_URL, VLM_MODEL
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(api_key=DASHSCOPE_API_KEY, base_url=DASHSCOPE_BASE_URL)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=VLM_MODEL,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": prompt},
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_data}"}}
|
||||
]
|
||||
}
|
||||
],
|
||||
max_tokens=200
|
||||
)
|
||||
|
||||
description = response.choices[0].message.content
|
||||
|
||||
# 缓存结果
|
||||
import hashlib
|
||||
img_hash = hashlib.md5(img_path.read_bytes()).hexdigest()
|
||||
cache_dir = Path('.data/cache/vlm')
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
(cache_dir / f'{img_hash}.txt').write_text(description, encoding='utf-8')
|
||||
|
||||
return description
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"生成图片描述失败: {e}")
|
||||
return ""
|
||||
|
||||
def update_image_descriptions(self, kb_name: str) -> dict:
|
||||
"""更新向量库中所有图片的描述"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
return {"success": False, "error": "向量库不存在"}
|
||||
|
||||
result = collection.get(where={"chunk_type": {"$in": ["image", "chart"]}})
|
||||
|
||||
if not result['ids']:
|
||||
return {"success": True, "updated": 0, "message": "没有图片切片"}
|
||||
|
||||
updated = 0
|
||||
failed = 0
|
||||
|
||||
for chunk_id, doc, meta in zip(
|
||||
result['ids'],
|
||||
result['documents'],
|
||||
result['metadatas']
|
||||
):
|
||||
image_path = meta.get('image_path', '')
|
||||
if not image_path:
|
||||
continue
|
||||
|
||||
description = self._generate_image_description(image_path, meta)
|
||||
if description:
|
||||
try:
|
||||
collection.update(
|
||||
ids=[chunk_id],
|
||||
documents=[description],
|
||||
metadatas=[{**meta, 'image_description': description}]
|
||||
)
|
||||
updated += 1
|
||||
except Exception as e:
|
||||
logger.warning(f"更新图片描述失败: {chunk_id}, {e}")
|
||||
failed += 1
|
||||
|
||||
# 重建 BM25 索引
|
||||
self.rebuild_bm25_index(kb_name)
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"updated": updated,
|
||||
"failed": failed,
|
||||
"total": len(result['ids'])
|
||||
}
|
||||
|
||||
# ==================== 文档生命周期 ====================
|
||||
|
||||
def mark_document_as_superseded(
|
||||
self,
|
||||
kb_name: str,
|
||||
old_filename: str,
|
||||
new_filename: str,
|
||||
reason: str = "版本更新"
|
||||
) -> Dict:
|
||||
"""标记文档为已替代版本"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
return {"success": False, "error": "向量库不存在"}
|
||||
|
||||
result = collection.get(where={"source": old_filename})
|
||||
|
||||
if not result['ids']:
|
||||
return {"success": False, "error": "旧文档不存在"}
|
||||
|
||||
from datetime import datetime
|
||||
superseded_date = datetime.now().isoformat()
|
||||
|
||||
updated_metadatas = [
|
||||
{
|
||||
**m,
|
||||
"status": "superseded",
|
||||
"superseded_by": new_filename,
|
||||
"superseded_date": superseded_date,
|
||||
"superseded_reason": reason
|
||||
}
|
||||
for m in result['metadatas']
|
||||
]
|
||||
|
||||
collection.update(
|
||||
ids=result['ids'],
|
||||
metadatas=updated_metadatas
|
||||
)
|
||||
|
||||
self.rebuild_bm25_index(kb_name)
|
||||
|
||||
logger.info(f"标记文档替代: {old_filename} -> {new_filename}")
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"superseded_chunks": len(result['ids']),
|
||||
"old_document": old_filename,
|
||||
"new_document": new_filename,
|
||||
"collection": kb_name
|
||||
}
|
||||
|
||||
# ==================== 辅助检索方法 ====================
|
||||
|
||||
def search_with_status_filter(
|
||||
self,
|
||||
kb_name: str,
|
||||
query_vector: List[float],
|
||||
top_k: int = 5,
|
||||
status_filter: str = None
|
||||
) -> Optional[SearchResult]:
|
||||
"""带状态过滤的检索"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
return None
|
||||
|
||||
where_filter = None
|
||||
if status_filter:
|
||||
where_filter = {"status": status_filter}
|
||||
|
||||
result = collection.query(
|
||||
query_embeddings=[query_vector],
|
||||
n_results=top_k,
|
||||
where=where_filter
|
||||
)
|
||||
|
||||
return SearchResult(
|
||||
ids=result['ids'][0] if result['ids'] else [],
|
||||
documents=result['documents'][0] if result['documents'] else [],
|
||||
metadatas=result['metadatas'][0] if result['metadatas'] else [],
|
||||
distances=result['distances'][0] if result['distances'] else [],
|
||||
collection_name=kb_name
|
||||
)
|
||||
|
||||
def _rebuild_bm25_index(self, kb_name: str):
|
||||
"""重建 BM25 索引（私有方法，兼容旧代码）"""
|
||||
return self.rebuild_bm25_index(kb_name)
|
||||
|
||||
|
||||
# ==================== 全局实例 ====================
|
||||
|
||||
_kb_manager: Optional[KnowledgeBaseManager] = None
|
||||
|
||||
|
||||
def get_kb_manager() -> KnowledgeBaseManager:
|
||||
"""获取全局知识库管理器实例"""
|
||||
global _kb_manager
|
||||
if _kb_manager is None:
|
||||
_kb_manager = KnowledgeBaseManager()
|
||||
return _kb_manager
|
||||
Reference in New Issue
Block a user