kangda_robot_backend/ruoyi-fastapi-backend/module_admin/service/ragflow_service.py

from sqlalchemy.ext.asyncio import AsyncSession
from utils.ragflow_asy_util import AsyncRAGFlowClient
from module_admin.entity.vo.ragflow_vo import RagflowListQueryModel, ListDocumentsQueryModel, UpdateFileModel, DeleteFileModel, CreateDatasetModel, DocumentIdsModel
from config.env import RAGFlowConfig
from typing import List

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
    # 创建数据集
    @classmethod
    async def create_dataset_services(cls, create_dataset_params: CreateDatasetModel):
        """创建数据集

        Args:
            create_dataset_params (CreateDatasetModel): 创建参数

        Returns:
            _type_: _description_
        """
        async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
            result = await client.create_dataset(
                **(create_dataset_params.model_dump())
            )
            
        return result
        
    # 更新数据集
    @classmethod
    async def update_dataset_services(
        cls,
        dataset_id: str,
        update_dataset_params: CreateDatasetModel,
    ):
        """更新数据集信息

        Args:
            dataset_id (str): 数据集id
            update_dataset_params (CreateDatasetModel): 更新参数

        Returns:
            _type_: _description_
        """
        async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
            result =  await client.update_dataset(
                dataset_id=dataset_id, **(update_dataset_params.model_dump())
            )
            
        return result
        
        
    # 获取数据集中文档列表
    @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
    
    # 上传文档到数据集
    @classmethod
    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
    # 开始解析文档
    @classmethod
    async def parse_documents_services(
        cls,
        dataset_id: str,
        parse_params: DocumentIdsModel,
    ):
        async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
            result = await client.parse_documents(dataset_id=dataset_id, document_ids=parse_params.documnet_ids)
            
        return result
    
    # 停止解析文档
    @classmethod
    async def stop_parse_documents_services(
        cls,
        dataset_id: str,
        parse_params: DocumentIdsModel,
    ):
        async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
            result = await client.stop_parsing_documents(dataset_id=dataset_id, document_ids=parse_params.documnet_ids)
        
        return result
    
    
    # 更新文档内容
    @classmethod
    async def update_file_dataset_services(
        cls,
        dataset_id: str,
        document_id: str,
        update_params: UpdateFileModel,
    ):
        async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
            result = await client.update_document(dataset_id=dataset_id, document_id=document_id, **(update_params.model_dump()))
             
        return result
    
    # 删除文档
    @classmethod
    async def delete_file_services(
        cls,
        dataset_id: str,
        delete_params: DeleteFileModel,
    ):
        async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
            result = await client.delete_documents(dataset_id=dataset_id, **(delete_params.model_dump()))
        return result 
    
    # 删除数据集
    @classmethod
    async def delete_datasets_services(
        cls,
        delete_params: DeleteFileModel,
    ):
        async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
            result = await client.delete_datasets(**(delete_params.model_dump()))
        return result
    
    
    # 查看聊天助手列表
    @classmethod
    async def get_chat_assistant_list_services(
        cls,
        query_params: RagflowListQueryModel,
    ):
        async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
            result = await client.list_chat_assistants(**(query_params.model_dump()))
        return result