Merge branch 'main' of https://github.com/ZIC143/grok2apifork

Author: ZIC143

app/api/v1/admin.py

@@ -1,7 +1,7 @@
-from fastapi import APIRouter, Depends, HTTPException, Request, Query, Body
+from fastapi import APIRouter, Depends, HTTPException, Request, Query, Body, WebSocket
 from fastapi.responses import HTMLResponse, RedirectResponse
 from pydantic import BaseModel
-from typing import Any
+from typing import Any, Optional
 
 from app.core.auth import verify_api_key
 from app.core.config import config, get_config
@@ -11,9 +11,21 @@
 import aiofiles
 import asyncio
 import json
+import time
+import uuid
+import orjson
+from starlette.websockets import WebSocketDisconnect, WebSocketState
 from app.core.logger import logger
 from app.services.register import get_auto_register_manager
 from app.services.api_keys import api_key_manager
+from app.services.grok.model import ModelService
+from app.services.grok.imagine_generation import (
+    collect_experimental_generation_images,
+    is_valid_image_value as is_valid_imagine_image_value,
+    resolve_aspect_ratio as resolve_imagine_aspect_ratio,
+)
+from app.services.token import get_token_manager
+from app.core.auth import _load_legacy_api_keys
 
 
 router = APIRouter()
@@ -82,6 +94,224 @@ async def admin_chat_page():
     """在线聊天页（后台入口）"""
     return await render_template("chat/chat_admin.html")
 
