rag-gateway/app/service/difyService.py

import json
from datetime import datetime

import httpx
from typing import Union, Dict, List
from fastapi import HTTPException
from starlette import status

from Log import logger
from app.config.config import settings
from app.utils.rsa_crypto import RagflowCrypto


class DifyService:
    def __init__(self, base_url: str):
        self.base_url = base_url

    def _handle_response(self, response: httpx.Response) -> Union[Dict, List]:
        if response.status_code != 200:
            if response.status_code == 201:
                return response.json()
            return {}

        data = response.json()
        ret_code = data.get("retcode")
        if ret_code == 401:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="登录过期",
            )
        # if ret_code != 0:
        #     return {}

        # 检查返回的数据类型
        if isinstance(data.get("data"), dict):
            return data.get("data", {})
        elif isinstance(data.get("data"), list):
            return data.get("data", [])
        else:
            return data

    async def register(self, username: str, password: str):
        password = RagflowCrypto(settings.PUBLIC_KEY, settings.PRIVATE_KEY).encrypt(password)
        async with httpx.AsyncClient() as client:
            response = await client.post(
                f"{self.base_url}/v1/user/register",
                headers={'Content-Type': 'application/json'},
                json={"nickname": username, "email":  f"{username}@example.com", "password": password}
            )
            if response.status_code != 200:
                raise Exception(f"Ragflow registration failed: {response.text}")
            return self._handle_response(response)

    async def login(self, username: str, password: str) -> str:
        password = RagflowCrypto(settings.PUBLIC_KEY, settings.PRIVATE_KEY).encrypt(password)
        async with httpx.AsyncClient() as client:
            response = await client.post(
                f"{self.base_url}/v1/user/login",
                headers={'Content-Type': 'application/json'},
                json={"email": f"{username}@example.com", "password": password}
            )
            if response.status_code != 200:
                raise Exception(f"Ragflow login failed: {response.text}")
            authorization = response.headers.get('Authorization')
            if not authorization:
                raise Exception("Authorization header not found in response")
            return authorization

    async def chat(self, token: str, user_id: int,  message: str, upload_file_id: str, conversation_id: str):

        target_url = f"{self.base_url}/v1/chat-messages"
        files = []
        if upload_file_id:
            files = [
                {
                    "type": "image",
                    "transfer_method": "local_file",
                    "url": "",
                    "upload_file_id": upload_file_id
                }
            ]
        data = {
            "inputs": {},
            "query": message,
            "response_mode": "streaming",
            "conversation_id": "",
            "user": str(user_id),
            "files": files
        }

        async with httpx.AsyncClient(timeout=300.0) as client:
            headers = {
                'Content-Type': 'application/json',
                'Authorization': f'Bearer {token}'
            }
            async with client.stream("POST", target_url, data=json.dumps(data), headers=headers) as response:
                if response.status_code == 200:
                    try:
                        async for answer in response.aiter_text():
                            # print(f"response of ragflow chat: {answer}")
                            yield answer
                    except GeneratorExit as e:
                        print(e)
                        return
                else:
                    yield f"Error: {response.status_code}"



    async def get_session_history(self, token: str, conversation_id: str, user: str):
        url = f"{self.base_url}/v1/messages"
        params = {
            'user': user,
            'conversation_id': conversation_id
        }
        headers = {"Authorization": f'Bearer {token}'}
        async with httpx.AsyncClient() as client:
            response = await client.get(url, params=params, headers=headers)
            # print(response.text)
            # print(response.status_code)
            # print(response.res)
            data = self._handle_response(response)
            # print("----------------data----------------------:", data)

            return data

    async def upload(self, token: str, filename: str, file: bytes, user_id) -> dict:
        url = f"{self.base_url}/v1/files/upload"
        headers = {
            # 'Content-Type': 'application/json',
            'Authorization': f'Bearer {token}'
        }
        data = {
            'user': str(user_id)
        }
        # 创建表单数据，包含文件
        files = {"file": (filename, file)}
        async with httpx.AsyncClient() as client:
            response = await client.post(url, headers=headers, files=files, data=data)
            data = self._handle_response(response)

            return data




if __name__ == "__main__":
    async def a():
        a = DifyService("http://192.168.20.119:11080")
        b = await a.get_knowledge_list("ImY3ZTZlZWQwYTY2NTExZWY5ZmFiMDI0MmFjMTMwMDA2Ig.Zzxwmw.uI_HAWzOkipQuga1aeQtoeIc3IM", 1,
                                 10)
        print(b)

    import asyncio

    asyncio.run(a())