app.py

import os
import json
import gradio as gr
import requests
from sseclient import SSEClient
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# API配置
API_URL = os.getenv("API_URL", "https://your-api-endpoint/v1/chat/completions")
API_KEY = os.getenv("API_KEY", "your-api-key")

def format_message(role, content):
    return {"role": role, "content": content}

class ChatBot:
    def __init__(self):
        self.headers = {
            "Authorization": f"Bearer {API_KEY}",
            "Content-Type": "application/json"
        }
    
    def generate_stream(self, messages):
        data = {
            "model": "gpt-3.5-turbo",  # 或其他模型
            "messages": messages,
            "stream": True,
            "temperature": 0.7
        }
        
        response = requests.post(
            API_URL,
            headers=self.headers,
            json=data,
            stream=True
        )
        
        client = SSEClient(response)
        return client

def chat_stream(message, history):
    chatbot = ChatBot()
    
    # 构建消息历史
    messages = []
    for human, assistant in history:
        messages.append(format_message("user", human))
        messages.append(format_message("assistant", assistant))
    messages.append(format_message("user", message))
    
    # 流式响应
    response_stream = chatbot.generate_stream(messages)
    partial_message = ""
    
    for event in response_stream:
        if event.data != "[DONE]":
            try:
                chunk = json.loads(event.data)
                if chunk and "choices" in chunk:
                    delta = chunk["choices"][0]["delta"]
                    if "content" in delta:
                        partial_message += delta["content"]
                        yield partial_message
            except json.JSONDecodeError:
                continue
    
    return partial_message

# Gradio界面配置
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    chatbot = gr.Chatbot(
        height=600,
        show_copy_button=True,
        avatar_images=["assets/user.png", "assets/assistant.png"],
    )
    msg = gr.Textbox(
        placeholder="在这里输入您的问题...",
        container=False,
        scale=7,
    )
    with gr.Row():
        submit = gr.Button("发送", scale=2, variant="primary")
        clear = gr.Button("清除对话", scale=1)
    
    # 事件处理
    msg.submit(
        chat_stream,
        [msg, chatbot],
        [chatbot],
        api_name="chat"
    ).then(
        lambda: "",
        None,
        [msg],
        api_name="clear_input"
    )
    
    submit.click(
        chat_stream,
        [msg, chatbot],
        [chatbot],
    ).then(
        lambda: "",
        None,
        [msg],
    )
    
    clear.click(lambda: None, None, chatbot)

# 启动应用
if __name__ == "__main__":
    demo.queue().launch()