feat: update image processing API to support response format customization and enhance request validation

.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/

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"
 

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={

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"