okooooooo/grok2api
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: improve token management and logging across services

Browse Source
This commit is contained in:
Chenyme committed 2026-02-07 13:11:23 +08:00
1 parent 2e201b5ebf
commit 5d369596c5
8 files changed
+176 -403
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

app/api/v1/admin.py
+79 -226
View File
@@ -6,6 +6,8 @@ from app.core.config import config, get_config
from app.core.batch_tasks import create_task, get_task, expire_task
from app.core.storage import get_storage, LocalStorage, RedisStorage, SQLStorage
from app.core.exceptions import AppException
from app.services.token.manager import get_token_manager
from app.services.grok.utils.batch import run_in_batches
import os
from pathlib import Path
import aiofiles
@@ -13,23 +15,41 @@ import asyncio
import orjson
from app.core.logger import logger
from app.services.grok.services.voice import VoiceService
from app.services.token import get_token_manager
DEFAULT_STORAGE_TYPE = ""
DEFAULT_USAGE_MAX_TOKENS = 1000
DEFAULT_USAGE_MAX_CONCURRENT = 25
DEFAULT_USAGE_BATCH_SIZE = 50
DEFAULT_NSFW_MAX_TOKENS = 1000
DEFAULT_NSFW_MAX_CONCURRENT = 10
DEFAULT_NSFW_BATCH_SIZE = 50
DEFAULT_ASSETS_MAX_CONCURRENT = 25
DEFAULT_ASSETS_BATCH_SIZE = 10
DEFAULT_ASSETS_MAX_TOKENS = 1000
TEMPLATE_DIR = Path(__file__).parent.parent.parent / "static"
router = APIRouter()
def _collect_tokens(data: dict) -> list[str]:
"""从请求数据中收集 token 列表"""
tokens = []
if isinstance(data.get("token"), str) and data["token"].strip():
tokens.append(data["token"].strip())
if isinstance(data.get("tokens"), list):
tokens.extend([str(t).strip() for t in data["tokens"] if str(t).strip()])
return tokens
def _truncate_tokens(tokens: list[str], max_tokens: int, operation: str = "operation") -> tuple[list[str], bool, int]:
"""去重并截断 token 列表，返回 (unique_tokens, truncated, original_count)"""
unique_tokens = list(dict.fromkeys(tokens))
original_count = len(unique_tokens)
truncated = False
if len(unique_tokens) > max_tokens:
unique_tokens = unique_tokens[:max_tokens]
truncated = True
logger.warning(f"{operation}: truncated from {original_count} to {max_tokens} tokens")
return unique_tokens, truncated, original_count
def _mask_token(token: str) -> str:
"""掩码 token 显示"""
return f"{token[:8]}...{token[-8:]}" if len(token) > 20 else token
async def render_template(filename: str):
"""渲染指定模板"""
template_path = TEMPLATE_DIR / filename
@@ -217,10 +237,9 @@ async def update_config_api(data: dict):
@router.get("/api/v1/admin/storage", dependencies=[Depends(verify_api_key)])
async def get_storage_info():
"""获取当前存储模式"""
storage_type = os.getenv("SERVER_STORAGE_TYPE", "local").lower()
logger.info(f"Storage type: {storage_type}")
storage_type = os.getenv("SERVER_STORAGE_TYPE", "").lower()
if not storage_type:
storage_type = str(get_config("storage.type", DEFAULT_STORAGE_TYPE)).lower()
storage_type = str(get_config("storage.type")).lower()
if not storage_type:
storage = get_storage()
if isinstance(storage, LocalStorage):
@@ -228,12 +247,10 @@ async def get_storage_info():
elif isinstance(storage, RedisStorage):
storage_type = "redis"
elif isinstance(storage, SQLStorage):
if storage.dialect in ("mysql", "mariadb"):
storage_type = "mysql"
elif storage.dialect in ("postgres", "postgresql", "pgsql"):
storage_type = "pgsql"
else:
storage_type = storage.dialect
storage_type = {
"mysql": "mysql", "mariadb": "mysql",
"postgres": "pgsql", "postgresql": "pgsql", "pgsql": "pgsql"
}.get(storage.dialect, storage.dialect)
return {"type": storage_type or "local"}
@@ -318,46 +335,20 @@ async def update_tokens_api(data: dict):
@router.post("/api/v1/admin/tokens/refresh", dependencies=[Depends(verify_api_key)])
async def refresh_tokens_api(data: dict):
"""刷新 Token 状态"""
from app.services.token.manager import get_token_manager
from app.services.grok.utils.batch import run_in_batches
try:
mgr = await get_token_manager()
tokens = []
if "token" in data:
tokens.append(data["token"])
if "tokens" in data and isinstance(data["tokens"], list):
tokens.extend(data["tokens"])
tokens = _collect_tokens(data)
if not tokens:
raise HTTPException(status_code=400, detail="No tokens provided")
# 去重并保持顺序
unique_tokens = list(dict.fromkeys(tokens))
# 最大数量限制
max_tokens = get_config("performance.usage_max_tokens", DEFAULT_USAGE_MAX_TOKENS)
try:
max_tokens = int(max_tokens)
except Exception:
max_tokens = 1000
truncated = False
original_count = len(unique_tokens)
if len(unique_tokens) > max_tokens:
unique_tokens = unique_tokens[:max_tokens]
truncated = True
logger.warning(
f"Usage refresh: truncated from {original_count} to {max_tokens} tokens"
)
# 去重并截断
max_tokens = int(get_config("performance.usage_max_tokens"))
unique_tokens, truncated, original_count = _truncate_tokens(tokens, max_tokens, "Usage refresh")
# 批量执行配置
max_concurrent = get_config(
"performance.usage_max_concurrent", DEFAULT_USAGE_MAX_CONCURRENT
)
batch_size = get_config(
"performance.usage_batch_size", DEFAULT_USAGE_BATCH_SIZE
)
max_concurrent = get_config("performance.usage_max_concurrent")
batch_size = get_config("performance.usage_batch_size")
async def _refresh_one(t):
return await mgr.sync_usage(
@@ -393,40 +384,20 @@ async def refresh_tokens_api(data: dict):
)
async def refresh_tokens_api_async(data: dict):
"""刷新 Token 状态（异步批量 + SSE 进度）"""
from app.services.token.manager import get_token_manager
from app.services.grok.utils.batch import run_in_batches
mgr = await get_token_manager()
tokens: list[str] = []
if isinstance(data.get("token"), str) and data["token"].strip():
tokens.append(data["token"].strip())
if isinstance(data.get("tokens"), list):
tokens.extend([str(t).strip() for t in data["tokens"] if str(t).strip()])
tokens = _collect_tokens(data)
if not tokens:
raise HTTPException(status_code=400, detail="No tokens provided")
unique_tokens = list(dict.fromkeys(tokens))
max_tokens = get_config("performance.usage_max_tokens", DEFAULT_USAGE_MAX_TOKENS)
try:
max_tokens = int(max_tokens)
except Exception:
max_tokens = 1000
truncated = False
original_count = len(unique_tokens)
if len(unique_tokens) > max_tokens:
unique_tokens = unique_tokens[:max_tokens]
truncated = True
logger.warning(
f"Usage refresh: truncated from {original_count} to {max_tokens} tokens"
)
# 去重并截断
max_tokens = int(get_config("performance.usage_max_tokens"))
unique_tokens, truncated, original_count = _truncate_tokens(tokens, max_tokens, "Usage refresh")
max_concurrent = get_config(
"performance.usage_max_concurrent", DEFAULT_USAGE_MAX_CONCURRENT
"performance.usage_max_concurrent"
)
batch_size = get_config("performance.usage_batch_size", DEFAULT_USAGE_BATCH_SIZE)
batch_size = get_config("performance.usage_batch_size")
task = create_task(len(unique_tokens))
@@ -500,19 +471,13 @@ async def refresh_tokens_api_async(data: dict):
async def enable_nsfw_api(data: dict):
"""批量开启 NSFW (Unhinged) 模式"""
from app.services.grok.services.nsfw import NSFWService
from app.services.grok.utils.batch import run_in_batches
from app.services.token.manager import get_token_manager
try:
mgr = await get_token_manager()
nsfw_service = NSFWService()
# 收集 token 列表
tokens: list[str] = []
if isinstance(data.get("token"), str) and data["token"].strip():
tokens.append(data["token"].strip())
if isinstance(data.get("tokens"), list):
tokens.extend([str(t).strip() for t in data["tokens"] if str(t).strip()])
tokens = _collect_tokens(data)
# 若未指定，则使用所有 pool 中的 token
if not tokens:
@@ -526,30 +491,15 @@ async def enable_nsfw_api(data: dict):
if not tokens:
raise HTTPException(status_code=400, detail="No tokens available")
# 去重并保持顺序
unique_tokens = list(dict.fromkeys(tokens))
# 限制最大数量（超出时截取前 N 个）
max_tokens = get_config("performance.nsfw_max_tokens", DEFAULT_NSFW_MAX_TOKENS)
try:
max_tokens = int(max_tokens)
except Exception:
max_tokens = 1000
truncated = False
original_count = len(unique_tokens)
if len(unique_tokens) > max_tokens:
unique_tokens = unique_tokens[:max_tokens]
truncated = True
logger.warning(
f"NSFW enable: truncated from {original_count} to {max_tokens} tokens"
)
# 去重并截断
max_tokens = int(get_config("performance.nsfw_max_tokens"))
unique_tokens, truncated, original_count = _truncate_tokens(tokens, max_tokens, "NSFW enable")
# 批量执行配置
max_concurrent = get_config(
"performance.nsfw_max_concurrent", DEFAULT_NSFW_MAX_CONCURRENT
"performance.nsfw_max_concurrent"
)
batch_size = get_config("performance.nsfw_batch_size", DEFAULT_NSFW_BATCH_SIZE)
batch_size = get_config("performance.nsfw_batch_size")
# 定义 worker
async def _enable(token: str):
@@ -576,7 +526,7 @@ async def enable_nsfw_api(data: dict):
fail_count = 0
for token, res in raw_results.items():
masked = f"{token[:8]}...{token[-8:]}" if len(token) > 20 else token
masked = _mask_token(token)
if res.get("ok") and res.get("data", {}).get("success"):
ok_count += 1
results[masked] = res.get("data", {})
@@ -615,17 +565,11 @@ async def enable_nsfw_api(data: dict):
async def enable_nsfw_api_async(data: dict):
"""批量开启 NSFW (Unhinged) 模式（异步批量 + SSE 进度）"""
from app.services.grok.services.nsfw import NSFWService
from app.services.grok.utils.batch import run_in_batches
from app.services.token.manager import get_token_manager
mgr = await get_token_manager()
nsfw_service = NSFWService()
tokens: list[str] = []
if isinstance(data.get("token"), str) and data["token"].strip():
tokens.append(data["token"].strip())
if isinstance(data.get("tokens"), list):
tokens.extend([str(t).strip() for t in data["tokens"] if str(t).strip()])
tokens = _collect_tokens(data)
if not tokens:
for pool_name, pool in mgr.pools.items():
@@ -636,27 +580,12 @@ async def enable_nsfw_api_async(data: dict):
if not tokens:
raise HTTPException(status_code=400, detail="No tokens available")
unique_tokens = list(dict.fromkeys(tokens))
# 去重并截断
max_tokens = int(get_config("performance.nsfw_max_tokens"))
unique_tokens, truncated, original_count = _truncate_tokens(tokens, max_tokens, "NSFW enable")
max_tokens = get_config("performance.nsfw_max_tokens", DEFAULT_NSFW_MAX_TOKENS)
try:
max_tokens = int(max_tokens)
except Exception:
max_tokens = 1000
truncated = False
original_count = len(unique_tokens)
if len(unique_tokens) > max_tokens:
unique_tokens = unique_tokens[:max_tokens]
truncated = True
logger.warning(
f"NSFW enable: truncated from {original_count} to {max_tokens} tokens"
)
max_concurrent = get_config(
"performance.nsfw_max_concurrent", DEFAULT_NSFW_MAX_CONCURRENT
)
batch_size = get_config("performance.nsfw_batch_size", DEFAULT_NSFW_BATCH_SIZE)
max_concurrent = get_config("performance.nsfw_max_concurrent")
batch_size = get_config("performance.nsfw_batch_size")
task = create_task(len(unique_tokens))
@@ -791,28 +720,9 @@ async def get_cache_stats_api(request: Request):
}
online_details = []
account_map = {a["token"]: a for a in accounts}
max_concurrent = get_config(
"performance.assets_max_concurrent", DEFAULT_ASSETS_MAX_CONCURRENT
)
batch_size = get_config(
"performance.assets_batch_size", DEFAULT_ASSETS_BATCH_SIZE
)
try:
max_concurrent = int(max_concurrent)
except Exception:
max_concurrent = 25
try:
batch_size = int(batch_size)
except Exception:
batch_size = 10
max_concurrent = max(1, max_concurrent)
batch_size = max(1, batch_size)
max_tokens = get_config("performance.assets_max_tokens", DEFAULT_ASSETS_MAX_TOKENS)
try:
max_tokens = int(max_tokens)
except Exception:
max_tokens = 1000
max_concurrent = max(1, int(get_config("performance.assets_max_concurrent")))
batch_size = max(1, int(get_config("performance.assets_batch_size")))
max_tokens = int(get_config("performance.assets_max_tokens"))
truncated = False
original_count = 0
@@ -855,11 +765,7 @@ async def get_cache_stats_api(request: Request):
}
if selected_tokens:
selected_tokens = list(dict.fromkeys(selected_tokens))
original_count = len(selected_tokens)
if len(selected_tokens) > max_tokens:
selected_tokens = selected_tokens[:max_tokens]
truncated = True
selected_tokens, truncated, original_count = _truncate_tokens(selected_tokens, max_tokens, "Assets fetch")
total = 0
raw_results = await run_in_batches(
selected_tokens,
@@ -1027,26 +933,11 @@ async def load_online_cache_api_async(data: dict):
else:
raise HTTPException(status_code=400, detail="No tokens provided")
selected_tokens = list(dict.fromkeys(selected_tokens))
max_tokens = int(get_config("performance.assets_max_tokens"))
selected_tokens, truncated, original_count = _truncate_tokens(selected_tokens, max_tokens, "Assets load")
max_tokens = get_config("performance.assets_max_tokens", DEFAULT_ASSETS_MAX_TOKENS)
try:
max_tokens = int(max_tokens)
except Exception:
max_tokens = 1000
truncated = False
original_count = len(selected_tokens)
if len(selected_tokens) > max_tokens:
selected_tokens = selected_tokens[:max_tokens]
truncated = True
max_concurrent = get_config(
"performance.assets_max_concurrent", DEFAULT_ASSETS_MAX_CONCURRENT
)
batch_size = get_config(
"performance.assets_batch_size", DEFAULT_ASSETS_BATCH_SIZE
)
max_concurrent = get_config("performance.assets_max_concurrent")
batch_size = get_config("performance.assets_batch_size")
task = create_task(len(selected_tokens))
@@ -1220,36 +1111,12 @@ async def clear_online_cache_api(data: dict):
token_list = list(dict.fromkeys(token_list))
# 最大数量限制
max_tokens = get_config(
"performance.assets_max_tokens", DEFAULT_ASSETS_MAX_TOKENS
)
try:
max_tokens = int(max_tokens)
except Exception:
max_tokens = 1000
truncated = False
original_count = len(token_list)
if len(token_list) > max_tokens:
token_list = token_list[:max_tokens]
truncated = True
max_tokens = int(get_config("performance.assets_max_tokens"))
token_list, truncated, original_count = _truncate_tokens(token_list, max_tokens, "Clear online cache")
results = {}
max_concurrent = get_config(
"performance.assets_max_concurrent", DEFAULT_ASSETS_MAX_CONCURRENT
)
batch_size = get_config(
"performance.assets_batch_size", DEFAULT_ASSETS_BATCH_SIZE
)
try:
max_concurrent = int(max_concurrent)
except Exception:
max_concurrent = 25
try:
batch_size = int(batch_size)
except Exception:
batch_size = 10
max_concurrent = max(1, max_concurrent)
batch_size = max(1, batch_size)
max_concurrent = max(1, int(get_config("performance.assets_max_concurrent")))
batch_size = max(1, int(get_config("performance.assets_batch_size")))
async def _clear_one(t: str):
try:
@@ -1312,25 +1179,11 @@ async def clear_online_cache_api_async(data: dict):
if not token_list:
raise HTTPException(status_code=400, detail="No tokens provided")
token_list = list(dict.fromkeys(token_list))
max_tokens = int(get_config("performance.assets_max_tokens"))
token_list, truncated, original_count = _truncate_tokens(token_list, max_tokens, "Clear online cache async")
max_tokens = get_config("performance.assets_max_tokens", DEFAULT_ASSETS_MAX_TOKENS)
try:
max_tokens = int(max_tokens)
except Exception:
max_tokens = 1000
truncated = False
original_count = len(token_list)
if len(token_list) > max_tokens:
token_list = token_list[:max_tokens]
truncated = True
max_concurrent = get_config(
"performance.assets_max_concurrent", DEFAULT_ASSETS_MAX_CONCURRENT
)
batch_size = get_config(
"performance.assets_batch_size", DEFAULT_ASSETS_BATCH_SIZE
)
max_concurrent = get_config("performance.assets_max_concurrent")
batch_size = get_config("performance.assets_batch_size")
task = create_task(len(token_list))
app/api/v1/image.py
+70 -165
View File
@@ -7,6 +7,7 @@ import base64
import math
import random
import re
import time
from pathlib import Path
from typing import List, Optional, Union
@@ -31,12 +32,6 @@ from app.core.config import get_config
from app.core.logger import logger
DEFAULT_IMAGE_WS = False
DEFAULT_IMAGE_FORMAT = "url"
DEFAULT_IMAGE_WS_NSFW = True
DEFAULT_GROK_TEMPORARY = True
router = APIRouter(tags=["Images"])
class ImageGenerationRequest(BaseModel):
@@ -97,15 +92,13 @@ def _validate_common_request(
if allow_ws_stream:
# WS 流式仅支持 b64_json (base64 视为同义)
if (
request.stream
and get_config("grok.image_ws", DEFAULT_IMAGE_WS)
request.stream
and get_config("grok.image_ws")
and request.response_format
and request.response_format not in {"b64_json", "base64"}
):
raise ValidationException(
message=(
"Streaming with image_ws only supports response_format=b64_json/base64"
),
message="Streaming with image_ws only supports response_format=b64_json/base64",
param="response_format",
code="invalid_response_format",
)
@@ -145,25 +138,22 @@ def validate_generation_request(request: ImageGenerationRequest):
def resolve_response_format(response_format: Optional[str]) -> str:
fmt = response_format or get_config("app.image_format", DEFAULT_IMAGE_FORMAT)
"""解析响应格式"""
fmt = response_format or get_config("app.image_format")
if isinstance(fmt, str):
fmt = fmt.lower()
if fmt in ("b64_json", "base64", "url"):
return fmt
allowed_formats = {"b64_json", "base64", "url"}
raise ValidationException(
message=f"response_format must be one of {sorted(allowed_formats)}",
message="response_format must be one of b64_json, base64, url",
param="response_format",
code="invalid_response_format",
)
def response_field_name(response_format: str) -> str:
if response_format == "url":
return "url"
if response_format == "base64":
return "base64"
return "b64_json"
"""获取响应字段名"""
return {"url": "url", "base64": "base64"}.get(response_format, "b64_json")
def resolve_aspect_ratio(size: str) -> str:
@@ -215,6 +205,48 @@ def validate_edit_request(request: ImageEditRequest, images: List[UploadFile]):
)
def _get_effort(model_info) -> EffortType:
"""获取模型消耗级别"""
return EffortType.HIGH if (model_info and model_info.cost.value == "high") else EffortType.LOW
async def _wrap_stream_with_usage(stream, token_mgr, token, model_info):
"""包装流式响应，成功完成时记录使用"""
success = False
try:
async for chunk in stream:
yield chunk
success = True
finally:
if success:
try:
await token_mgr.consume(token, _get_effort(model_info))
except Exception as e:
logger.warning(f"Failed to consume token: {e}")
async def _get_token(model: str):
"""获取可用 token"""
token_mgr = await get_token_manager()
await token_mgr.reload_if_stale()
token = None
for pool_name in ModelService.pool_candidates_for_model(model):
token = token_mgr.get_token(pool_name)
if token:
break
if not token:
raise AppException(
message="No available tokens. Please try again later.",
error_type=ErrorType.RATE_LIMIT.value,
code="rate_limit_exceeded",
status_code=429,
)
return token_mgr, token
async def call_grok(
token_mgr,
token: str,
@@ -237,7 +269,6 @@ async def call_grok(
file_attachments=file_attachments,
)
# 收集图片
processor = ImageCollectProcessor(
model_info.model_id, token, response_format=response_format
)
@@ -249,15 +280,9 @@ async def call_grok(
logger.error(f"Grok image call failed: {e}")
return []
finally:
# 只在成功时记录使用，失败时不扣费（避免清零 fail_count）
if success:
try:
effort = (
EffortType.HIGH
if (model_info and model_info.cost.value == "high")
else EffortType.LOW
)
await token_mgr.consume(token, effort)
await token_mgr.consume(token, _get_effort(model_info))
except Exception as e:
logger.warning(f"Failed to consume token: {e}")
@@ -291,40 +316,16 @@ async def create_image(request: ImageGenerationRequest):
response_format = resolve_response_format(request.response_format)
response_field = response_field_name(response_format)
# 获取 token
try:
token_mgr = await get_token_manager()
await token_mgr.reload_if_stale()
token = None
for pool_name in ModelService.pool_candidates_for_model(request.model):
token = token_mgr.get_token(pool_name)
if token:
break
except Exception as e:
logger.error(f"Failed to get token: {e}")
raise AppException(
message="Internal service error obtaining token",
error_type=ErrorType.SERVER.value,
code="internal_error",
)
if not token:
raise AppException(
message="No available tokens. Please try again later.",
error_type=ErrorType.RATE_LIMIT.value,
code="rate_limit_exceeded",
status_code=429,
)
# 获取模型信息
# 获取 token 和模型信息
token_mgr, token = await _get_token(request.model)
model_info = ModelService.get(request.model)
use_ws = bool(get_config("grok.image_ws", DEFAULT_IMAGE_WS))
use_ws = bool(get_config("grok.image_ws"))
# 流式模式
if request.stream:
if use_ws:
aspect_ratio = resolve_aspect_ratio(request.size)
enable_nsfw = bool(get_config("grok.image_ws_nsfw", DEFAULT_IMAGE_WS_NSFW))
enable_nsfw = bool(get_config("grok.image_ws_nsfw"))
upstream = image_service.stream(
token=token,
prompt=request.prompt,
@@ -340,26 +341,8 @@ async def create_image(request: ImageGenerationRequest):
size=request.size,
)
async def _wrap_stream(stream):
success = False
try:
async for chunk in stream:
yield chunk
success = True
finally:
if success:
try:
effort = (
EffortType.HIGH
if (model_info and model_info.cost.value == "high")
else EffortType.LOW
)
await token_mgr.consume(token, effort)
except Exception as e:
logger.warning(f"Failed to consume token: {e}")
return StreamingResponse(
_wrap_stream(processor.process(upstream)),
_wrap_stream_with_usage(processor.process(upstream), token_mgr, token, model_info),
media_type="text/event-stream",
headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
)
@@ -377,28 +360,8 @@ async def create_image(request: ImageGenerationRequest):
model_info.model_id, token, n=request.n, response_format=response_format
)
# 包装流式响应，在成功完成时记录使用
async def _wrap_stream(stream):
success = False
try:
async for chunk in stream:
yield chunk
success = True
finally:
# 只在成功完成时扣费
if success:
try:
effort = (
EffortType.HIGH
if (model_info and model_info.cost.value == "high")
else EffortType.LOW
)
await token_mgr.consume(token, effort)
except Exception as e:
logger.warning(f"Failed to consume token: {e}")
return StreamingResponse(
_wrap_stream(processor.process(response)),
_wrap_stream_with_usage(processor.process(response), token_mgr, token, model_info),
media_type="text/event-stream",
headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
)
@@ -409,7 +372,7 @@ async def create_image(request: ImageGenerationRequest):
usage_override = None
if use_ws:
aspect_ratio = resolve_aspect_ratio(request.size)
enable_nsfw = bool(get_config("grok.image_ws_nsfw", DEFAULT_IMAGE_WS_NSFW))
enable_nsfw = bool(get_config("grok.image_ws_nsfw"))
all_images = []
seen = set()
expected_per_call = 6
@@ -452,12 +415,7 @@ async def create_image(request: ImageGenerationRequest):
if len(all_images) >= n:
break
try:
effort = (
EffortType.HIGH
if (model_info and model_info.cost.value == "high")
else EffortType.LOW
)
await token_mgr.consume(token, effort)
await token_mgr.consume(token, _get_effort(model_info))
except Exception as e:
logger.warning(f"Failed to consume token: {e}")
usage_override = {
@@ -510,16 +468,11 @@ async def create_image(request: ImageGenerationRequest):
selected_images.append("error")
# 构建响应
import time
data = [{response_field: img} for img in selected_images]
usage = usage_override or {
"total_tokens": 0
* len([img for img in selected_images if img != "error"]),
"total_tokens": 0,
"input_tokens": 0,
"output_tokens": 0
* len([img for img in selected_images if img != "error"]),
"output_tokens": 0,
"input_tokens_details": {"text_tokens": 0, "image_tokens": 0},
}
@@ -627,32 +580,8 @@ async def edit_image(
b64 = base64.b64encode(content).decode()
images.append(f"data:{mime};base64,{b64}")
# 获取 token
try:
token_mgr = await get_token_manager()
await token_mgr.reload_if_stale()
token = None
for pool_name in ModelService.pool_candidates_for_model(edit_request.model):
token = token_mgr.get_token(pool_name)
if token:
break
except Exception as e:
logger.error(f"Failed to get token: {e}")
raise AppException(
message="Internal service error obtaining token",
error_type=ErrorType.SERVER.value,
code="internal_error",
)
if not token:
raise AppException(
message="No available tokens. Please try again later.",
error_type=ErrorType.RATE_LIMIT.value,
code="rate_limit_exceeded",
status_code=429,
)
# 获取模型信息
# 获取 token 和模型信息
token_mgr, token = await _get_token(edit_request.model)
model_info = ModelService.get(edit_request.model)
# 上传图片
@@ -714,7 +643,7 @@ async def edit_image(
] = parent_post_id
raw_payload = {
"temporary": bool(get_config("grok.temporary", DEFAULT_GROK_TEMPORARY)),
"temporary": bool(get_config("grok.temporary")),
"modelName": model_info.grok_model,
"message": edit_request.prompt,
"enableImageGeneration": True,
@@ -749,26 +678,8 @@ async def edit_image(
model_info.model_id, token, n=edit_request.n, response_format=response_format
)
async def _wrap_stream(stream):
success = False
try:
async for chunk in stream:
yield chunk
success = True
finally:
if success:
try:
effort = (
EffortType.HIGH
if (model_info and model_info.cost.value == "high")
else EffortType.LOW
)
await token_mgr.consume(token, effort)
except Exception as e:
logger.warning(f"Failed to consume token: {e}")
return StreamingResponse(
_wrap_stream(processor.process(response)),
_wrap_stream_with_usage(processor.process(response), token_mgr, token, model_info),
media_type="text/event-stream",
headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
)
@@ -795,10 +706,7 @@ async def edit_image(
if calls_needed == 1:
all_images = await _call_edit()
else:
tasks = [
_call_edit()
for _ in range(calls_needed)
]
tasks = [_call_edit() for _ in range(calls_needed)]
results = await asyncio.gather(*tasks, return_exceptions=True)
all_images = []
@@ -808,6 +716,7 @@ async def edit_image(
elif isinstance(result, list):
all_images.extend(result)
# 选择图片
if len(all_images) >= n:
selected_images = random.sample(all_images, n)
else:
@@ -815,8 +724,6 @@ async def edit_image(
while len(selected_images) < n:
selected_images.append("error")
import time
data = [{response_field: img} for img in selected_images]
return JSONResponse(
@@ -824,11 +731,9 @@ async def edit_image(
"created": int(time.time()),
"data": data,
"usage": {
"total_tokens": 0
* len([img for img in selected_images if img != "error"]),
"total_tokens": 0,
"input_tokens": 0,
"output_tokens": 0
* len([img for img in selected_images if img != "error"]),
"output_tokens": 0,
"input_tokens_details": {"text_tokens": 0, "image_tokens": 0},
},
}
app/services/grok/services/assets.py
+3
View File
@@ -246,6 +246,8 @@ class UploadService(BaseService):
else:
filename, b64, mime = self.parse_b64(file_input)
logger.debug(f"Upload prepare: filename={filename}, type={mime}, size={len(b64)}")
if not b64:
raise ValidationException("Invalid file input: empty content")
@@ -353,6 +355,7 @@ class ListService(BaseService):
total = 0
async for page_assets in self.iter_assets(token):
total += len(page_assets)
logger.debug(f"Asset count: {total}")
return total
app/services/grok/services/chat.py
+9 -4
View File
@@ -197,7 +197,7 @@ class GrokChatService:
proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None
timeout = get_config("grok.timeout")
logger.debug(f"Payload: {payload}")
logger.debug(f"Chat request: model={model}, mode={mode}, stream={stream}, attachments={len(file_attachments or [])}")
# 建立连接
async def establish_connection():
@@ -220,8 +220,7 @@ class GrokChatService:
except Exception:
pass
logger.error(f"Chat failed: {response.status_code}",
extra={"status": response.status_code, "token": token[:10] + "..."})
logger.error(f"Chat failed: status={response.status_code}, token={token[:10]}...")
await session.close()
raise UpstreamException(
@@ -229,6 +228,7 @@ class GrokChatService:
details={"status": response.status_code, "body": content},
)
logger.info(f"Chat connected: model={model}, stream={stream}")
return session, response
except UpstreamException:
@@ -287,6 +287,7 @@ class GrokChatService:
# 提取消息和附件
try:
message, attachments = MessageExtractor.extract(request.messages, is_video=is_video)
logger.debug(f"Extracted message length={len(message)}, attachments={len(attachments)}")
except ValueError as e:
raise ValidationException(str(e))
@@ -295,9 +296,10 @@ class GrokChatService:
if attachments:
upload_service = UploadService()
try:
for _, attach_data in attachments:
for attach_type, attach_data in attachments:
file_id, _ = await upload_service.upload(attach_data, token)
file_ids.append(file_id)
logger.debug(f"Attachment uploaded: type={attach_type}, file_id={file_id}")
finally:
await upload_service.close()
@@ -349,15 +351,18 @@ class ChatService:
# 处理响应
if is_stream:
logger.debug(f"Processing stream response: model={model}")
processor = StreamProcessor(model_name, token, think)
return wrap_stream_with_usage(processor.process(response), token_mgr, token, model)
# 非流式
logger.debug(f"Processing non-stream response: model={model}")
result = await CollectProcessor(model_name, token).process(response)
try:
model_info = ModelService.get(model)
effort = EffortType.HIGH if (model_info and model_info.cost.value == "high") else EffortType.LOW
await token_mgr.consume(token, effort)
logger.info(f"Chat completed: model={model}, effort={effort.value}")
except Exception as e:
logger.warning(f"Failed to record usage: {e}")
return result
app/services/grok/services/image.py
+2
View File
@@ -93,6 +93,7 @@ class ImageService:
max_retries: int = None,
) -> AsyncGenerator[Dict[str, object], None]:
retries = max(1, max_retries if max_retries is not None else 1)
logger.info(f"Image generation: prompt='{prompt[:50]}...', n={n}, ratio={aspect_ratio}, nsfw={enable_nsfw}")
for attempt in range(retries):
try:
@@ -199,6 +200,7 @@ class ImageService:
if info["is_final"] and not existing.get("is_final"):
completed += 1
logger.debug(f"Final image received: id={image_id}, size={info['blob_size']}")
images[image_id] = {"is_final": info["is_final"] or existing.get("is_final")}
yield info
app/services/grok/services/media.py
+5 -1
View File
@@ -183,12 +183,14 @@ class VideoService:
)
if response.status_code != 200:
logger.error(f"Video generation failed: {response.status_code}")
logger.error(f"Video generation failed: status={response.status_code}, post_id={post_id}")
raise UpstreamException(
message=f"Video generation failed: {response.status_code}",
details={"status": response.status_code},
)
logger.info(f"Video generation started: post_id={post_id}")
async def stream_response():
try:
async for line in response.aiter_lines():
@@ -217,6 +219,7 @@ class VideoService:
preset: str = "normal"
) -> AsyncGenerator[bytes, None]:
"""生成视频"""
logger.info(f"Video generation: prompt='{prompt[:50]}...', ratio={aspect_ratio}, length={video_length}s, preset={preset}")
async with _get_semaphore():
post_id = await self.create_post(token, prompt)
return await self._generate_internal(token, post_id, prompt, aspect_ratio, video_length, resolution_name, preset)
@@ -229,6 +232,7 @@ class VideoService:
preset: str = "normal"
) -> AsyncGenerator[bytes, None]:
"""从图片生成视频"""
logger.info(f"Image to video: prompt='{prompt[:50]}...', image={image_url[:80]}")
async with _get_semaphore():
post_id = await self.create_image_post(token, image_url)
return await self._generate_internal(token, post_id, prompt, aspect_ratio, video_length, resolution, preset)
app/services/grok/services/usage.py
+2 -2
View File
@@ -112,10 +112,10 @@ class UsageService:
if response.status_code == 200:
data = response.json()
remaining = data.get("remainingTokens", 0)
logger.info(f"Usage: quota {remaining} remaining")
logger.info(f"Usage sync success: remaining={remaining}, token={token[:10]}...")
return data
logger.error(f"Usage failed: {response.status_code}")
logger.error(f"Usage sync failed: status={response.status_code}, token={token[:10]}...")
raise UpstreamException(
message=f"Failed to get usage stats: {response.status_code}",
app/services/grok/services/voice.py
+6 -5
View File
@@ -35,6 +35,7 @@ class VoiceService:
Returns:
Dict containing token and livekitUrl
"""
logger.debug(f"Voice token request: voice={voice}, personality={personality}, speed={speed}")
headers = self._build_headers(token)
payload = self._build_payload(voice, personality, speed)
@@ -53,16 +54,16 @@ class VoiceService:
)
if response.status_code != 200:
logger.error(
f"Voice token failed: {response.status_code}",
extra={"response": response.text[:200]}
)
body = response.text[:200]
logger.error(f"Voice token failed: status={response.status_code}, body={body}")
raise UpstreamException(
message=f"Failed to get voice token: {response.status_code}",
details={"status": response.status_code, "body": response.text}
)
return response.json()
result = response.json()
logger.info(f"Voice token obtained: voice={voice}")
return result
except Exception as e:
logger.error(f"Voice service error: {e}")