diff --git a/app/api/v1/admin.py b/app/api/v1/admin.py index c2dc480..8c05b05 100644 --- a/app/api/v1/admin.py +++ b/app/api/v1/admin.py @@ -16,7 +16,11 @@ 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 +from app.services.grok.batch_services import ( + BatchUsageService, + BatchNSFWService, + BatchAssetsService, +) import os import time import uuid @@ -27,10 +31,8 @@ import orjson from app.core.logger import logger from app.api.v1.image import resolve_aspect_ratio from app.services.grok.services.voice import VoiceService -from app.services.grok.services.image import image_service +from app.services.grok.services.image import ImageGenerationService from app.services.grok.models.model import ModelService -from app.services.grok.processors.image_ws_processors import ImageWSCollectProcessor -from app.services.token import EffortType TEMPLATE_DIR = Path(__file__).parent.parent.parent / "static" @@ -363,7 +365,6 @@ async def admin_imagine_ws(websocket: WebSocket): return token_mgr = await get_token_manager() - enable_nsfw = bool(get_config("image.image_ws_nsfw", True)) sequence = 0 run_id = uuid.uuid4().hex @@ -399,26 +400,23 @@ async def admin_imagine_ws(websocket: WebSocket): await asyncio.sleep(2) continue - upstream = image_service.stream( + start_at = time.time() + result = await ImageGenerationService().generate( + token_mgr=token_mgr, token=token, + model_info=model_info, prompt=prompt, - aspect_ratio=aspect_ratio, - n=6, - enable_nsfw=enable_nsfw, - ) - - processor = ImageWSCollectProcessor( - model_info.model_id, - token, n=6, response_format="b64_json", + size="1024x1024", + aspect_ratio=aspect_ratio, + stream=False, + use_ws=True, ) - - start_at = time.time() - images = await processor.process(upstream) elapsed_ms = int((time.time() - start_at) * 1000) - if images and all(img and img != "error" for img in images): + images = [img for img in result.data if img and img != "error"] + if images: # 一次发送所有 6 张图片 for img_b64 in images: sequence += 1 @@ -433,17 +431,6 @@ async def admin_imagine_ws(websocket: WebSocket): "run_id": run_id, } ) - - # 消耗 token（6 张图片按高成本计算） - 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}") else: await _send( { @@ -602,7 +589,6 @@ async def admin_imagine_sse( return token_mgr = await get_token_manager() - enable_nsfw = bool(get_config("image.image_ws_nsfw", True)) sequence = 0 run_id = uuid.uuid4().hex @@ -645,26 +631,23 @@ async def admin_imagine_sse( await asyncio.sleep(2) continue - upstream = image_service.stream( + start_at = time.time() + result = await ImageGenerationService().generate( + token_mgr=token_mgr, token=token, + model_info=model_info, prompt=prompt, - aspect_ratio=ratio, - n=6, - enable_nsfw=enable_nsfw, - ) - - processor = ImageWSCollectProcessor( - model_info.model_id, - token, n=6, response_format="b64_json", + size="1024x1024", + aspect_ratio=ratio, + stream=False, + use_ws=True, ) - - start_at = time.time() - images = await processor.process(upstream) elapsed_ms = int((time.time() - start_at) * 1000) - if images and all(img and img != "error" for img in images): + images = [img for img in result.data if img and img != "error"] + if images: for img_b64 in images: sequence += 1 yield _sse_event( @@ -678,16 +661,6 @@ async def admin_imagine_sse( "run_id": run_id, } ) - - 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}") else: yield _sse_event( { @@ -863,12 +836,9 @@ async def refresh_tokens_api(data: dict): 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(t, consume_on_fail=False, is_usage=False) - - raw_results = await run_in_batches( + raw_results = await BatchUsageService.refresh( unique_tokens, - _refresh_one, + mgr, max_concurrent=max_concurrent, batch_size=batch_size, ) @@ -915,15 +885,12 @@ async def refresh_tokens_api_async(data: dict): async def _run(): try: - async def _refresh_one(t: str): - return await mgr.sync_usage(t, consume_on_fail=False, is_usage=False) - async def _on_item(item: str, res: dict): task.record(bool(res.get("ok"))) - raw_results = await run_in_batches( + raw_results = await BatchUsageService.refresh( unique_tokens, - _refresh_one, + mgr, max_concurrent=max_concurrent, batch_size=batch_size, on_item=_on_item, @@ -979,11 +946,8 @@ async def refresh_tokens_api_async(data: dict): @router.post("/api/v1/admin/tokens/nsfw/enable", dependencies=[Depends(verify_api_key)]) async def enable_nsfw_api(data: dict): """批量开启 NSFW (Unhinged) 模式""" - from app.services.grok.services.nsfw import NSFWService - try: mgr = await get_token_manager() - nsfw_service = NSFWService() # 收集 token 列表 tokens = _collect_tokens(data) @@ -1010,23 +974,11 @@ async def enable_nsfw_api(data: dict): max_concurrent = get_config("performance.nsfw_max_concurrent") batch_size = get_config("performance.nsfw_batch_size") - # 定义 worker - async def _enable(token: str): - result = await nsfw_service.enable(token) - # 成功后添加 nsfw tag - if result.success: - await mgr.add_tag(token, "nsfw") - return { - "success": result.success, - "http_status": result.http_status, - "grpc_status": result.grpc_status, - "grpc_message": result.grpc_message, - "error": result.error, - } - - # 执行批量操作 - raw_results = await run_in_batches( - unique_tokens, _enable, max_concurrent=max_concurrent, batch_size=batch_size + raw_results = await BatchNSFWService.enable( + unique_tokens, + mgr, + max_concurrent=max_concurrent, + batch_size=batch_size, ) # 构造返回结果（mask token） @@ -1073,10 +1025,7 @@ 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 - mgr = await get_token_manager() - nsfw_service = NSFWService() tokens = _collect_tokens(data) @@ -1103,25 +1052,13 @@ async def enable_nsfw_api_async(data: dict): async def _run(): try: - async def _enable(token: str): - result = await nsfw_service.enable(token) - if result.success: - await mgr.add_tag(token, "nsfw") - return { - "success": result.success, - "http_status": result.http_status, - "grpc_status": result.grpc_status, - "grpc_message": result.grpc_message, - "error": result.error, - } - async def _on_item(item: str, res: dict): ok = bool(res.get("ok") and res.get("data", {}).get("success")) task.record(ok) - raw_results = await run_in_batches( + raw_results = await BatchNSFWService.enable( unique_tokens, - _enable, + mgr, max_concurrent=max_concurrent, batch_size=batch_size, on_item=_on_item, @@ -1184,9 +1121,8 @@ async def admin_cache_page(): @router.get("/api/v1/admin/cache", dependencies=[Depends(verify_api_key)]) async def get_cache_stats_api(request: Request): """获取缓存统计""" - from app.services.grok.services.assets import DownloadService, ListService + from app.services.grok.utils.download import DownloadService from app.services.token.manager import get_token_manager - from app.services.grok.utils.batch import run_in_batches try: dl_service = DownloadService() @@ -1238,51 +1174,14 @@ async def get_cache_stats_api(request: Request): truncated = False original_count = 0 - async def _fetch_assets(token: str): - list_service = ListService() - try: - return await list_service.count(token) - finally: - await list_service.close() - - async def _fetch_detail(token: str): - account = account_map.get(token) - try: - count = await _fetch_assets(token) - return { - "detail": { - "token": token, - "token_masked": account["token_masked"] if account else token, - "count": count, - "status": "ok", - "last_asset_clear_at": account["last_asset_clear_at"] - if account - else None, - }, - "count": count, - } - except Exception as e: - return { - "detail": { - "token": token, - "token_masked": account["token_masked"] if account else token, - "count": 0, - "status": f"error: {str(e)}", - "last_asset_clear_at": account["last_asset_clear_at"] - if account - else None, - }, - "count": 0, - } - if selected_tokens: selected_tokens, truncated, original_count = _truncate_tokens( selected_tokens, max_tokens, "Assets fetch" ) total = 0 - raw_results = await run_in_batches( + raw_results = await BatchAssetsService.fetch_details( selected_tokens, - _fetch_detail, + account_map, max_concurrent=max_concurrent, batch_size=batch_size, ) @@ -1318,9 +1217,9 @@ async def get_cache_stats_api(request: Request): if len(tokens) > max_tokens: tokens = tokens[:max_tokens] truncated = True - raw_results = await run_in_batches( + raw_results = await BatchAssetsService.fetch_details( tokens, - _fetch_detail, + account_map, max_concurrent=max_concurrent, batch_size=batch_size, ) @@ -1351,23 +1250,28 @@ async def get_cache_stats_api(request: Request): else: token = selected_token if token: - try: - count = await _fetch_assets(token) - match = next((a for a in accounts if a["token"] == token), None) + raw_results = await BatchAssetsService.fetch_details( + [token], + account_map, + max_concurrent=1, + batch_size=1, + ) + res = raw_results.get(token, {}) + data = res.get("data", {}) + detail = data.get("detail") if res.get("ok") else None + if detail: online_stats = { - "count": count, - "status": "ok", - "token": token, - "token_masked": match["token_masked"] if match else token, - "last_asset_clear_at": match["last_asset_clear_at"] - if match - else None, + "count": data.get("count", 0), + "status": detail.get("status", "ok"), + "token": detail.get("token"), + "token_masked": detail.get("token_masked"), + "last_asset_clear_at": detail.get("last_asset_clear_at"), } - except Exception as e: + else: match = next((a for a in accounts if a["token"] == token), None) online_stats = { "count": 0, - "status": f"error: {str(e)}", + "status": f"error: {res.get('error')}", "token": token, "token_masked": match["token_masked"] if match else token, "last_asset_clear_at": match["last_asset_clear_at"] @@ -1404,9 +1308,8 @@ async def get_cache_stats_api(request: Request): ) async def load_online_cache_api_async(data: dict): """在线资产统计（异步批量 + SSE 进度）""" - from app.services.grok.services.assets import DownloadService, ListService + from app.services.grok.utils.download import DownloadService from app.services.token.manager import get_token_manager - from app.services.grok.utils.batch import run_in_batches mgr = await get_token_manager() @@ -1462,44 +1365,16 @@ async def load_online_cache_api_async(data: dict): image_stats = dl_service.get_stats("image") video_stats = dl_service.get_stats("video") - async def _fetch_detail(token: str): - account = account_map.get(token) - list_service = ListService() - try: - count = await list_service.count(token) - detail = { - "token": token, - "token_masked": account["token_masked"] if account else token, - "count": count, - "status": "ok", - "last_asset_clear_at": account["last_asset_clear_at"] - if account - else None, - } - return {"ok": True, "detail": detail, "count": count} - except Exception as e: - detail = { - "token": token, - "token_masked": account["token_masked"] if account else token, - "count": 0, - "status": f"error: {str(e)}", - "last_asset_clear_at": account["last_asset_clear_at"] - if account - else None, - } - return {"ok": False, "detail": detail, "count": 0} - finally: - await list_service.close() - async def _on_item(item: str, res: dict): ok = bool(res.get("data", {}).get("ok")) task.record(ok) - raw_results = await run_in_batches( + raw_results = await BatchAssetsService.fetch_details( selected_tokens, - _fetch_detail, + account_map, max_concurrent=max_concurrent, batch_size=batch_size, + include_ok=True, on_item=_on_item, should_cancel=lambda: task.cancelled, ) @@ -1555,7 +1430,7 @@ async def load_online_cache_api_async(data: dict): @router.post("/api/v1/admin/cache/clear", dependencies=[Depends(verify_api_key)]) async def clear_local_cache_api(data: dict): """清理本地缓存""" - from app.services.grok.services.assets import DownloadService + from app.services.grok.utils.download import DownloadService cache_type = data.get("type", "image") @@ -1575,7 +1450,7 @@ async def list_local_cache_api( page_size: int = 1000, ): """列出本地缓存文件""" - from app.services.grok.services.assets import DownloadService + from app.services.grok.utils.download import DownloadService try: if type_: @@ -1590,7 +1465,7 @@ async def list_local_cache_api( @router.post("/api/v1/admin/cache/item/delete", dependencies=[Depends(verify_api_key)]) async def delete_local_cache_item_api(data: dict): """删除单个本地缓存文件""" - from app.services.grok.services.assets import DownloadService + from app.services.grok.utils.download import DownloadService cache_type = data.get("type", "image") name = data.get("name") @@ -1607,15 +1482,10 @@ async def delete_local_cache_item_api(data: dict): @router.post("/api/v1/admin/cache/online/clear", dependencies=[Depends(verify_api_key)]) async def clear_online_cache_api(data: dict): """清理在线缓存""" - from app.services.grok.services.assets import DeleteService from app.services.token.manager import get_token_manager - from app.services.grok.utils.batch import run_in_batches - - delete_service = None try: mgr = await get_token_manager() tokens = data.get("tokens") - delete_service = DeleteService() if isinstance(tokens, list): token_list = [t.strip() for t in tokens if isinstance(t, str) and t.strip()] @@ -1637,17 +1507,9 @@ async def clear_online_cache_api(data: dict): ) batch_size = max(1, int(get_config("performance.assets_batch_size"))) - async def _clear_one(t: str): - try: - result = await delete_service.delete_all(t) - await mgr.mark_asset_clear(t) - return {"status": "success", "result": result} - except Exception as e: - return {"status": "error", "error": str(e)} - - raw_results = await run_in_batches( + raw_results = await BatchAssetsService.clear_online( token_list, - _clear_one, + mgr, max_concurrent=max_concurrent, batch_size=batch_size, ) @@ -1670,14 +1532,19 @@ async def clear_online_cache_api(data: dict): status_code=400, detail="No available token to perform cleanup" ) - result = await delete_service.delete_all(token) - await mgr.mark_asset_clear(token) - return {"status": "success", "result": result} + raw_results = await BatchAssetsService.clear_online( + [token], + mgr, + max_concurrent=1, + batch_size=1, + ) + res = raw_results.get(token, {}) + data = res.get("data", {}) + if res.get("ok") and data.get("status") == "success": + return {"status": "success", "result": data.get("result")} + return {"status": "error", "error": data.get("error") or res.get("error")} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) - finally: - if delete_service: - await delete_service.close() @router.post( @@ -1685,9 +1552,7 @@ async def clear_online_cache_api(data: dict): ) async def clear_online_cache_api_async(data: dict): """清理在线缓存（异步批量 + SSE 进度）""" - from app.services.grok.services.assets import DeleteService 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 = data.get("tokens") @@ -1709,26 +1574,17 @@ async def clear_online_cache_api_async(data: dict): task = create_task(len(token_list)) async def _run(): - delete_service = DeleteService() try: - - async def _clear_one(t: str): - try: - result = await delete_service.delete_all(t) - await mgr.mark_asset_clear(t) - return {"ok": True, "result": result} - except Exception as e: - return {"ok": False, "error": str(e)} - async def _on_item(item: str, res: dict): ok = bool(res.get("data", {}).get("ok")) task.record(ok) - raw_results = await run_in_batches( + raw_results = await BatchAssetsService.clear_online( token_list, - _clear_one, + mgr, max_concurrent=max_concurrent, batch_size=batch_size, + include_ok=True, on_item=_on_item, should_cancel=lambda: task.cancelled, ) @@ -1767,7 +1623,6 @@ async def clear_online_cache_api_async(data: dict): except Exception as e: task.fail_task(str(e)) finally: - await delete_service.close() asyncio.create_task(expire_task(task.id, 300)) asyncio.create_task(_run()) diff --git a/app/api/v1/chat.py b/app/api/v1/chat.py index 6c42056..ba5a820 100644 --- a/app/api/v1/chat.py +++ b/app/api/v1/chat.py @@ -260,7 +260,7 @@ async def chat_completions(request: ChatCompletionRequest): # 检测视频模型 model_info = ModelService.get(request.model) if model_info and model_info.is_video: - from app.services.grok.services.media import VideoService + from app.services.grok.services.video import VideoService # 提取视频配置 (默认值在 Pydantic 模型中处理) v_conf = request.video_config or VideoConfig() diff --git a/app/api/v1/image.py b/app/api/v1/image.py index bc47b16..6e8ab67 100644 --- a/app/api/v1/image.py +++ b/app/api/v1/image.py @@ -2,11 +2,7 @@ Image Generation API 路由 """ -import asyncio import base64 -import math -import random -import re import time from pathlib import Path from typing import List, Optional, Union @@ -15,21 +11,12 @@ from fastapi import APIRouter, File, Form, UploadFile from fastapi.responses import StreamingResponse, JSONResponse from pydantic import BaseModel, Field, ValidationError -from app.services.grok.services.chat import GrokChatService -from app.services.grok.services.image import image_service -from app.services.grok.services.assets import UploadService -from app.services.grok.services.media import VideoService +from app.services.grok.services.image import ImageGenerationService +from app.services.grok.services.image_edit import ImageEditService from app.services.grok.models.model import ModelService -from app.services.grok.processors import ( - ImageStreamProcessor, - ImageCollectProcessor, - ImageWSStreamProcessor, - ImageWSCollectProcessor, -) -from app.services.token import get_token_manager, EffortType +from app.services.token import get_token_manager from app.core.exceptions import ValidationException, AppException, ErrorType from app.core.config import get_config -from app.core.logger import logger router = APIRouter(tags=["Images"]) @@ -202,30 +189,6 @@ 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() @@ -248,46 +211,6 @@ async def _get_token(model: str): return token_mgr, token -async def call_grok( - token_mgr, - token: str, - prompt: str, - model_info, - file_attachments: Optional[List[str]] = None, - response_format: str = "b64_json", -) -> List[str]: - """调用 Grok 获取图片，返回 base64 列表""" - chat_service = GrokChatService() - success = False - - try: - response = await chat_service.chat( - token=token, - message=prompt, - model=model_info.grok_model, - mode=model_info.model_mode, - stream=True, - file_attachments=file_attachments, - ) - - processor = ImageCollectProcessor( - model_info.model_id, token, response_format=response_format - ) - images = await processor.process(response) - success = True - return images - - except Exception as e: - logger.error(f"Grok image call failed: {e}") - return [] - 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}") - - @router.post("/images/generations") async def create_image(request: ImageGenerationRequest): """ @@ -321,160 +244,30 @@ async def create_image(request: ImageGenerationRequest): token_mgr, token = await _get_token(request.model) model_info = ModelService.get(request.model) use_ws = bool(get_config("image.image_ws")) + aspect_ratio = resolve_aspect_ratio(request.size) - # 流式模式 - if request.stream: - if use_ws: - aspect_ratio = resolve_aspect_ratio(request.size) - enable_nsfw = bool(get_config("image.image_ws_nsfw")) - upstream = image_service.stream( - token=token, - prompt=request.prompt, - aspect_ratio=aspect_ratio, - n=request.n, - enable_nsfw=enable_nsfw, - ) - processor = ImageWSStreamProcessor( - model_info.model_id, - token, - n=request.n, - response_format=response_format, - size=request.size, - ) - - return StreamingResponse( - _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"}, - ) - - chat_service = GrokChatService() - response = await chat_service.chat( - token=token, - message=f"Image Generation: {request.prompt}", - model=model_info.grok_model, - mode=model_info.model_mode, - stream=True, - ) - - processor = ImageStreamProcessor( - model_info.model_id, token, n=request.n, response_format=response_format - ) + result = await ImageGenerationService().generate( + token_mgr=token_mgr, + token=token, + model_info=model_info, + prompt=request.prompt, + n=request.n, + response_format=response_format, + size=request.size, + aspect_ratio=aspect_ratio, + stream=bool(request.stream), + use_ws=use_ws, + ) + if result.stream: return StreamingResponse( - _wrap_stream_with_usage( - processor.process(response), token_mgr, token, model_info - ), + result.data, media_type="text/event-stream", headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, ) - # 非流式模式 - n = request.n - - usage_override = None - if use_ws: - aspect_ratio = resolve_aspect_ratio(request.size) - enable_nsfw = bool(get_config("image.image_ws_nsfw")) - all_images = [] - seen = set() - expected_per_call = 6 - calls_needed = max(1, math.ceil(n / expected_per_call)) - calls_needed = min(calls_needed, n) - - async def _fetch_batch(call_target: int): - upstream = image_service.stream( - token=token, - prompt=request.prompt, - aspect_ratio=aspect_ratio, - n=call_target, - enable_nsfw=enable_nsfw, - ) - processor = ImageWSCollectProcessor( - model_info.model_id, - token, - n=call_target, - response_format=response_format, - ) - return await processor.process(upstream) - - tasks = [] - for i in range(calls_needed): - remaining = n - (i * expected_per_call) - call_target = min(expected_per_call, remaining) - tasks.append(_fetch_batch(call_target)) - - results = await asyncio.gather(*tasks, return_exceptions=True) - for batch in results: - if isinstance(batch, Exception): - logger.warning(f"WS batch failed: {batch}") - continue - for img in batch: - if img not in seen: - seen.add(img) - all_images.append(img) - if len(all_images) >= n: - break - if len(all_images) >= n: - break - try: - await token_mgr.consume(token, _get_effort(model_info)) - except Exception as e: - logger.warning(f"Failed to consume token: {e}") - usage_override = { - "total_tokens": 0, - "input_tokens": 0, - "output_tokens": 0, - "input_tokens_details": {"text_tokens": 0, "image_tokens": 0}, - } - else: - calls_needed = (n + 1) // 2 - - if calls_needed == 1: - # 单次调用 - all_images = await call_grok( - token_mgr, - token, - f"Image Generation: {request.prompt}", - model_info, - response_format=response_format, - ) - else: - # 并发调用 - tasks = [ - call_grok( - token_mgr, - token, - f"Image Generation: {request.prompt}", - model_info, - response_format=response_format, - ) - for _ in range(calls_needed) - ] - results = await asyncio.gather(*tasks, return_exceptions=True) - - # 收集成功的图片 - all_images = [] - for result in results: - if isinstance(result, Exception): - logger.error(f"Concurrent call failed: {result}") - elif isinstance(result, list): - all_images.extend(result) - - # 随机选取 n 张图片 - if len(all_images) >= n: - selected_images = random.sample(all_images, n) - else: - # 全部返回，error 填充缺失 - selected_images = all_images.copy() - while len(selected_images) < n: - selected_images.append("error") - - # 构建响应 - data = [{response_field: img} for img in selected_images] - usage = usage_override or { + data = [{response_field: img} for img in result.data] + usage = result.usage_override or { "total_tokens": 0, "input_tokens": 0, "output_tokens": 0, @@ -588,134 +381,25 @@ async def edit_image( token_mgr, token = await _get_token(edit_request.model) model_info = ModelService.get(edit_request.model) - # 上传图片 - image_urls: List[str] = [] - upload_service = UploadService() - try: - for image in images: - file_id, file_uri = await upload_service.upload(image, token) - if file_uri: - if file_uri.startswith("http"): - image_urls.append(file_uri) - else: - image_urls.append(f"https://assets.grok.com/{file_uri.lstrip('/')}") - finally: - await upload_service.close() - - if not image_urls: - raise AppException( - message="Image upload failed", - error_type=ErrorType.SERVER.value, - code="upload_failed", - ) - - parent_post_id = None - try: - media_service = VideoService() - parent_post_id = await media_service.create_image_post(token, image_urls[0]) - logger.debug(f"Parent post ID: {parent_post_id}") - except Exception as e: - logger.warning(f"Create image post failed: {e}") - - if not parent_post_id: - for url in image_urls: - match = re.search(r"/generated/([a-f0-9-]+)/", url) - if match: - parent_post_id = match.group(1) - logger.debug(f"Parent post ID: {parent_post_id}") - break - match = re.search(r"/users/[^/]+/([a-f0-9-]+)/content", url) - if match: - parent_post_id = match.group(1) - logger.debug(f"Parent post ID: {parent_post_id}") - break - - model_config_override = { - "modelMap": { - "imageEditModel": "imagine", - "imageEditModelConfig": { - "imageReferences": image_urls, - }, - } - } - - if parent_post_id: - model_config_override["modelMap"]["imageEditModelConfig"]["parentPostId"] = ( - parent_post_id - ) - - tool_overrides = {"imageGen": True} - - # 流式模式 - if edit_request.stream: - chat_service = GrokChatService() - response = await chat_service.chat( - token=token, - message=edit_request.prompt, - model=model_info.grok_model, - mode=None, - stream=True, - tool_overrides=tool_overrides, - model_config_override=model_config_override, - ) - - processor = ImageStreamProcessor( - model_info.model_id, - token, - n=edit_request.n, - response_format=response_format, - ) + result = await ImageEditService().edit( + token_mgr=token_mgr, + token=token, + model_info=model_info, + prompt=edit_request.prompt, + images=images, + n=edit_request.n, + response_format=response_format, + stream=bool(edit_request.stream), + ) + if result.stream: return StreamingResponse( - _wrap_stream_with_usage( - processor.process(response), token_mgr, token, model_info - ), + result.data, media_type="text/event-stream", headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, ) - # 非流式模式 - n = edit_request.n - calls_needed = (n + 1) // 2 - - async def _call_edit(): - chat_service = GrokChatService() - response = await chat_service.chat( - token=token, - message=edit_request.prompt, - model=model_info.grok_model, - mode=None, - stream=True, - tool_overrides=tool_overrides, - model_config_override=model_config_override, - ) - processor = ImageCollectProcessor( - model_info.model_id, token, response_format=response_format - ) - return await processor.process(response) - - if calls_needed == 1: - all_images = await _call_edit() - else: - tasks = [_call_edit() for _ in range(calls_needed)] - results = await asyncio.gather(*tasks, return_exceptions=True) - - all_images = [] - for result in results: - if isinstance(result, Exception): - logger.error(f"Concurrent call failed: {result}") - elif isinstance(result, list): - all_images.extend(result) - - # 选择图片 - if len(all_images) >= n: - selected_images = random.sample(all_images, n) - else: - selected_images = all_images.copy() - while len(selected_images) < n: - selected_images.append("error") - - data = [{response_field: img} for img in selected_images] + data = [{response_field: img} for img in result.data] return JSONResponse( content={ diff --git a/app/services/grok/batch_services/__init__.py b/app/services/grok/batch_services/__init__.py new file mode 100644 index 0000000..4979637 --- /dev/null +++ b/app/services/grok/batch_services/__init__.py @@ -0,0 +1,7 @@ +"""Batch services.""" + +from .usage import BatchUsageService +from .nsfw import BatchNSFWService +from .assets import BatchAssetsService + +__all__ = ["BatchUsageService", "BatchNSFWService", "BatchAssetsService"] diff --git a/app/services/grok/batch_services/assets.py b/app/services/grok/batch_services/assets.py new file mode 100644 index 0000000..3ad213e --- /dev/null +++ b/app/services/grok/batch_services/assets.py @@ -0,0 +1,107 @@ +""" +Batch assets service. +""" + +from typing import Callable, Awaitable, Dict, Any, Optional + +from app.services.grok.services.assets import ListService, DeleteService +from app.services.grok.utils.batch import run_in_batches + + +class BatchAssetsService: + """Batch assets orchestration.""" + + @staticmethod + async def fetch_details( + tokens: list[str], + account_map: Dict[str, Dict[str, Any]], + *, + max_concurrent: int, + batch_size: int, + include_ok: bool = False, + on_item: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None, + should_cancel: Optional[Callable[[], bool]] = None, + ) -> Dict[str, Dict[str, Any]]: + account_map = account_map or {} + + async def _fetch_detail(token: str): + account = account_map.get(token) + list_service = ListService() + try: + count = await list_service.count(token) + detail = { + "token": token, + "token_masked": account["token_masked"] if account else token, + "count": count, + "status": "ok", + "last_asset_clear_at": account["last_asset_clear_at"] + if account + else None, + } + if include_ok: + return {"ok": True, "detail": detail, "count": count} + return {"detail": detail, "count": count} + except Exception as e: + detail = { + "token": token, + "token_masked": account["token_masked"] if account else token, + "count": 0, + "status": f"error: {str(e)}", + "last_asset_clear_at": account["last_asset_clear_at"] + if account + else None, + } + if include_ok: + return {"ok": False, "detail": detail, "count": 0} + return {"detail": detail, "count": 0} + finally: + await list_service.close() + + return await run_in_batches( + tokens, + _fetch_detail, + max_concurrent=max_concurrent, + batch_size=batch_size, + on_item=on_item, + should_cancel=should_cancel, + ) + + @staticmethod + async def clear_online( + tokens: list[str], + mgr, + *, + max_concurrent: int, + batch_size: int, + include_ok: bool = False, + on_item: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None, + should_cancel: Optional[Callable[[], bool]] = None, + ) -> Dict[str, Dict[str, Any]]: + delete_service = DeleteService() + + async def _clear_one(token: str): + try: + result = await delete_service.delete_all(token) + await mgr.mark_asset_clear(token) + if include_ok: + return {"ok": True, "result": result} + return {"status": "success", "result": result} + except Exception as e: + if include_ok: + return {"ok": False, "error": str(e)} + return {"status": "error", "error": str(e)} + + try: + return await run_in_batches( + tokens, + _clear_one, + max_concurrent=max_concurrent, + batch_size=batch_size, + on_item=on_item, + should_cancel=should_cancel, + ) + finally: + await delete_service.close() + + +__all__ = ["BatchAssetsService"] diff --git a/app/services/grok/batch_services/nsfw.py b/app/services/grok/batch_services/nsfw.py new file mode 100644 index 0000000..3c8ae06 --- /dev/null +++ b/app/services/grok/batch_services/nsfw.py @@ -0,0 +1,48 @@ +""" +Batch NSFW service. +""" + +from typing import Callable, Awaitable, Dict, Any, Optional + +from app.services.grok.services.nsfw import NSFWService +from app.services.grok.utils.batch import run_in_batches + + +class BatchNSFWService: + """Batch NSFW orchestration.""" + + @staticmethod + async def enable( + tokens: list[str], + mgr, + *, + max_concurrent: int, + batch_size: int, + on_item: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None, + should_cancel: Optional[Callable[[], bool]] = None, + ) -> Dict[str, Dict[str, Any]]: + nsfw_service = NSFWService() + + async def _enable(token: str): + result = await nsfw_service.enable(token) + if result.success: + await mgr.add_tag(token, "nsfw") + return { + "success": result.success, + "http_status": result.http_status, + "grpc_status": result.grpc_status, + "grpc_message": result.grpc_message, + "error": result.error, + } + + return await run_in_batches( + tokens, + _enable, + max_concurrent=max_concurrent, + batch_size=batch_size, + on_item=on_item, + should_cancel=should_cancel, + ) + + +__all__ = ["BatchNSFWService"] diff --git a/app/services/grok/batch_services/usage.py b/app/services/grok/batch_services/usage.py new file mode 100644 index 0000000..29e1b4d --- /dev/null +++ b/app/services/grok/batch_services/usage.py @@ -0,0 +1,36 @@ +""" +Batch usage service. +""" + +from typing import Callable, Awaitable, Dict, Any, Optional + +from app.services.grok.utils.batch import run_in_batches + + +class BatchUsageService: + """Batch usage orchestration.""" + + @staticmethod + async def refresh( + tokens: list[str], + mgr, + *, + max_concurrent: int, + batch_size: int, + on_item: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None, + should_cancel: Optional[Callable[[], bool]] = None, + ) -> Dict[str, Dict[str, Any]]: + async def _refresh_one(t: str): + return await mgr.sync_usage(t, consume_on_fail=False, is_usage=False) + + return await run_in_batches( + tokens, + _refresh_one, + max_concurrent=max_concurrent, + batch_size=batch_size, + on_item=on_item, + should_cancel=should_cancel, + ) + + +__all__ = ["BatchUsageService"] diff --git a/app/services/grok/processors/__init__.py b/app/services/grok/processors/__init__.py index 04773f6..72fb3d5 100644 --- a/app/services/grok/processors/__init__.py +++ b/app/services/grok/processors/__init__.py @@ -3,10 +3,14 @@ OpenAI 响应格式处理器 """ from .base import BaseProcessor, StreamIdleTimeoutError -from .chat_processors import StreamProcessor, CollectProcessor -from .video_processors import VideoStreamProcessor, VideoCollectProcessor -from .image_processors import ImageStreamProcessor, ImageCollectProcessor -from .image_ws_processors import ImageWSStreamProcessor, ImageWSCollectProcessor +from .chat import StreamProcessor, CollectProcessor +from .video import VideoStreamProcessor, VideoCollectProcessor +from .image import ( + ImageStreamProcessor, + ImageCollectProcessor, + ImageWSStreamProcessor, + ImageWSCollectProcessor, +) __all__ = [ "BaseProcessor", diff --git a/app/services/grok/processors/base.py b/app/services/grok/processors/base.py index 76b4838..f1c208e 100644 --- a/app/services/grok/processors/base.py +++ b/app/services/grok/processors/base.py @@ -8,7 +8,7 @@ from typing import Any, AsyncGenerator, Optional, AsyncIterable, List, TypeVar from app.core.config import get_config from app.core.logger import logger -from app.services.grok.services.assets import DownloadService +from app.services.grok.utils.download import DownloadService ASSET_URL = "https://assets.grok.com/" diff --git a/app/services/grok/processors/chat_processors.py b/app/services/grok/processors/chat.py similarity index 96% rename from app/services/grok/processors/chat_processors.py rename to app/services/grok/processors/chat.py index 000e09b..f55f853 100644 --- a/app/services/grok/processors/chat_processors.py +++ b/app/services/grok/processors/chat.py @@ -1,5 +1,5 @@ """ -聊天响应处理器 +Chat response processors. """ import asyncio @@ -24,7 +24,7 @@ from .base import ( class StreamProcessor(BaseProcessor): - """流式响应处理器""" + """Stream response processor.""" def __init__(self, model: str, token: str = "", think: bool = None): super().__init__(model, token) @@ -43,7 +43,7 @@ class StreamProcessor(BaseProcessor): self.show_think = think def _filter_token(self, token: str) -> str: - """过滤 token 中的特殊标签（如 ...），支持跨 token 的标签过滤""" + """Filter special tags (supports cross-token tag filtering).""" if not self.filter_tags: return token @@ -92,7 +92,7 @@ class StreamProcessor(BaseProcessor): return "".join(result) def _sse(self, content: str = "", role: str = None, finish: str = None) -> str: - """构建 SSE 响应""" + """Build SSE response.""" delta = {} if role: delta["role"] = role @@ -115,7 +115,7 @@ class StreamProcessor(BaseProcessor): async def process( self, response: AsyncIterable[bytes] ) -> AsyncGenerator[str, None]: - """处理流式响应""" + """Process stream response.""" idle_timeout = get_config("timeout.stream_idle_timeout") try: @@ -139,7 +139,7 @@ class StreamProcessor(BaseProcessor): yield self._sse(role="assistant") self.role_sent = True - # 图像生成进度 + # Image generation progress if img := resp.get("streamingImageGenerationResponse"): if self.show_think: if not self.think_opened: @@ -160,7 +160,7 @@ class StreamProcessor(BaseProcessor): yield self._sse("\n") self.think_opened = False - # 处理生成的图片 + # Handle generated images for url in _collect_image_urls(mr): parts = url.split("/") img_id = parts[-2] if len(parts) >= 2 else "image" @@ -194,7 +194,7 @@ class StreamProcessor(BaseProcessor): self.fingerprint = meta["llm_info"]["modelHash"] continue - # 普通 token + # Normal token if (token := resp.get("token")) is not None: if token: filtered = self._filter_token(token) @@ -242,7 +242,7 @@ class StreamProcessor(BaseProcessor): class CollectProcessor(BaseProcessor): - """非流式响应处理器""" + """Non-stream response processor.""" def __init__(self, model: str, token: str = ""): super().__init__(model, token) @@ -250,7 +250,7 @@ class CollectProcessor(BaseProcessor): self.filter_tags = get_config("chat.filter_tags") def _filter_content(self, content: str) -> str: - """过滤内容中的特殊标签""" + """Filter special tags in content.""" if not content or not self.filter_tags: return content @@ -262,7 +262,7 @@ class CollectProcessor(BaseProcessor): return result async def process(self, response: AsyncIterable[bytes]) -> dict[str, Any]: - """处理并收集完整响应""" + """Process and collect full response.""" response_id = "" fingerprint = "" content = "" diff --git a/app/services/grok/processors/image.py b/app/services/grok/processors/image.py new file mode 100644 index 0000000..a53b94a --- /dev/null +++ b/app/services/grok/processors/image.py @@ -0,0 +1,506 @@ +""" +Image response processors (HTTP + WebSocket). +""" + +import asyncio +import base64 +import random +import time +from pathlib import Path +from typing import AsyncGenerator, AsyncIterable, List, Dict, Optional + +import orjson +from curl_cffi.requests.errors import RequestsError + +from app.core.config import get_config +from app.core.logger import logger +from app.core.storage import DATA_DIR +from app.core.exceptions import UpstreamException +from .base import ( + BaseProcessor, + StreamIdleTimeoutError, + _with_idle_timeout, + _normalize_stream_line, + _collect_image_urls, + _is_http2_stream_error, +) + + +class ImageStreamProcessor(BaseProcessor): + """HTTP image stream processor.""" + + def __init__( + self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json" + ): + super().__init__(model, token) + self.partial_index = 0 + self.n = n + self.target_index = random.randint(0, 1) if n == 1 else None + self.response_format = response_format + if response_format == "url": + self.response_field = "url" + elif response_format == "base64": + self.response_field = "base64" + else: + self.response_field = "b64_json" + + def _sse(self, event: str, data: dict) -> str: + """Build SSE response.""" + return f"event: {event}\ndata: {orjson.dumps(data).decode()}\n\n" + + async def process( + self, response: AsyncIterable[bytes] + ) -> AsyncGenerator[str, None]: + """Process stream response.""" + final_images = [] + idle_timeout = get_config("timeout.stream_idle_timeout") + + try: + async for line in _with_idle_timeout(response, idle_timeout, self.model): + line = _normalize_stream_line(line) + if not line: + continue + try: + data = orjson.loads(line) + except orjson.JSONDecodeError: + continue + + resp = data.get("result", {}).get("response", {}) + + # Image generation progress + if img := resp.get("streamingImageGenerationResponse"): + image_index = img.get("imageIndex", 0) + progress = img.get("progress", 0) + + if self.n == 1 and image_index != self.target_index: + continue + + out_index = 0 if self.n == 1 else image_index + + yield self._sse( + "image_generation.partial_image", + { + "type": "image_generation.partial_image", + self.response_field: "", + "index": out_index, + "progress": progress, + }, + ) + continue + + # modelResponse + if mr := resp.get("modelResponse"): + if urls := _collect_image_urls(mr): + for url in urls: + if self.response_format == "url": + processed = await self.process_url(url, "image") + if processed: + final_images.append(processed) + continue + try: + dl_service = self._get_dl() + base64_data = await dl_service.to_base64( + url, self.token, "image" + ) + if base64_data: + if "," in base64_data: + b64 = base64_data.split(",", 1)[1] + else: + b64 = base64_data + final_images.append(b64) + except Exception as e: + logger.warning( + f"Failed to convert image to base64, falling back to URL: {e}" + ) + processed = await self.process_url(url, "image") + if processed: + final_images.append(processed) + continue + + for index, b64 in enumerate(final_images): + if self.n == 1: + if index != self.target_index: + continue + out_index = 0 + else: + out_index = index + + yield self._sse( + "image_generation.completed", + { + "type": "image_generation.completed", + self.response_field: b64, + "index": out_index, + "usage": { + "total_tokens": 0, + "input_tokens": 0, + "output_tokens": 0, + "input_tokens_details": { + "text_tokens": 0, + "image_tokens": 0, + }, + }, + }, + ) + except asyncio.CancelledError: + logger.debug("Image stream cancelled by client") + except StreamIdleTimeoutError as e: + raise UpstreamException( + message=f"Image stream idle timeout after {e.idle_seconds}s", + status_code=504, + details={ + "error": str(e), + "type": "stream_idle_timeout", + "idle_seconds": e.idle_seconds, + }, + ) + except RequestsError as e: + if _is_http2_stream_error(e): + logger.warning(f"HTTP/2 stream error in image: {e}") + raise UpstreamException( + message="Upstream connection closed unexpectedly", + status_code=502, + details={"error": str(e), "type": "http2_stream_error"}, + ) + logger.error(f"Image stream request error: {e}") + raise UpstreamException( + message=f"Upstream request failed: {e}", + status_code=502, + details={"error": str(e)}, + ) + except Exception as e: + logger.error( + f"Image stream processing error: {e}", + extra={"error_type": type(e).__name__}, + ) + raise + finally: + await self.close() + + +class ImageCollectProcessor(BaseProcessor): + """HTTP image non-stream processor.""" + + def __init__(self, model: str, token: str = "", response_format: str = "b64_json"): + super().__init__(model, token) + self.response_format = response_format + + async def process(self, response: AsyncIterable[bytes]) -> List[str]: + """Process and collect images.""" + images = [] + idle_timeout = get_config("timeout.stream_idle_timeout") + + try: + async for line in _with_idle_timeout(response, idle_timeout, self.model): + line = _normalize_stream_line(line) + if not line: + continue + try: + data = orjson.loads(line) + except orjson.JSONDecodeError: + continue + + resp = data.get("result", {}).get("response", {}) + + if mr := resp.get("modelResponse"): + if urls := _collect_image_urls(mr): + for url in urls: + if self.response_format == "url": + processed = await self.process_url(url, "image") + if processed: + images.append(processed) + continue + try: + dl_service = self._get_dl() + base64_data = await dl_service.to_base64( + url, self.token, "image" + ) + if base64_data: + if "," in base64_data: + b64 = base64_data.split(",", 1)[1] + else: + b64 = base64_data + images.append(b64) + except Exception as e: + logger.warning( + f"Failed to convert image to base64, falling back to URL: {e}" + ) + processed = await self.process_url(url, "image") + if processed: + images.append(processed) + + except asyncio.CancelledError: + logger.debug("Image collect cancelled by client") + except StreamIdleTimeoutError as e: + logger.warning(f"Image collect idle timeout: {e}") + except RequestsError as e: + if _is_http2_stream_error(e): + logger.warning(f"HTTP/2 stream error in image collect: {e}") + else: + logger.error(f"Image collect request error: {e}") + except Exception as e: + logger.error( + f"Image collect processing error: {e}", + extra={"error_type": type(e).__name__}, + ) + finally: + await self.close() + + return images + + +class ImageWSBaseProcessor(BaseProcessor): + """WebSocket image processor base.""" + + def __init__(self, model: str, token: str = "", response_format: str = "b64_json"): + super().__init__(model, token) + self.response_format = response_format + if response_format == "url": + self.response_field = "url" + elif response_format == "base64": + self.response_field = "base64" + else: + self.response_field = "b64_json" + self._image_dir: Optional[Path] = None + + def _ensure_image_dir(self) -> Path: + if self._image_dir is None: + base_dir = DATA_DIR / "tmp" / "image" + base_dir.mkdir(parents=True, exist_ok=True) + self._image_dir = base_dir + return self._image_dir + + def _strip_base64(self, blob: str) -> str: + if not blob: + return "" + if "," in blob and "base64" in blob.split(",", 1)[0]: + return blob.split(",", 1)[1] + return blob + + def _filename(self, image_id: str, is_final: bool) -> str: + ext = "jpg" if is_final else "png" + return f"{image_id}.{ext}" + + def _build_file_url(self, filename: str) -> str: + app_url = get_config("app.app_url") + if app_url: + return f"{app_url.rstrip('/')}/v1/files/image/{filename}" + return f"/v1/files/image/{filename}" + + def _save_blob(self, image_id: str, blob: str, is_final: bool) -> str: + data = self._strip_base64(blob) + if not data: + return "" + image_dir = self._ensure_image_dir() + filename = self._filename(image_id, is_final) + filepath = image_dir / filename + with open(filepath, "wb") as f: + f.write(base64.b64decode(data)) + return self._build_file_url(filename) + + def _pick_best(self, existing: Optional[Dict], incoming: Dict) -> Dict: + if not existing: + return incoming + if incoming.get("is_final") and not existing.get("is_final"): + return incoming + if existing.get("is_final") and not incoming.get("is_final"): + return existing + if incoming.get("blob_size", 0) > existing.get("blob_size", 0): + return incoming + return existing + + def _to_output(self, image_id: str, item: Dict) -> str: + try: + if self.response_format == "url": + return self._save_blob( + image_id, item.get("blob", ""), item.get("is_final", False) + ) + return self._strip_base64(item.get("blob", "")) + except Exception as e: + logger.warning(f"Image output failed: {e}") + return "" + + +class ImageWSStreamProcessor(ImageWSBaseProcessor): + """WebSocket image stream processor.""" + + def __init__( + self, + model: str, + token: str = "", + n: int = 1, + response_format: str = "b64_json", + size: str = "1024x1024", + ): + super().__init__(model, token, "b64_json") + self.n = n + self.size = size + self._target_id: Optional[str] = None + self._index_map: Dict[str, int] = {} + self._partial_map: Dict[str, int] = {} + + def _assign_index(self, image_id: str) -> Optional[int]: + if image_id in self._index_map: + return self._index_map[image_id] + if len(self._index_map) >= self.n: + return None + self._index_map[image_id] = len(self._index_map) + return self._index_map[image_id] + + def _sse(self, event: str, data: dict) -> str: + return f"event: {event}\ndata: {orjson.dumps(data).decode()}\n\n" + + async def process(self, response: AsyncIterable[dict]) -> AsyncGenerator[str, None]: + images: Dict[str, Dict] = {} + + async for item in response: + if item.get("type") == "error": + message = item.get("error") or "Upstream error" + code = item.get("error_code") or "upstream_error" + yield self._sse( + "error", + { + "error": { + "message": message, + "type": "server_error", + "code": code, + } + }, + ) + return + if item.get("type") != "image": + continue + + image_id = item.get("image_id") + if not image_id: + continue + + if self.n == 1: + if self._target_id is None: + self._target_id = image_id + index = 0 if image_id == self._target_id else None + else: + index = self._assign_index(image_id) + + images[image_id] = self._pick_best(images.get(image_id), item) + + if index is None: + continue + + if item.get("stage") != "final": + partial_b64 = self._strip_base64(item.get("blob", "")) + if not partial_b64: + continue + partial_index = self._partial_map.get(image_id, 0) + if item.get("stage") == "medium": + partial_index = max(partial_index, 1) + self._partial_map[image_id] = partial_index + yield self._sse( + "image_generation.partial_image", + { + "type": "image_generation.partial_image", + "b64_json": partial_b64, + "created_at": int(time.time()), + "size": self.size, + "index": index, + "partial_image_index": partial_index, + }, + ) + + if self.n == 1: + if self._target_id and self._target_id in images: + selected = [(self._target_id, images[self._target_id])] + else: + selected = ( + [ + max( + images.items(), + key=lambda x: ( + x[1].get("is_final", False), + x[1].get("blob_size", 0), + ), + ) + ] + if images + else [] + ) + else: + selected = [ + (image_id, images[image_id]) + for image_id in self._index_map + if image_id in images + ] + + for image_id, item in selected: + output = self._strip_base64(item.get("blob", "")) + if not output: + continue + + if self.n == 1: + index = 0 + else: + index = self._index_map.get(image_id, 0) + yield self._sse( + "image_generation.completed", + { + "type": "image_generation.completed", + "b64_json": output, + "created_at": int(time.time()), + "size": self.size, + "index": index, + "usage": { + "total_tokens": 0, + "input_tokens": 0, + "output_tokens": 0, + "input_tokens_details": {"text_tokens": 0, "image_tokens": 0}, + }, + }, + ) + + +class ImageWSCollectProcessor(ImageWSBaseProcessor): + """WebSocket image non-stream processor.""" + + def __init__( + self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json" + ): + super().__init__(model, token, response_format) + self.n = n + + async def process(self, response: AsyncIterable[dict]) -> List[str]: + images: Dict[str, Dict] = {} + + async for item in response: + if item.get("type") == "error": + message = item.get("error") or "Upstream error" + raise UpstreamException(message, details=item) + if item.get("type") != "image": + continue + image_id = item.get("image_id") + if not image_id: + continue + images[image_id] = self._pick_best(images.get(image_id), item) + + selected = sorted( + images.values(), + key=lambda x: (x.get("is_final", False), x.get("blob_size", 0)), + reverse=True, + ) + if self.n: + selected = selected[: self.n] + + results: List[str] = [] + for item in selected: + output = self._to_output(item.get("image_id", ""), item) + if output: + results.append(output) + + return results + + +__all__ = [ + "ImageStreamProcessor", + "ImageCollectProcessor", + "ImageWSStreamProcessor", + "ImageWSCollectProcessor", +] diff --git a/app/services/grok/processors/image_processors.py b/app/services/grok/processors/image_processors.py deleted file mode 100644 index 8f78ac3..0000000 --- a/app/services/grok/processors/image_processors.py +++ /dev/null @@ -1,248 +0,0 @@ -""" -图片生成响应处理器（HTTP） -""" - -import asyncio -import random -from typing import AsyncGenerator, AsyncIterable, List - -import orjson -from curl_cffi.requests.errors import RequestsError - -from app.core.config import get_config -from app.core.logger import logger -from app.core.exceptions import UpstreamException -from .base import ( - BaseProcessor, - StreamIdleTimeoutError, - _with_idle_timeout, - _normalize_stream_line, - _collect_image_urls, - _is_http2_stream_error, -) - - -class ImageStreamProcessor(BaseProcessor): - """图片生成流式响应处理器""" - - def __init__( - self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json" - ): - super().__init__(model, token) - self.partial_index = 0 - self.n = n - self.target_index = random.randint(0, 1) if n == 1 else None - self.response_format = response_format - if response_format == "url": - self.response_field = "url" - elif response_format == "base64": - self.response_field = "base64" - else: - self.response_field = "b64_json" - - def _sse(self, event: str, data: dict) -> str: - """构建 SSE 响应""" - return f"event: {event}\ndata: {orjson.dumps(data).decode()}\n\n" - - async def process( - self, response: AsyncIterable[bytes] - ) -> AsyncGenerator[str, None]: - """处理流式响应""" - final_images = [] - idle_timeout = get_config("timeout.stream_idle_timeout") - - try: - async for line in _with_idle_timeout(response, idle_timeout, self.model): - line = _normalize_stream_line(line) - if not line: - continue - try: - data = orjson.loads(line) - except orjson.JSONDecodeError: - continue - - resp = data.get("result", {}).get("response", {}) - - # 图片生成进度 - if img := resp.get("streamingImageGenerationResponse"): - image_index = img.get("imageIndex", 0) - progress = img.get("progress", 0) - - if self.n == 1 and image_index != self.target_index: - continue - - out_index = 0 if self.n == 1 else image_index - - yield self._sse( - "image_generation.partial_image", - { - "type": "image_generation.partial_image", - self.response_field: "", - "index": out_index, - "progress": progress, - }, - ) - continue - - # modelResponse - if mr := resp.get("modelResponse"): - if urls := _collect_image_urls(mr): - for url in urls: - if self.response_format == "url": - processed = await self.process_url(url, "image") - if processed: - final_images.append(processed) - continue - try: - dl_service = self._get_dl() - base64_data = await dl_service.to_base64( - url, self.token, "image" - ) - if base64_data: - if "," in base64_data: - b64 = base64_data.split(",", 1)[1] - else: - b64 = base64_data - final_images.append(b64) - except Exception as e: - logger.warning( - f"Failed to convert image to base64, falling back to URL: {e}" - ) - processed = await self.process_url(url, "image") - if processed: - final_images.append(processed) - continue - - for index, b64 in enumerate(final_images): - if self.n == 1: - if index != self.target_index: - continue - out_index = 0 - else: - out_index = index - - yield self._sse( - "image_generation.completed", - { - "type": "image_generation.completed", - self.response_field: b64, - "index": out_index, - "usage": { - "total_tokens": 0, - "input_tokens": 0, - "output_tokens": 0, - "input_tokens_details": { - "text_tokens": 0, - "image_tokens": 0, - }, - }, - }, - ) - except asyncio.CancelledError: - logger.debug("Image stream cancelled by client") - except StreamIdleTimeoutError as e: - raise UpstreamException( - message=f"Image stream idle timeout after {e.idle_seconds}s", - status_code=504, - details={ - "error": str(e), - "type": "stream_idle_timeout", - "idle_seconds": e.idle_seconds, - }, - ) - except RequestsError as e: - if _is_http2_stream_error(e): - logger.warning(f"HTTP/2 stream error in image: {e}") - raise UpstreamException( - message="Upstream connection closed unexpectedly", - status_code=502, - details={"error": str(e), "type": "http2_stream_error"}, - ) - logger.error(f"Image stream request error: {e}") - raise UpstreamException( - message=f"Upstream request failed: {e}", - status_code=502, - details={"error": str(e)}, - ) - except Exception as e: - logger.error( - f"Image stream processing error: {e}", - extra={"error_type": type(e).__name__}, - ) - raise - finally: - await self.close() - - -class ImageCollectProcessor(BaseProcessor): - """图片生成非流式响应处理器""" - - def __init__(self, model: str, token: str = "", response_format: str = "b64_json"): - super().__init__(model, token) - self.response_format = response_format - - async def process(self, response: AsyncIterable[bytes]) -> List[str]: - """处理并收集图片""" - images = [] - idle_timeout = get_config("timeout.stream_idle_timeout") - - try: - async for line in _with_idle_timeout(response, idle_timeout, self.model): - line = _normalize_stream_line(line) - if not line: - continue - try: - data = orjson.loads(line) - except orjson.JSONDecodeError: - continue - - resp = data.get("result", {}).get("response", {}) - - if mr := resp.get("modelResponse"): - if urls := _collect_image_urls(mr): - for url in urls: - if self.response_format == "url": - processed = await self.process_url(url, "image") - if processed: - images.append(processed) - continue - try: - dl_service = self._get_dl() - base64_data = await dl_service.to_base64( - url, self.token, "image" - ) - if base64_data: - if "," in base64_data: - b64 = base64_data.split(",", 1)[1] - else: - b64 = base64_data - images.append(b64) - except Exception as e: - logger.warning( - f"Failed to convert image to base64, falling back to URL: {e}" - ) - processed = await self.process_url(url, "image") - if processed: - images.append(processed) - - except asyncio.CancelledError: - logger.debug("Image collect cancelled by client") - except StreamIdleTimeoutError as e: - logger.warning(f"Image collect idle timeout: {e}") - except RequestsError as e: - if _is_http2_stream_error(e): - logger.warning(f"HTTP/2 stream error in image collect: {e}") - else: - logger.error(f"Image collect request error: {e}") - except Exception as e: - logger.error( - f"Image collect processing error: {e}", - extra={"error_type": type(e).__name__}, - ) - finally: - await self.close() - - return images - - -__all__ = ["ImageStreamProcessor", "ImageCollectProcessor"] diff --git a/app/services/grok/processors/image_ws_processors.py b/app/services/grok/processors/image_ws_processors.py deleted file mode 100644 index 788a442..0000000 --- a/app/services/grok/processors/image_ws_processors.py +++ /dev/null @@ -1,268 +0,0 @@ -""" -图片生成响应处理器（WebSocket） -""" - -import base64 -import time -from pathlib import Path -from typing import AsyncGenerator, AsyncIterable, List, Dict, Optional - -import orjson - -from app.core.config import get_config -from app.core.logger import logger -from app.core.storage import DATA_DIR -from app.core.exceptions import UpstreamException -from .base import BaseProcessor - - -class ImageWSBaseProcessor(BaseProcessor): - """WebSocket 图片处理基类""" - - def __init__(self, model: str, token: str = "", response_format: str = "b64_json"): - super().__init__(model, token) - self.response_format = response_format - if response_format == "url": - self.response_field = "url" - elif response_format == "base64": - self.response_field = "base64" - else: - self.response_field = "b64_json" - self._image_dir: Optional[Path] = None - - def _ensure_image_dir(self) -> Path: - if self._image_dir is None: - base_dir = DATA_DIR / "tmp" / "image" - base_dir.mkdir(parents=True, exist_ok=True) - self._image_dir = base_dir - return self._image_dir - - def _strip_base64(self, blob: str) -> str: - if not blob: - return "" - if "," in blob and "base64" in blob.split(",", 1)[0]: - return blob.split(",", 1)[1] - return blob - - def _filename(self, image_id: str, is_final: bool) -> str: - ext = "jpg" if is_final else "png" - return f"{image_id}.{ext}" - - def _build_file_url(self, filename: str) -> str: - app_url = get_config("app.app_url") - if app_url: - return f"{app_url.rstrip('/')}/v1/files/image/{filename}" - return f"/v1/files/image/{filename}" - - def _save_blob(self, image_id: str, blob: str, is_final: bool) -> str: - data = self._strip_base64(blob) - if not data: - return "" - image_dir = self._ensure_image_dir() - filename = self._filename(image_id, is_final) - filepath = image_dir / filename - with open(filepath, "wb") as f: - f.write(base64.b64decode(data)) - return self._build_file_url(filename) - - def _pick_best(self, existing: Optional[Dict], incoming: Dict) -> Dict: - if not existing: - return incoming - if incoming.get("is_final") and not existing.get("is_final"): - return incoming - if existing.get("is_final") and not incoming.get("is_final"): - return existing - if incoming.get("blob_size", 0) > existing.get("blob_size", 0): - return incoming - return existing - - def _to_output(self, image_id: str, item: Dict) -> str: - try: - if self.response_format == "url": - return self._save_blob( - image_id, item.get("blob", ""), item.get("is_final", False) - ) - return self._strip_base64(item.get("blob", "")) - except Exception as e: - logger.warning(f"Image output failed: {e}") - return "" - - -class ImageWSStreamProcessor(ImageWSBaseProcessor): - """WebSocket 图片流式响应处理器""" - - def __init__( - self, - model: str, - token: str = "", - n: int = 1, - response_format: str = "b64_json", - size: str = "1024x1024", - ): - super().__init__(model, token, "b64_json") - self.n = n - self.size = size - self._target_id: Optional[str] = None - self._index_map: Dict[str, int] = {} - self._partial_map: Dict[str, int] = {} - - def _assign_index(self, image_id: str) -> Optional[int]: - if image_id in self._index_map: - return self._index_map[image_id] - if len(self._index_map) >= self.n: - return None - self._index_map[image_id] = len(self._index_map) - return self._index_map[image_id] - - def _sse(self, event: str, data: dict) -> str: - return f"event: {event}\ndata: {orjson.dumps(data).decode()}\n\n" - - async def process(self, response: AsyncIterable[dict]) -> AsyncGenerator[str, None]: - images: Dict[str, Dict] = {} - - async for item in response: - if item.get("type") == "error": - message = item.get("error") or "Upstream error" - code = item.get("error_code") or "upstream_error" - yield self._sse( - "error", - { - "error": { - "message": message, - "type": "server_error", - "code": code, - } - }, - ) - return - if item.get("type") != "image": - continue - - image_id = item.get("image_id") - if not image_id: - continue - - if self.n == 1: - if self._target_id is None: - self._target_id = image_id - index = 0 if image_id == self._target_id else None - else: - index = self._assign_index(image_id) - - images[image_id] = self._pick_best(images.get(image_id), item) - - if index is None: - continue - - if item.get("stage") != "final": - partial_b64 = self._strip_base64(item.get("blob", "")) - if not partial_b64: - continue - partial_index = self._partial_map.get(image_id, 0) - if item.get("stage") == "medium": - partial_index = max(partial_index, 1) - self._partial_map[image_id] = partial_index - yield self._sse( - "image_generation.partial_image", - { - "type": "image_generation.partial_image", - "b64_json": partial_b64, - "created_at": int(time.time()), - "size": self.size, - "index": index, - "partial_image_index": partial_index, - }, - ) - - if self.n == 1: - if self._target_id and self._target_id in images: - selected = [(self._target_id, images[self._target_id])] - else: - selected = ( - [ - max( - images.items(), - key=lambda x: ( - x[1].get("is_final", False), - x[1].get("blob_size", 0), - ), - ) - ] - if images - else [] - ) - else: - selected = [ - (image_id, images[image_id]) - for image_id in self._index_map - if image_id in images - ] - - for image_id, item in selected: - output = self._strip_base64(item.get("blob", "")) - if not output: - continue - - if self.n == 1: - index = 0 - else: - index = self._index_map.get(image_id, 0) - yield self._sse( - "image_generation.completed", - { - "type": "image_generation.completed", - "b64_json": output, - "created_at": int(time.time()), - "size": self.size, - "index": index, - "usage": { - "total_tokens": 0, - "input_tokens": 0, - "output_tokens": 0, - "input_tokens_details": {"text_tokens": 0, "image_tokens": 0}, - }, - }, - ) - - -class ImageWSCollectProcessor(ImageWSBaseProcessor): - """WebSocket 图片非流式响应处理器""" - - def __init__( - self, model: str, token: str = "", n: int = 1, response_format: str = "b64_json" - ): - super().__init__(model, token, response_format) - self.n = n - - async def process(self, response: AsyncIterable[dict]) -> List[str]: - images: Dict[str, Dict] = {} - - async for item in response: - if item.get("type") == "error": - message = item.get("error") or "Upstream error" - raise UpstreamException(message, details=item) - if item.get("type") != "image": - continue - image_id = item.get("image_id") - if not image_id: - continue - images[image_id] = self._pick_best(images.get(image_id), item) - - selected = sorted( - images.values(), - key=lambda x: (x.get("is_final", False), x.get("blob_size", 0)), - reverse=True, - ) - if self.n: - selected = selected[: self.n] - - results: List[str] = [] - for item in selected: - output = self._to_output(item.get("image_id", ""), item) - if output: - results.append(output) - - return results - - -__all__ = ["ImageWSStreamProcessor", "ImageWSCollectProcessor"] diff --git a/app/services/grok/processors/video_processors.py b/app/services/grok/processors/video.py similarity index 96% rename from app/services/grok/processors/video_processors.py rename to app/services/grok/processors/video.py index a0ead8c..89521b9 100644 --- a/app/services/grok/processors/video_processors.py +++ b/app/services/grok/processors/video.py @@ -1,5 +1,5 @@ """ -视频响应处理器 +Video response processors. """ import asyncio @@ -22,7 +22,7 @@ from .base import ( class VideoStreamProcessor(BaseProcessor): - """视频流式响应处理器""" + """Video stream response processor.""" def __init__(self, model: str, token: str = "", think: bool = None): super().__init__(model, token) @@ -37,7 +37,7 @@ class VideoStreamProcessor(BaseProcessor): self.show_think = think def _sse(self, content: str = "", role: str = None, finish: str = None) -> str: - """构建 SSE 响应""" + """Build SSE response.""" delta = {} if role: delta["role"] = role @@ -57,7 +57,7 @@ class VideoStreamProcessor(BaseProcessor): return f"data: {orjson.dumps(chunk).decode()}\n\n" def _build_video_html(self, video_url: str, thumbnail_url: str = "") -> str: - """构建视频 HTML 标签""" + """Build video HTML tag.""" import html safe_video_url = html.escape(video_url) @@ -70,7 +70,7 @@ class VideoStreamProcessor(BaseProcessor): async def process( self, response: AsyncIterable[bytes] ) -> AsyncGenerator[str, None]: - """处理视频流式响应""" + """Process video stream response.""" idle_timeout = get_config("timeout.video_idle_timeout") try: @@ -92,7 +92,7 @@ class VideoStreamProcessor(BaseProcessor): yield self._sse(role="assistant") self.role_sent = True - # 视频生成进度 + # Video generation progress if video_resp := resp.get("streamingVideoGenerationResponse"): progress = video_resp.get("progress", 0) @@ -175,7 +175,7 @@ class VideoStreamProcessor(BaseProcessor): class VideoCollectProcessor(BaseProcessor): - """视频非流式响应处理器""" + """Video non-stream response processor.""" def __init__(self, model: str, token: str = ""): super().__init__(model, token) @@ -188,7 +188,7 @@ class VideoCollectProcessor(BaseProcessor): ''' async def process(self, response: AsyncIterable[bytes]) -> dict[str, Any]: - """处理并收集视频响应""" + """Process and collect video response.""" response_id = "" content = "" idle_timeout = get_config("timeout.video_idle_timeout") diff --git a/app/services/grok/services/image.py b/app/services/grok/services/image.py index 7313976..4bba8d9 100644 --- a/app/services/grok/services/image.py +++ b/app/services/grok/services/image.py @@ -1,11 +1,307 @@ """ -Grok Imagine WebSocket image service. +Grok image services. """ +import asyncio +import math +import random +from dataclasses import dataclass +from typing import Any, AsyncGenerator, List, Optional, Union + +from app.core.config import get_config +from app.core.logger import logger +from app.services.grok.processors import ( + ImageStreamProcessor, + ImageCollectProcessor, + ImageWSStreamProcessor, + ImageWSCollectProcessor, +) +from app.services.grok.services.chat import GrokChatService +from app.services.grok.utils.stream import wrap_stream_with_usage +from app.services.token import EffortType from app.services.reverse.ws_imagine import ImagineWebSocketReverse ImageService = ImagineWebSocketReverse image_service = ImagineWebSocketReverse() -__all__ = ["image_service", "ImageService"] + +@dataclass +class ImageGenerationResult: + stream: bool + data: Union[AsyncGenerator[str, None], List[str]] + usage_override: Optional[dict] = None + + +class ImageGenerationService: + """Image generation orchestration service.""" + + async def generate( + self, + *, + token_mgr: Any, + token: str, + model_info: Any, + prompt: str, + n: int, + response_format: str, + size: str, + aspect_ratio: str, + stream: bool, + use_ws: bool, + ) -> ImageGenerationResult: + if stream: + if use_ws: + return await self._stream_ws( + token_mgr=token_mgr, + token=token, + model_info=model_info, + prompt=prompt, + n=n, + response_format=response_format, + size=size, + aspect_ratio=aspect_ratio, + ) + return await self._stream_http( + token_mgr=token_mgr, + token=token, + model_info=model_info, + prompt=prompt, + n=n, + response_format=response_format, + ) + + if use_ws: + return await self._collect_ws( + token_mgr=token_mgr, + token=token, + model_info=model_info, + prompt=prompt, + n=n, + response_format=response_format, + aspect_ratio=aspect_ratio, + ) + + return await self._collect_http( + token_mgr=token_mgr, + token=token, + model_info=model_info, + prompt=prompt, + n=n, + response_format=response_format, + ) + + async def _stream_ws( + self, + *, + token_mgr: Any, + token: str, + model_info: Any, + prompt: str, + n: int, + response_format: str, + size: str, + aspect_ratio: str, + ) -> ImageGenerationResult: + enable_nsfw = bool(get_config("image.image_ws_nsfw")) + upstream = image_service.stream( + token=token, + prompt=prompt, + aspect_ratio=aspect_ratio, + n=n, + enable_nsfw=enable_nsfw, + ) + processor = ImageWSStreamProcessor( + model_info.model_id, + token, + n=n, + response_format=response_format, + size=size, + ) + stream = wrap_stream_with_usage( + processor.process(upstream), + token_mgr, + token, + model_info.model_id, + ) + return ImageGenerationResult(stream=True, data=stream) + + async def _stream_http( + self, + *, + token_mgr: Any, + token: str, + model_info: Any, + prompt: str, + n: int, + response_format: str, + ) -> ImageGenerationResult: + response = await GrokChatService().chat( + token=token, + message=f"Image Generation: {prompt}", + model=model_info.grok_model, + mode=model_info.model_mode, + stream=True, + ) + processor = ImageStreamProcessor( + model_info.model_id, + token, + n=n, + response_format=response_format, + ) + stream = wrap_stream_with_usage( + processor.process(response), + token_mgr, + token, + model_info.model_id, + ) + return ImageGenerationResult(stream=True, data=stream) + + async def _collect_ws( + self, + *, + token_mgr: Any, + token: str, + model_info: Any, + prompt: str, + n: int, + response_format: str, + aspect_ratio: str, + ) -> ImageGenerationResult: + enable_nsfw = bool(get_config("image.image_ws_nsfw")) + all_images: List[str] = [] + seen = set() + expected_per_call = 6 + calls_needed = max(1, int(math.ceil(n / expected_per_call))) + calls_needed = min(calls_needed, n) + + async def _fetch_batch(call_target: int): + upstream = image_service.stream( + token=token, + prompt=prompt, + aspect_ratio=aspect_ratio, + n=call_target, + enable_nsfw=enable_nsfw, + ) + processor = ImageWSCollectProcessor( + model_info.model_id, + token, + n=call_target, + response_format=response_format, + ) + return await processor.process(upstream) + + tasks = [] + for i in range(calls_needed): + remaining = n - (i * expected_per_call) + call_target = min(expected_per_call, remaining) + tasks.append(_fetch_batch(call_target)) + + results = await asyncio.gather(*tasks, return_exceptions=True) + for batch in results: + if isinstance(batch, Exception): + logger.warning(f"WS batch failed: {batch}") + continue + for img in batch: + if img not in seen: + seen.add(img) + all_images.append(img) + if len(all_images) >= n: + break + if len(all_images) >= n: + break + + try: + await token_mgr.consume(token, self._get_effort(model_info)) + except Exception as e: + logger.warning(f"Failed to consume token: {e}") + + selected = self._select_images(all_images, n) + usage_override = { + "total_tokens": 0, + "input_tokens": 0, + "output_tokens": 0, + "input_tokens_details": {"text_tokens": 0, "image_tokens": 0}, + } + return ImageGenerationResult( + stream=False, data=selected, usage_override=usage_override + ) + + async def _collect_http( + self, + *, + token_mgr: Any, + token: str, + model_info: Any, + prompt: str, + n: int, + response_format: str, + ) -> ImageGenerationResult: + calls_needed = (n + 1) // 2 + + async def _call_grok(): + success = False + try: + response = await GrokChatService().chat( + token=token, + message=f"Image Generation: {prompt}", + model=model_info.grok_model, + mode=model_info.model_mode, + stream=True, + ) + processor = ImageCollectProcessor( + model_info.model_id, token, response_format=response_format + ) + images = await processor.process(response) + success = True + return images + except Exception as e: + logger.error(f"Grok image call failed: {e}") + return [] + finally: + if success: + try: + await token_mgr.consume(token, self._get_effort(model_info)) + except Exception as e: + logger.warning(f"Failed to consume token: {e}") + + if calls_needed == 1: + all_images = await _call_grok() + else: + tasks = [_call_grok() for _ in range(calls_needed)] + results = await asyncio.gather(*tasks, return_exceptions=True) + + all_images: List[str] = [] + for result in results: + if isinstance(result, Exception): + logger.error(f"Concurrent call failed: {result}") + elif isinstance(result, list): + all_images.extend(result) + + selected = self._select_images(all_images, n) + return ImageGenerationResult(stream=False, data=selected) + + @staticmethod + def _get_effort(model_info: Any) -> EffortType: + return ( + EffortType.HIGH + if (model_info and model_info.cost.value == "high") + else EffortType.LOW + ) + + @staticmethod + def _select_images(images: List[str], n: int) -> List[str]: + if len(images) >= n: + return random.sample(images, n) + selected = images.copy() + while len(selected) < n: + selected.append("error") + return selected + + +__all__ = [ + "image_service", + "ImageService", + "ImageGenerationService", + "ImageGenerationResult", +] diff --git a/app/services/grok/services/image_edit.py b/app/services/grok/services/image_edit.py new file mode 100644 index 0000000..ca460bc --- /dev/null +++ b/app/services/grok/services/image_edit.py @@ -0,0 +1,197 @@ +""" +Grok image edit service. +""" + +import asyncio +import random +import re +from dataclasses import dataclass +from typing import AsyncGenerator, List, Union, Any + +from app.core.exceptions import AppException, ErrorType +from app.core.logger import logger +from app.services.grok.processors import ImageCollectProcessor, ImageStreamProcessor +from app.services.grok.services.assets import UploadService +from app.services.grok.services.chat import GrokChatService +from app.services.grok.services.video import VideoService +from app.services.grok.utils.stream import wrap_stream_with_usage + + +@dataclass +class ImageEditResult: + stream: bool + data: Union[AsyncGenerator[str, None], List[str]] + + +class ImageEditService: + """Image edit orchestration service.""" + + async def edit( + self, + *, + token_mgr: Any, + token: str, + model_info: Any, + prompt: str, + images: List[str], + n: int, + response_format: str, + stream: bool, + ) -> ImageEditResult: + image_urls = await self._upload_images(images, token) + parent_post_id = await self._get_parent_post_id(token, image_urls) + + model_config_override = { + "modelMap": { + "imageEditModel": "imagine", + "imageEditModelConfig": { + "imageReferences": image_urls, + }, + } + } + if parent_post_id: + model_config_override["modelMap"]["imageEditModelConfig"][ + "parentPostId" + ] = parent_post_id + + tool_overrides = {"imageGen": True} + + if stream: + response = await GrokChatService().chat( + token=token, + message=prompt, + model=model_info.grok_model, + mode=None, + stream=True, + tool_overrides=tool_overrides, + model_config_override=model_config_override, + ) + processor = ImageStreamProcessor( + model_info.model_id, + token, + n=n, + response_format=response_format, + ) + return ImageEditResult( + stream=True, + data=wrap_stream_with_usage( + processor.process(response), + token_mgr, + token, + model_info.model_id, + ), + ) + + images_out = await self._collect_images( + token=token, + prompt=prompt, + model_info=model_info, + n=n, + response_format=response_format, + tool_overrides=tool_overrides, + model_config_override=model_config_override, + ) + return ImageEditResult(stream=False, data=images_out) + + async def _upload_images(self, images: List[str], token: str) -> List[str]: + image_urls: List[str] = [] + upload_service = UploadService() + try: + for image in images: + _, file_uri = await upload_service.upload(image, token) + if file_uri: + if file_uri.startswith("http"): + image_urls.append(file_uri) + else: + image_urls.append( + f"https://assets.grok.com/{file_uri.lstrip('/')}" + ) + finally: + await upload_service.close() + + if not image_urls: + raise AppException( + message="Image upload failed", + error_type=ErrorType.SERVER.value, + code="upload_failed", + ) + + return image_urls + + async def _get_parent_post_id(self, token: str, image_urls: List[str]) -> str: + parent_post_id = None + try: + media_service = VideoService() + parent_post_id = await media_service.create_image_post(token, image_urls[0]) + logger.debug(f"Parent post ID: {parent_post_id}") + except Exception as e: + logger.warning(f"Create image post failed: {e}") + + if parent_post_id: + return parent_post_id + + for url in image_urls: + match = re.search(r"/generated/([a-f0-9-]+)/", url) + if match: + parent_post_id = match.group(1) + logger.debug(f"Parent post ID: {parent_post_id}") + break + match = re.search(r"/users/[^/]+/([a-f0-9-]+)/content", url) + if match: + parent_post_id = match.group(1) + logger.debug(f"Parent post ID: {parent_post_id}") + break + + return parent_post_id or "" + + async def _collect_images( + self, + *, + token: str, + prompt: str, + model_info: Any, + n: int, + response_format: str, + tool_overrides: dict, + model_config_override: dict, + ) -> List[str]: + calls_needed = (n + 1) // 2 + + async def _call_edit(): + response = await GrokChatService().chat( + token=token, + message=prompt, + model=model_info.grok_model, + mode=None, + stream=True, + tool_overrides=tool_overrides, + model_config_override=model_config_override, + ) + processor = ImageCollectProcessor( + model_info.model_id, token, response_format=response_format + ) + return await processor.process(response) + + if calls_needed == 1: + all_images = await _call_edit() + else: + tasks = [_call_edit() for _ in range(calls_needed)] + results = await asyncio.gather(*tasks, return_exceptions=True) + + all_images: List[str] = [] + for result in results: + if isinstance(result, Exception): + logger.error(f"Concurrent call failed: {result}") + elif isinstance(result, list): + all_images.extend(result) + + if len(all_images) >= n: + return random.sample(all_images, n) + + selected_images = all_images.copy() + while len(selected_images) < n: + selected_images.append("error") + return selected_images + + +__all__ = ["ImageEditService", "ImageEditResult"] diff --git a/app/services/grok/services/media.py b/app/services/grok/services/video.py similarity index 93% rename from app/services/grok/services/media.py rename to app/services/grok/services/video.py index 50aa025..bb53e53 100644 --- a/app/services/grok/services/media.py +++ b/app/services/grok/services/video.py @@ -1,5 +1,5 @@ """ -Grok 视频生成服务 +Grok video generation service. """ import asyncio @@ -25,7 +25,7 @@ _MEDIA_SEM_VALUE = 0 def _get_semaphore() -> asyncio.Semaphore: - """获取或更新信号量""" + """Get or refresh the semaphore.""" global _MEDIA_SEMAPHORE, _MEDIA_SEM_VALUE value = max(1, int(get_config("performance.media_max_concurrent"))) if value != _MEDIA_SEM_VALUE: @@ -35,7 +35,7 @@ def _get_semaphore() -> asyncio.Semaphore: class VideoService: - """视频生成服务""" + """Video generation service.""" def __init__(self): self.timeout = get_config("network.timeout") @@ -47,7 +47,7 @@ class VideoService: media_type: str = "MEDIA_POST_TYPE_VIDEO", media_url: str = None, ) -> str: - """创建媒体帖子，返回 post ID""" + """Create media post and return post ID.""" try: if media_type == "MEDIA_POST_TYPE_IMAGE" and not media_url: raise ValidationException("media_url is required for image posts") @@ -74,7 +74,7 @@ class VideoService: raise UpstreamException(f"Create post error: {str(e)}") async def create_image_post(self, token: str, image_url: str) -> str: - """创建图片帖子，返回 post ID""" + """Create image post and return post ID.""" return await self.create_post( token, prompt="", media_type="MEDIA_POST_TYPE_IMAGE", media_url=image_url ) @@ -88,7 +88,7 @@ class VideoService: resolution_name: str = "480p", preset: str = "normal", ) -> dict: - """构建视频生成载荷""" + """Build video generation payload.""" mode_map = { "fun": "--mode=extremely-crazy", "normal": "--mode=normal", @@ -139,7 +139,7 @@ class VideoService: resolution_name: str, preset: str, ) -> AsyncGenerator[bytes, None]: - """内部生成逻辑""" + """Internal generation logic.""" session = None try: payload = self._build_payload( @@ -182,7 +182,7 @@ class VideoService: resolution_name: str = "480p", preset: str = "normal", ) -> AsyncGenerator[bytes, None]: - """生成视频""" + """Generate video.""" logger.info( f"Video generation: prompt='{prompt[:50]}...', ratio={aspect_ratio}, length={video_length}s, preset={preset}" ) @@ -208,7 +208,7 @@ class VideoService: resolution: str = "480p", preset: str = "normal", ) -> AsyncGenerator[bytes, None]: - """从图片生成视频""" + """Generate video from image.""" logger.info( f"Image to video: prompt='{prompt[:50]}...', image={image_url[:80]}" ) @@ -229,12 +229,12 @@ class VideoService: resolution: str = "480p", preset: str = "normal", ): - """视频生成入口""" - # 获取 token（使用智能路由） + """Video generation entrypoint.""" + # Get token via intelligent routing. token_mgr = await get_token_manager() await token_mgr.reload_if_stale() - # 使用智能路由选择 token（根据视频需求与候选池） + # Select token based on video requirements and pool candidates. pool_candidates = ModelService.pool_candidates_for_model(model) token_info = token_mgr.get_token_for_video( resolution=resolution, @@ -250,7 +250,7 @@ class VideoService: status_code=429, ) - # 从 TokenInfo 对象中提取 token 字符串 + # Extract token string from TokenInfo. token = token_info.token if token.startswith("sso="): token = token[4:] @@ -258,7 +258,7 @@ class VideoService: think = {"enabled": True, "disabled": False}.get(thinking) is_stream = stream if stream is not None else get_config("chat.stream") - # 提取内容 + # Extract content. from app.services.grok.services.chat import MessageExtractor from app.services.grok.services.assets import UploadService @@ -267,7 +267,7 @@ class VideoService: except ValueError as e: raise ValidationException(str(e)) - # 处理图片附件 + # Handle image attachments. image_url = None if attachments: upload_service = UploadService() @@ -281,7 +281,7 @@ class VideoService: finally: await upload_service.close() - # 生成视频 + # Generate video. service = VideoService() if image_url: response = await service.generate_from_image( @@ -292,7 +292,7 @@ class VideoService: token, prompt, aspect_ratio, video_length, resolution, preset ) - # 处理响应 + # Process response. if is_stream: processor = VideoStreamProcessor(model, token, think) return wrap_stream_with_usage( diff --git a/app/services/grok/utils/download.py b/app/services/grok/utils/download.py new file mode 100644 index 0000000..b3b930f --- /dev/null +++ b/app/services/grok/utils/download.py @@ -0,0 +1,5 @@ +"""Download service (compat wrapper).""" + +from app.services.grok.services.assets import DownloadService + +__all__ = ["DownloadService"]