feat: restore rag dataset management and linkage
This commit is contained in:
@@ -2,7 +2,9 @@ from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from fastapi import Depends, Query
|
||||
import json
|
||||
|
||||
from fastapi import Depends, File, Form, Query, UploadFile
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
|
||||
from fastapi_common.fastapi_common_security.security import verify_access_token
|
||||
@@ -14,6 +16,7 @@ from fastapi_modules.fastapi_leaudit.domian.Dto.ragChatDto import (
|
||||
RagChatSendMessageDTO,
|
||||
RagMessageFeedbackDTO,
|
||||
)
|
||||
from fastapi_modules.fastapi_leaudit.domian.Dto.ragDatasetDto import RagDatasetUpdateDTO
|
||||
from fastapi_modules.fastapi_leaudit.domian.vo.ragChatVo import (
|
||||
RagAppParametersVO,
|
||||
RagChatAppListVO,
|
||||
@@ -23,7 +26,16 @@ from fastapi_modules.fastapi_leaudit.domian.vo.ragChatVo import (
|
||||
RagMessagePageVO,
|
||||
RagOperationResultVO,
|
||||
)
|
||||
from fastapi_modules.fastapi_leaudit.domian.vo.ragDatasetVo import RagDatasetPageVO
|
||||
from fastapi_modules.fastapi_leaudit.domian.vo.ragDatasetVo import (
|
||||
RagDatasetDetailVO,
|
||||
RagDatasetDocumentItemVO,
|
||||
RagDatasetDocumentPageVO,
|
||||
RagDatasetPageVO,
|
||||
RagDatasetRetrieveResponseVO,
|
||||
RagDatasetSegmentItemVO,
|
||||
RagDatasetSegmentPageVO,
|
||||
RagDatasetUploadDocumentVO,
|
||||
)
|
||||
from fastapi_modules.fastapi_leaudit.services.impl.permissionServiceImpl import PermissionServiceImpl
|
||||
from fastapi_modules.fastapi_leaudit.services.impl.ragChatServiceImpl import RagChatServiceImpl
|
||||
from fastapi_modules.fastapi_leaudit.services.impl.ragDatasetServiceImpl import RagDatasetServiceImpl
|
||||
@@ -82,6 +94,332 @@ class RagChatController(BaseController):
|
||||
)
|
||||
return Result.success(data=data)
|
||||
|
||||
@self.router.get("/datasets/admin", response_model=Result[RagDatasetPageVO])
|
||||
async def GetAdminDatasets(
|
||||
area: str | None = Query(None),
|
||||
onlyEnabled: bool | None = Query(None),
|
||||
page: int = Query(1, ge=1),
|
||||
pageSize: int = Query(20, ge=1, le=200),
|
||||
payload: dict[str, Any] = Depends(verify_access_token),
|
||||
):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有管理知识库权限", "data": None})
|
||||
data = await self.RagDatasetService.GetAdminDatasets(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
Area=area,
|
||||
OnlyEnabled=onlyEnabled,
|
||||
Page=page,
|
||||
PageSize=pageSize,
|
||||
)
|
||||
return Result.success(data=data)
|
||||
|
||||
@self.router.post("/datasets/admin", response_model=Result[RagDatasetDetailVO])
|
||||
async def CreateAdminDataset(Body: dict[str, Any], payload: dict[str, Any] = Depends(verify_access_token)):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有创建知识库权限", "data": None})
|
||||
data = await self.RagDatasetService.CreateAdminDataset(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
Body=Body,
|
||||
)
|
||||
return Result.success(data=data)
|
||||
|
||||
@self.router.put("/datasets/admin/{DatasetId}", response_model=Result[RagDatasetDetailVO | None])
|
||||
async def UpdateAdminDataset(DatasetId: int, Body: dict[str, Any], payload: dict[str, Any] = Depends(verify_access_token)):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有更新知识库权限", "data": None})
|
||||
data = await self.RagDatasetService.UpdateAdminDataset(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
DatasetId=DatasetId,
|
||||
Body=Body,
|
||||
)
|
||||
return Result.success(data=data)
|
||||
|
||||
@self.router.delete("/datasets/admin/{DatasetId}", response_model=Result[RagOperationResultVO])
|
||||
async def DeleteAdminDataset(DatasetId: int, payload: dict[str, Any] = Depends(verify_access_token)):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有删除知识库权限", "data": None})
|
||||
data = await self.RagDatasetService.DeleteAdminDataset(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
DatasetId=DatasetId,
|
||||
)
|
||||
return Result.success(data=data)
|
||||
|
||||
@self.router.get("/datasets/{DatasetId}", response_model=Result[RagDatasetDetailVO | None])
|
||||
async def GetDatasetDetail(DatasetId: int, payload: dict[str, Any] = Depends(verify_access_token)):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有查看知识库权限", "data": None})
|
||||
data = await self.RagDatasetService.GetDatasetDetail(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
DatasetId=DatasetId,
|
||||
)
|
||||
return Result.success(data=data)
|
||||
|
||||
@self.router.patch("/datasets/{DatasetId}", response_model=Result[RagDatasetDetailVO | None])
|
||||
async def UpdateDataset(DatasetId: int, Body: RagDatasetUpdateDTO, payload: dict[str, Any] = Depends(verify_access_token)):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有修改知识库权限", "data": None})
|
||||
data = await self.RagDatasetService.UpdateDataset(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
DatasetId=DatasetId,
|
||||
Body=Body,
|
||||
)
|
||||
return Result.success(data=data)
|
||||
|
||||
@self.router.get("/datasets/{DatasetId}/documents", response_model=Result[RagDatasetDocumentPageVO])
|
||||
async def GetDatasetDocuments(
|
||||
DatasetId: int,
|
||||
page: int = Query(1, ge=1),
|
||||
limit: int = Query(20, ge=1, le=100),
|
||||
keyword: str | None = Query(None),
|
||||
payload: dict[str, Any] = Depends(verify_access_token),
|
||||
):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有查看知识库文档权限", "data": None})
|
||||
data = await self.RagDatasetService.GetDatasetDocuments(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
DatasetId=DatasetId,
|
||||
Page=page,
|
||||
Limit=limit,
|
||||
Keyword=keyword,
|
||||
)
|
||||
return Result.success(data=data)
|
||||
|
||||
@self.router.get("/datasets/{DatasetId}/documents/{DocumentId}", response_model=Result[RagDatasetDocumentItemVO | None])
|
||||
async def GetDatasetDocumentDetail(
|
||||
DatasetId: int,
|
||||
DocumentId: int,
|
||||
payload: dict[str, Any] = Depends(verify_access_token),
|
||||
):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有查看知识库文档权限", "data": None})
|
||||
data = await self.RagDatasetService.GetDatasetDocumentDetail(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
DatasetId=DatasetId,
|
||||
DocumentId=DocumentId,
|
||||
)
|
||||
return Result.success(data=data)
|
||||
|
||||
@self.router.post("/datasets/{DatasetId}/documents", response_model=Result[RagDatasetUploadDocumentVO])
|
||||
async def UploadDatasetDocument(
|
||||
DatasetId: int,
|
||||
file: UploadFile = File(...),
|
||||
data: str | None = Form(None),
|
||||
payload: dict[str, Any] = Depends(verify_access_token),
|
||||
):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有上传知识库文档权限", "data": None})
|
||||
process_config = json.loads(data) if data else None
|
||||
file_bytes = await file.read()
|
||||
result = await self.RagDatasetService.UploadDatasetDocument(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
DatasetId=DatasetId,
|
||||
FileName=file.filename or "document",
|
||||
ContentType=file.content_type,
|
||||
Content=file_bytes,
|
||||
ProcessConfig=process_config,
|
||||
)
|
||||
return Result.success(data=result)
|
||||
|
||||
@self.router.post("/datasets/{DatasetId}/documents/{DocumentId}/update-by-file", response_model=Result[RagDatasetUploadDocumentVO])
|
||||
async def UpdateDatasetDocumentByFile(
|
||||
DatasetId: int,
|
||||
DocumentId: int,
|
||||
file: UploadFile = File(...),
|
||||
data: str | None = Form(None),
|
||||
payload: dict[str, Any] = Depends(verify_access_token),
|
||||
):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有重处理知识库文档权限", "data": None})
|
||||
process_config = json.loads(data) if data else None
|
||||
file_bytes = await file.read()
|
||||
result = await self.RagDatasetService.UpdateDatasetDocumentByFile(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
DatasetId=DatasetId,
|
||||
DocumentId=DocumentId,
|
||||
FileName=file.filename or "document",
|
||||
ContentType=file.content_type,
|
||||
Content=file_bytes,
|
||||
ProcessConfig=process_config,
|
||||
)
|
||||
return Result.success(data=result)
|
||||
|
||||
@self.router.get("/datasets/{DatasetId}/documents/{DocumentId}/indexing-status")
|
||||
async def GetDatasetDocumentIndexingStatus(
|
||||
DatasetId: int,
|
||||
DocumentId: int,
|
||||
payload: dict[str, Any] = Depends(verify_access_token),
|
||||
):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有查看知识库处理进度权限", "data": None})
|
||||
result = await self.RagDatasetService.GetDatasetDocumentIndexingStatus(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
DatasetId=DatasetId,
|
||||
DocumentId=DocumentId,
|
||||
)
|
||||
return Result.success(data=result["data"])
|
||||
|
||||
@self.router.patch("/datasets/{DatasetId}/documents/status/{Action}", response_model=Result[RagOperationResultVO])
|
||||
async def BatchUpdateDatasetDocumentStatus(
|
||||
DatasetId: int,
|
||||
Action: str,
|
||||
Body: dict[str, Any],
|
||||
payload: dict[str, Any] = Depends(verify_access_token),
|
||||
):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有修改知识库文档状态权限", "data": None})
|
||||
enabled = Action == "enable"
|
||||
if Action not in {"enable", "disable"}:
|
||||
return JSONResponse(status_code=400, content={"code": 400, "msg": "当前仅支持启用和禁用", "data": None})
|
||||
document_ids = Body.get("document_ids") or []
|
||||
result = await self.RagDatasetService.BatchUpdateDatasetDocumentStatus(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
DatasetId=DatasetId,
|
||||
DocumentIds=[int(item) for item in document_ids],
|
||||
Enabled=enabled,
|
||||
)
|
||||
return Result.success(data=result)
|
||||
|
||||
@self.router.get("/datasets/{DatasetId}/documents/{DocumentId}/segments", response_model=Result[RagDatasetSegmentPageVO])
|
||||
async def GetDatasetDocumentSegments(
|
||||
DatasetId: int,
|
||||
DocumentId: int,
|
||||
page: int = Query(1, ge=1),
|
||||
limit: int = Query(20, ge=1, le=200),
|
||||
keyword: str | None = Query(None),
|
||||
payload: dict[str, Any] = Depends(verify_access_token),
|
||||
):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有查看知识库分段权限", "data": None})
|
||||
result = await self.RagDatasetService.GetDatasetDocumentSegments(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
DatasetId=DatasetId,
|
||||
DocumentId=DocumentId,
|
||||
Page=page,
|
||||
Limit=limit,
|
||||
Keyword=keyword,
|
||||
)
|
||||
return Result.success(data=result)
|
||||
|
||||
@self.router.delete("/datasets/{DatasetId}/documents/{DocumentId}", response_model=Result[RagOperationResultVO])
|
||||
async def DeleteDatasetDocument(
|
||||
DatasetId: int,
|
||||
DocumentId: int,
|
||||
payload: dict[str, Any] = Depends(verify_access_token),
|
||||
):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有删除知识库文档权限", "data": None})
|
||||
result = await self.RagDatasetService.DeleteDatasetDocument(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
DatasetId=DatasetId,
|
||||
DocumentId=DocumentId,
|
||||
)
|
||||
return Result.success(data=result)
|
||||
|
||||
@self.router.post("/datasets/{DatasetId}/retrieve", response_model=Result[RagDatasetRetrieveResponseVO])
|
||||
async def RetrieveDataset(
|
||||
DatasetId: int,
|
||||
Body: dict[str, Any],
|
||||
payload: dict[str, Any] = Depends(verify_access_token),
|
||||
):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有检索知识库权限", "data": None})
|
||||
result = await self.RagDatasetService.RetrieveDataset(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
DatasetId=DatasetId,
|
||||
Query=str(Body.get("query") or ""),
|
||||
RetrievalModel=Body.get("retrieval_model") if isinstance(Body.get("retrieval_model"), dict) else None,
|
||||
)
|
||||
return Result.success(data=result)
|
||||
|
||||
@self.router.get("/datasets/{DatasetId}/documents/{DocumentId}/segments/{SegmentId}", response_model=Result[RagDatasetSegmentItemVO | None])
|
||||
async def GetDatasetDocumentSegmentDetail(
|
||||
DatasetId: int,
|
||||
DocumentId: int,
|
||||
SegmentId: str,
|
||||
payload: dict[str, Any] = Depends(verify_access_token),
|
||||
):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有查看知识库分段权限", "data": None})
|
||||
result = await self.RagDatasetService.GetDatasetDocumentSegmentDetail(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
DatasetId=DatasetId,
|
||||
DocumentId=DocumentId,
|
||||
SegmentId=SegmentId,
|
||||
)
|
||||
return Result.success(data=result)
|
||||
|
||||
@self.router.post("/datasets/{DatasetId}/documents/{DocumentId}/segments/{SegmentId}", response_model=Result[RagDatasetSegmentItemVO | None])
|
||||
async def UpdateDatasetDocumentSegment(
|
||||
DatasetId: int,
|
||||
DocumentId: int,
|
||||
SegmentId: str,
|
||||
Body: dict[str, Any],
|
||||
payload: dict[str, Any] = Depends(verify_access_token),
|
||||
):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有修改知识库分段权限", "data": None})
|
||||
result = await self.RagDatasetService.UpdateDatasetDocumentSegment(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
DatasetId=DatasetId,
|
||||
DocumentId=DocumentId,
|
||||
SegmentId=SegmentId,
|
||||
Body=Body,
|
||||
)
|
||||
return Result.success(data=result)
|
||||
|
||||
@self.router.delete("/datasets/{DatasetId}/documents/{DocumentId}/segments/{SegmentId}", response_model=Result[RagOperationResultVO])
|
||||
async def DeleteDatasetDocumentSegment(
|
||||
DatasetId: int,
|
||||
DocumentId: int,
|
||||
SegmentId: str,
|
||||
payload: dict[str, Any] = Depends(verify_access_token),
|
||||
):
|
||||
if not await self._check_permission(int(payload["user_id"]), [self._PERMISSIONS["dataset_read"]]):
|
||||
return JSONResponse(status_code=403, content={"code": 403, "msg": "当前用户没有删除知识库分段权限", "data": None})
|
||||
result = await self.RagDatasetService.DeleteDatasetDocumentSegment(
|
||||
CurrentUserId=int(payload["user_id"]),
|
||||
UserArea=payload.get("area"),
|
||||
UserRole=payload.get("user_role"),
|
||||
DatasetId=DatasetId,
|
||||
DocumentId=DocumentId,
|
||||
SegmentId=SegmentId,
|
||||
)
|
||||
return Result.success(data=result)
|
||||
|
||||
@self.router.get("/chat/parameters", response_model=Result[RagAppParametersVO])
|
||||
async def GetAppParameters(
|
||||
appId: int | None = Query(None, description="聊天应用ID"),
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class RagDatasetUpdateDTO(BaseModel):
|
||||
name: str | None = Field(None, min_length=1, max_length=255)
|
||||
retrieval_model: dict | None = Field(None)
|
||||
@@ -11,8 +11,114 @@ class RagDatasetItemVO(BaseModel):
|
||||
documentCount: int = Field(0)
|
||||
totalChunks: int = Field(0)
|
||||
status: int = Field(1)
|
||||
sortOrder: int = Field(0)
|
||||
createdAt: int = Field(0)
|
||||
updatedAt: int = Field(0)
|
||||
appId: int | None = Field(default=None)
|
||||
appName: str = Field("")
|
||||
appIsDefault: bool = Field(False)
|
||||
|
||||
|
||||
class RagDatasetPageVO(BaseModel):
|
||||
data: list[RagDatasetItemVO] = Field(default_factory=list)
|
||||
total: int = Field(0)
|
||||
|
||||
|
||||
class RagDatasetDetailVO(BaseModel):
|
||||
id: int = Field(...)
|
||||
name: str = Field(...)
|
||||
description: str = Field("")
|
||||
area: str = Field("")
|
||||
isPublic: bool = Field(False)
|
||||
isDefault: bool = Field(False)
|
||||
status: int = Field(1)
|
||||
documentCount: int = Field(0)
|
||||
totalChunks: int = Field(0)
|
||||
chunkMaxSize: int = Field(800)
|
||||
chunkMinSize: int = Field(20)
|
||||
sortOrder: int = Field(0)
|
||||
retrievalModel: dict = Field(default_factory=dict)
|
||||
createdAt: int = Field(0)
|
||||
updatedAt: int = Field(0)
|
||||
appId: int | None = Field(default=None)
|
||||
appName: str = Field("")
|
||||
appIsDefault: bool = Field(False)
|
||||
|
||||
|
||||
class RagDatasetDocumentItemVO(BaseModel):
|
||||
id: int = Field(...)
|
||||
datasetId: int = Field(...)
|
||||
name: str = Field(...)
|
||||
fileType: str = Field("")
|
||||
fileSize: int = Field(0)
|
||||
chunkCount: int = Field(0)
|
||||
indexingStatus: str = Field("waiting")
|
||||
error: str = Field("")
|
||||
enabled: bool = Field(True)
|
||||
hitCount: int = Field(0)
|
||||
createdBy: int | None = Field(None)
|
||||
createdAt: int = Field(0)
|
||||
updatedAt: int = Field(0)
|
||||
|
||||
|
||||
class RagDatasetDocumentPageVO(BaseModel):
|
||||
data: list[RagDatasetDocumentItemVO] = Field(default_factory=list)
|
||||
total: int = Field(0)
|
||||
page: int = Field(1)
|
||||
limit: int = Field(20)
|
||||
hasMore: bool = Field(False)
|
||||
|
||||
|
||||
class RagDatasetUploadDocumentVO(BaseModel):
|
||||
document: dict = Field(default_factory=dict)
|
||||
batch: str = Field("")
|
||||
|
||||
|
||||
class RagDatasetSegmentItemVO(BaseModel):
|
||||
id: str = Field(...)
|
||||
position: int = Field(0)
|
||||
documentId: str = Field("")
|
||||
content: str = Field("")
|
||||
wordCount: int = Field(0)
|
||||
hitCount: int = Field(0)
|
||||
enabled: bool = Field(True)
|
||||
status: str = Field("completed")
|
||||
createdAt: int = Field(0)
|
||||
|
||||
|
||||
class RagDatasetSegmentPageVO(BaseModel):
|
||||
data: list[RagDatasetSegmentItemVO] = Field(default_factory=list)
|
||||
total: int = Field(0)
|
||||
limit: int = Field(20)
|
||||
hasMore: bool = Field(False)
|
||||
|
||||
|
||||
class RagDatasetRetrieveDocumentVO(BaseModel):
|
||||
id: str = Field("")
|
||||
dataSourceType: str = Field("upload_file")
|
||||
name: str = Field("")
|
||||
docType: str | None = Field(default=None)
|
||||
|
||||
|
||||
class RagDatasetRetrieveSegmentVO(BaseModel):
|
||||
id: str = Field(...)
|
||||
position: int = Field(0)
|
||||
documentId: str = Field("")
|
||||
content: str = Field("")
|
||||
answer: str = Field("")
|
||||
wordCount: int = Field(0)
|
||||
hitCount: int = Field(0)
|
||||
enabled: bool = Field(True)
|
||||
status: str = Field("completed")
|
||||
createdAt: int = Field(0)
|
||||
document: RagDatasetRetrieveDocumentVO | None = Field(default=None)
|
||||
|
||||
|
||||
class RagDatasetRetrieveRecordVO(BaseModel):
|
||||
segment: RagDatasetRetrieveSegmentVO = Field(...)
|
||||
score: float = Field(0.0)
|
||||
|
||||
|
||||
class RagDatasetRetrieveResponseVO(BaseModel):
|
||||
query: dict = Field(default_factory=dict)
|
||||
records: list[RagDatasetRetrieveRecordVO] = Field(default_factory=list)
|
||||
|
||||
@@ -80,7 +80,7 @@ class RagChatServiceImpl(IRagChatService):
|
||||
|
||||
context_chunks, dataset_name = await self._retrieve_context(app.get("dataset_id"), Query)
|
||||
collected_answer = ""
|
||||
held_message_end: bytes | None = None
|
||||
held_message_end: dict | None = None
|
||||
|
||||
async for chunk in generate_stream(
|
||||
query=Query,
|
||||
@@ -94,17 +94,21 @@ class RagChatServiceImpl(IRagChatService):
|
||||
dataset_name=dataset_name,
|
||||
):
|
||||
chunk_bytes = chunk.encode("utf-8")
|
||||
for line in chunk.strip().split("\n"):
|
||||
if not line.startswith("data: "):
|
||||
continue
|
||||
data = json.loads(line[6:])
|
||||
if data.get("event") == "message":
|
||||
collected_answer += data.get("answer", "")
|
||||
elif data.get("event") == "message_end":
|
||||
held_message_end = chunk_bytes
|
||||
continue
|
||||
if held_message_end is None:
|
||||
data = self._parse_sse_event(chunk)
|
||||
if not data:
|
||||
yield chunk_bytes
|
||||
continue
|
||||
|
||||
if data.get("event") == "message":
|
||||
collected_answer += data.get("answer", "")
|
||||
yield chunk_bytes
|
||||
continue
|
||||
|
||||
if data.get("event") == "message_end":
|
||||
held_message_end = data
|
||||
continue
|
||||
|
||||
yield chunk_bytes
|
||||
|
||||
followups: list[str] = []
|
||||
try:
|
||||
@@ -114,15 +118,10 @@ class RagChatServiceImpl(IRagChatService):
|
||||
|
||||
if held_message_end:
|
||||
try:
|
||||
for line in held_message_end.decode("utf-8").strip().split("\n"):
|
||||
if not line.startswith("data: "):
|
||||
continue
|
||||
end_data = json.loads(line[6:])
|
||||
if end_data.get("event") == "message_end":
|
||||
end_data.setdefault("metadata", {})["suggested_questions"] = followups
|
||||
yield f"data: {json.dumps(end_data, ensure_ascii=False)}\\n\\n".encode("utf-8")
|
||||
held_message_end.setdefault("metadata", {})["suggested_questions"] = followups
|
||||
yield f"data: {json.dumps(held_message_end, ensure_ascii=False)}\n\n".encode("utf-8")
|
||||
except Exception:
|
||||
yield held_message_end
|
||||
yield f"data: {json.dumps(held_message_end, ensure_ascii=False)}\n\n".encode("utf-8")
|
||||
|
||||
async with GetAsyncSession() as session:
|
||||
async with session.begin():
|
||||
@@ -158,7 +157,7 @@ class RagChatServiceImpl(IRagChatService):
|
||||
FROM rag_conversation
|
||||
WHERE user_id = :user_id
|
||||
AND deleted_at IS NULL
|
||||
AND (:app_id IS NULL OR app_id = :app_id)
|
||||
AND (CAST(:app_id AS BIGINT) IS NULL OR app_id = CAST(:app_id AS BIGINT))
|
||||
ORDER BY updated_at DESC
|
||||
OFFSET :offset LIMIT :limit
|
||||
"""
|
||||
@@ -277,11 +276,10 @@ class RagChatServiceImpl(IRagChatService):
|
||||
raise LeauditException(StatusCodeEnum.HTTP_404_NOT_FOUND, "消息不存在")
|
||||
if int(owner) != CurrentUserId:
|
||||
raise LeauditException(StatusCodeEnum.HTTP_403_FORBIDDEN, "当前用户无权修改该消息反馈")
|
||||
async with session.begin():
|
||||
await session.execute(
|
||||
text("UPDATE rag_message SET feedback = :feedback WHERE message_id = :message_id"),
|
||||
{"feedback": Body.rating, "message_id": MessageId},
|
||||
)
|
||||
await session.execute(
|
||||
text("UPDATE rag_message SET feedback = :feedback WHERE message_id = :message_id"),
|
||||
{"feedback": Body.rating, "message_id": MessageId},
|
||||
)
|
||||
return RagOperationResultVO(result="success")
|
||||
|
||||
async def GetAppParameters(
|
||||
@@ -587,3 +585,26 @@ class RagChatServiceImpl(IRagChatService):
|
||||
)
|
||||
|
||||
return visible_chunks
|
||||
|
||||
def _parse_sse_event(self, chunk: str) -> dict | None:
|
||||
data_lines: list[str] = []
|
||||
for line in chunk.splitlines():
|
||||
if line.startswith("data: "):
|
||||
data_lines.append(line[6:])
|
||||
elif line.startswith("data:"):
|
||||
data_lines.append(line[5:].lstrip())
|
||||
|
||||
if not data_lines:
|
||||
return None
|
||||
|
||||
payload = "\n".join(part for part in data_lines if part.strip()).strip()
|
||||
if not payload or payload == "[DONE]":
|
||||
return None
|
||||
payload = payload.removesuffix("\\n\\n").removesuffix("\\n").strip()
|
||||
|
||||
try:
|
||||
data = json.loads(payload)
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
|
||||
return data if isinstance(data, dict) else None
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -2,9 +2,203 @@ from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from fastapi_modules.fastapi_leaudit.domian.vo.ragDatasetVo import RagDatasetPageVO
|
||||
from fastapi_modules.fastapi_leaudit.domian.Dto.ragDatasetDto import RagDatasetUpdateDTO
|
||||
from fastapi_modules.fastapi_leaudit.domian.vo.ragDatasetVo import (
|
||||
RagDatasetDetailVO,
|
||||
RagDatasetDocumentItemVO,
|
||||
RagDatasetDocumentPageVO,
|
||||
RagDatasetPageVO,
|
||||
RagDatasetRetrieveResponseVO,
|
||||
RagDatasetSegmentPageVO,
|
||||
RagDatasetUploadDocumentVO,
|
||||
)
|
||||
from fastapi_modules.fastapi_leaudit.domian.vo.ragChatVo import RagOperationResultVO
|
||||
|
||||
|
||||
class IRagDatasetService(ABC):
|
||||
@abstractmethod
|
||||
async def GetAdminDatasets(
|
||||
self,
|
||||
CurrentUserId: int,
|
||||
UserArea: str | None,
|
||||
UserRole: str | None,
|
||||
Area: str | None,
|
||||
OnlyEnabled: bool | None,
|
||||
Page: int,
|
||||
PageSize: int,
|
||||
) -> RagDatasetPageVO: ...
|
||||
|
||||
@abstractmethod
|
||||
async def CreateAdminDataset(
|
||||
self,
|
||||
CurrentUserId: int,
|
||||
UserArea: str | None,
|
||||
UserRole: str | None,
|
||||
Body: dict,
|
||||
) -> RagDatasetDetailVO: ...
|
||||
|
||||
@abstractmethod
|
||||
async def UpdateAdminDataset(
|
||||
self,
|
||||
CurrentUserId: int,
|
||||
UserArea: str | None,
|
||||
UserRole: str | None,
|
||||
DatasetId: int,
|
||||
Body: dict,
|
||||
) -> RagDatasetDetailVO | None: ...
|
||||
|
||||
@abstractmethod
|
||||
async def DeleteAdminDataset(
|
||||
self,
|
||||
CurrentUserId: int,
|
||||
UserArea: str | None,
|
||||
UserRole: str | None,
|
||||
DatasetId: int,
|
||||
) -> RagOperationResultVO: ...
|
||||
|
||||
@abstractmethod
|
||||
async def GetMyDatasets(self, CurrentUserId: int, UserArea: str | None, UserRole: str | None) -> RagDatasetPageVO: ...
|
||||
|
||||
@abstractmethod
|
||||
async def GetDatasetDetail(self, CurrentUserId: int, UserArea: str | None, UserRole: str | None, DatasetId: int) -> RagDatasetDetailVO | None: ...
|
||||
|
||||
@abstractmethod
|
||||
async def UpdateDataset(self, CurrentUserId: int, UserArea: str | None, UserRole: str | None, DatasetId: int, Body: RagDatasetUpdateDTO) -> RagDatasetDetailVO | None: ...
|
||||
|
||||
@abstractmethod
|
||||
async def GetDatasetDocuments(
|
||||
self,
|
||||
CurrentUserId: int,
|
||||
UserArea: str | None,
|
||||
UserRole: str | None,
|
||||
DatasetId: int,
|
||||
Page: int,
|
||||
Limit: int,
|
||||
Keyword: str | None,
|
||||
) -> RagDatasetDocumentPageVO: ...
|
||||
|
||||
@abstractmethod
|
||||
async def GetDatasetDocumentDetail(
|
||||
self,
|
||||
CurrentUserId: int,
|
||||
UserArea: str | None,
|
||||
UserRole: str | None,
|
||||
DatasetId: int,
|
||||
DocumentId: int,
|
||||
) -> RagDatasetDocumentItemVO | None: ...
|
||||
|
||||
@abstractmethod
|
||||
async def UploadDatasetDocument(
|
||||
self,
|
||||
CurrentUserId: int,
|
||||
UserArea: str | None,
|
||||
UserRole: str | None,
|
||||
DatasetId: int,
|
||||
FileName: str,
|
||||
ContentType: str | None,
|
||||
Content: bytes,
|
||||
ProcessConfig: dict | None,
|
||||
) -> RagDatasetUploadDocumentVO: ...
|
||||
|
||||
@abstractmethod
|
||||
async def GetDatasetDocumentSegments(
|
||||
self,
|
||||
CurrentUserId: int,
|
||||
UserArea: str | None,
|
||||
UserRole: str | None,
|
||||
DatasetId: int,
|
||||
DocumentId: int,
|
||||
Page: int,
|
||||
Limit: int,
|
||||
Keyword: str | None,
|
||||
) -> RagDatasetSegmentPageVO: ...
|
||||
|
||||
@abstractmethod
|
||||
async def DeleteDatasetDocument(
|
||||
self,
|
||||
CurrentUserId: int,
|
||||
UserArea: str | None,
|
||||
UserRole: str | None,
|
||||
DatasetId: int,
|
||||
DocumentId: int,
|
||||
) -> RagOperationResultVO: ...
|
||||
|
||||
@abstractmethod
|
||||
async def RetrieveDataset(
|
||||
self,
|
||||
CurrentUserId: int,
|
||||
UserArea: str | None,
|
||||
UserRole: str | None,
|
||||
DatasetId: int,
|
||||
Query: str,
|
||||
RetrievalModel: dict | None,
|
||||
) -> RagDatasetRetrieveResponseVO: ...
|
||||
|
||||
@abstractmethod
|
||||
async def GetDatasetDocumentIndexingStatus(
|
||||
self,
|
||||
CurrentUserId: int,
|
||||
UserArea: str | None,
|
||||
UserRole: str | None,
|
||||
DatasetId: int,
|
||||
DocumentId: int,
|
||||
) -> dict: ...
|
||||
|
||||
@abstractmethod
|
||||
async def UpdateDatasetDocumentByFile(
|
||||
self,
|
||||
CurrentUserId: int,
|
||||
UserArea: str | None,
|
||||
UserRole: str | None,
|
||||
DatasetId: int,
|
||||
DocumentId: int,
|
||||
FileName: str,
|
||||
ContentType: str | None,
|
||||
Content: bytes,
|
||||
ProcessConfig: dict | None,
|
||||
) -> RagDatasetUploadDocumentVO: ...
|
||||
|
||||
@abstractmethod
|
||||
async def BatchUpdateDatasetDocumentStatus(
|
||||
self,
|
||||
CurrentUserId: int,
|
||||
UserArea: str | None,
|
||||
UserRole: str | None,
|
||||
DatasetId: int,
|
||||
DocumentIds: list[int],
|
||||
Enabled: bool,
|
||||
) -> RagOperationResultVO: ...
|
||||
|
||||
@abstractmethod
|
||||
async def GetDatasetDocumentSegmentDetail(
|
||||
self,
|
||||
CurrentUserId: int,
|
||||
UserArea: str | None,
|
||||
UserRole: str | None,
|
||||
DatasetId: int,
|
||||
DocumentId: int,
|
||||
SegmentId: str,
|
||||
) -> RagDatasetSegmentItemVO | None: ...
|
||||
|
||||
@abstractmethod
|
||||
async def UpdateDatasetDocumentSegment(
|
||||
self,
|
||||
CurrentUserId: int,
|
||||
UserArea: str | None,
|
||||
UserRole: str | None,
|
||||
DatasetId: int,
|
||||
DocumentId: int,
|
||||
SegmentId: str,
|
||||
Body: dict,
|
||||
) -> RagDatasetSegmentItemVO | None: ...
|
||||
|
||||
@abstractmethod
|
||||
async def DeleteDatasetDocumentSegment(
|
||||
self,
|
||||
CurrentUserId: int,
|
||||
UserArea: str | None,
|
||||
UserRole: str | None,
|
||||
DatasetId: int,
|
||||
DocumentId: int,
|
||||
SegmentId: str,
|
||||
) -> RagOperationResultVO: ...
|
||||
|
||||
Reference in New Issue
Block a user