AI 大模型应用进阶系列（三）：大模型流式输出

import requests
import json
import sys

def stream_chat(api_key, api_url, model, prompt, history=None):
    """
    流式调用大模型API并实时输出结果
    :param api_key: API密钥
    :param api_url: API地址
    :param model: 模型名称
    :param prompt: 用户当前输入
    :param history: 历史对话列表
    """
    # 初始化历史对话
    history = history or []
    history.append({"role": "user", "content": prompt})
    
    # 构建请求参数
    payload = {
        "model": model,
        "messages": history,
        "stream": True,  # 启用流式输出
        "temperature": 0.7,
        "max_tokens": 1024
    }
    
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    
    try:
        # 发送流式请求
        with requests.post(
            api_url,
            headers=headers,
            data=json.dumps(payload),
            stream=True  # 保持连接并流式接收响应
        ) as response:
            response.raise_for_status()  # 检查HTTP错误状态
            
            print("\nAI: ", end="", flush=True)
            full_response = ""
            # 逐行处理流式响应
            for line in response.iter_lines():
                if line:
                    # 处理SSE格式（移除前缀"data: "）
                    line = line.decode('utf-8').lstrip('data: ').strip()
                    if line == "[DONE]":  # 结束标志
                        break
                    
                    try:
                        chunk = json.loads(line)
                        # 提取内容片段
                        content = chunk.get("choices", [{}])[0].get("delta", {}).get("content", "")
                        if content:
                            print(content, end="", flush=True)  # 实时输出
                            full_response += content
                            sys.stdout.flush()
                    except json.JSONDecodeError:
                        continue
            
            print("\n")  # 结束后换行
            # 更新历史对话
            history.append({"role": "assistant", "content": full_response})
            return history
            
    except requests.exceptions.RequestException as e:
        print(f"\n请求发生错误: {str(e)}")
        return history

def main():
    # ==================== 请替换为你的API信息 ====================
    API_KEY = "your_api_key_here"
    API_URL = "https://api.openai.com/v1/chat/completions"  # 示例地址，根据实际模型修改
    MODEL = "gpt-3.5-turbo"  # 模型名称，根据实际模型修改
    # ==========================================================
    
    print("大模型流式输出演示（输入'exit'退出）")
    history = []
    
    while True:
        user_input = input("\n你: ")
        if user_input.lower() == "exit":
            print("再见！")
            break
        # 调用流式接口
        history = stream_chat(API_KEY, API_URL, MODEL, user_input, history)

if __name__ == "__main__":
    main()