+
+async def _verify_ws_api_key(websocket: WebSocket) -> bool:
+    api_key = str(get_config("app.api_key", "") or "").strip()
+    legacy_keys = await _load_legacy_api_keys()
+    if not api_key and not legacy_keys:
+        return True
+    token = str(websocket.query_params.get("api_key") or "").strip()
+    if not token:
+        return False
+    if (api_key and token == api_key) or token in legacy_keys:
+        return True
+    try:
+        await api_key_manager.init()
+        if api_key_manager.validate_key(token):
+            return True
+    except Exception as e:
+        logger.warning(f"Imagine ws api_key validation fallback failed: {e}")
+    return False
+
+
+async def _collect_imagine_batch(token: str, prompt: str, aspect_ratio: str) -> list[str]:
+    return await collect_experimental_generation_images(
+        token=token,
+        prompt=prompt,
+        n=6,
+        response_format="b64_json",
+        aspect_ratio=aspect_ratio,
+        concurrency=1,
+    )
+
+
+@router.websocket("/api/v1/admin/imagine/ws")
+async def admin_imagine_ws(websocket: WebSocket):
+    if not await _verify_ws_api_key(websocket):
+        await websocket.close(code=1008)
+        return
+
+    await websocket.accept()
+    stop_event = asyncio.Event()
+    run_task: Optional[asyncio.Task] = None
+
+    async def _send(payload: dict) -> bool:
+        try:
+            await websocket.send_text(orjson.dumps(payload).decode())
+            return True
+        except Exception:
+            return False
+
+    async def _stop_run():
+        nonlocal run_task
+        stop_event.set()
+        if run_task and not run_task.done():
+            run_task.cancel()
+            try:
+                await run_task
+            except Exception:
+                pass
+        run_task = None
+        stop_event.clear()
+
+    async def _run(prompt: str, aspect_ratio: str):
+        model_id = "grok-imagine-1.0"
+        model_info = ModelService.get(model_id)
+        if not model_info or not model_info.is_image:
+            await _send(
+                {
+                    "type": "error",
+                    "message": "Image model is not available.",
+                    "code": "model_not_supported",
+                }
+            )
+            return
+
+        token_mgr = await get_token_manager()
+        sequence = 0
+        run_id = uuid.uuid4().hex
+        await _send(
+            {
+                "type": "status",
+                "status": "running",
+                "prompt": prompt,
+                "aspect_ratio": aspect_ratio,
+                "run_id": run_id,
+            }
+        )
+
+        while not stop_event.is_set():
+            try:
+                await token_mgr.reload_if_stale()
+                token = token_mgr.get_token_for_model(model_info.model_id)
+                if not token:
+                    await _send(
+                        {
+                            "type": "error",
+                            "message": "No available tokens. Please try again later.",
+                            "code": "rate_limit_exceeded",
+                        }
+                    )
+                    await asyncio.sleep(2)
+                    continue
+
+                start_at = time.time()
+                images = await _collect_imagine_batch(token, prompt, aspect_ratio)
+                elapsed_ms = int((time.time() - start_at) * 1000)
+
+                sent_any = False
+                for image_b64 in images:
+                    if not is_valid_imagine_image_value(image_b64):
+                        continue
+                    sent_any = True
+                    sequence += 1
+                    ok = await _send(
+                        {
+                            "type": "image",
+                            "b64_json": image_b64,
+                            "sequence": sequence,
+                            "created_at": int(time.time() * 1000),
+                            "elapsed_ms": elapsed_ms,
+                            "aspect_ratio": aspect_ratio,
+                            "run_id": run_id,
+                        }
+                    )
+                    if not ok:
+                        stop_event.set()
+                        break
+
+                if sent_any:
+                    try:
+                        await token_mgr.sync_usage(
+                            token,
+                            model_info.model_id,
+                            consume_on_fail=True,
+                            is_usage=True,
+                        )
+                    except Exception as e:
+                        logger.warning(f"Imagine ws token sync failed: {e}")
+                else:
+                    await _send(
+                        {
+                            "type": "error",
+                            "message": "Image generation returned empty data.",
+                            "code": "empty_image",
+                        }
+                    )
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.warning(f"Imagine stream error: {e}")
+                await _send(
+                    {
+                        "type": "error",
+                        "message": str(e),
+                        "code": "internal_error",
+                    }
+                )
+                await asyncio.sleep(1.5)
+
+        await _send({"type": "status", "status": "stopped", "run_id": run_id})
+
+    try:
+        while True:
+            try:
+                raw = await websocket.receive_text()
+            except (RuntimeError, WebSocketDisconnect):
+                break
+
+            try:
+                payload = orjson.loads(raw)
+            except Exception:
+                await _send(
+                    {
+                        "type": "error",
+                        "message": "Invalid message format.",
+                        "code": "invalid_payload",
+                    }
+                )
+                continue
+
+            msg_type = payload.get("type")
+            if msg_type == "start":
+                prompt = str(payload.get("prompt") or "").strip()
+                if not prompt:
+                    await _send(
+                        {
+                            "type": "error",
+                            "message": "Prompt cannot be empty.",
+                            "code": "empty_prompt",
+                        }
+                    )
+                    continue
+                ratio = resolve_imagine_aspect_ratio(str(payload.get("aspect_ratio") or "2:3").strip())
+                await _stop_run()
+                run_task = asyncio.create_task(_run(prompt, ratio))
+            elif msg_type == "stop":
+                await _stop_run()
+            elif msg_type == "ping":
+                await _send({"type": "pong"})
+            else:
+                await _send(
+                    {
+                        "type": "error",
+                        "message": "Unknown command.",
+                        "code": "unknown_command",
+                    }
+                )
+    except WebSocketDisconnect:
+        logger.debug("WebSocket disconnected by client")
+    except Exception as e:
+        logger.warning(f"WebSocket error: {e}")
+    finally:
+        await _stop_run()
+        try:
+            if websocket.client_state == WebSocketState.CONNECTED:
+                await websocket.close(code=1000, reason="Server closing connection")
+        except Exception as e:
+            logger.debug(f"WebSocket close ignored: {e}")
+
+
 @router.post("/api/v1/admin/login")
 async def admin_login_api(request: Request, body: AdminLoginBody | None = Body(default=None)):
     """管理后台登录验证（用户名+密码）

app/api/v1/image.py

@@ -25,6 +25,13 @@
     ImagineExperimentalService,
     resolve_image_generation_method,
 )
+from app.services.grok.imagine_generation import (
+    call_experimental_generation_once,
+    collect_experimental_generation_images,
+    dedupe_images as dedupe_imagine_images,
+    is_valid_image_value as is_valid_imagine_image_value,
+    resolve_aspect_ratio as resolve_imagine_aspect_ratio,
+)
 from app.services.grok.model import ModelService
 from app.services.grok.processor import ImageCollectProcessor, ImageStreamProcessor
 from app.services.quota import enforce_daily_quota
@@ -236,44 +243,15 @@ def _image_generation_method() -> str:
 
 
 def resolve_aspect_ratio(size: Optional[str]) -> str:
-    value = str(size or "").strip().lower()
-    if value in {"16:9", "9:16", "1:1", "2:3", "3:2"}:
-        return value
-
-    mapping = {
-        "1024x1024": "1:1",
-        "512x512": "1:1",
-        "1024x576": "16:9",
-        "1280x720": "16:9",
-        "1536x864": "16:9",
-        "576x1024": "9:16",
-        "720x1280": "9:16",
-        "864x1536": "9:16",
-        "1024x1536": "2:3",
-        "512x768": "2:3",
-        "768x1024": "2:3",
-        "1536x1024": "3:2",
-        "768x512": "3:2",
-        "1024x768": "3:2",
-    }
-    return mapping.get(value, "2:3")
+    return resolve_imagine_aspect_ratio(size)
 
 
 def _is_valid_image_value(value: Any) -> bool:
-    return isinstance(value, str) and bool(value) and value != "error"
+    return is_valid_imagine_image_value(value)
 
 
 def _dedupe_images(images: List[str]) -> List[str]:
-    out: List[str] = []
-    seen: set[str] = set()
-    for image in images:
-        if not isinstance(image, str):
-            continue
-        if image in seen:
-            continue
-        seen.add(image)
-        out.append(image)
-    return out
+    return dedupe_imagine_images(images)
 
 
 async def _gather_limited(
@@ -330,14 +308,13 @@ async def call_grok_experimental_ws(
     n: int = 4,
     aspect_ratio: str = "2:3",
 ) -> List[str]:
-    service = ImagineExperimentalService()
-    raw_urls = await service.generate_ws(
+    return await call_experimental_generation_once(
         token=token,
         prompt=prompt,
+        response_format=response_format,
         n=n,
         aspect_ratio=aspect_ratio,
     )
-    return await service.convert_urls(token=token, urls=raw_urls, response_format=response_format)
 
 
 async def call_grok_experimental_edit(
@@ -365,36 +342,14 @@ async def _collect_experimental_generation_images(
     aspect_ratio: str,
     concurrency: int,
 ) -> List[str]:
-    calls_needed = max(1, (n + 3) // 4)
-    task_factories: List[Callable[[], Awaitable[List[str]]]] = []
-    remain = n
-    for _ in range(calls_needed):
-        target_n = max(1, min(4, remain))
-        remain -= target_n
-        task_factories.append(
-            lambda target_n=target_n: call_grok_experimental_ws(
-                token,
-                prompt,
-                response_format=response_format,
-                n=target_n,
-                aspect_ratio=aspect_ratio,
-            )
-        )
-    results = await _gather_limited(
-        task_factories,
-        max_concurrency=min(calls_needed, max(1, int(concurrency or 1))),
+    return await collect_experimental_generation_images(
+        token=token,
+        prompt=prompt,
+        n=n,
+        response_format=response_format,
+        aspect_ratio=aspect_ratio,
+        concurrency=concurrency,
     )
-    all_images: List[str] = []
-    for result in results:
-        if isinstance(result, Exception):
-            logger.warning(f"Experimental imagine websocket call failed: {result}")
-            continue
-        if isinstance(result, list):
-            all_images.extend(result)
-    all_images = _dedupe_images(all_images)
-    if not any(_is_valid_image_value(item) for item in all_images):
-        raise UpstreamException("Experimental imagine websocket returned no images")
-    return all_images
 
 
 async def _experimental_stream_generation(