import os import uuid import httpx from fastapi import APIRouter, HTTPException, Depends from typing import Dict, Any import logging from config import settings from api.ai_schemas import AIGenerationRequest, AITaskResponse, AITaskResult, AITemplateGenerationRequest from api.prompts import get_prompt router = APIRouter() logger = logging.getLogger(__name__) DASHSCOPE_API_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation" DASHSCOPE_TASK_URL = "https://dashscope.aliyuncs.com/api/v1/tasks" async def _download_image(url: str) -> str: """ 下载图片并保存到media目录 :param url: 图片URL :return: 本地文件相对路径 (e.g., "/media/xxx.png") """ try: async with httpx.AsyncClient() as client: response = await client.get(url, timeout=30.0) if response.status_code != 200: logger.error(f"Failed to download image: {url}, status: {response.status_code}") return url # 下载失败返回原URL # 生成文件名 filename = f"{uuid.uuid4()}.png" filepath = os.path.join(settings.media_dir, filename) # 写入文件 with open(filepath, "wb") as f: f.write(response.content) # 返回可访问的URL路径 return f"/media/{filename}" except Exception as e: logger.error(f"Error downloading image: {str(e)}") return url async def _submit_dashscope_task(prompt: str, negative_prompt: str, size: str, n: int, model: str): if not settings.dashscope_api_key: raise HTTPException(status_code=500, detail="DashScope API Key not configured") headers = { "Content-Type": "application/json", "Authorization": f"Bearer {settings.dashscope_api_key}" # "X-DashScope-Async": "enable" # 移除强制异步，因为某些账号/模型不支持 } # 构建请求体 payload = { "model": model, "input": { "messages": [ { "role": "user", "content": [ { "text": prompt } ] } ] }, "parameters": { "prompt_extend": True, "watermark": False, "n": n, "size": size } } if negative_prompt: payload["parameters"]["negative_prompt"] = negative_prompt try: async with httpx.AsyncClient() as client: response = await client.post( DASHSCOPE_API_URL, headers=headers, json=payload, timeout=60.0 # 增加超时时间，同步请求可能较慢 ) if response.status_code != 200: logger.error(f"DashScope API error: {response.text}") try: error_detail = response.json() except: error_detail = {"message": response.text} raise HTTPException(status_code=response.status_code, detail=error_detail) result = response.json() # 检查是异步任务返回还是同步结果返回 task_id = None if "output" in result and "task_id" in result["output"]: task_id = result["output"]["task_id"] elif "task_id" in result: task_id = result["task_id"] # 如果没有task_id，检查是否有直接结果（同步模式） if not task_id: if "output" in result and "choices" in result["output"]: # 同步返回成功 # 提取结果以便直接返回 choices = result["output"]["choices"] results = [] for choice in choices: msg_content = choice.get("message", {}).get("content", []) for item in msg_content: if "image" in item: # 下载图片到本地 local_url = await _download_image(item["image"]) results.append({"url": local_url, "origin_url": item["image"]}) return AITaskResponse( request_id=result.get("request_id"), status="SUCCEEDED", results=results ) # 既没有task_id也没有choices，可能是其他结构或错误 logger.warning(f"Unexpected response structure: {result}") # 尝试从output.task_id再找一次 task_id = result.get("output", {}).get("task_id") if not task_id: # 依然找不到，报错 raise HTTPException(status_code=500, detail="Failed to retrieve task_id or results from DashScope response") return AITaskResponse( task_id=task_id, request_id=result.get("request_id"), status="PENDING" # 异步任务初始状态 ) except httpx.RequestError as e: logger.error(f"Request error: {str(e)}") raise HTTPException(status_code=500, detail=f"Request failed: {str(e)}") @router.post("/generate", response_model=AITaskResponse, summary="提交AI图片生成任务") async def generate_image(request: AIGenerationRequest): """ 提交AI图片生成任务，使用阿里云DashScope服务 """ return await _submit_dashscope_task( prompt=request.prompt, negative_prompt=request.negative_prompt, size=request.size, n=request.n, model=request.model ) @router.post("/generate_from_template", response_model=AITaskResponse, summary="使用模板提交AI图片生成任务") async def generate_image_from_template(request: AITemplateGenerationRequest): """ 使用预设模板提交AI图片生成任务 """ try: prompt = get_prompt(request.template_id, **request.params) if prompt is None: raise HTTPException(status_code=404, detail=f"未找到模板ID: {request.template_id}") except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) return await _submit_dashscope_task( prompt=prompt, negative_prompt=request.negative_prompt, size=request.size, n=request.n, model=request.model ) @router.get("/tasks/{task_id}", response_model=AITaskResult, summary="查询AI任务结果") async def get_task_result(task_id: str): """ 查询AI生成任务的结果 """ if not settings.dashscope_api_key: raise HTTPException(status_code=500, detail="DashScope API Key not configured") headers = { "Authorization": f"Bearer {settings.dashscope_api_key}" } try: async with httpx.AsyncClient() as client: response = await client.get( f"{DASHSCOPE_TASK_URL}/{task_id}", headers=headers, timeout=30.0 ) if response.status_code != 200: logger.error(f"DashScope Task API error: {response.text}") raise HTTPException(status_code=response.status_code, detail="Failed to fetch task status") result = response.json() # 构建返回结果 # DashScope 任务查询返回结构： # { # "request_id": "...", # "output": { # "task_id": "...", # "task_status": "SUCCEEDED", # "results": [...] # }, # "usage": ... # } task_status = result.get("output", {}).get("task_status", "UNKNOWN") return AITaskResult( task_id=task_id, status=task_status, code=result.get("code"), message=result.get("message"), output=result.get("output"), usage=result.get("usage") ) except httpx.RequestError as e: logger.error(f"Request error: {str(e)}") raise HTTPException(status_code=500, detail=f"Request failed: {str(e)}")