kangda-robot-backend/ruoyi-fastapi-backend/module_admin/service/ragflow_service.py

from module_admin.entity.vo.ragflow_vo import (
    RagflowListQueryModel, 
    ListDocumentsQueryModel, 
    UpdateFileModel, 
    DeleteFileModel, 
    CreateDatasetModel, 
    DocumentIdsModel, 
    UpdateChatAssistantModel,
    CreateSessionWithChatModel, 
    ConverseWithChatAssistantModel
)
from utils.ragflow_util import RAGFlowClient
from config.env import RAGFlowConfig
from utils.ragflow_client_manager import get_ragflow_client


class RAGFlowService:
    """ 
        RAGFlow服务 - 简化版本，使用同步操作
    """
    
    # 获取数据集列表
    @classmethod
    def get_ragflow_dataset_list_services(cls, query_db, rage_flow_query: RagflowListQueryModel):
        """
        获取数据集列表 - 同步版本
        """
        client = RAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY)
        result = client.list_datasets(**(rage_flow_query.model_dump()))
        return result
        
    # 创建数据集
    @classmethod
    def create_dataset_services(cls, create_dataset_params: CreateDatasetModel):
        """创建数据集 - 同步版本"""
        client = RAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY)
        result = client.create_dataset(**(create_dataset_params.model_dump()))
        return result
        
    # 更新数据集
    @classmethod
    def update_dataset_services(cls, dataset_id: str, update_dataset_params: CreateDatasetModel):
        """更新数据集信息 - 同步版本"""
        client = RAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY)
        result = client.update_dataset(dataset_id=dataset_id, **(update_dataset_params.model_dump()))
        return result
        
    # 获取数据集中文档列表
    @classmethod
    def list_documents_services(cls, query_db, dataset_id: str, list_documents_query: ListDocumentsQueryModel):
        """获取文档列表 - 同步版本"""
        client = RAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY)
        result = client.list_documents(dataset_id=dataset_id, **(list_documents_query.model_dump()))
        return result
    
    # 上传文档到数据集
    @classmethod
    def upload_file_dataset_services(cls, dataset_id: str, files):
        """上传文档 - 同步版本"""
        client = RAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY)
        result = client.upload_documents_bytes(dataset_id=dataset_id, file_bytes=files)
        return result
    
    # 开始解析文档
    @classmethod
    def parse_documents_services(cls, dataset_id: str, parse_params: DocumentIdsModel):
        """解析文档 - 同步版本"""
        client = RAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY)
        result = client.parse_documents(dataset_id=dataset_id, document_ids=parse_params.documnet_ids)
        return result
    
    # 停止解析文档
    @classmethod
    def stop_parse_documents_services(cls, dataset_id: str, parse_params: DocumentIdsModel):
        """停止解析文档 - 同步版本"""
        client = RAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY)
        result = client.stop_parsing_documents(dataset_id=dataset_id, document_ids=parse_params.documnet_ids)
        return result
    
    # 更新文档内容
    @classmethod
    def update_file_dataset_services(cls, dataset_id: str, document_id: str, update_params: UpdateFileModel):
        """更新文档 - 同步版本"""
        client = RAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY)
        result = client.update_document(dataset_id=dataset_id, document_id=document_id, **(update_params.model_dump()))
        return result
    
    # 删除文档
    @classmethod
    def delete_file_services(cls, dataset_id: str, delete_params: DeleteFileModel):
        """删除文档 - 同步版本"""
        client = RAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY)
        result = client.delete_documents(dataset_id=dataset_id, **(delete_params.model_dump()))
        return result 
    
    # 删除数据集
    @classmethod
    def delete_datasets_services(cls, delete_params: DeleteFileModel):
        """删除数据集 - 同步版本"""
        client = RAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY)
        result = client.delete_datasets(**(delete_params.model_dump()))
        return result
    
    # 查看聊天助手列表
    @classmethod
    def get_chat_assistant_list_services(cls, query_params: RagflowListQueryModel):
        """获取聊天助手列表 - 同步版本"""
        client = RAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY)
        result = client.list_chat_assistants(**(query_params.model_dump()))
        return result
    
    # 修改聊天助手
    @classmethod
    def update_chat_assistant_services(cls, update_params: UpdateChatAssistantModel):
        """更新聊天助手 - 同步版本"""
        client = RAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY)
        result = client.update_chat_assistant(**(update_params.model_dump()))
        return result
    
    # 创建助手会话
    @classmethod
    def create_session_with_chat_services(cls, create_params: CreateSessionWithChatModel):
        """创建会话 - 同步版本"""
        client = RAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY)
        result = client.create_session_with_chat(**(create_params.model_dump()))
        return result
    
    # 与助手聊天 - 核心方法
    @classmethod
    def converse_with_chat_assistant_services(cls, converse_params: ConverseWithChatAssistantModel):
        """与聊天助手对话 - 同步版本，返回Generator"""
        client = RAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY)
        # 直接返回Generator，支持流式响应
        return client.converse_with_chat_assistant(
            chat_id=converse_params.chat_id,
            question=converse_params.question,
            stream=converse_params.stream,
            session_id=converse_params.session_id
        )

    # 与助手聊天 (OpenAI Compatible) - 同步版本
    @classmethod
    def converse_with_chat_assistant_services_openai(cls, converse_params: ConverseWithChatAssistantModel):
        """OpenAI兼容格式 - 同步版本"""
        client = RAGFlowClient(RAGFlowConfig.RAGFLOW_BASE_URL, RAGFlowConfig.RAGFLOW_API_KEY)
        # Construct messages list for OpenAI format
        messages = [{"role": "user", "content": converse_params.question}]
        # Uses defaults for model name as per user indication "server will parse this automatically"
        return client.create_chat_completion(
            chat_id=converse_params.chat_id,
            model="ragflow", 
            messages=messages,
            stream=converse_params.stream
        )