deepdata: core product code (backend, frontend, runtime, deepflow module)
This commit is contained in:
figmar
@@ -0,0 +1,12 @@
|
||||
"""
|
||||
RAG 仓储层导出
|
||||
|
||||
集中导出所有仓储类
|
||||
"""
|
||||
from .knowledge_base_repository import KnowledgeBaseRepository
|
||||
from .file_repository import RagFileRepository
|
||||
|
||||
__all__ = [
|
||||
"KnowledgeBaseRepository",
|
||||
"RagFileRepository",
|
||||
]
|
||||
@@ -0,0 +1,346 @@
|
||||
"""
|
||||
RAG 文件仓储层
|
||||
|
||||
提供 RAG 文件数据访问操作
|
||||
使用 SQLAlchemy 异步 session 进行数据库操作
|
||||
"""
|
||||
from typing import List, Optional, Tuple
|
||||
from sqlalchemy import select, func, and_, or_
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from app.db.models.knowledge_gen import RagFile, FileStatus
|
||||
from app.core.exception import BusinessError, ErrorCodes
|
||||
|
||||
|
||||
class RagFileRepository:
|
||||
"""RAG 文件仓储类
|
||||
|
||||
对应 Java: com.deepdata.rag.indexer.domain.repository.RagFileRepository
|
||||
提供 RAG 文件的 CRUD 操作和查询功能
|
||||
"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
"""初始化仓储
|
||||
|
||||
Args:
|
||||
db: SQLAlchemy 异步 session
|
||||
"""
|
||||
self.db = db
|
||||
|
||||
async def create(self, rag_file: RagFile) -> RagFile:
|
||||
"""创建 RAG 文件
|
||||
|
||||
Args:
|
||||
rag_file: RAG 文件实体
|
||||
|
||||
Returns:
|
||||
创建的 RAG 文件实体
|
||||
"""
|
||||
self.db.add(rag_file)
|
||||
await self.db.flush()
|
||||
return rag_file
|
||||
|
||||
async def batch_create(self, rag_files: List[RagFile]) -> List[RagFile]:
|
||||
"""批量创建 RAG 文件
|
||||
|
||||
Args:
|
||||
rag_files: RAG 文件实体列表
|
||||
|
||||
Returns:
|
||||
创建的 RAG 文件实体列表
|
||||
"""
|
||||
self.db.add_all(rag_files)
|
||||
await self.db.flush()
|
||||
return rag_files
|
||||
|
||||
async def update(self, rag_file: RagFile) -> RagFile:
|
||||
"""更新 RAG 文件
|
||||
|
||||
Args:
|
||||
rag_file: RAG 文件实体（必须包含 id）
|
||||
|
||||
Returns:
|
||||
更新后的 RAG 文件实体
|
||||
|
||||
Raises:
|
||||
BusinessError: 文件不存在
|
||||
"""
|
||||
existing = await self.get_by_id(rag_file.id)
|
||||
if not existing:
|
||||
raise BusinessError(ErrorCodes.RAG_FILE_NOT_FOUND)
|
||||
|
||||
# 更新字段
|
||||
if rag_file.chunk_count is not None:
|
||||
existing.chunk_count = rag_file.chunk_count
|
||||
if rag_file.metadata is not None:
|
||||
existing.metadata = rag_file.metadata
|
||||
if rag_file.status is not None:
|
||||
existing.status = rag_file.status
|
||||
if rag_file.err_msg is not None:
|
||||
existing.err_msg = rag_file.err_msg
|
||||
|
||||
await self.db.flush()
|
||||
return existing
|
||||
|
||||
async def delete(self, rag_file_id: str) -> None:
|
||||
"""删除 RAG 文件
|
||||
|
||||
Args:
|
||||
rag_file_id: RAG 文件 ID
|
||||
|
||||
Raises:
|
||||
BusinessError: 文件不存在
|
||||
"""
|
||||
rag_file = await self.get_by_id(rag_file_id)
|
||||
if not rag_file:
|
||||
raise BusinessError(ErrorCodes.RAG_FILE_NOT_FOUND)
|
||||
|
||||
await self.db.delete(rag_file)
|
||||
await self.db.flush()
|
||||
|
||||
async def batch_delete(self, rag_file_ids: List[str]) -> None:
|
||||
"""批量删除 RAG 文件
|
||||
|
||||
Args:
|
||||
rag_file_ids: RAG 文件 ID 列表
|
||||
"""
|
||||
if not rag_file_ids:
|
||||
return
|
||||
|
||||
await self.db.execute(
|
||||
select(RagFile).where(RagFile.id.in_(rag_file_ids))
|
||||
)
|
||||
# 注意：实际删除需要在查询后进行，这里简化处理
|
||||
# 在实际使用时，应该先查询再删除
|
||||
|
||||
async def delete_by_knowledge_base(
|
||||
self,
|
||||
knowledge_base_id: str
|
||||
) -> int:
|
||||
"""删除知识库的所有文件
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
|
||||
Returns:
|
||||
删除的文件数量
|
||||
"""
|
||||
result = await self.db.execute(
|
||||
select(RagFile).where(
|
||||
RagFile.knowledge_base_id == knowledge_base_id
|
||||
)
|
||||
)
|
||||
files = result.scalars().all()
|
||||
|
||||
count = len(files)
|
||||
for file in files:
|
||||
await self.db.delete(file)
|
||||
|
||||
await self.db.flush()
|
||||
return count
|
||||
|
||||
async def get_by_id(self, rag_file_id: str) -> Optional[RagFile]:
|
||||
"""根据 ID 获取 RAG 文件
|
||||
|
||||
Args:
|
||||
rag_file_id: RAG 文件 ID
|
||||
|
||||
Returns:
|
||||
RAG 文件实体，不存在则返回 None
|
||||
"""
|
||||
result = await self.db.execute(
|
||||
select(RagFile).where(RagFile.id == rag_file_id)
|
||||
)
|
||||
return result.scalars().first()
|
||||
|
||||
async def get_by_file_id(self, file_id: str) -> Optional[RagFile]:
|
||||
"""根据原始文件 ID 获取 RAG 文件
|
||||
|
||||
Args:
|
||||
file_id: 原始文件 ID
|
||||
|
||||
Returns:
|
||||
RAG 文件实体，不存在则返回 None
|
||||
"""
|
||||
result = await self.db.execute(
|
||||
select(RagFile).where(RagFile.file_id == file_id)
|
||||
)
|
||||
return result.scalars().first()
|
||||
|
||||
async def list_by_knowledge_base(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
keyword: Optional[str] = None,
|
||||
status: Optional[FileStatus] = None,
|
||||
page: int = 1,
|
||||
page_size: int = 10
|
||||
) -> Tuple[List[RagFile], int]:
|
||||
"""分页查询知识库的文件列表
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
keyword: 搜索关键词（模糊匹配文件名）
|
||||
status: 文件状态筛选
|
||||
page: 页码（从 1 开始）
|
||||
page_size: 每页数量
|
||||
|
||||
Returns:
|
||||
(RAG 文件列表, 总记录数)
|
||||
"""
|
||||
# 构建查询条件
|
||||
conditions = [RagFile.knowledge_base_id == knowledge_base_id]
|
||||
|
||||
if keyword:
|
||||
conditions.append(RagFile.file_name.like(f"%{keyword}%"))
|
||||
|
||||
if status:
|
||||
conditions.append(RagFile.status == status)
|
||||
|
||||
# 查询总数
|
||||
count_query = select(func.count()).select_from(RagFile).where(
|
||||
and_(*conditions)
|
||||
)
|
||||
total_result = await self.db.execute(count_query)
|
||||
total = total_result.scalar() or 0
|
||||
|
||||
# 分页查询
|
||||
query = select(RagFile).where(and_(*conditions))
|
||||
query = query.order_by(RagFile.created_at.desc())
|
||||
query = query.offset((page - 1) * page_size).limit(page_size)
|
||||
|
||||
result = await self.db.execute(query)
|
||||
items = result.scalars().all()
|
||||
|
||||
return list(items), total
|
||||
|
||||
async def get_unprocessed_files(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
limit: int = 100
|
||||
) -> List[RagFile]:
|
||||
"""获取待处理的文件
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
limit: 最大返回数量
|
||||
|
||||
Returns:
|
||||
待处理的 RAG 文件列表
|
||||
"""
|
||||
result = await self.db.execute(
|
||||
select(RagFile).where(
|
||||
and_(
|
||||
RagFile.knowledge_base_id == knowledge_base_id,
|
||||
RagFile.status == FileStatus.UNPROCESSED
|
||||
)
|
||||
).limit(limit)
|
||||
)
|
||||
return list(result.scalars().all())
|
||||
|
||||
async def update_status(
|
||||
self,
|
||||
rag_file_id: str,
|
||||
status: FileStatus,
|
||||
err_msg: Optional[str] = None
|
||||
) -> None:
|
||||
"""更新文件状态
|
||||
|
||||
Args:
|
||||
rag_file_id: RAG 文件 ID
|
||||
status: 新状态
|
||||
err_msg: 错误信息（可选）
|
||||
|
||||
Raises:
|
||||
BusinessError: 文件不存在
|
||||
"""
|
||||
rag_file = await self.get_by_id(rag_file_id)
|
||||
if not rag_file:
|
||||
raise BusinessError(ErrorCodes.RAG_FILE_NOT_FOUND)
|
||||
|
||||
rag_file.status = status
|
||||
if err_msg is not None:
|
||||
rag_file.err_msg = err_msg
|
||||
|
||||
await self.db.flush()
|
||||
|
||||
async def update_chunk_count(
|
||||
self,
|
||||
rag_file_id: str,
|
||||
chunk_count: int
|
||||
) -> None:
|
||||
"""更新文件分块数量
|
||||
|
||||
Args:
|
||||
rag_file_id: RAG 文件 ID
|
||||
chunk_count: 分块数量
|
||||
|
||||
Raises:
|
||||
BusinessError: 文件不存在
|
||||
"""
|
||||
rag_file = await self.get_by_id(rag_file_id)
|
||||
if not rag_file:
|
||||
raise BusinessError(ErrorCodes.RAG_FILE_NOT_FOUND)
|
||||
|
||||
rag_file.chunk_count = chunk_count
|
||||
await self.db.flush()
|
||||
|
||||
async def update_progress(
|
||||
self,
|
||||
rag_file_id: str,
|
||||
progress: int
|
||||
) -> None:
|
||||
"""更新文件处理进度
|
||||
|
||||
Args:
|
||||
rag_file_id: RAG 文件 ID
|
||||
progress: 进度值 (0-100)
|
||||
|
||||
Raises:
|
||||
BusinessError: 文件不存在
|
||||
"""
|
||||
rag_file = await self.get_by_id(rag_file_id)
|
||||
if not rag_file:
|
||||
raise BusinessError(ErrorCodes.RAG_FILE_NOT_FOUND)
|
||||
|
||||
rag_file.progress = max(0, min(100, progress))
|
||||
await self.db.flush()
|
||||
|
||||
async def count_by_knowledge_base(
|
||||
self,
|
||||
knowledge_base_id: str
|
||||
) -> int:
|
||||
"""统计知识库的文件数量
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
|
||||
Returns:
|
||||
文件数量
|
||||
"""
|
||||
result = await self.db.execute(
|
||||
select(func.count()).select_from(RagFile).where(
|
||||
RagFile.knowledge_base_id == knowledge_base_id
|
||||
)
|
||||
)
|
||||
return result.scalar() or 0
|
||||
|
||||
async def count_chunks_by_knowledge_base(
|
||||
self,
|
||||
knowledge_base_id: str
|
||||
) -> int:
|
||||
"""统计知识库的总分块数量
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
|
||||
Returns:
|
||||
总分块数量
|
||||
"""
|
||||
result = await self.db.execute(
|
||||
select(func.sum(RagFile.chunk_count)).where(
|
||||
and_(
|
||||
RagFile.knowledge_base_id == knowledge_base_id,
|
||||
RagFile.chunk_count.isnot(None)
|
||||
)
|
||||
)
|
||||
)
|
||||
return result.scalar() or 0
|
||||
@@ -0,0 +1,203 @@
|
||||
"""
|
||||
知识库仓储层
|
||||
|
||||
提供知识库数据访问操作
|
||||
使用 SQLAlchemy 异步 session 进行数据库操作
|
||||
"""
|
||||
from typing import List, Optional
|
||||
from sqlalchemy import select, func, and_, or_
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from app.db.models.knowledge_gen import KnowledgeBase, RagType
|
||||
from app.core.exception import BusinessError, ErrorCodes
|
||||
|
||||
|
||||
class KnowledgeBaseRepository:
|
||||
"""知识库仓储类
|
||||
|
||||
对应 Java: com.deepdata.rag.indexer.domain.repository.KnowledgeBaseRepository
|
||||
提供知识库的 CRUD 操作和查询功能
|
||||
"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
"""初始化仓储
|
||||
|
||||
Args:
|
||||
db: SQLAlchemy 异步 session
|
||||
"""
|
||||
self.db = db
|
||||
|
||||
async def create(self, knowledge_base: KnowledgeBase) -> KnowledgeBase:
|
||||
"""创建知识库
|
||||
|
||||
Args:
|
||||
knowledge_base: 知识库实体
|
||||
|
||||
Returns:
|
||||
创建的知识库实体
|
||||
|
||||
Raises:
|
||||
BusinessError: 知识库名称已存在
|
||||
"""
|
||||
# 检查名称是否已存在
|
||||
existing = await self.get_by_name(knowledge_base.name)
|
||||
if existing:
|
||||
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_ALREADY_EXISTS)
|
||||
|
||||
self.db.add(knowledge_base)
|
||||
await self.db.flush()
|
||||
return knowledge_base
|
||||
|
||||
async def update(self, knowledge_base: KnowledgeBase) -> KnowledgeBase:
|
||||
"""更新知识库
|
||||
|
||||
Args:
|
||||
knowledge_base: 知识库实体（必须包含 id）
|
||||
|
||||
Returns:
|
||||
更新后的知识库实体
|
||||
|
||||
Raises:
|
||||
BusinessError: 知识库不存在
|
||||
"""
|
||||
existing = await self.get_by_id(knowledge_base.id)
|
||||
if not existing:
|
||||
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
||||
|
||||
# 如果名称变更，检查新名称是否已存在
|
||||
if existing.name != knowledge_base.name:
|
||||
name_exists = await self.get_by_name(knowledge_base.name)
|
||||
if name_exists:
|
||||
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_ALREADY_EXISTS)
|
||||
|
||||
# 更新字段
|
||||
existing.name = knowledge_base.name
|
||||
existing.description = knowledge_base.description
|
||||
|
||||
await self.db.flush()
|
||||
return existing
|
||||
|
||||
async def delete(self, knowledge_base_id: str) -> None:
|
||||
"""删除知识库
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库ID
|
||||
|
||||
Raises:
|
||||
BusinessError: 知识库不存在
|
||||
"""
|
||||
knowledge_base = await self.get_by_id(knowledge_base_id)
|
||||
if not knowledge_base:
|
||||
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
||||
|
||||
await self.db.delete(knowledge_base)
|
||||
await self.db.flush()
|
||||
|
||||
async def get_by_id(self, knowledge_base_id: str) -> Optional[KnowledgeBase]:
|
||||
"""根据 ID 获取知识库
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库ID
|
||||
|
||||
Returns:
|
||||
知识库实体，不存在则返回 None
|
||||
"""
|
||||
result = await self.db.execute(
|
||||
select(KnowledgeBase).where(KnowledgeBase.id == knowledge_base_id)
|
||||
)
|
||||
return result.scalars().first()
|
||||
|
||||
async def get_by_name(self, name: str) -> Optional[KnowledgeBase]:
|
||||
"""根据名称获取知识库
|
||||
|
||||
Args:
|
||||
name: 知识库名称
|
||||
|
||||
Returns:
|
||||
知识库实体，不存在则返回 None
|
||||
"""
|
||||
result = await self.db.execute(
|
||||
select(KnowledgeBase).where(KnowledgeBase.name == name)
|
||||
)
|
||||
return result.scalars().first()
|
||||
|
||||
async def list(
|
||||
self,
|
||||
keyword: Optional[str] = None,
|
||||
rag_type: Optional[RagType] = None,
|
||||
page: int = 1,
|
||||
page_size: int = 10
|
||||
) -> tuple[List[KnowledgeBase], int]:
|
||||
"""分页查询知识库列表
|
||||
|
||||
Args:
|
||||
keyword: 搜索关键词（模糊匹配名称或描述）
|
||||
rag_type: RAG 类型筛选
|
||||
page: 页码（从 1 开始）
|
||||
page_size: 每页数量
|
||||
|
||||
Returns:
|
||||
(知识库列表, 总记录数)
|
||||
"""
|
||||
# 构建查询条件
|
||||
conditions = []
|
||||
|
||||
if keyword:
|
||||
conditions.append(
|
||||
or_(
|
||||
KnowledgeBase.name.like(f"%{keyword}%"),
|
||||
KnowledgeBase.description.like(f"%{keyword}%")
|
||||
)
|
||||
)
|
||||
|
||||
if rag_type:
|
||||
conditions.append(KnowledgeBase.type == rag_type)
|
||||
|
||||
# 构建查询
|
||||
query = select(KnowledgeBase)
|
||||
if conditions:
|
||||
query = query.where(and_(*conditions))
|
||||
|
||||
# 查询总数
|
||||
count_query = select(func.count()).select_from(KnowledgeBase)
|
||||
if conditions:
|
||||
count_query = count_query.where(and_(*conditions))
|
||||
|
||||
total_result = await self.db.execute(count_query)
|
||||
total = total_result.scalar() or 0
|
||||
|
||||
# 分页查询
|
||||
query = query.order_by(KnowledgeBase.created_at.desc())
|
||||
query = query.offset((page - 1) * page_size).limit(page_size)
|
||||
|
||||
result = await self.db.execute(query)
|
||||
items = result.scalars().all()
|
||||
|
||||
return list(items), total
|
||||
|
||||
async def exists_by_name(self, name: str, exclude_id: Optional[str] = None) -> bool:
|
||||
"""检查知识库名称是否存在
|
||||
|
||||
Args:
|
||||
name: 知识库名称
|
||||
exclude_id: 排除的知识库ID（用于更新时检查）
|
||||
|
||||
Returns:
|
||||
True 表示名称已存在，False 表示名称可用
|
||||
"""
|
||||
query = select(KnowledgeBase).where(KnowledgeBase.name == name)
|
||||
if exclude_id:
|
||||
query = query.where(KnowledgeBase.id != exclude_id)
|
||||
|
||||
result = await self.db.execute(query)
|
||||
return result.scalars().first() is not None
|
||||
|
||||
async def get_all_ids(self) -> List[str]:
|
||||
"""获取所有知识库 ID
|
||||
|
||||
Returns:
|
||||
知识库 ID 列表
|
||||
"""
|
||||
result = await self.db.execute(
|
||||
select(KnowledgeBase.id).order_by(KnowledgeBase.created_at.desc())
|
||||
)
|
||||
return list(result.scalars().all())
|
||||
Reference in New Issue
Block a user