File size: 8,820 Bytes
1a9e2c2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 | """图片上传管理器 - 支持Base64和URL图片上传"""
import asyncio
import base64
import re
from typing import Tuple, Optional
from urllib.parse import urlparse
from curl_cffi.requests import AsyncSession
from app.services.grok.statsig import get_dynamic_headers
from app.core.exception import GrokApiException
from app.core.config import setting
from app.core.logger import logger
# 常量
UPLOAD_API = "https://grok.com/rest/app-chat/upload-file"
TIMEOUT = 30
BROWSER = "chrome133a"
# MIME类型
MIME_TYPES = {
'.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.png': 'image/png',
'.gif': 'image/gif', '.webp': 'image/webp', '.bmp': 'image/bmp',
}
DEFAULT_MIME = "image/jpeg"
DEFAULT_EXT = "jpg"
class ImageUploadManager:
"""图片上传管理器"""
@staticmethod
async def upload(image_input: str, auth_token: str) -> Tuple[str, str]:
"""上传图片（支持Base64或URL）
Returns:
(file_id, file_uri) 元组
"""
try:
# 判断类型并处理
if ImageUploadManager._is_url(image_input):
buffer, mime = await ImageUploadManager._download(image_input)
filename, _ = ImageUploadManager._get_info("", mime)
else:
buffer = image_input.split(",")[1] if "data:image" in image_input else image_input
filename, mime = ImageUploadManager._get_info(image_input)
# 构建数据
data = {
"fileName": filename,
"fileMimeType": mime,
"content": buffer,
}
if not auth_token:
raise GrokApiException("认证令牌缺失", "NO_AUTH_TOKEN")
# 外层重试：可配置状态码（401/429等）
retry_codes = setting.grok_config.get("retry_status_codes", [401, 429])
MAX_OUTER_RETRY = 3
for outer_retry in range(MAX_OUTER_RETRY + 1): # +1 确保实际重试3次
try:
# 内层重试：403代理池重试
max_403_retries = 5
retry_403_count = 0
while retry_403_count <= max_403_retries:
# 请求配置
cf = setting.grok_config.get("cf_clearance", "")
headers = {
**get_dynamic_headers("/rest/app-chat/upload-file"),
"Cookie": f"{auth_token};{cf}" if cf else auth_token,
}
# 异步获取代理（支持代理池）
from app.core.proxy_pool import proxy_pool
# 如果是403重试且使用代理池，强制刷新代理
if retry_403_count > 0 and proxy_pool._enabled:
logger.info(f"[Upload] 403重试 {retry_403_count}/{max_403_retries}，刷新代理...")
proxy = await proxy_pool.force_refresh()
else:
proxy = await setting.get_proxy_async("service")
proxies = {"http": proxy, "https": proxy} if proxy else None
# 上传
async with AsyncSession() as session:
response = await session.post(
UPLOAD_API,
headers=headers,
json=data,
impersonate=BROWSER,
timeout=TIMEOUT,
proxies=proxies,
)
# 内层403重试：仅当有代理池时触发
if response.status_code == 403 and proxy_pool._enabled:
retry_403_count += 1
if retry_403_count <= max_403_retries:
logger.warning(f"[Upload] 遇到403错误，正在重试 ({retry_403_count}/{max_403_retries})...")
await asyncio.sleep(0.5)
continue
# 内层重试全部失败
logger.error(f"[Upload] 403错误，已重试{retry_403_count-1}次，放弃")
# 检查可配置状态码错误 - 外层重试
if response.status_code in retry_codes:
if outer_retry < MAX_OUTER_RETRY:
delay = (outer_retry + 1) * 0.1 # 渐进延迟：0.1s, 0.2s, 0.3s
logger.warning(f"[Upload] 遇到{response.status_code}错误，外层重试 ({outer_retry+1}/{MAX_OUTER_RETRY})，等待{delay}s...")
await asyncio.sleep(delay)
break # 跳出内层循环，进入外层重试
else:
logger.error(f"[Upload] {response.status_code}错误，已重试{outer_retry}次，放弃")
return "", ""
if response.status_code == 200:
result = response.json()
file_id = result.get("fileMetadataId", "")
file_uri = result.get("fileUri", "")
if outer_retry > 0 or retry_403_count > 0:
logger.info(f"[Upload] 重试成功！")
logger.debug(f"[Upload] 成功，ID: {file_id}")
return file_id, file_uri
# 其他错误直接返回
logger.error(f"[Upload] 失败，状态码: {response._status_code}")
return "", ""
# 内层循环正常结束（非break），说明403重试全部失败
return "", ""
except Exception as e:
if outer_retry < MAX_OUTER_RETRY - 1:
logger.warning(f"[Upload] 异常: {e}，外层重试 ({outer_retry+1}/{MAX_OUTER_RETRY})...")
await asyncio.sleep(0.5)
continue
logger.warning(f"[Upload] 失败: {e}")
return "", ""
return "", ""
except Exception as e:
logger.warning(f"[Upload] 失败: {e}")
return "", ""
@staticmethod
def _is_url(input_str: str) -> bool:
"""检查是否为URL"""
try:
result = urlparse(input_str)
return all([result.scheme, result.netloc]) and result.scheme in ['http', 'https']
except:
return False
@staticmethod
async def _download(url: str) -> Tuple[str, str]:
"""下载图片并转Base64
Returns:
(base64_string, mime_type) 元组
"""
try:
async with AsyncSession() as session:
response = await session.get(url, timeout=5)
response.raise_for_status()
content_type = response.headers.get('content-type', DEFAULT_MIME)
if not content_type.startswith('image/'):
content_type = DEFAULT_MIME
b64 = base64.b64encode(response.content).decode()
return b64, content_type
except Exception as e:
logger.warning(f"[Upload] 下载失败: {e}")
return "", ""
@staticmethod
def _get_info(image_data: str, mime_type: Optional[str] = None) -> Tuple[str, str]:
"""获取文件名和MIME类型
Returns:
(file_name, mime_type) 元组
"""
# 已提供MIME类型
if mime_type:
ext = mime_type.split("/")[1] if "/" in mime_type else DEFAULT_EXT
return f"image.{ext}", mime_type
# 从Base64提取
mime = DEFAULT_MIME
ext = DEFAULT_EXT
if "data:image" in image_data:
if match := re.search(r"data:([a-zA-Z0-9]+/[a-zA-Z0-9-.+]+);base64,", image_data):
mime = match.group(1)
ext = mime.split("/")[1]
return f"image.{ext}", mime |