abc / api /image_tasks.py
xiaolongmr
init
85270f9
Raw History Blame Contribute Delete
18.7 kB
from __future__ import annotations
from decimal import Decimal, ROUND_HALF_UP
from fastapi import APIRouter, Header, HTTPException, Query, Request
from fastapi.concurrency import run_in_threadpool
from pydantic import BaseModel, Field
from api.image_inputs import parse_image_edit_request, read_image_sources
from api.support import require_identity, resolve_image_base_url
from services.content_filter import check_request
from services.config import config
from services.auth_service import DAILY_GIFT_ACCOUNT_POOL_MODEL_ID, auth_service
from services.image_task_service import image_task_service
from services.log_service import LoggedCall
from services.pool_pricing_service import pool_pricing_service
from services.share_evidence_service import share_evidence_service
from services.third_party_image_service import third_party_image_service
def _credit_amount(value: object) -> int | float:
try:
amount = max(Decimal("0"), Decimal(str(value or 0))).quantize(Decimal("0.01"), rounding=ROUND_HALF_UP)
except Exception:
amount = Decimal("0")
return int(amount) if amount == amount.to_integral_value() else float(amount)
def _third_party_activity_quote(identity: dict[str, object], *, group_id: str, base_unit_price: object) -> dict[str, object]:
pricing_config = config.normal_user_third_party_image_pricing
if str(identity.get("role") or "").strip().lower() != "normal" or not bool(pricing_config.get("enabled")):
base = _credit_amount(base_unit_price)
return {
"applies": False,
"target_type": "third_party_group",
"third_party_group_id": group_id,
"base_unit_price": base,
"pay_percentage": 100,
"unit_price": base,
"quantity": 1,
"estimated_total": base,
"rule_id": "",
"rule_name": "",
}
return pool_pricing_service.quote(
identity,
quantity=1,
pricing_config=pricing_config,
target_type="third_party_group",
third_party_group_id=group_id,
base_unit_price=base_unit_price,
)
class ImageGenerationTaskRequest(BaseModel):
client_task_id: str = Field(..., min_length=1)
prompt: str = Field(..., min_length=1)
model: str = "gpt-image-2"
size: str | None = None
quality: str | None = None
resolution: str | None = None
image_resolution: str | None = None
third_party_image_group_id: str | None = None
account_pool_strategy: str = "own_first"
memory: bool = False
memory_conversation_id: str = ""
memory_reset: bool = False
class ResumePollRequest(BaseModel):
extra_timeout_secs: float | None = Field(default=None, ge=5.0, le=120.0)
class ShareEvidenceRequest(BaseModel):
conversation_id: str = Field(..., min_length=1)
def _parse_task_ids(value: str) -> list[str]:
return [item.strip() for item in value.split(",") if item.strip()]
def _require_owned_task(identity: dict[str, object], conversation_id: str) -> dict[str, object]:
"""校验这条上游会话确实属于当前用户,返回内部任务字典。
没有这层校验的话,任何人拿到一个别人的 conversation_id 就能把它变成公开分享链接。
⚠️ 生图任务记录可能已被清理,但官网取证台账里还有这条会话 → 也算「属于你」。
否则用户点「查看官网」时只会看到一句像越权警告的 404,而真实原因是
「这条会话已过保留期,程序已按策略把官网对话删了」。
"""
task = image_task_service.find_owned_conversation(identity, conversation_id)
if task is None:
raise HTTPException(status_code=404, detail={
"error": "这条会话的官网记录已不存在:可能已过保留期被自动清理,或不属于当前账号",
"reason": "conversation_not_found",
})
return task
def _internal_model_limit_kwargs(identity: dict[str, object], model: object) -> dict[str, object]:
if str(identity.get("role") or "").strip().lower() != "user":
return {}
model_id = str(model or "gpt-image-2").strip() or "gpt-image-2"
return {
"internal_model_limit_user_id": str(identity.get("id") or "").strip(),
"internal_model_limit_model_id": model_id,
"internal_model_limit_amount": 1,
}
def _credit_kwargs(identity: dict[str, object], model: object) -> dict[str, object]:
if str(identity.get("role") or "").strip().lower() != "normal":
return _internal_model_limit_kwargs(identity, model)
pricing = pool_pricing_service.quote(identity, 1)
return {
"credit_user_id": str(identity.get("id") or "").strip(),
"credit_amount": pricing["estimated_total"],
"credit_model_id": DAILY_GIFT_ACCOUNT_POOL_MODEL_ID,
"credit_pricing": pricing,
}
def _third_party_wallet_remaining(identity: dict[str, object], group_id: str, model_id: str) -> int | float:
try:
summary = auth_service.credit_summary(identity)
except Exception:
return 0
wallets = summary.get("third_party_image_wallets") if isinstance(summary, dict) else []
wallet_key = f"{group_id}:{model_id}"
wallet_items = wallets if isinstance(wallets, list) else []
balances: list[int | float] = []
for wallet in wallet_items:
if not isinstance(wallet, dict):
continue
if str(wallet.get("id") or "") == wallet_key or (
str(wallet.get("group_id") or "") == group_id
and str(wallet.get("model_id") or "").strip() == model_id
):
try:
balances.append(_credit_amount(wallet.get("remaining")))
except (TypeError, ValueError):
balances.append(0)
return max(balances) if balances else 0
def _image_api_billing_kwargs(
identity: dict[str, object],
*,
model: object,
quality: object = "",
resolution: object = "",
image_resolution: object = "",
third_party_image_group_id: object = "",
) -> dict[str, object]:
resolution = str(resolution or image_resolution or "").strip().lower()
group_id = str(third_party_image_group_id or "").strip()
explicit_third_party = resolution in {"1k", "2k", "4k"} or bool(group_id)
route = None
if not explicit_third_party:
find_named_route = getattr(third_party_image_service, "find_named_route", None)
route = find_named_route(model, advance=False) if callable(find_named_route) else None
if not explicit_third_party and route is None:
return _credit_kwargs(identity, model)
require_4k = resolution == "4k"
if route is None and str(identity.get("role") or "").strip().lower() == "normal":
find_wallet_route = getattr(third_party_image_service, "find_wallet_route_for_identity", None)
if callable(find_wallet_route):
route = find_wallet_route(
identity,
model,
group_id=group_id,
require_4k=require_4k,
resolution=resolution,
advance=False,
)
if route is None:
route = third_party_image_service.find_route(
model,
group_id=group_id,
require_4k=require_4k,
resolution=resolution or "1k",
advance=False,
)
if route is None:
raise ValueError("No enabled image API model route")
try:
third_party_image_service.assert_access(identity, route)
except PermissionError as exc:
raise PermissionError(str(exc)) from exc
cost = third_party_image_service.cost_for_route(route, quality, 1)
common = {
"resolution": str(route.model.get("resolution") or resolution or "1k"),
"third_party_image_group_id": str(route.group.get("id") or ""),
}
route_model_id = str(route.model.get("id") or "")
if str(identity.get("role") or "").strip().lower() != "normal":
return {
**common,
**_internal_model_limit_kwargs(identity, route_model_id),
}
billing_mode = str(route.group.get("billing_mode") or "")
route_group_id = str(route.group.get("id") or "")
if billing_mode == "general_multiplier" and _third_party_wallet_remaining(identity, route_group_id, route_model_id) > 0:
# Keep existing model-specific dedicated balances usable even if an
# older/stale group setting still says general-only.
billing_mode = "both"
wallet_pricing = _third_party_activity_quote(identity, group_id=route_group_id, base_unit_price=1)
wallet_cost = _credit_amount(wallet_pricing.get("estimated_total", 1))
if billing_mode == "dedicated_wallet" or (
billing_mode == "both" and cost > 0 and _third_party_wallet_remaining(identity, route_group_id, route_model_id) >= wallet_cost
):
wallet_kwargs: dict[str, object] = {
**common,
"third_party_wallet_user_id": str(identity.get("id") or "").strip(),
"third_party_wallet_group_id": route_group_id,
"third_party_wallet_model_id": route_model_id,
"third_party_wallet_amount": wallet_cost,
"third_party_wallet_cost_per_image": _credit_amount(wallet_pricing.get("unit_price", 1)),
"credit_pricing": wallet_pricing,
}
if billing_mode == "both" and cost > 0:
fallback_pricing = _third_party_activity_quote(identity, group_id=route_group_id, base_unit_price=cost)
wallet_kwargs["third_party_wallet_fallback_credit_user_id"] = str(identity.get("id") or "").strip()
wallet_kwargs["third_party_wallet_fallback_credit_amount"] = fallback_pricing["estimated_total"]
wallet_kwargs["third_party_wallet_fallback_credit_model_id"] = f"{route_group_id}:{route_model_id}"
wallet_kwargs["third_party_wallet_fallback_credit_pricing"] = fallback_pricing
return wallet_kwargs
pricing = _third_party_activity_quote(identity, group_id=route_group_id, base_unit_price=cost)
return {
**common,
"credit_user_id": str(identity.get("id") or "").strip(),
"credit_amount": pricing["estimated_total"],
"credit_model_id": f"{route_group_id}:{route_model_id}",
"credit_pricing": pricing,
}
def _generation_billing_kwargs(identity: dict[str, object], body: ImageGenerationTaskRequest) -> dict[str, object]:
return _image_api_billing_kwargs(
identity,
model=body.model,
quality=body.quality,
resolution=body.resolution,
image_resolution=body.image_resolution,
third_party_image_group_id=body.third_party_image_group_id,
)
def _effective_account_pool_strategy(identity: dict[str, object], value: object) -> str:
if identity.get("role") == "admin" or bool(identity.get("account_pool_enabled")):
return str(value or "own_first")
return "admin_only"
async def filter_or_log(call: LoggedCall, text: str) -> None:
try:
await run_in_threadpool(check_request, text)
except HTTPException as exc:
call.log("调用失败", status="failed", error=str(exc.detail))
raise
def create_router() -> APIRouter:
router = APIRouter()
@router.get("/api/image-tasks")
async def list_image_tasks(
ids: str = Query(default=""),
authorization: str | None = Header(default=None),
):
identity = require_identity(authorization)
return await run_in_threadpool(image_task_service.list_tasks, identity, _parse_task_ids(ids))
@router.post("/api/image-tasks/generations")
async def create_generation_task(
body: ImageGenerationTaskRequest,
request: Request,
authorization: str | None = Header(default=None),
):
identity = require_identity(authorization)
await filter_or_log(LoggedCall(identity, "/api/image-tasks/generations", body.model, "文生图任务", request_text=body.prompt), body.prompt)
try:
return await run_in_threadpool(
image_task_service.submit_generation,
identity,
client_task_id=body.client_task_id,
prompt=body.prompt,
model=body.model,
size=body.size,
quality=str(body.quality or ""),
base_url=resolve_image_base_url(request),
memory_enabled=body.memory,
memory_conversation_id=body.memory_conversation_id,
memory_reset=body.memory_reset,
account_pool_strategy=_effective_account_pool_strategy(identity, body.account_pool_strategy),
**_generation_billing_kwargs(identity, body),
)
except PermissionError as exc:
raise HTTPException(status_code=403, detail={"error": str(exc)}) from exc
except ValueError as exc:
error_text = str(exc)
status_code = 402 if "次数不足" in error_text or "not enough" in error_text.lower() or "balance" in error_text.lower() else 400
raise HTTPException(status_code=status_code, detail={"error": str(exc)}) from exc
@router.post("/api/image-tasks/edits")
async def create_edit_task(
request: Request,
authorization: str | None = Header(default=None),
):
identity = require_identity(authorization)
payload, image_sources, mask_sources = await parse_image_edit_request(request)
client_task_id = str(payload.get("client_task_id") or "").strip()
if not client_task_id:
raise HTTPException(status_code=400, detail={"error": "client_task_id is required"})
prompt = str(payload["prompt"])
model = str(payload["model"])
await filter_or_log(LoggedCall(identity, "/api/image-tasks/edits", model, "图生图任务", request_text=prompt), prompt)
images = await read_image_sources(image_sources)
masks = await read_image_sources(mask_sources) if mask_sources else None
try:
return await run_in_threadpool(
image_task_service.submit_edit,
identity,
client_task_id=client_task_id,
prompt=prompt,
model=model,
size=payload["size"],
base_url=resolve_image_base_url(request),
images=images,
masks=masks,
memory_enabled=bool(payload.get("memory")),
memory_conversation_id=str(payload.get("memory_conversation_id") or ""),
memory_reset=bool(payload.get("memory_reset")),
account_pool_strategy=_effective_account_pool_strategy(identity, payload.get("account_pool_strategy") or "own_first"),
quality=str(payload.get("quality") or ""),
**_image_api_billing_kwargs(
identity,
model=model,
quality=payload.get("quality"),
resolution=payload.get("resolution"),
image_resolution=payload.get("image_resolution"),
third_party_image_group_id=payload.get("third_party_image_group_id"),
),
)
except PermissionError as exc:
raise HTTPException(status_code=403, detail={"error": str(exc)}) from exc
except ValueError as exc:
status_code = 402 if "次数不足" in str(exc) else 400
raise HTTPException(status_code=status_code, detail={"error": str(exc)}) from exc
@router.post("/api/image-tasks/{task_id}/resume-poll")
async def resume_image_poll(
task_id: str,
body: ResumePollRequest,
authorization: str | None = Header(default=None),
):
identity = require_identity(authorization)
try:
return await run_in_threadpool(
image_task_service.resume_poll,
identity,
task_id,
body.extra_timeout_secs or config.image_timeout_retry_secs,
)
except ValueError as exc:
raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc
@router.get("/api/image-tasks/share-evidence")
async def share_evidence_status(
conversation_id: str = Query(..., min_length=1),
authorization: str | None = Header(default=None),
):
"""查询某条会话的「官网来源」取证状态(链接、剩余有效秒数)。"""
identity = require_identity(authorization)
_require_owned_task(identity, conversation_id)
return share_evidence_service.describe(conversation_id)
@router.post("/api/image-tasks/share-evidence")
async def create_share_evidence(
body: ShareEvidenceRequest,
authorization: str | None = Header(default=None),
):
"""把一条会话取成为 ChatGPT 官方匿名分享链接(幂等)。"""
identity = require_identity(authorization)
if not config.image_share_evidence_enabled:
raise HTTPException(status_code=403, detail={"error": "后台未开启「官网来源取证」,无法生成分享链接"})
task = _require_owned_task(identity, body.conversation_id)
try:
await run_in_threadpool(
share_evidence_service.create_share,
conversation_id=body.conversation_id,
owner_id=str(identity.get("id") or ""),
account_email=str(task.get("account_email") or ""),
)
except ValueError as exc:
raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc
except RuntimeError as exc:
raise HTTPException(status_code=502, detail={"error": str(exc)}) from exc
return share_evidence_service.describe(body.conversation_id)
@router.delete("/api/image-tasks/share-evidence")
async def revoke_share_evidence(
conversation_id: str = Query(..., min_length=1),
authorization: str | None = Header(default=None),
):
"""提前撤销分享链接并回收会话。"""
identity = require_identity(authorization)
_require_owned_task(identity, conversation_id)
try:
await run_in_threadpool(share_evidence_service.revoke, conversation_id)
except ValueError as exc:
raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc
except RuntimeError as exc:
raise HTTPException(status_code=502, detail={"error": str(exc)}) from exc
return share_evidence_service.describe(conversation_id)
return router