rag-gateway/app/service/v2/app_driver/chat_agent.py

import json

from Log import logger
# from Log import logger
from app.models.v2.session_model import ChatData
from app.service.v2.app_driver.chat_base import ChatBase


class ChatAgent(ChatBase):

    async def chat_completions(self, url, data, headers):
        complete_response = ""
        # print(data)
        async for line in self.http_stream(url, data, headers):
            # logger.error(line)
            if line.startswith("data:"):
                complete_response = line.strip("data:").strip()
            elif line.startswith("Error: "):
                yield {"event": "error", "message": line}
            else:
                complete_response += line.strip()
            try:
                json_data = json.loads(complete_response)
                # 处理 JSON 数据
                # print(json_data)
                complete_response = ""
                yield json_data

            except json.JSONDecodeError as e:
                logger.info("Invalid JSON data------------------")
                # print(e)

    @staticmethod
    async def request_data(query: str, conversation_id: str, user:str, chat_data: ChatData) -> dict:
        inputs = []
        files = []
        if hasattr(chat_data,  "inputs"):
            inputs = chat_data.inputs
        if hasattr(chat_data,  "files"):
            files = chat_data.files

        return {
            "inputs":inputs,
            "query": query,
            "response_mode": "streaming",
            "conversation_id": conversation_id,
            "user": user,
            "files": files
        }

    @staticmethod
    async def complex_request_data(query: str, conversation_id: str, user: str, files: list=None, inputs: dict=None) -> dict:
        if not files:
            files = []
        if not inputs:
            inputs = {}
        return {
            "inputs": inputs,
            "query": query,
            "response_mode": "streaming",
            "conversation_id": conversation_id,
            "user": user,
            "files": files
        }


if __name__ == "__main__":
    async def aa():
        chat_id = "16954f6d-c1e6-4a0b-b371-363c28e8a48b"
        token = "app-79ndndjNAFSV3qTuDAjDwuSO"
        base_url = "http://192.168.20.116"
        url = f"{base_url}/v1/chat-messages"
        chat = ChatAgent()
        data = {
            "inputs":{},
            "query": "你好，你能做什么？",
            "response_mode": "streaming",
            "conversation_id": "",
            "user": "1",
            "files": []
        }
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f"Bearer {token}"
        }
        async for ans in chat.chat_completions(url, data, headers):
            print(ans)


    import asyncio

    asyncio.run(aa())