deepdata: core product code (backend, frontend, runtime, deepflow module)
This commit is contained in:
figmar
@@ -0,0 +1,23 @@
|
||||
"""
|
||||
RAG 基础设施层
|
||||
|
||||
提供文档处理、向量存储、嵌入模型和后台任务功能。
|
||||
|
||||
使用示例:
|
||||
from app.module.rag.infra.document import ingest_file_to_chunks
|
||||
from app.module.rag.infra.vectorstore import VectorStoreFactory
|
||||
from app.module.rag.infra.task import get_global_pool
|
||||
"""
|
||||
from app.module.rag.infra.document import (
|
||||
SplitOptions,
|
||||
default_split_options,
|
||||
ingest_file_to_chunks,
|
||||
load_and_split,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"load_and_split",
|
||||
"ingest_file_to_chunks",
|
||||
"SplitOptions",
|
||||
"default_split_options",
|
||||
]
|
||||
@@ -0,0 +1,42 @@
|
||||
"""
|
||||
文档处理模块
|
||||
|
||||
提供文档加载、分块和处理管道功能。
|
||||
|
||||
使用示例:
|
||||
from app.module.rag.infra.document import (
|
||||
ingest_file_to_chunks,
|
||||
SplitOptions,
|
||||
DocumentChunk,
|
||||
)
|
||||
|
||||
chunks = await ingest_file_to_chunks(
|
||||
"/path/to/doc.pdf",
|
||||
chunk_size=500,
|
||||
overlap_size=50,
|
||||
)
|
||||
"""
|
||||
from app.module.rag.infra.document.processor import (
|
||||
SplitOptions,
|
||||
default_split_options,
|
||||
ingest_file_to_chunks,
|
||||
load_and_split,
|
||||
)
|
||||
from app.module.rag.infra.document.types import (
|
||||
DocumentChunk,
|
||||
ParsedDocument,
|
||||
langchain_documents_to_parsed,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# 处理管道入口
|
||||
"load_and_split",
|
||||
"ingest_file_to_chunks",
|
||||
# 选项
|
||||
"SplitOptions",
|
||||
"default_split_options",
|
||||
# 类型
|
||||
"DocumentChunk",
|
||||
"ParsedDocument",
|
||||
"langchain_documents_to_parsed",
|
||||
]
|
||||
@@ -0,0 +1,26 @@
|
||||
"""
|
||||
文档加载器
|
||||
|
||||
封装 UniversalDocLoader，提供统一的文档加载接口。
|
||||
"""
|
||||
import asyncio
|
||||
from typing import List
|
||||
|
||||
from langchain_core.documents import Document
|
||||
|
||||
from app.module.shared.common.document_loaders import UniversalDocLoader
|
||||
|
||||
|
||||
async def load_document(file_path: str) -> List[Document]:
|
||||
"""加载文档（异步封装）
|
||||
|
||||
使用 UniversalDocLoader 加载文档，支持多种格式。
|
||||
|
||||
Args:
|
||||
file_path: 文件绝对路径
|
||||
|
||||
Returns:
|
||||
LangChain Document 列表
|
||||
"""
|
||||
loader = UniversalDocLoader(file_path)
|
||||
return await asyncio.to_thread(loader.load)
|
||||
@@ -0,0 +1,125 @@
|
||||
"""
|
||||
文档处理管道
|
||||
|
||||
提供统一的文档加载和分块入口，合并原有的 pipeline.py 和 options.py。
|
||||
"""
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from app.module.rag.infra.document.loader import load_document
|
||||
from app.module.rag.infra.document.splitter import DocumentSplitterFactory
|
||||
from app.module.rag.infra.document.types import (
|
||||
DocumentChunk,
|
||||
ParsedDocument,
|
||||
langchain_documents_to_parsed,
|
||||
)
|
||||
from app.module.rag.schema.enums import ProcessType
|
||||
|
||||
|
||||
@dataclass
|
||||
class SplitOptions:
|
||||
"""文档分片选项
|
||||
|
||||
Attributes:
|
||||
process_type: 分片策略
|
||||
chunk_size: 块大小（字符）
|
||||
overlap_size: 块间重叠
|
||||
delimiter: 仅 CUSTOM_SEPARATOR_CHUNK 时有效
|
||||
"""
|
||||
process_type: ProcessType = ProcessType.DEFAULT_CHUNK
|
||||
chunk_size: int = 500
|
||||
overlap_size: int = 50
|
||||
delimiter: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""转换为字典"""
|
||||
return {
|
||||
"process_type": self.process_type,
|
||||
"chunk_size": self.chunk_size,
|
||||
"overlap_size": self.overlap_size,
|
||||
"delimiter": self.delimiter,
|
||||
}
|
||||
|
||||
|
||||
def default_split_options() -> SplitOptions:
|
||||
"""默认分片选项：递归分块 500/50"""
|
||||
return SplitOptions()
|
||||
|
||||
|
||||
async def load_and_split(
|
||||
file_path: str,
|
||||
split_options: Optional[SplitOptions] = None,
|
||||
**chunk_metadata: Any,
|
||||
) -> List[DocumentChunk]:
|
||||
"""加载文档并分块
|
||||
|
||||
使用 UniversalDocLoader 加载文档，然后按指定策略分块。
|
||||
|
||||
Args:
|
||||
file_path: 文件绝对路径
|
||||
split_options: 分片选项，None 表示使用默认
|
||||
**chunk_metadata: 写入每个 chunk.metadata 的额外字段
|
||||
|
||||
Returns:
|
||||
分块列表
|
||||
"""
|
||||
documents = await load_document(file_path)
|
||||
|
||||
parser_metadata = {}
|
||||
for key in ["original_file_id", "rag_file_id", "file_name"]:
|
||||
if key in chunk_metadata:
|
||||
parser_metadata[key] = chunk_metadata[key]
|
||||
|
||||
parsed = langchain_documents_to_parsed(documents, file_path, **parser_metadata)
|
||||
|
||||
options = split_options or default_split_options()
|
||||
|
||||
base_chunk_metadata = {
|
||||
"file_name": parsed.metadata.get("file_name", ""),
|
||||
"file_extension": parsed.metadata.get("file_extension", ""),
|
||||
"absolute_directory_path": parsed.metadata.get("absolute_directory_path", ""),
|
||||
"original_file_id": parsed.metadata.get("original_file_id", ""),
|
||||
"rag_file_id": parsed.metadata.get("rag_file_id", ""),
|
||||
}
|
||||
base_chunk_metadata.update(chunk_metadata)
|
||||
|
||||
splitter = DocumentSplitterFactory.create_splitter(
|
||||
options.process_type,
|
||||
chunk_size=options.chunk_size,
|
||||
overlap_size=options.overlap_size,
|
||||
delimiter=options.delimiter,
|
||||
)
|
||||
|
||||
return await splitter.split(parsed.text, **base_chunk_metadata)
|
||||
|
||||
|
||||
async def ingest_file_to_chunks(
|
||||
file_path: str,
|
||||
process_type: ProcessType = ProcessType.DEFAULT_CHUNK,
|
||||
chunk_size: int = 500,
|
||||
overlap_size: int = 50,
|
||||
delimiter: Optional[str] = None,
|
||||
**chunk_metadata: Any,
|
||||
) -> List[DocumentChunk]:
|
||||
"""从本地文件加载文档并分块（便捷入口）
|
||||
|
||||
可被 ETL、URL 抓取、S3 等场景复用。
|
||||
|
||||
Args:
|
||||
file_path: 文件绝对路径
|
||||
process_type: 分块策略
|
||||
chunk_size: 块大小
|
||||
overlap_size: 重叠大小
|
||||
delimiter: 自定义分隔符
|
||||
**chunk_metadata: 写入每个 chunk.metadata 的额外字段
|
||||
|
||||
Returns:
|
||||
分块列表
|
||||
"""
|
||||
split_options = SplitOptions(
|
||||
process_type=process_type,
|
||||
chunk_size=chunk_size,
|
||||
overlap_size=overlap_size,
|
||||
delimiter=delimiter,
|
||||
)
|
||||
return await load_and_split(file_path, split_options=split_options, **chunk_metadata)
|
||||
@@ -0,0 +1,148 @@
|
||||
"""
|
||||
文档分块器
|
||||
|
||||
包含 DocumentSplitter 基类、工厂和 LangChain 实现。
|
||||
根据 ProcessType 创建对应的分块策略。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from langchain_text_splitters import (
|
||||
CharacterTextSplitter,
|
||||
RecursiveCharacterTextSplitter,
|
||||
)
|
||||
|
||||
from app.module.rag.infra.document.types import DocumentChunk
|
||||
from app.module.rag.schema.enums import ProcessType
|
||||
|
||||
|
||||
# 各 ProcessType 对应的分隔符配置（优先保持较大语义块）
|
||||
SEPARATORS_BY_PROCESS_TYPE = {
|
||||
ProcessType.PARAGRAPH_CHUNK: ["\n\n", "\n", " ", ""],
|
||||
ProcessType.SENTENCE_CHUNK: ["\n\n", "\n", "。", "！", "？", ". ", "! ", "? ", " ", ""],
|
||||
ProcessType.DEFAULT_CHUNK: ["\n\n", "\n", " ", ""],
|
||||
ProcessType.CUSTOM_SEPARATOR_CHUNK: None,
|
||||
}
|
||||
|
||||
|
||||
class DocumentSplitter(ABC):
|
||||
"""文档分块器基类
|
||||
|
||||
所有具体的分块器都需要继承此类并实现 split 方法。
|
||||
"""
|
||||
|
||||
def __init__(self, chunk_size: int = 500, overlap_size: int = 50):
|
||||
"""初始化分块器
|
||||
|
||||
Args:
|
||||
chunk_size: 分块大小
|
||||
overlap_size: 重叠大小
|
||||
"""
|
||||
self.chunk_size = chunk_size
|
||||
self.overlap_size = overlap_size
|
||||
|
||||
@abstractmethod
|
||||
async def split(self, text: str, **metadata: Any) -> List[DocumentChunk]:
|
||||
"""分割文档
|
||||
|
||||
Args:
|
||||
text: 文档文本
|
||||
**metadata: 额外的元数据
|
||||
|
||||
Returns:
|
||||
分块列表
|
||||
"""
|
||||
pass
|
||||
|
||||
def _create_chunk(self, text: str, chunk_index: int, **metadata: Any) -> DocumentChunk:
|
||||
"""创建文档分块"""
|
||||
chunk_metadata = {"chunk_index": chunk_index, **metadata}
|
||||
return DocumentChunk(text=text, metadata=chunk_metadata)
|
||||
|
||||
|
||||
class LangChainDocumentSplitter(DocumentSplitter):
|
||||
"""基于 LangChain 的分块器实现
|
||||
|
||||
根据 ProcessType 选择 RecursiveCharacterTextSplitter 或 CharacterTextSplitter。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
process_type: ProcessType,
|
||||
chunk_size: int = 500,
|
||||
overlap_size: int = 50,
|
||||
delimiter: Optional[str] = None,
|
||||
):
|
||||
super().__init__(chunk_size=chunk_size, overlap_size=overlap_size)
|
||||
self._process_type = process_type
|
||||
self._delimiter = delimiter or "\n\n"
|
||||
self._splitter = self._create_splitter()
|
||||
|
||||
def _create_splitter(self) -> RecursiveCharacterTextSplitter | CharacterTextSplitter:
|
||||
"""创建 LangChain 分块器实例"""
|
||||
if self._process_type == ProcessType.LENGTH_CHUNK:
|
||||
return CharacterTextSplitter(
|
||||
chunk_size=self.chunk_size,
|
||||
chunk_overlap=self.overlap_size,
|
||||
length_function=len,
|
||||
)
|
||||
|
||||
separators = SEPARATORS_BY_PROCESS_TYPE.get(self._process_type)
|
||||
if self._process_type == ProcessType.CUSTOM_SEPARATOR_CHUNK:
|
||||
separators = [self._delimiter, "\n", " ", ""]
|
||||
if separators is None:
|
||||
separators = ["\n\n", "\n", " ", ""]
|
||||
|
||||
return RecursiveCharacterTextSplitter(
|
||||
chunk_size=self.chunk_size,
|
||||
chunk_overlap=self.overlap_size,
|
||||
separators=separators,
|
||||
length_function=len,
|
||||
)
|
||||
|
||||
async def split(self, text: str, **metadata: Any) -> List[DocumentChunk]:
|
||||
"""分割文档（异步）"""
|
||||
if not text or not text.strip():
|
||||
return []
|
||||
|
||||
texts = await asyncio.to_thread(self._splitter.split_text, text)
|
||||
return [
|
||||
DocumentChunk(text=t, metadata={**metadata, "chunk_index": i})
|
||||
for i, t in enumerate(texts)
|
||||
]
|
||||
|
||||
|
||||
class DocumentSplitterFactory:
|
||||
"""文档分块器工厂
|
||||
|
||||
根据处理类型创建对应的分块器实例。
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def create_splitter(
|
||||
cls,
|
||||
process_type: ProcessType,
|
||||
chunk_size: int = 500,
|
||||
overlap_size: int = 50,
|
||||
delimiter: Optional[str] = None,
|
||||
) -> DocumentSplitter:
|
||||
"""创建分块器
|
||||
|
||||
Args:
|
||||
process_type: 处理类型
|
||||
chunk_size: 分块大小
|
||||
overlap_size: 重叠大小
|
||||
delimiter: 自定义分隔符（仅用于 CUSTOM_SEPARATOR_CHUNK）
|
||||
|
||||
Returns:
|
||||
DocumentSplitter 实例
|
||||
"""
|
||||
return LangChainDocumentSplitter(
|
||||
process_type=process_type,
|
||||
chunk_size=chunk_size,
|
||||
overlap_size=overlap_size,
|
||||
delimiter=delimiter,
|
||||
)
|
||||
@@ -0,0 +1,116 @@
|
||||
"""
|
||||
文档处理类型定义
|
||||
|
||||
包含 ParsedDocument（解析后文档）和 DocumentChunk（文档分块）。
|
||||
"""
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, Any, List, Optional
|
||||
from pathlib import Path
|
||||
|
||||
from langchain_core.documents import Document
|
||||
|
||||
|
||||
@dataclass
|
||||
class DocumentChunk:
|
||||
"""文档分块
|
||||
|
||||
包含分块的文本内容和元数据。
|
||||
|
||||
Attributes:
|
||||
text: 分块文本内容
|
||||
metadata: 分块元数据（包含文件信息、分块索引等）
|
||||
"""
|
||||
text: str
|
||||
metadata: dict
|
||||
|
||||
def __repr__(self):
|
||||
return f"<DocumentChunk(text_length={len(self.text)}, metadata_keys={list(self.metadata.keys())})>"
|
||||
|
||||
|
||||
class ParsedDocument:
|
||||
"""解析后的文档
|
||||
|
||||
包含文档的完整文本内容和元数据，由文档加载器生成，供分块器使用。
|
||||
|
||||
Attributes:
|
||||
text: 文档完整文本内容
|
||||
metadata: 文档元数据（如文件名、扩展名、路径等）
|
||||
file_name: 文件名
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
text: str,
|
||||
metadata: Dict[str, Any],
|
||||
file_name: str
|
||||
):
|
||||
"""初始化解析后的文档
|
||||
|
||||
Args:
|
||||
text: 文档文本内容
|
||||
metadata: 文档元数据（如作者、创建时间等）
|
||||
file_name: 文件名
|
||||
"""
|
||||
self.text = text
|
||||
self.metadata = metadata
|
||||
self.file_name = file_name
|
||||
|
||||
def __repr__(self):
|
||||
return f"<ParsedDocument(file_name={self.file_name}, text_length={len(self.text)})>"
|
||||
|
||||
|
||||
def langchain_documents_to_parsed(
|
||||
documents: List[Document],
|
||||
file_path: str,
|
||||
file_name: Optional[str] = None,
|
||||
**extra_metadata: Any,
|
||||
) -> ParsedDocument:
|
||||
"""将 LangChain Document 列表转换为 ParsedDocument
|
||||
|
||||
多页/多段结果合并为一个文档，用于后续分块处理。
|
||||
|
||||
Args:
|
||||
documents: LangChain 加载器返回的 Document 列表
|
||||
file_path: 源文件路径
|
||||
file_name: 文件名，若提供则优先使用
|
||||
**extra_metadata: 额外的元数据字段（会合并到返回的 metadata 中）
|
||||
|
||||
Returns:
|
||||
ParsedDocument: 合并后的文档对象
|
||||
"""
|
||||
path = Path(file_path)
|
||||
name = file_name or path.name
|
||||
|
||||
if not documents:
|
||||
base_metadata = {
|
||||
"file_name": name,
|
||||
"file_extension": path.suffix.lower(),
|
||||
"file_size": path.stat().st_size if path.exists() else 0,
|
||||
}
|
||||
base_metadata.update(extra_metadata)
|
||||
return ParsedDocument(
|
||||
text="",
|
||||
metadata=base_metadata,
|
||||
file_name=name,
|
||||
)
|
||||
|
||||
texts = [d.page_content for d in documents if d.page_content]
|
||||
merged_text = "\n\n".join(texts)
|
||||
|
||||
meta: Dict[str, Any] = {
|
||||
"file_name": name,
|
||||
"file_extension": path.suffix.lower(),
|
||||
"file_size": path.stat().st_size if path.exists() else 0,
|
||||
"absolute_directory_path": str(path.parent),
|
||||
"file_path": str(path),
|
||||
}
|
||||
|
||||
meta.update(extra_metadata)
|
||||
|
||||
if documents and isinstance(documents[0].metadata, dict):
|
||||
first_meta = documents[0].metadata
|
||||
for k, v in first_meta.items():
|
||||
if k not in meta and v is not None:
|
||||
meta[k] = v
|
||||
|
||||
return ParsedDocument(text=merged_text, metadata=meta, file_name=name)
|
||||
@@ -0,0 +1,97 @@
|
||||
"""
|
||||
LangChain Embeddings 封装
|
||||
|
||||
直接使用 LangChain 的 embeddings 功能，支持多种提供商：
|
||||
- OpenAI: langchain-openai
|
||||
- Ollama: langchain-community
|
||||
- 其他: 通过 LangChain 生态
|
||||
"""
|
||||
from typing import Optional, Any
|
||||
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from langchain_openai import OpenAIEmbeddings
|
||||
|
||||
|
||||
class EmbeddingFactory:
|
||||
"""LangChain Embeddings 工厂类"""
|
||||
|
||||
@staticmethod
|
||||
def create_embeddings(
|
||||
model_name: str,
|
||||
base_url: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> Embeddings:
|
||||
"""
|
||||
创建 LangChain Embeddings 实例
|
||||
|
||||
Args:
|
||||
model_name: 模型名称（如 text-embedding-3-small）
|
||||
base_url: API 基础 URL
|
||||
api_key: API 密钥
|
||||
**kwargs: 其他参数
|
||||
|
||||
Returns:
|
||||
LangChain Embeddings 实例
|
||||
"""
|
||||
# ========== 新增修复开始 ==========
|
||||
# DashScope兼容模式不支持encoding_format=base64，强制删除该key
|
||||
if "encoding_format" in kwargs:
|
||||
del kwargs["encoding_format"]
|
||||
# ========== 新增修复结束 ==========
|
||||
|
||||
# OpenAI / OpenAI 兼容接口
|
||||
if "openai" in model_name.lower() or model_name.startswith("text-embedding"):
|
||||
return OpenAIEmbeddings(
|
||||
model=model_name,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
check_embedding_ctx_length=False,
|
||||
**kwargs,
|
||||
)
|
||||
# Ollama
|
||||
if base_url and "ollama" in base_url.lower():
|
||||
from langchain_community.embeddings.ollama import OllamaEmbeddings
|
||||
|
||||
return OllamaEmbeddings(
|
||||
model=model_name,
|
||||
base_url=base_url,
|
||||
**kwargs,
|
||||
)
|
||||
# 默认使用 OpenAI 兼容
|
||||
return OpenAIEmbeddings(
|
||||
model=model_name,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# xuhuanqing
|
||||
@staticmethod
|
||||
async def aembed_texts(
|
||||
texts: list[str],
|
||||
model_name: str,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
**kwargs,
|
||||
) -> list[list[float]]:
|
||||
"""
|
||||
LangChain Embeddings 的异步包装
|
||||
专门给 LightRAG / 非 LangChain 调用方使用
|
||||
"""
|
||||
embeddings = EmbeddingFactory.create_embeddings(
|
||||
model_name=model_name,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# LangChain 的 embed_documents 是同步的
|
||||
# 在 async 环境中用 run_in_executor 包装
|
||||
import asyncio
|
||||
loop = asyncio.get_running_loop()
|
||||
return await loop.run_in_executor(
|
||||
None,
|
||||
embeddings.embed_documents,
|
||||
texts,
|
||||
)
|
||||
@@ -0,0 +1,8 @@
|
||||
"""
|
||||
异步任务处理模块
|
||||
|
||||
提供工作池和异步任务处理功能。
|
||||
"""
|
||||
from app.module.rag.infra.task.worker_pool import WorkerPool, get_global_pool
|
||||
|
||||
__all__ = ["WorkerPool", "get_global_pool"]
|
||||
@@ -0,0 +1,137 @@
|
||||
"""
|
||||
工作协程池
|
||||
|
||||
使用 asyncio.Semaphore 控制并发数，替代 Java 的虚拟线程 + 信号量。
|
||||
提供全局单例，确保所有文件处理共享同一个并发池。
|
||||
"""
|
||||
import asyncio
|
||||
from typing import Callable, Any, Coroutine, Optional
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 全局单例
|
||||
_global_pool: Optional["WorkerPool"] = None
|
||||
|
||||
|
||||
def get_global_pool(max_workers: int = 10) -> "WorkerPool":
|
||||
"""获取全局 WorkerPool 单例
|
||||
|
||||
所有文件处理任务共享同一个并发池，确保最多 10 个文件并行处理。
|
||||
|
||||
Args:
|
||||
max_workers: 最大并发数（仅在首次创建时生效）
|
||||
|
||||
Returns:
|
||||
全局 WorkerPool 实例
|
||||
"""
|
||||
global _global_pool
|
||||
if _global_pool is None:
|
||||
_global_pool = WorkerPool(max_workers)
|
||||
logger.info("创建全局 WorkerPool，最大并发数: %d", max_workers)
|
||||
return _global_pool
|
||||
|
||||
|
||||
class WorkerPool:
|
||||
"""工作协程池
|
||||
|
||||
对应 Java 的虚拟线程 + 信号量方案
|
||||
|
||||
使用 asyncio.Semaphore 控制并发数，避免资源耗尽
|
||||
|
||||
使用示例：
|
||||
pool = WorkerPool(max_workers=10)
|
||||
|
||||
async def task_func(item):
|
||||
# 处理任务
|
||||
return result
|
||||
|
||||
# 并发执行多个任务
|
||||
tasks = [pool.submit(task_func, item) for item in items]
|
||||
results = await asyncio.gather(*tasks)
|
||||
"""
|
||||
|
||||
def __init__(self, max_workers: int = 10):
|
||||
"""初始化工作协程池
|
||||
|
||||
Args:
|
||||
max_workers: 最大并发数（默认 10）
|
||||
"""
|
||||
self.semaphore = asyncio.Semaphore(max_workers)
|
||||
self.max_workers = max_workers
|
||||
self._lock = asyncio.Lock()
|
||||
self._active_count = 0
|
||||
|
||||
async def submit(
|
||||
self,
|
||||
coro: Callable[..., Coroutine],
|
||||
*args: Any,
|
||||
**kwargs: Any
|
||||
) -> Any:
|
||||
"""提交异步任务并等待完成
|
||||
|
||||
Args:
|
||||
coro: 异步协程函数
|
||||
*args: 位置参数
|
||||
**kwargs: 关键字参数
|
||||
|
||||
Returns:
|
||||
协程的返回值
|
||||
"""
|
||||
async with self.semaphore:
|
||||
async with self._lock:
|
||||
self._active_count += 1
|
||||
|
||||
try:
|
||||
result = await coro(*args, **kwargs)
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"任务执行失败: {e}")
|
||||
raise
|
||||
finally:
|
||||
async with self._lock:
|
||||
self._active_count -= 1
|
||||
|
||||
async def submit_batch(
|
||||
self,
|
||||
coro: Callable[..., Coroutine],
|
||||
items: list[Any],
|
||||
*args: Any,
|
||||
**kwargs: Any
|
||||
) -> list[Any]:
|
||||
"""批量提交异步任务
|
||||
|
||||
Args:
|
||||
coro: 异步协程函数
|
||||
items: 要处理的项目列表
|
||||
*args: 额外的位置参数
|
||||
**kwargs: 额外的关键字参数
|
||||
|
||||
Returns:
|
||||
结果列表
|
||||
"""
|
||||
tasks = [
|
||||
self.submit(coro, item, *args, **kwargs)
|
||||
for item in items
|
||||
]
|
||||
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
# 检查是否有任务失败
|
||||
for i, result in enumerate(results):
|
||||
if isinstance(result, Exception):
|
||||
logger.error(f"批次任务 {i} 失败: {result}")
|
||||
|
||||
return [r for r in results if not isinstance(r, Exception)]
|
||||
|
||||
async def get_status(self) -> dict:
|
||||
"""获取工作池状态
|
||||
|
||||
Returns:
|
||||
状态字典，包含 max_workers, active_count, available
|
||||
"""
|
||||
return {
|
||||
"max_workers": self.max_workers,
|
||||
"active_count": self._active_count,
|
||||
"available": self.max_workers - self._active_count,
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
"""
|
||||
向量存储模块
|
||||
|
||||
提供 Milvus 向量存储的创建、管理和数据操作功能。
|
||||
"""
|
||||
from app.module.rag.infra.vectorstore.factory import VectorStoreFactory
|
||||
from app.module.rag.infra.vectorstore.milvus_client import get_milvus_client
|
||||
from app.module.rag.infra.vectorstore.store import (
|
||||
chunks_to_documents,
|
||||
create_collection,
|
||||
delete_chunks_by_rag_file_ids,
|
||||
drop_collection,
|
||||
get_vector_dimension,
|
||||
rename_collection,
|
||||
update_chunk_by_id,
|
||||
delete_chunk_by_id,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"VectorStoreFactory",
|
||||
"get_milvus_client",
|
||||
"create_collection",
|
||||
"drop_collection",
|
||||
"rename_collection",
|
||||
"get_vector_dimension",
|
||||
"delete_chunks_by_rag_file_ids",
|
||||
"chunks_to_documents",
|
||||
"update_chunk_by_id",
|
||||
"delete_chunk_by_id",
|
||||
]
|
||||
@@ -0,0 +1,120 @@
|
||||
"""
|
||||
向量存储工厂
|
||||
|
||||
使用 LangChain Milvus 创建向量存储实例，支持混合检索（向量 + BM25）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from langchain_core.embeddings import Embeddings
|
||||
|
||||
from app.core.config import settings
|
||||
from app.module.rag.infra.vectorstore.store import (
|
||||
create_collection,
|
||||
drop_collection,
|
||||
get_vector_dimension,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class VectorStoreFactory:
|
||||
"""LangChain Milvus 向量存储工厂
|
||||
|
||||
使用单例模式缓存 Milvus 实例，确保每个 collection 只有一个实例。
|
||||
"""
|
||||
|
||||
_instances: Dict[str, Any] = {}
|
||||
_lock = threading.Lock()
|
||||
|
||||
@staticmethod
|
||||
def get_connection_args() -> dict:
|
||||
"""获取 Milvus 连接参数"""
|
||||
args: dict = {"uri": settings.milvus_uri}
|
||||
if getattr(settings, "milvus_token", None):
|
||||
args["token"] = settings.milvus_token
|
||||
return args
|
||||
|
||||
@classmethod
|
||||
def create(
|
||||
cls,
|
||||
collection_name: str,
|
||||
embedding: Embeddings,
|
||||
*,
|
||||
drop_old: bool = False,
|
||||
consistency_level: str = "Strong",
|
||||
force_new: bool = False,
|
||||
) -> Any:
|
||||
"""创建或获取 Milvus 向量存储实例（支持混合检索）
|
||||
|
||||
Args:
|
||||
collection_name: 集合名称（知识库名称）
|
||||
embedding: LangChain Embeddings 实例
|
||||
drop_old: 是否删除已存在同名集合（默认 False，避免数据丢失）
|
||||
consistency_level: 一致性级别
|
||||
force_new: 是否强制创建新实例（默认 False，优先使用缓存）
|
||||
|
||||
Returns:
|
||||
langchain_milvus.Milvus 实例
|
||||
"""
|
||||
if drop_old:
|
||||
drop_collection(collection_name)
|
||||
with cls._lock:
|
||||
cls._instances.pop(collection_name, None)
|
||||
|
||||
if not force_new and collection_name in cls._instances:
|
||||
logger.debug("使用缓存的 Milvus 实例: %s", collection_name)
|
||||
return cls._instances[collection_name]
|
||||
|
||||
from langchain_milvus import BM25BuiltInFunction, Milvus
|
||||
|
||||
dimension = get_vector_dimension(
|
||||
embedding_model="",
|
||||
embedding_instance=embedding,
|
||||
)
|
||||
|
||||
create_collection(
|
||||
collection_name=collection_name,
|
||||
dimension=dimension,
|
||||
consistency_level=consistency_level,
|
||||
)
|
||||
|
||||
instance = Milvus(
|
||||
embedding_function=embedding,
|
||||
collection_name=collection_name,
|
||||
connection_args=cls.get_connection_args(),
|
||||
builtin_function=BM25BuiltInFunction(),
|
||||
text_field="text",
|
||||
vector_field=["vector"],
|
||||
drop_old=False,
|
||||
consistency_level=consistency_level,
|
||||
auto_id=False,
|
||||
primary_field="id",
|
||||
metadata_field="metadata",
|
||||
enable_dynamic_field=False,
|
||||
metadata_schema={"metadata": "JSON"},
|
||||
)
|
||||
|
||||
with cls._lock:
|
||||
cls._instances[collection_name] = instance
|
||||
logger.info("创建并缓存 Milvus 实例: %s", collection_name)
|
||||
|
||||
return instance
|
||||
|
||||
@classmethod
|
||||
def clear_cache(cls, collection_name: Optional[str] = None) -> None:
|
||||
"""清除缓存
|
||||
|
||||
Args:
|
||||
collection_name: 集合名称，如果为 None 则清除所有缓存
|
||||
"""
|
||||
with cls._lock:
|
||||
if collection_name:
|
||||
cls._instances.pop(collection_name, None)
|
||||
logger.info("清除 Milvus 实例缓存: %s", collection_name)
|
||||
else:
|
||||
cls._instances.clear()
|
||||
logger.info("清除所有 Milvus 实例缓存")
|
||||
@@ -0,0 +1,83 @@
|
||||
"""
|
||||
Milvus 客户端单例管理器
|
||||
|
||||
确保 MilvusClient 在全局范围内只创建一个实例，避免重复连接。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from typing import Optional
|
||||
|
||||
from pymilvus import MilvusClient
|
||||
|
||||
from app.core.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MilvusClientManager:
|
||||
"""Milvus 客户端单例管理器
|
||||
|
||||
使用线程安全的单例模式，确保全局只有一个 MilvusClient 实例。
|
||||
"""
|
||||
|
||||
_instance: Optional[MilvusClientManager] = None
|
||||
_lock = threading.Lock()
|
||||
_client: Optional[MilvusClient] = None
|
||||
|
||||
def __new__(cls) -> MilvusClientManager:
|
||||
"""单例模式：确保只有一个管理器实例"""
|
||||
if cls._instance is None:
|
||||
with cls._lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
@staticmethod
|
||||
def _get_connection_args() -> dict:
|
||||
"""获取 Milvus 连接参数"""
|
||||
args: dict = {"uri": settings.milvus_uri}
|
||||
token = getattr(settings, "milvus_token", None)
|
||||
if token:
|
||||
args["token"] = token
|
||||
return args
|
||||
|
||||
def get_client(self) -> MilvusClient:
|
||||
"""获取 Milvus 客户端实例（单例）
|
||||
|
||||
Returns:
|
||||
MilvusClient 实例
|
||||
"""
|
||||
if self._client is None:
|
||||
with self._lock:
|
||||
if self._client is None:
|
||||
conn_args = self._get_connection_args()
|
||||
self._client = MilvusClient(
|
||||
uri=conn_args["uri"],
|
||||
token=conn_args.get("token", "")
|
||||
)
|
||||
logger.info("创建 Milvus 客户端单例: uri=%s", conn_args["uri"])
|
||||
return self._client
|
||||
|
||||
def close(self) -> None:
|
||||
"""关闭 Milvus 客户端连接"""
|
||||
if self._client is not None:
|
||||
with self._lock:
|
||||
if self._client is not None:
|
||||
try:
|
||||
self._client.close()
|
||||
logger.info("关闭 Milvus 客户端连接")
|
||||
except Exception as e:
|
||||
logger.warning("关闭 Milvus 客户端时出错: %s", e)
|
||||
finally:
|
||||
self._client = None
|
||||
|
||||
|
||||
def get_milvus_client() -> MilvusClient:
|
||||
"""获取 Milvus 客户端实例（全局单例）
|
||||
|
||||
Returns:
|
||||
MilvusClient 实例
|
||||
"""
|
||||
return MilvusClientManager().get_client()
|
||||
@@ -0,0 +1,466 @@
|
||||
"""
|
||||
向量存储管理
|
||||
|
||||
提供 Milvus 集合的创建、删除、重命名和数据操作功能。
|
||||
使用 Milvus 2.6+ 的 BM25 内置函数实现混合检索。
|
||||
|
||||
字段结构：
|
||||
- id (VarChar, 主键)
|
||||
- text (VarChar, with analyzer for BM25)
|
||||
- metadata (JSON)
|
||||
- vector (FloatVector, 密集向量)
|
||||
- sparse (SparseFloatVector, BM25 稀疏向量)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
from langchain_core.documents import Document
|
||||
from pymilvus import DataType, FunctionType, CollectionSchema, FieldSchema, Function
|
||||
|
||||
from app.core.exception import BusinessError, ErrorCodes
|
||||
from app.module.rag.infra.document.types import DocumentChunk
|
||||
from app.module.rag.infra.embeddings import EmbeddingFactory
|
||||
from app.module.rag.infra.vectorstore.milvus_client import get_milvus_client
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
BATCH_DELETE_SIZE = 500
|
||||
|
||||
|
||||
def _delete_chunks_by_rag_file_id_batched(client, collection_name: str, rag_file_id: str) -> int:
|
||||
"""分批删除指定 rag_file_id 的所有 chunks
|
||||
|
||||
Args:
|
||||
client: Milvus 客户端
|
||||
collection_name: 集合名称
|
||||
rag_file_id: RAG 文件 ID
|
||||
|
||||
Returns:
|
||||
删除的总数量
|
||||
"""
|
||||
filter_expr = f'metadata["rag_file_id"] == "{rag_file_id}"'
|
||||
total_deleted = 0
|
||||
|
||||
while True:
|
||||
try:
|
||||
results = client.query(
|
||||
collection_name=collection_name,
|
||||
filter=filter_expr,
|
||||
output_fields=["id"],
|
||||
limit=BATCH_DELETE_SIZE,
|
||||
)
|
||||
if not results:
|
||||
break
|
||||
|
||||
chunk_ids = [r["id"] for r in results]
|
||||
id_filter = ' || '.join([f'id == "{cid}"' for cid in chunk_ids])
|
||||
client.delete(collection_name=collection_name, filter=f"({id_filter})")
|
||||
total_deleted += len(chunk_ids)
|
||||
|
||||
if len(chunk_ids) < BATCH_DELETE_SIZE:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.warning("分批删除失败: collection=%s rag_file_id=%s error=%s", collection_name, rag_file_id, e)
|
||||
break
|
||||
|
||||
return total_deleted
|
||||
|
||||
|
||||
def drop_collection(collection_name: str) -> None:
|
||||
"""删除 Milvus 集合
|
||||
|
||||
Args:
|
||||
collection_name: 集合名称
|
||||
"""
|
||||
try:
|
||||
client = get_milvus_client()
|
||||
if client.has_collection(collection_name):
|
||||
client.drop_collection(collection_name)
|
||||
logger.info("成功删除集合: %s", collection_name)
|
||||
except Exception as e:
|
||||
logger.error("删除集合失败: %s", e)
|
||||
raise BusinessError(ErrorCodes.RAG_MILVUS_ERROR, f"删除集合失败: {str(e)}") from e
|
||||
|
||||
|
||||
def rename_collection(old_name: str, new_name: str) -> None:
|
||||
"""重命名 Milvus 集合
|
||||
|
||||
Args:
|
||||
old_name: 原集合名称
|
||||
new_name: 新集合名称
|
||||
"""
|
||||
from pymilvus import utility, connections
|
||||
from app.core.config import settings
|
||||
|
||||
try:
|
||||
uri = settings.milvus_uri
|
||||
token = getattr(settings, "milvus_token", None) or ""
|
||||
connections.connect(
|
||||
alias="default",
|
||||
uri=uri,
|
||||
token=token,
|
||||
)
|
||||
if utility.has_collection(old_name, using="default"):
|
||||
utility.rename_collection(old_name, new_name, using="default")
|
||||
logger.info("成功重命名集合: %s -> %s", old_name, new_name)
|
||||
except Exception as e:
|
||||
logger.error("重命名集合失败: %s", e)
|
||||
raise BusinessError(ErrorCodes.RAG_MILVUS_ERROR, f"重命名集合失败: {str(e)}") from e
|
||||
|
||||
#xuhuanqing
|
||||
def create_collection(
|
||||
collection_name: str,
|
||||
dimension: int,
|
||||
consistency_level: str = "Strong",
|
||||
) -> None:
|
||||
"""创建 Milvus 集合（Milvus 2.6 + Woodpecker + BM25 稳定版）"""
|
||||
BM25_INDEX_PARAMS = {
|
||||
"inverted_index_algo": "DAAT_MAXSCORE",
|
||||
"bm25_k1": 1.2,
|
||||
"bm25_b": 0.75,
|
||||
}
|
||||
|
||||
try:
|
||||
client = get_milvus_client()
|
||||
|
||||
if client.has_collection(collection_name):
|
||||
logger.info("集合 %s 已存在，跳过创建", collection_name)
|
||||
return
|
||||
|
||||
fields = [
|
||||
FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=36, is_primary=True, auto_id=False),
|
||||
FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535, enable_analyzer=True),
|
||||
FieldSchema(name="metadata", dtype=DataType.JSON),
|
||||
FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dimension),
|
||||
FieldSchema(name="sparse", dtype=DataType.SPARSE_FLOAT_VECTOR),
|
||||
]
|
||||
|
||||
bm25_function = Function(
|
||||
name="text_bm25_emb",
|
||||
function_type=FunctionType.BM25,
|
||||
input_field_names=["text"],
|
||||
output_field_names=["sparse"],
|
||||
# ❌ 不写 params
|
||||
)
|
||||
|
||||
schema = CollectionSchema(
|
||||
fields=fields,
|
||||
functions=[bm25_function],
|
||||
description="Knowledge base collection",
|
||||
enable_dynamic_field=False,
|
||||
)
|
||||
|
||||
index_params = client.prepare_index_params()
|
||||
|
||||
# ✅ BM25 稀疏索引（唯一合法写法）
|
||||
index_params.add_index(
|
||||
field_name="sparse",
|
||||
index_type="SPARSE_INVERTED_INDEX",
|
||||
metric_type="BM25",
|
||||
params=BM25_INDEX_PARAMS,
|
||||
)
|
||||
|
||||
# ✅ 稠密向量索引
|
||||
index_params.add_index(
|
||||
field_name="vector",
|
||||
index_type="FLAT",
|
||||
metric_type="COSINE",
|
||||
params={},
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"开始创建集合: %s (维度: %d, consistency: %s)",
|
||||
collection_name, dimension, consistency_level,
|
||||
)
|
||||
|
||||
client.create_collection(
|
||||
collection_name=collection_name,
|
||||
schema=schema,
|
||||
index_params=index_params,
|
||||
consistency_level=consistency_level,
|
||||
timeout=None, # ✅ 防止首次 BM25 建库超时
|
||||
)
|
||||
|
||||
logger.info("成功创建集合: %s (维度: %d)，索引创建完成", collection_name, dimension)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("创建集合失败: %s", e)
|
||||
raise BusinessError(
|
||||
ErrorCodes.RAG_MILVUS_ERROR,
|
||||
f"创建集合失败: {str(e)}",
|
||||
) from e
|
||||
|
||||
'''
|
||||
def create_collection(
|
||||
collection_name: str,
|
||||
dimension: int,
|
||||
consistency_level: str = "Strong",
|
||||
) -> None:
|
||||
"""创建 Milvus 集合
|
||||
|
||||
使用标准的5字段结构：id、text、metadata、vector、sparse
|
||||
|
||||
Args:
|
||||
collection_name: 集合名称
|
||||
dimension: 向量维度
|
||||
consistency_level: 一致性级别
|
||||
"""
|
||||
try:
|
||||
client = get_milvus_client()
|
||||
|
||||
if client.has_collection(collection_name):
|
||||
logger.info("集合 %s 已存在，跳过创建", collection_name)
|
||||
return
|
||||
|
||||
# 创建字段
|
||||
fields = [
|
||||
FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=36, is_primary=True, auto_id=False),
|
||||
FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535, enable_analyzer=True),
|
||||
FieldSchema(name="metadata", dtype=DataType.JSON),
|
||||
FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dimension),
|
||||
FieldSchema(name="sparse", dtype=DataType.SPARSE_FLOAT_VECTOR),
|
||||
]
|
||||
|
||||
# 创建 BM25 函数
|
||||
bm25_function = Function(
|
||||
name="text_bm25_emb",
|
||||
function_type=FunctionType.BM25,
|
||||
input_field_names=["text"],
|
||||
output_field_names=["sparse"],
|
||||
)
|
||||
|
||||
schema = CollectionSchema(
|
||||
fields=fields,
|
||||
functions=[bm25_function],
|
||||
description="Knowledge base collection",
|
||||
enable_dynamic_field=False,
|
||||
)
|
||||
|
||||
# BM25 索引参数
|
||||
sparse_index_params = {
|
||||
"inverted_index_algo": "DAAT_MAXSCORE",
|
||||
"bm25_k1": 1.2,
|
||||
"bm25_b": 0.75,
|
||||
}
|
||||
|
||||
index_params = client.prepare_index_params()
|
||||
index_params.add_index(
|
||||
field_name="sparse",
|
||||
index_type="SPARSE_INVERTED_INDEX",
|
||||
metric_type="BM25",
|
||||
params=sparse_index_params
|
||||
)
|
||||
index_params.add_index(
|
||||
field_name="vector",
|
||||
index_type="FLAT",
|
||||
metric_type="COSINE",
|
||||
params={}
|
||||
)
|
||||
|
||||
client.create_collection(
|
||||
collection_name=collection_name,
|
||||
schema=schema,
|
||||
index_params=index_params,
|
||||
consistency_level=consistency_level,
|
||||
)
|
||||
|
||||
logger.info("成功创建集合: %s (维度: %d)", collection_name, dimension)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("创建集合失败: %s", e)
|
||||
raise BusinessError(ErrorCodes.RAG_MILVUS_ERROR, f"创建集合失败: {str(e)}") from e
|
||||
'''
|
||||
|
||||
def get_vector_dimension(
|
||||
embedding_model: str = "",
|
||||
base_url: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
embedding_instance=None,
|
||||
) -> int:
|
||||
"""获取嵌入模型的向量维度
|
||||
|
||||
Args:
|
||||
embedding_model: 模型名称
|
||||
base_url: API 基础 URL
|
||||
api_key: API 密钥
|
||||
embedding_instance: 已有的 Embeddings 实例（优先使用）
|
||||
|
||||
Returns:
|
||||
向量维度
|
||||
"""
|
||||
try:
|
||||
if embedding_instance:
|
||||
embedding = embedding_instance
|
||||
else:
|
||||
embedding = EmbeddingFactory.create_embeddings(
|
||||
model_name=embedding_model,
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
test_vector = embedding.embed_query("test")
|
||||
dimension = len(test_vector)
|
||||
logger.info("获取向量维度: %d", dimension)
|
||||
return dimension
|
||||
|
||||
except Exception as e:
|
||||
logger.error("获取向量维度失败: %s", e)
|
||||
raise BusinessError(ErrorCodes.RAG_EMBEDDING_FAILED, f"获取向量维度失败: {str(e)}") from e
|
||||
|
||||
def delete_chunks_by_rag_file_ids(collection_name: str, rag_file_ids: List[str]) -> None:
|
||||
"""按 RAG 文件 ID 列表删除 Milvus 中的分块
|
||||
|
||||
Args:
|
||||
collection_name: 集合名称
|
||||
rag_file_ids: RAG 文件 ID 列表
|
||||
"""
|
||||
if not rag_file_ids:
|
||||
return
|
||||
|
||||
try:
|
||||
client = get_milvus_client()
|
||||
|
||||
for rid in rag_file_ids:
|
||||
deleted = _delete_chunks_by_rag_file_id_batched(client, collection_name, rid)
|
||||
logger.info("删除文件分块: collection=%s rag_file_id=%s deleted=%d", collection_name, rid, deleted)
|
||||
|
||||
logger.info("已按 rag_file_id 删除集合 %s 中的分块: %s", collection_name, rag_file_ids)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("删除 Milvus 分块失败: %s", e)
|
||||
raise BusinessError(ErrorCodes.RAG_MILVUS_ERROR, f"删除分块失败: {str(e)}") from e
|
||||
|
||||
|
||||
def chunks_to_documents(
|
||||
chunks: List[DocumentChunk],
|
||||
ids: Optional[List[str]] = None,
|
||||
) -> tuple[List[Document], List[str]]:
|
||||
"""将 DocumentChunk 转换为 LangChain Document 格式
|
||||
|
||||
Args:
|
||||
chunks: DocumentChunk 列表
|
||||
ids: 可选的 ID 列表
|
||||
|
||||
Returns:
|
||||
(documents, ids): LangChain Document 列表和对应的 ID 列表
|
||||
"""
|
||||
if ids is None:
|
||||
import uuid
|
||||
ids = [str(uuid.uuid4()) for _ in chunks]
|
||||
|
||||
documents = []
|
||||
for chunk, chunk_id in zip(chunks, ids):
|
||||
doc = Document(page_content=chunk.text, metadata=chunk.metadata)
|
||||
documents.append(doc)
|
||||
|
||||
return documents, ids
|
||||
|
||||
|
||||
def update_chunk_by_id(
|
||||
collection_name: str,
|
||||
chunk_id: str,
|
||||
text: str,
|
||||
metadata: Optional[dict] = None,
|
||||
embedding_instance=None,
|
||||
) -> None:
|
||||
"""更新指定 ID 的分块
|
||||
|
||||
Args:
|
||||
collection_name: 集合名称
|
||||
chunk_id: 分块 ID
|
||||
text: 新的文本内容
|
||||
metadata: 新的元数据（可选）
|
||||
embedding_instance: Embeddings 实例
|
||||
"""
|
||||
try:
|
||||
client = get_milvus_client()
|
||||
|
||||
filter_expr = f'id == "{chunk_id}"'
|
||||
existing = client.query(
|
||||
collection_name=collection_name,
|
||||
filter=filter_expr,
|
||||
output_fields=["metadata"],
|
||||
)
|
||||
|
||||
if not existing:
|
||||
raise BusinessError(
|
||||
ErrorCodes.RAG_CHUNK_NOT_FOUND,
|
||||
f"Chunk not found: {chunk_id}"
|
||||
)
|
||||
|
||||
existing_metadata = existing[0].get("metadata", {})
|
||||
|
||||
if metadata is None:
|
||||
metadata = existing_metadata
|
||||
else:
|
||||
# 确保保留原有的 rag_file_id 字段，防止用户修改时丢失
|
||||
if "rag_file_id" in existing_metadata and "rag_file_id" not in metadata:
|
||||
metadata = {**metadata, "rag_file_id": existing_metadata["rag_file_id"]}
|
||||
|
||||
if embedding_instance:
|
||||
embedding = embedding_instance
|
||||
else:
|
||||
from app.module.rag.infra.embeddings import EmbeddingFactory
|
||||
embedding = EmbeddingFactory.create_embeddings()
|
||||
|
||||
vector = embedding.embed_query(text)
|
||||
|
||||
client.delete(collection_name=collection_name, filter=filter_expr)
|
||||
|
||||
client.insert(
|
||||
collection_name=collection_name,
|
||||
data=[{
|
||||
"id": chunk_id,
|
||||
"text": text,
|
||||
"metadata": metadata,
|
||||
"vector": vector,
|
||||
}]
|
||||
)
|
||||
|
||||
logger.info("成功更新分块: collection=%s chunk_id=%s", collection_name, chunk_id)
|
||||
|
||||
except BusinessError:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error("更新分块失败: %s", e)
|
||||
raise BusinessError(ErrorCodes.RAG_MILVUS_ERROR, f"更新分块失败: {str(e)}") from e
|
||||
|
||||
|
||||
def delete_chunk_by_id(collection_name: str, chunk_id: str) -> Optional[str]:
|
||||
"""删除指定 ID 的分块
|
||||
|
||||
Args:
|
||||
collection_name: 集合名称
|
||||
chunk_id: 分块 ID
|
||||
|
||||
Returns:
|
||||
被删除分块对应的 rag_file_id（如果存在），否则返回 None
|
||||
"""
|
||||
try:
|
||||
client = get_milvus_client()
|
||||
|
||||
filter_expr = f'id == "{chunk_id}"'
|
||||
|
||||
# 先查询 chunk 的 metadata 获取 rag_file_id
|
||||
existing = client.query(
|
||||
collection_name=collection_name,
|
||||
filter=filter_expr,
|
||||
output_fields=["metadata"],
|
||||
)
|
||||
|
||||
rag_file_id = None
|
||||
if existing:
|
||||
metadata = existing[0].get("metadata", {})
|
||||
rag_file_id = metadata.get("rag_file_id")
|
||||
|
||||
client.delete(collection_name=collection_name, filter=filter_expr)
|
||||
|
||||
logger.info("成功删除分块: collection=%s chunk_id=%s rag_file_id=%s", collection_name, chunk_id, rag_file_id)
|
||||
|
||||
return rag_file_id
|
||||
|
||||
except Exception as e:
|
||||
logger.error("删除分块失败: %s", e)
|
||||
raise BusinessError(ErrorCodes.RAG_MILVUS_ERROR, f"删除分块失败: {str(e)}") from e
|
||||
Reference in New Issue
Block a user