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