From 9d6af73f35fec1fb785a2d5b93066b61329342a0 Mon Sep 17 00:00:00 2001 From: Chenyme <118253778+chenyme@users.noreply.github.com> Date: Fri, 6 Feb 2026 15:36:12 +0800 Subject: [PATCH] feat: update image processing API to support response format customization and enhance request validation --- .gitignore | 8 +- app/api/v1/files.py | 2 +- app/api/v1/image.py | 366 ++++++++++++++++++++-- app/services/grok/processors/processor.py | 30 +- 4 files changed, 379 insertions(+), 27 deletions(-) diff --git a/.gitignore b/.gitignore index fc9be0f..7b897ce 100644 --- a/.gitignore +++ b/.gitignore @@ -46,9 +46,11 @@ logs/ # Data data/*.json -data/token.json -data/tmp/ -data/.locks/ +data/tmp/* +data/.locks/* +app/data/*.json +app/data/tmp/* +app/data/.locks/* # Testing .pytest_cache/ diff --git a/app/api/v1/files.py b/app/api/v1/files.py index 46e7213..a9a4477 100644 --- a/app/api/v1/files.py +++ b/app/api/v1/files.py @@ -12,7 +12,7 @@ from app.core.logger import logger router = APIRouter(tags=["Files"]) # 缓存根目录 -BASE_DIR = Path(__file__).parent.parent.parent.parent / "data" / "tmp" +BASE_DIR = Path(__file__).parent.parent.parent / "data" / "tmp" IMAGE_DIR = BASE_DIR / "image" VIDEO_DIR = BASE_DIR / "video" diff --git a/app/api/v1/image.py b/app/api/v1/image.py index a4124cd..e874b8e 100644 --- a/app/api/v1/image.py +++ b/app/api/v1/image.py @@ -3,18 +3,22 @@ Image Generation API 路由 """ import asyncio +import base64 import random -from typing import List, Optional +from pathlib import Path +from typing import List, Optional, Union -from fastapi import APIRouter +from fastapi import APIRouter, File, Form, UploadFile from fastapi.responses import StreamingResponse, JSONResponse -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, ValidationError from app.services.grok.services.chat import GrokChatService +from app.services.grok.services.assets import UploadService from app.services.grok.models.model import ModelService from app.services.grok.processors.processor import ImageStreamProcessor, ImageCollectProcessor from app.services.token import get_token_manager, EffortType from app.core.exceptions import ValidationException, AppException, ErrorType +from app.core.config import get_config from app.core.logger import logger @@ -29,12 +33,28 @@ class ImageGenerationRequest(BaseModel): n: Optional[int] = Field(1, ge=1, le=10, description="生成数量 (1-10)") size: Optional[str] = Field("1024x1024", description="图片尺寸 (暂不支持)") quality: Optional[str] = Field("standard", description="图片质量 (暂不支持)") - response_format: Optional[str] = Field("b64_json", description="响应格式") + response_format: Optional[str] = Field(None, description="响应格式") style: Optional[str] = Field(None, description="风格 (暂不支持)") stream: Optional[bool] = Field(False, description="是否流式输出") -def validate_request(request: ImageGenerationRequest): + +class ImageEditRequest(BaseModel): + """图片编辑请求 - OpenAI 兼容""" + + prompt: str = Field(..., description="编辑描述") + model: Optional[str] = Field("grok-imagine-1.0", description="模型名称") + image: Optional[Union[str, List[str]]] = Field(None, description="待编辑图片文件") + n: Optional[int] = Field(1, ge=1, le=10, description="生成数量 (1-10)") + size: Optional[str] = Field("1024x1024", description="图片尺寸 (暂不支持)") + quality: Optional[str] = Field("standard", description="图片质量 (暂不支持)") + response_format: Optional[str] = Field(None, description="响应格式") + style: Optional[str] = Field(None, description="风格 (暂不支持)") + stream: Optional[bool] = Field(False, description="是否流式输出") + + + +def validate_generation_request(request: ImageGenerationRequest): """验证请求参数""" # 验证模型 - 通过 is_image 检查 model_info = ModelService.get(request.model) @@ -67,26 +87,81 @@ def validate_request(request: ImageGenerationRequest): code="invalid_stream_n", ) + if request.response_format: + allowed_formats = {"b64_json", "base64", "url"} + if request.response_format not in allowed_formats: + raise ValidationException( + message=f"response_format must be one of {sorted(allowed_formats)}", + param="response_format", + code="invalid_response_format", + ) -async def call_grok(token_mgr, token: str, prompt: str, model_info) -> List[str]: - """ - 调用 Grok 获取图片，返回 base64 列表 - """ + +def resolve_response_format(response_format: Optional[str]) -> str: + fmt = response_format or get_config("app.image_format", "url") + 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)}", + 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" + + +def validate_edit_request(request: ImageEditRequest, images: List[UploadFile]): + """验证图片编辑请求参数""" + validate_generation_request(request) + if not images: + raise ValidationException( + message="Image is required", + param="image", + code="missing_image", + ) + if len(images) > 16: + raise ValidationException( + message="Too many images. Maximum is 16.", + param="image", + code="invalid_image_count", + ) + + +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=f"Image Generation:{prompt}", + message=prompt, model=model_info.grok_model, mode=model_info.model_mode, - think=False, stream=True, + file_attachments=file_attachments, ) # 收集图片 - processor = ImageCollectProcessor(model_info.model_id, token) + processor = ImageCollectProcessor( + model_info.model_id, token, response_format=response_format + ) images = await processor.process(response) success = True return images @@ -124,8 +199,14 @@ async def create_image(request: ImageGenerationRequest): if request.stream is None: request.stream = False + if request.response_format is None: + request.response_format = resolve_response_format(None) + # 参数验证 - validate_request(request) + validate_generation_request(request) + + response_format = resolve_response_format(request.response_format) + response_field = response_field_name(response_format) # 获取 token try: @@ -160,14 +241,15 @@ async def create_image(request: ImageGenerationRequest): chat_service = GrokChatService() response = await chat_service.chat( token=token, - message=f"Image Generation:{request.prompt}", + message=f"Image Generation: {request.prompt}", model=model_info.grok_model, mode=model_info.model_mode, - think=False, stream=True, ) - processor = ImageStreamProcessor(model_info.model_id, token, n=request.n) + processor = ImageStreamProcessor( + model_info.model_id, token, n=request.n, response_format=response_format + ) # 包装流式响应，在成功完成时记录使用 async def _wrap_stream(stream): @@ -202,11 +284,23 @@ async def create_image(request: ImageGenerationRequest): if calls_needed == 1: # 单次调用 - all_images = await call_grok(token_mgr, token, request.prompt, model_info) + 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, request.prompt, model_info) + 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) @@ -231,7 +325,241 @@ async def create_image(request: ImageGenerationRequest): # 构建响应 import time - data = [{"b64_json": img} for img in selected_images] + data = [{response_field: img} for img in selected_images] + + return JSONResponse( + content={ + "created": int(time.time()), + "data": data, + "usage": { + "total_tokens": 0 + * len([img for img in selected_images if img != "error"]), + "input_tokens": 0, + "output_tokens": 0 + * len([img for img in selected_images if img != "error"]), + "input_tokens_details": {"text_tokens": 0, "image_tokens": 0}, + }, + } + ) + + + +@router.post("/images/edits") +async def edit_image( + prompt: str = Form(...), + image: List[UploadFile] = File(...), + model: Optional[str] = Form("grok-imagine-1.0"), + n: int = Form(1), + size: str = Form("1024x1024"), + quality: str = Form("standard"), + response_format: Optional[str] = Form(None), + style: Optional[str] = Form(None), + stream: Optional[bool] = Form(False), +): + """ + Image Edits API + + 同官方 API 格式，仅支持 multipart/form-data 文件上传 + """ + if response_format is None: + response_format = resolve_response_format(None) + + try: + edit_request = ImageEditRequest( + prompt=prompt, + model=model, + n=n, + size=size, + quality=quality, + response_format=response_format, + style=style, + stream=stream, + ) + except ValidationError as exc: + errors = exc.errors() + if errors: + first = errors[0] + loc = first.get("loc", []) + msg = first.get("msg", "Invalid request") + code = first.get("type", "invalid_value") + param_parts = [ + str(x) for x in loc if not (isinstance(x, int) or str(x).isdigit()) + ] + param = ".".join(param_parts) if param_parts else None + raise ValidationException(message=msg, param=param, code=code) + raise ValidationException(message="Invalid request", code="invalid_value") + + if edit_request.stream is None: + edit_request.stream = False + + response_format = resolve_response_format(edit_request.response_format) + edit_request.response_format = response_format + response_field = response_field_name(response_format) + + # 参数验证 + validate_edit_request(edit_request, image) + + max_image_bytes = 50 * 1024 * 1024 + allowed_types = {"image/png", "image/jpeg", "image/webp", "image/jpg"} + + images: List[str] = [] + for item in image: + content = await item.read() + await item.close() + if not content: + raise ValidationException( + message="File content is empty", + param="image", + code="empty_file", + ) + if len(content) > max_image_bytes: + raise ValidationException( + message="Image file too large. Maximum is 50MB.", + param="image", + code="file_too_large", + ) + mime = (item.content_type or "").lower() + if mime == "image/jpg": + mime = "image/jpeg" + ext = Path(item.filename or "").suffix.lower() + if mime not in allowed_types: + if ext in (".jpg", ".jpeg"): + mime = "image/jpeg" + elif ext == ".png": + mime = "image/png" + elif ext == ".webp": + mime = "image/webp" + else: + raise ValidationException( + message="Unsupported image type. Supported: png, jpg, webp.", + param="image", + code="invalid_image_type", + ) + 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, + ) + + # 获取模型信息 + model_info = ModelService.get(edit_request.model) + + # 上传图片 + file_ids: List[str] = [] + upload_service = UploadService() + try: + for image in images: + file_id, _ = await upload_service.upload(image, token) + file_ids.append(file_id) + finally: + await upload_service.close() + + # 流式模式 + if edit_request.stream: + chat_service = GrokChatService() + response = await chat_service.chat( + token=token, + message=f"Image Edit: {edit_request.prompt}", + model=model_info.grok_model, + mode=model_info.model_mode, + stream=True, + file_attachments=file_ids, + ) + + processor = ImageStreamProcessor( + 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)), + media_type="text/event-stream", + headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, + ) + + # 非流式模式 + n = edit_request.n + calls_needed = (n + 1) // 2 + + if calls_needed == 1: + all_images = await call_grok( + token_mgr, + token, + f"Image Edit: {edit_request.prompt}", + model_info, + file_attachments=file_ids, + response_format=response_format, + ) + else: + tasks = [ + call_grok( + token_mgr, + token, + f"Image Edit: {edit_request.prompt}", + model_info, + file_attachments=file_ids, + 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) + + 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") + + import time + + data = [{response_field: img} for img in selected_images] return JSONResponse( content={ diff --git a/app/services/grok/processors/processor.py b/app/services/grok/processors/processor.py index a351672..b7ca639 100644 --- a/app/services/grok/processors/processor.py +++ b/app/services/grok/processors/processor.py @@ -726,11 +726,20 @@ class VideoCollectProcessor(BaseProcessor): class ImageStreamProcessor(BaseProcessor): """图片生成流式响应处理器""" - def __init__(self, model: str, token: str = "", n: int = 1): + 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 响应 (覆盖基类)""" @@ -769,7 +778,7 @@ class ImageStreamProcessor(BaseProcessor): "image_generation.partial_image", { "type": "image_generation.partial_image", - "b64_json": "", + self.response_field: "", "index": out_index, "progress": progress, }, @@ -780,6 +789,11 @@ class ImageStreamProcessor(BaseProcessor): if mr := resp.get("modelResponse"): if urls := mr.get("generatedImageUrls"): for url in urls: + if self.response_format == "url": + processed = await self.process_url(url, "image") + if processed: + final_images.append(processed) + continue dl_service = self._get_dl() base64_data = await dl_service.to_base64( url, self.token, "image" @@ -804,7 +818,7 @@ class ImageStreamProcessor(BaseProcessor): "image_generation.completed", { "type": "image_generation.completed", - "b64_json": b64, + self.response_field: b64, "index": out_index, "usage": { "total_tokens": 50, @@ -856,8 +870,11 @@ class ImageStreamProcessor(BaseProcessor): class ImageCollectProcessor(BaseProcessor): """图片生成非流式响应处理器""" - def __init__(self, model: str, token: str = ""): + 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]: """处理并收集图片""" @@ -879,6 +896,11 @@ class ImageCollectProcessor(BaseProcessor): if mr := resp.get("modelResponse"): if urls := mr.get("generatedImageUrls"): for url in urls: + if self.response_format == "url": + processed = await self.process_url(url, "image") + if processed: + images.append(processed) + continue dl_service = self._get_dl() base64_data = await dl_service.to_base64( url, self.token, "image"