init: RAG 知识库服务初始提交
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
This commit is contained in:
lacerate551
369
knowledge/chunk.py
Normal file
369
knowledge/chunk.py
Normal file
@@ -0,0 +1,369 @@
|
||||
"""
|
||||
知识库管理器 - 切片管理 Mixin
|
||||
|
||||
提供文档切片（Chunk）的 CRUD 操作，支持：
|
||||
- 新增切片：手动添加文本切片到向量库
|
||||
- 修改切片：更新切片内容和元数据
|
||||
- 删除切片：从向量库移除切片
|
||||
- 查询切片：分页获取切片列表
|
||||
|
||||
切片是向量检索的基本单位，每个切片包含：
|
||||
- 文档内容（用于向量化）
|
||||
- 元数据（来源、页码、状态等）
|
||||
- 向量嵌入（由 embedding 模型生成）
|
||||
|
||||
主要方法：
|
||||
- add_chunk: 新增切片
|
||||
- update_chunk: 修改切片
|
||||
- delete_chunk: 删除切片
|
||||
- list_chunks: 查询切片列表
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from typing import List, Dict, Optional, Tuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ChunkMixin:
|
||||
"""
|
||||
切片管理 Mixin
|
||||
|
||||
提供切片级别的 CRUD 操作，同时维护向量索引和 BM25 索引。
|
||||
|
||||
依赖属性（需由主类提供）：
|
||||
- self.get_collection: 获取向量库集合的方法
|
||||
- self.get_bm25_index: 获取 BM25 索引的方法
|
||||
- self.save_bm25_index: 保存 BM25 索引的方法
|
||||
"""
|
||||
|
||||
def add_chunk(
|
||||
self,
|
||||
kb_name: str,
|
||||
content: str,
|
||||
metadata: dict = None
|
||||
) -> str:
|
||||
"""
|
||||
新增单个切片
|
||||
|
||||
将文本内容切片添加到指定向量库，自动生成向量嵌入和 BM25 索引。
|
||||
|
||||
Args:
|
||||
kb_name: 目标向量库名称
|
||||
content: 切片文本内容
|
||||
metadata: 可选的元数据字典，可包含：
|
||||
- source: 来源文件名
|
||||
- page: 页码
|
||||
- chunk_type: 切片类型 ('text' | 'table' | 'image')
|
||||
- status: 状态 ('active' | 'deprecated')
|
||||
- 其他自定义字段
|
||||
|
||||
Returns:
|
||||
新创建的切片 ID，格式为: manual_{时间戳}_{内容哈希}
|
||||
|
||||
Raises:
|
||||
ValueError: 向量库不存在
|
||||
Exception: 向量生成失败
|
||||
|
||||
Example:
|
||||
>>> chunk_id = kb_manager.add_chunk(
|
||||
... "public_kb",
|
||||
... "这是要添加的文本内容",
|
||||
... {"source": "manual", "page": 1}
|
||||
... )
|
||||
"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
raise ValueError(f"向量库 '{kb_name}' 不存在")
|
||||
|
||||
content_hash = hashlib.md5(content.encode('utf-8')).hexdigest()[:8]
|
||||
chunk_id = f"manual_{datetime.now().strftime('%Y%m%d%H%M%S')}_{content_hash}"
|
||||
|
||||
chunk_metadata = {
|
||||
"chunk_id": chunk_id,
|
||||
"chunk_type": "text",
|
||||
"collection": kb_name,
|
||||
"source": "manual",
|
||||
"status": "active",
|
||||
"version": "v1",
|
||||
"change_time": datetime.now().isoformat(),
|
||||
}
|
||||
if metadata:
|
||||
chunk_metadata.update(metadata)
|
||||
|
||||
try:
|
||||
from core.engine import get_engine
|
||||
engine = get_engine()
|
||||
if not engine._initialized:
|
||||
engine.initialize()
|
||||
embedding = engine.embedding_model.encode(content).tolist()
|
||||
except Exception as e:
|
||||
logger.error(f"生成向量失败: {e}")
|
||||
raise
|
||||
|
||||
collection.add(
|
||||
ids=[chunk_id],
|
||||
documents=[content],
|
||||
metadatas=[chunk_metadata],
|
||||
embeddings=[embedding]
|
||||
)
|
||||
|
||||
bm25 = self.get_bm25_index(kb_name)
|
||||
if bm25:
|
||||
bm25.add_documents([chunk_id], [content], [chunk_metadata])
|
||||
self.save_bm25_index(kb_name)
|
||||
|
||||
logger.info(f"新增切片: {chunk_id} -> {kb_name}")
|
||||
return chunk_id
|
||||
|
||||
def update_chunk(
|
||||
self,
|
||||
kb_name: str,
|
||||
chunk_id: str,
|
||||
content: str = None,
|
||||
metadata: dict = None
|
||||
) -> bool:
|
||||
"""
|
||||
修改切片
|
||||
|
||||
更新切片的内容和/或元数据。如果更新内容，会重新生成向量嵌入。
|
||||
|
||||
Args:
|
||||
kb_name: 向量库名称
|
||||
chunk_id: 切片 ID
|
||||
content: 新的文本内容（None 表示不更新）
|
||||
metadata: 要更新的元数据字段（None 表示不更新）
|
||||
|
||||
Returns:
|
||||
更新成功返回 True，切片不存在或更新失败返回 False
|
||||
|
||||
Example:
|
||||
>>> kb_manager.update_chunk(
|
||||
... "public_kb",
|
||||
... "manual_20260517_abc123",
|
||||
... content="更新后的内容",
|
||||
... metadata={"status": "updated"}
|
||||
... )
|
||||
True
|
||||
"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
return False
|
||||
|
||||
existing = collection.get(ids=[chunk_id])
|
||||
if not existing['ids']:
|
||||
logger.warning(f"切片不存在: {chunk_id}")
|
||||
return False
|
||||
|
||||
update_kwargs = {"ids": [chunk_id]}
|
||||
|
||||
if content is not None:
|
||||
update_kwargs["documents"] = [content]
|
||||
try:
|
||||
from core.engine import get_engine
|
||||
engine = get_engine()
|
||||
if not engine._initialized:
|
||||
engine.initialize()
|
||||
embedding = engine.embedding_model.encode(content).tolist()
|
||||
update_kwargs["embeddings"] = [embedding]
|
||||
except Exception as e:
|
||||
logger.error(f"生成向量失败: {e}")
|
||||
return False
|
||||
|
||||
if metadata is not None:
|
||||
existing_metadata = existing['metadatas'][0] if existing['metadatas'] else {}
|
||||
existing_metadata.update(metadata)
|
||||
update_kwargs["metadatas"] = [existing_metadata]
|
||||
|
||||
collection.update(**update_kwargs)
|
||||
|
||||
if content is not None:
|
||||
bm25 = self.get_bm25_index(kb_name)
|
||||
if bm25:
|
||||
bm25.add_documents([chunk_id], [content], [metadata or {}])
|
||||
self.save_bm25_index(kb_name)
|
||||
|
||||
logger.info(f"更新切片: {chunk_id}")
|
||||
return True
|
||||
|
||||
def get_chunk(self, kb_name: str, chunk_id: str) -> Optional[Dict]:
|
||||
"""
|
||||
获取单个切片信息
|
||||
|
||||
Args:
|
||||
kb_name: 向量库名称
|
||||
chunk_id: 切片 ID
|
||||
|
||||
Returns:
|
||||
切片信息字典，包含 id, document, metadata 等字段
|
||||
切片不存在时返回 None
|
||||
"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
return None
|
||||
|
||||
result = collection.get(ids=[chunk_id], include=["documents", "metadatas"])
|
||||
if not result['ids']:
|
||||
return None
|
||||
|
||||
return {
|
||||
'id': result['ids'][0],
|
||||
'document': result['documents'][0] if result['documents'] else '',
|
||||
'metadata': result['metadatas'][0] if result['metadatas'] else {}
|
||||
}
|
||||
|
||||
def delete_chunk(self, kb_name: str, chunk_id: str) -> Tuple[bool, Optional[str]]:
|
||||
"""
|
||||
删除切片
|
||||
|
||||
从向量库中移除指定的切片。
|
||||
|
||||
Args:
|
||||
kb_name: 向量库名称
|
||||
chunk_id: 切片 ID
|
||||
|
||||
Returns:
|
||||
元组 (success, source_file)：
|
||||
- success: 删除是否成功
|
||||
- source_file: 被删除切片的来源文件名（用于清理哈希记录）
|
||||
|
||||
Note:
|
||||
删除后 BM25 索引不会立即更新，需要手动调用 rebuild_bm25_index。
|
||||
"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
return False, None
|
||||
|
||||
# 获取切片信息（删除前）
|
||||
existing = collection.get(ids=[chunk_id], include=["metadatas"])
|
||||
if not existing['ids']:
|
||||
logger.warning(f"切片不存在: {chunk_id}")
|
||||
return False, None
|
||||
|
||||
source_file = None
|
||||
if existing['metadatas'] and existing['metadatas'][0]:
|
||||
source_file = existing['metadatas'][0].get('source')
|
||||
|
||||
collection.delete(ids=[chunk_id])
|
||||
logger.info(f"删除切片: {chunk_id}")
|
||||
return True, source_file
|
||||
|
||||
def delete_chunks_by_source(self, kb_name: str, source: str) -> int:
|
||||
"""
|
||||
批量删除指定文件的所有切片
|
||||
|
||||
Args:
|
||||
kb_name: 向量库名称
|
||||
source: 文件名（source 字段值）
|
||||
|
||||
Returns:
|
||||
删除的切片数量
|
||||
|
||||
Note:
|
||||
删除后会重建 BM25 索引。
|
||||
"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
return 0
|
||||
|
||||
# 获取所有切片，找到匹配 source 的
|
||||
all_chunks = collection.get(include=["metadatas"])
|
||||
|
||||
# 找到要删除的切片 ID
|
||||
chunk_ids_to_delete = []
|
||||
for i, chunk_id in enumerate(all_chunks['ids']):
|
||||
metadata = all_chunks['metadatas'][i] if all_chunks['metadatas'] else {}
|
||||
if metadata.get('source') == source:
|
||||
chunk_ids_to_delete.append(chunk_id)
|
||||
|
||||
if not chunk_ids_to_delete:
|
||||
logger.warning(f"未找到文件 {source} 的切片")
|
||||
return 0
|
||||
|
||||
# 批量删除
|
||||
collection.delete(ids=chunk_ids_to_delete)
|
||||
logger.info(f"批量删除切片: {source}, 共 {len(chunk_ids_to_delete)} 个")
|
||||
|
||||
# 重建 BM25 索引
|
||||
try:
|
||||
self._bm25_indexes.pop(kb_name, None)
|
||||
bm25 = self.get_bm25_index(kb_name)
|
||||
if bm25:
|
||||
remaining = collection.get(include=["documents", "metadatas"])
|
||||
if remaining['ids']:
|
||||
bm25.add_documents(
|
||||
remaining['ids'],
|
||||
remaining['documents'] or [],
|
||||
remaining['metadatas'] or []
|
||||
)
|
||||
self.save_bm25_index(kb_name)
|
||||
except Exception as e:
|
||||
logger.warning(f"重建 BM25 索引失败: {e}")
|
||||
|
||||
return len(chunk_ids_to_delete)
|
||||
|
||||
def list_chunks(
|
||||
self,
|
||||
kb_name: str,
|
||||
document_id: str = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0
|
||||
) -> List[Dict]:
|
||||
"""
|
||||
获取向量库中的切片列表
|
||||
|
||||
支持分页和按文档过滤。
|
||||
|
||||
Args:
|
||||
kb_name: 向量库名称
|
||||
document_id: 文档 ID（source 字段），用于过滤特定文档的切片
|
||||
limit: 返回数量限制（默认 100）
|
||||
offset: 偏移量（默认 0）
|
||||
|
||||
Returns:
|
||||
切片字典列表，每个字典包含：
|
||||
- id: 切片 ID
|
||||
- document: 切片内容
|
||||
- metadata: 元数据字典
|
||||
- status: 状态
|
||||
- version: 版本号
|
||||
|
||||
Example:
|
||||
>>> chunks = kb_manager.list_chunks("public_kb", limit=10)
|
||||
>>> for chunk in chunks:
|
||||
... print(chunk["id"], chunk["status"])
|
||||
"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
return []
|
||||
|
||||
where_filter = {}
|
||||
if document_id:
|
||||
where_filter["source"] = document_id
|
||||
|
||||
try:
|
||||
result = collection.get(
|
||||
where=where_filter if where_filter else None,
|
||||
limit=limit,
|
||||
offset=offset
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"获取切片失败: {e}")
|
||||
return []
|
||||
|
||||
return [
|
||||
{
|
||||
"id": id,
|
||||
"document": doc,
|
||||
"metadata": meta,
|
||||
"status": meta.get("status", "active") if meta else "active",
|
||||
"version": meta.get("version", "v1") if meta else "v1"
|
||||
}
|
||||
for id, doc, meta in zip(
|
||||
result['ids'],
|
||||
result.get('documents', []),
|
||||
result.get('metadatas', [])
|
||||
)
|
||||
]
|
||||
Reference in New Issue
Block a user