ai_v/services/generation_service.py
24024 0adabaa235 ```
feat(api): 新增gpt-image-2编辑模型支持

- 添加gpt-image-2模型的专门接口配置和API密钥
- 实现gpt-image-2编辑模式的完整处理流程,包括Base64图片解码、
  图片验证和MinIO上传功能
- 更新前端界面以支持gpt-image-2的特定参数(如尺寸选项)
- 修复prompt不能为空的校验逻辑
- 优化图片上传处理,支持编辑模式下的参考图必需校验
- 更新视频生成任务状态消息文案
```
2026-04-24 23:17:10 +08:00

240 lines
8.0 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import json
import threading
import uuid
import requests
from config import Config
from extensions import db, redis_client
from models import GenerationRecord, SystemDict, User
from services.logger import system_logger
from services.task_service import process_image_generation, process_video_generation
from utils import get_api_candidates, get_proxied_url, should_switch_to_backup
def get_model_cost(model_value, is_video=False):
"""获取模型消耗积分。"""
dict_type = "video_model" if is_video else "ai_model"
model_dict = SystemDict.query.filter_by(dict_type=dict_type, value=model_value).first()
if model_dict:
return model_dict.cost
# 默认计费兜底
if is_video:
return 15 if "pro" in model_value.lower() or "3.1" in model_value else 10
return 1
def validate_generation_request(user, data):
"""校验生图请求并返回 (api_key, target_api, cost, use_trial, error)。"""
mode = data.get("mode", "trial")
is_premium = data.get("is_premium", False)
input_key = data.get("apiKey")
model_value = data.get("model")
is_gemini_flash_image_preview = model_value == "gemini-3.1-flash-image-preview"
is_gpt_image_edit = model_value == "gpt-image-2"
if isinstance(is_premium, str):
is_premium = is_premium.lower() in ("1", "true", "yes", "on")
# gpt-image-2 走独立的 edits 接口,其余模型沿用现有 generations 接口
target_api = Config.GPT_IMAGE_EDIT_API if is_gpt_image_edit else Config.AI_API
api_key = None
use_trial = False
if mode == "key":
api_key = input_key or user.api_key
if not api_key:
return None, None, 0, False, "请先输入您的 API 密钥"
# 用户手动输入了新 key 时,同步保存到账号
if input_key and input_key != user.api_key:
user.api_key = input_key
db.session.commit()
else:
if user.points <= 0:
return None, None, 0, False, "可用积分已耗尽,请充值或切换至自定义 Key 模式"
# 积分模式下gpt-image-2 使用固定的专用 key
if is_gpt_image_edit:
api_key = Config.GPT_IMAGE_EDIT_KEY
elif is_premium:
api_key = (
Config.GEMINI_FLASH_IMAGE_PREMIUM_KEY
if is_gemini_flash_image_preview
else Config.PREMIUM_KEY
)
elif is_gemini_flash_image_preview:
api_key = Config.PREMIUM_KEY
else:
api_key = Config.TRIAL_KEY
target_api = Config.GPT_IMAGE_EDIT_API if is_gpt_image_edit else Config.TRIAL_API
use_trial = True
# 计算本次任务消耗
cost = get_model_cost(model_value, is_video=False)
if use_trial and is_premium:
cost *= 2
if use_trial and user.points < cost:
return None, None, cost, True, f"可用积分不足,本次需要 {cost} 积分"
return api_key, target_api, cost, use_trial, None
def deduct_points(user_id, cost):
"""原子扣减积分。"""
user = db.session.query(User).filter_by(id=user_id).populate_existing().with_for_update().first()
if user:
user.points -= cost
user.has_used_points = True
db.session.commit()
def refund_points(user_id, cost):
"""原子退回积分。"""
try:
user = db.session.query(User).filter_by(id=user_id).populate_existing().with_for_update().first()
if user:
user.points += cost
db.session.commit()
except Exception:
db.session.rollback()
def handle_chat_generation_sync(user_id, api_key, model_value, prompt, use_trial, cost):
"""同步处理聊天类模型。"""
chat_payload = {
"model": model_value,
"messages": [{"role": "user", "content": prompt}],
}
try:
resp = None
last_error = None
candidates = get_api_candidates(Config.CHAT_API, api_key)
for idx, candidate in enumerate(candidates):
headers = {
"Authorization": f"Bearer {candidate['api_key']}",
"Content-Type": "application/json",
}
try:
resp = requests.post(
get_proxied_url(candidate["url"]),
json=chat_payload,
headers=headers,
timeout=Config.PROXY_TIMEOUT_LONG,
)
if resp.status_code == 200:
break
last_error = resp.text
has_backup = idx < len(candidates) - 1
if has_backup and should_switch_to_backup(resp.status_code):
system_logger.warning(
"聊天主线路失败,切换备用线路",
user_id=user_id,
model=model_value,
status_code=resp.status_code,
)
continue
break
except requests.RequestException as e:
last_error = str(e)
if idx < len(candidates) - 1:
system_logger.warning(
"聊天请求异常,切换备用线路",
user_id=user_id,
model=model_value,
error=last_error,
)
continue
raise
if not resp or resp.status_code != 200:
if use_trial:
refund_points(user_id, cost)
return {"error": last_error or "聊天请求失败"}, resp.status_code if resp else 500
api_result = resp.json()
content = api_result["choices"][0]["message"]["content"]
# 非验光单解读的文本生成写入历史
if prompt != "解读验光单":
new_record = GenerationRecord(
user_id=user_id,
prompt=prompt,
model=model_value,
cost=cost,
image_urls=json.dumps([{"type": "text", "content": content}]),
)
db.session.add(new_record)
db.session.commit()
return {
"data": [{"content": content, "type": "text"}],
"message": "生成成功",
}, 200
except Exception as e:
if use_trial:
refund_points(user_id, cost)
return {"error": str(e)}, 500
def start_async_image_task(app, user_id, payload, api_key, target_api, cost, mode, model_value, use_trial=False):
"""启动异步生图任务。"""
task_id = str(uuid.uuid4())
log_msg = "用户发起验光单解读" if payload.get("prompt") == "解读验光单" else "用户发起生图任务"
system_logger.info(log_msg, model=model_value, mode=mode)
redis_client.setex(
f"task:{task_id}",
3600,
json.dumps({"status": "queued", "message": "任务已提交,等待处理..."}),
)
threading.Thread(
target=process_image_generation,
args=(app, user_id, task_id, payload, api_key, target_api, cost, use_trial),
).start()
return task_id
def validate_video_request(user, data):
"""校验视频生成请求。"""
if user.points <= 0:
return None, 0, "可用积分不足,请先充值"
model_value = data.get("model", "veo3.1")
cost = get_model_cost(model_value, is_video=True)
if user.points < cost:
return None, cost, f"积分不足,生成该视频需要 {cost} 积分"
return model_value, cost, None
def start_async_video_task(app, user_id, payload, cost, model_value):
"""启动异步视频任务。"""
api_key = Config.VIDEO_KEY
task_id = str(uuid.uuid4())
system_logger.info("用户发起视频生成任务 (积分模式)", model=model_value, cost=cost)
redis_client.setex(
f"task:{task_id}",
3600,
json.dumps({"status": "queued", "message": "视频任务已提交,准备开始处理..."}),
)
threading.Thread(
target=process_video_generation,
args=(app, user_id, task_id, payload, api_key, cost, True),
).start()
return task_id