import json import time import asyncio import logging from typing import AsyncGenerator, List, Dict, Any from ..sora_integration import SoraClient from ..config import Config from ..utils import localize_image_urls from .image_service import format_think_block logger = logging.getLogger("sora-api.streaming") async def generate_streaming_response( sora_client: SoraClient, prompt: str, n_images: int = 1 ) -> AsyncGenerator[str, None]: """ 文本到图像的流式响应生成器 Args: sora_client: Sora客户端 prompt: 提示词 n_images: 生成图像数量 Yields: SSE格式的响应数据 """ request_id = f"chatcmpl-stream-{time.time()}-{hash(prompt) % 10000}" # 发送开始事件 yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {'role': 'assistant'}, 'finish_reason': None}]})}\n\n" # 发送处理中的消息（放在代码块中） start_msg = "```think\n正在生成图像，请稍候...\n" yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {'content': start_msg}, 'finish_reason': None}]})}\n\n" # 创建一个后台任务来生成图像 logger.info(f"[流式响应 {request_id}] 开始生成图像, 提示词: {prompt}") generation_task = asyncio.create_task(sora_client.generate_image( prompt=prompt, num_images=n_images, width=720, height=480 )) # 每5秒发送一条"仍在生成中"的消息，防止连接超时 progress_messages = [ "正在处理您的请求...", "仍在生成图像中，请继续等待...", "Sora正在创作您的图像作品...", "图像生成需要一点时间，感谢您的耐心等待...", "我们正在努力为您创作高质量图像..." ] i = 0 while not generation_task.done(): # 每5秒发送一次进度消息 await asyncio.sleep(5) progress_msg = progress_messages[i % len(progress_messages)] i += 1 content = "\n" + progress_msg + "\n" yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {'content': content}, 'finish_reason': None}]})}\n\n" try: # 获取生成结果 image_urls = await generation_task logger.info(f"[流式响应 {request_id}] 图像生成完成，获取到 {len(image_urls) if isinstance(image_urls, list) else '非列表'} 个URL") # 本地化图片URL if Config.IMAGE_LOCALIZATION and isinstance(image_urls, list) and image_urls: logger.info(f"[流式响应 {request_id}] 准备进行图片本地化处理") try: localized_urls = await localize_image_urls(image_urls) image_urls = localized_urls logger.info(f"[流式响应 {request_id}] 图片本地化处理完成") except Exception as e: logger.error(f"[流式响应 {request_id}] 图片本地化过程中发生错误: {str(e)}", exc_info=True) logger.info(f"[流式响应 {request_id}] 由于错误，将使用原始URL") elif not Config.IMAGE_LOCALIZATION: logger.info(f"[流式响应 {request_id}] 图片本地化功能未启用") # 结束代码块 content_str = "\n```\n\n" yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {'content': content_str}, 'finish_reason': None}]})}\n\n" # 添加生成的图片URLs for i, url in enumerate(image_urls): if i > 0: content_str = "\n\n" yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {'content': content_str}, 'finish_reason': None}]})}\n\n" image_markdown = f"![Generated Image]({url})" yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {'content': image_markdown}, 'finish_reason': None}]})}\n\n" # 发送完成事件 yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]})}\n\n" # 发送结束标志 yield "data: [DONE]\n\n" except Exception as e: error_msg = f"图像生成失败: {str(e)}" logger.error(f"[流式响应 {request_id}] 错误: {error_msg}", exc_info=True) error_content = f"\n{error_msg}\n```" yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {'content': error_content}, 'finish_reason': 'error'}]})}\n\n" yield "data: [DONE]\n\n" async def generate_streaming_remix_response( sora_client: SoraClient, prompt: str, image_data: str, n_images: int = 1 ) -> AsyncGenerator[str, None]: """ 图像到图像的流式响应生成器 Args: sora_client: Sora客户端 prompt: 提示词 image_data: Base64编码的图像数据 n_images: 生成图像数量 Yields: SSE格式的响应数据 """ import os import tempfile import uuid import base64 request_id = f"chatcmpl-stream-remix-{time.time()}-{hash(prompt) % 10000}" # 发送开始事件 yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {'role': 'assistant'}, 'finish_reason': None}]})}\n\n" try: # 保存base64图片到临时文件 temp_dir = tempfile.mkdtemp() temp_image_path = os.path.join(temp_dir, f"upload_{uuid.uuid4()}.png") try: # 解码并保存图片 with open(temp_image_path, "wb") as f: f.write(base64.b64decode(image_data)) # 上传图片 upload_msg = "```think\n上传图片中...\n" yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {'content': upload_msg}, 'finish_reason': None}]})}\n\n" logger.info(f"[流式响应Remix {request_id}] 上传图片中") upload_result = await sora_client.upload_image(temp_image_path) media_id = upload_result['id'] # 发送生成中消息 generate_msg = "\n基于图片生成新图像中...\n" yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {'content': generate_msg}, 'finish_reason': None}]})}\n\n" # 创建后台任务生成图像 logger.info(f"[流式响应Remix {request_id}] 开始生成图像，提示词: {prompt}") generation_task = asyncio.create_task(sora_client.generate_image_remix( prompt=prompt, media_id=media_id, num_images=n_images )) # 每5秒发送一条"仍在生成中"的消息 progress_messages = [ "正在处理您的请求...", "仍在生成图像中，请继续等待...", "Sora正在基于您的图片创作新作品...", "图像生成需要一点时间，感谢您的耐心等待...", "正在融合您的风格和提示词，打造专属图像..." ] i = 0 while not generation_task.done(): await asyncio.sleep(5) progress_msg = progress_messages[i % len(progress_messages)] i += 1 content = "\n" + progress_msg + "\n" yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {'content': content}, 'finish_reason': None}]})}\n\n" # 获取生成结果 image_urls = await generation_task logger.info(f"[流式响应Remix {request_id}] 图像生成完成") # 本地化图片URL if Config.IMAGE_LOCALIZATION: logger.info(f"[流式响应Remix {request_id}] 进行图片本地化处理") localized_urls = await localize_image_urls(image_urls) image_urls = localized_urls logger.info(f"[流式响应Remix {request_id}] 图片本地化处理完成") else: logger.info(f"[流式响应Remix {request_id}] 图片本地化功能未启用") # 结束代码块 content_str = "\n```\n\n" yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {'content': content_str}, 'finish_reason': None}]})}\n\n" # 发送图片URL作为Markdown for i, url in enumerate(image_urls): if i > 0: newline_str = "\n\n" yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {'content': newline_str}, 'finish_reason': None}]})}\n\n" image_markdown = f"![Generated Image]({url})" yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {'content': image_markdown}, 'finish_reason': None}]})}\n\n" # 发送完成事件 yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]})}\n\n" # 发送结束标志 yield "data: [DONE]\n\n" finally: # 清理临时文件 if os.path.exists(temp_image_path): os.remove(temp_image_path) if os.path.exists(temp_dir): os.rmdir(temp_dir) except Exception as e: error_msg = f"图像Remix失败: {str(e)}" logger.error(f"[流式响应Remix {request_id}] 错误: {error_msg}", exc_info=True) error_content = f"\n{error_msg}\n```" yield f"data: {json.dumps({'id': request_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': 'sora-1.0', 'choices': [{'index': 0, 'delta': {'content': error_content}, 'finish_reason': 'error'}]})}\n\n" # 结束流 yield "data: [DONE]\n\n"