Tool_Agent/tools/local/gpt_image_2/main.py

import os
import sys
import re
from typing import Optional
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from dotenv import load_dotenv
from openai import OpenAI

# 加载环境变量
load_dotenv()

app = FastAPI(title="GPT Image 2 API", version="1.0.0")

# 初始化 OpenAI 客户端，使用 apiyi 平台
api_key = os.getenv("APIYI_KEY")
if not api_key:
    raise ValueError("APIYI_KEY environment variable is required")

client = OpenAI(
    api_key=api_key,
    base_url="https://api.apiyi.com/v1"
)


def extract_url_from_content(text: str) -> Optional[str]:
    """从 API 返回内容中提取图片 URL（可能是 markdown 格式）"""
    # 匹配 ![image](url) 或 ![...](url) 格式
    match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', text)
    if match:
        return match.group(1)
    # 如果不是 markdown 格式，尝试直接作为 URL
    stripped = text.strip()
    if stripped.startswith("http"):
        return stripped
    return None


class ImageRequest(BaseModel):
    prompt: str = Field(..., description="图像生成提示词，可包含尺寸描述如'1024x1024 方图'")


class ImageResponse(BaseModel):
    status: str
    image_url: Optional[str] = None
    error: Optional[str] = None


@app.post("/generate", response_model=ImageResponse)
async def generate_image(request: ImageRequest):
    """生成图像 - 使用 gpt-image-2-all 模型"""
    try:
        # 使用 chat/completions 接口调用 gpt-image-2-all
        response = client.chat.completions.create(
            model="gpt-image-2-all",
            messages=[
                {"role": "user", "content": request.prompt}
            ],
            extra_body={"response_format": {"type": "url"}}
        )

        if response.choices and len(response.choices) > 0:
            content = response.choices[0].message.content
            image_url = extract_url_from_content(content)

            if image_url:
                return ImageResponse(status="success", image_url=image_url)
            else:
                return ImageResponse(
                    status="error",
                    error=f"Could not extract URL from response: {content}"
                )
        else:
            return ImageResponse(status="error", error="No content returned from API")

    except Exception as e:
        return ImageResponse(status="error", error=str(e))


@app.get("/health")
async def health_check():
    return {"status": "ok", "service": "gpt-image-2"}


if __name__ == "__main__":
    import uvicorn

    port = 5000
    if "--port" in sys.argv:
        try:
            port_idx = sys.argv.index("--port")
            port = int(sys.argv[port_idx + 1])
        except (IndexError, ValueError):
            pass

    uvicorn.run(app, host="0.0.0.0", port=port)