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, UpdateChatAssistantModel \
    ,CreateSessionWithChatModel, ConverseWithChatAssistantModel
from config.env import RAGFlowConfig
from typing import List
import asyncio
import string
import re

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
    
    # 修改聊天助手
    @classmethod
    async def update_chat_assistant_services(cls, update_params: UpdateChatAssistantModel):
        async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
            result = await client.update_chat_assistant(**(update_params.model_dump()))
        return result
    
    # 创建助手会话
    @classmethod
    async def create_session_with_chat_services(cls, create_params: CreateSessionWithChatModel):
        async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
            result = await client.create_session_with_chat(**(create_params.model_dump()))
        return result
    
    
    # 与助手聊天
    @classmethod
    async def converse_with_chat_assistant_services(cls, converse_params: ConverseWithChatAssistantModel):
        if converse_params.stream:
            async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
                result = await client.converse_with_chat_assistant(**(converse_params.model_dump()))
                
                i = 0
                async for t in result:
                    try:
                        answer = t["data"].get("answer", "")
                        answer = cls.clean_text(answer)
                        print(repr(answer[i:]))
                        i = len(answer)
                    except Exception as e:
                        print(e)
                        print(t)
                    await asyncio.sleep(0.2)
        else:
            async with AsyncRAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY) as client:
                result = await client.converse_with_chat_assistant(**(converse_params.model_dump()))
                return result
                
    
    
    @classmethod
    def clean_text(cls,text):
        # 只保留 数字、中文、英文、常见标点
        translator = str.maketrans('', '', string.whitespace)
        clean = text.translate(translator)  # 删除所有空白字符（包括 \n, space 等）
        return clean