46 lines
1.7 KiB
Python
46 lines
1.7 KiB
Python
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from utils.ragflow_asy_util import AsyncRAGFlowClient
|
|
from module_admin.entity.vo.ragflow_vo import RagflowListQueryModel, ListDocumentsQueryModel
|
|
from config.env import RAGFlowConfig
|
|
|
|
class RAGFlowService:
|
|
"""
|
|
RAGFlow服务
|
|
"""
|
|
|
|
# 获取数据集列表
|
|
@classmethod
|
|
async def get_ragflow_dataset_list_services(cls, query_db: AsyncSession, rage_flow_query: RagflowListQueryModel):
|
|
"""
|
|
获取数据集列表
|
|
"""
|
|
|
|
async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
|
|
result = await client.list_datasets(**(rage_flow_query.model_dump()))
|
|
|
|
# 获取分页数据
|
|
return result.get('data', None)
|
|
|
|
# 获取数据集中文档列表
|
|
@classmethod
|
|
async def list_documents_services(
|
|
cls,
|
|
query_db: AsyncSession,
|
|
dataset_id: str,
|
|
list_documents_query: ListDocumentsQueryModel,
|
|
):
|
|
async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
|
|
result = await client.list_documents(dataset_id=dataset_id, **(list_documents_query.model_dump()))
|
|
|
|
return result.get('data', None)
|
|
|
|
# 上传文档到数据集
|
|
async def upload_file_dataset_services(
|
|
cls,
|
|
dataset_id: str,
|
|
files,
|
|
):
|
|
async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
|
|
result = await client.upload_documents_bytes(dataset_id=dataset_id, file_bytes=files)
|
|
|
|
return result.get('data', None) |