Tool_Agent/tools/local/midjourney/main.py

"""Midjourney 本地代理 — 三个 POST 与上游 JSON 对齐。

环境变量：
  MIDJOURNEY_API_BASE  必填，例如 https://your-host（后接 /submit_job 等）

接口（与 Router 注册一致）：
  GET  /health
  POST /submit_job        cookie, prompt, user_id, mode ∈ relaxed|fast
  POST /query_job_status  cookie, job_id
  POST /get_image_urls    job_id
"""

from __future__ import annotations

import argparse
from typing import Any, Literal

import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

from midjourney_client import forward_post

app = FastAPI(title="Midjourney API Proxy")


class SubmitJobRequest(BaseModel):
    cookie: str = Field(..., description="Midjourney 会话 cookie")
    prompt: str = Field(..., description="提示词")
    user_id: str = Field(..., description="用户 ID")
    mode: Literal["relaxed", "fast"] = Field(..., description="relaxed 或 fast")


class QueryJobStatusRequest(BaseModel):
    cookie: str = Field(..., description="Midjourney 会话 cookie")
    job_id: str = Field(..., description="submit_job 返回的任务 ID")


class GetImageUrlsRequest(BaseModel):
    job_id: str = Field(..., description="任务 ID")


@app.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@app.post("/submit_job")
def submit_job(req: SubmitJobRequest) -> Any:
    try:
        return forward_post(
            "/submit_job",
            {
                "cookie": req.cookie,
                "prompt": req.prompt,
                "user_id": req.user_id,
                "mode": req.mode,
            },
        )
    except ValueError as e:
        raise HTTPException(status_code=503, detail=str(e)) from e
    except RuntimeError as e:
        raise HTTPException(status_code=502, detail=str(e)) from e
    except Exception as e:
        raise HTTPException(status_code=502, detail=str(e)) from e


@app.post("/query_job_status")
def query_job_status(req: QueryJobStatusRequest) -> Any:
    try:
        return forward_post(
            "/query_job_status",
            {"cookie": req.cookie, "job_id": req.job_id},
        )
    except ValueError as e:
        raise HTTPException(status_code=503, detail=str(e)) from e
    except RuntimeError as e:
        raise HTTPException(status_code=502, detail=str(e)) from e
    except Exception as e:
        raise HTTPException(status_code=502, detail=str(e)) from e


@app.post("/get_image_urls")
def get_image_urls(req: GetImageUrlsRequest) -> Any:
    try:
        return forward_post("/get_image_urls", {"job_id": req.job_id})
    except ValueError as e:
        raise HTTPException(status_code=503, detail=str(e)) from e
    except RuntimeError as e:
        raise HTTPException(status_code=502, detail=str(e)) from e
    except Exception as e:
        raise HTTPException(status_code=502, detail=str(e)) from e


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=8001)
    args = parser.parse_args()
    uvicorn.run(app, host="0.0.0.0", port=args.port)