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

refactor: remove gRPC-Web protocol implementation and integrate WebSocket reverse handlers for image service

Browse Source
This commit is contained in:
Chenyme committed 2026-02-10 11:46:22 +08:00
1 parent 4573a1928e
commit 06357ad6c5
9 files changed
+345 -576
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

app/services/grok/protocols/grpc_web.py
-157
View File
@@ -1,157 +0,0 @@
"""
gRPC-Web 协议工具
提供 framing 编码/解码、trailer 解析等通用功能。
支持 application/grpc-web+proto 和 application/grpc-web-text (base64) 两种格式。
"""
from __future__ import annotations
import base64
import re
import struct
from dataclasses import dataclass
from typing import Dict, List, Mapping, Tuple
from urllib.parse import unquote
_B64_RE = re.compile(rb"^[A-Za-z0-9+/=\r\n]+$")
def encode_grpc_web_payload(data: bytes) -> bytes:
"""
编码 gRPC-Web data frame
Frame format:
1-byte flags + 4-byte big-endian length + message bytes
"""
return b"\x00" + struct.pack(">I", len(data)) + data
def _maybe_decode_grpc_web_text(body: bytes, content_type: str | None) -> bytes:
"""处理 grpc-web-text 模式的 base64 解码"""
ct = (content_type or "").lower()
if "grpc-web-text" in ct:
compact = b"".join(body.split())
return base64.b64decode(compact, validate=False)
# 启发式：body 仅包含 base64 字符才尝试解码
head = body[: min(len(body), 2048)]
if head and _B64_RE.fullmatch(head):
compact = b"".join(body.split())
try:
return base64.b64decode(compact, validate=True)
except Exception:
return body
return body
def _parse_trailer_block(payload: bytes) -> Dict[str, str]:
"""解析 trailer frame 内容"""
text = payload.decode("utf-8", errors="replace")
lines = [ln for ln in re.split(r"\r\n|\n", text) if ln]
trailers: Dict[str, str] = {}
for ln in lines:
if ":" not in ln:
continue
k, v = ln.split(":", 1)
trailers[k.strip().lower()] = v.strip()
# grpc-message 可能是 percent-encoding
if "grpc-message" in trailers:
trailers["grpc-message"] = unquote(trailers["grpc-message"])
return trailers
def parse_grpc_web_response(
body: bytes,
content_type: str | None = None,
headers: Mapping[str, str] | None = None,
) -> Tuple[List[bytes], Dict[str, str]]:
"""
解析 gRPC-Web 响应
Returns:
(messages, trailers): data frames 列表和合并后的 trailers
"""
decoded = _maybe_decode_grpc_web_text(body, content_type)
messages: List[bytes] = []
trailers: Dict[str, str] = {}
i = 0
n = len(decoded)
while i < n:
if n - i < 5:
break
flag = decoded[i]
length = int.from_bytes(decoded[i + 1 : i + 5], "big")
i += 5
if n - i < length:
break
payload = decoded[i : i + length]
i += length
if flag & 0x80: # trailer frame
trailers.update(_parse_trailer_block(payload))
elif flag & 0x01: # compressed (不支持)
raise ValueError("grpc-web compressed flag not supported")
else:
messages.append(payload)
# 兼容：grpc-status 可能在 response headers 中
if headers:
lower = {k.lower(): v for k, v in headers.items()}
if "grpc-status" in lower and "grpc-status" not in trailers:
trailers["grpc-status"] = str(lower["grpc-status"]).strip()
if "grpc-message" in lower and "grpc-message" not in trailers:
trailers["grpc-message"] = unquote(str(lower["grpc-message"]).strip())
return messages, trailers
@dataclass(frozen=True)
class GrpcStatus:
code: int
message: str = ""
@property
def ok(self) -> bool:
return self.code == 0
@property
def http_equiv(self) -> int:
"""映射到类 HTTP 状态码"""
mapping = {
0: 200, # OK
16: 401, # UNAUTHENTICATED
7: 403, # PERMISSION_DENIED
8: 429, # RESOURCE_EXHAUSTED
4: 504, # DEADLINE_EXCEEDED
14: 503, # UNAVAILABLE
}
return mapping.get(self.code, 502)
def get_grpc_status(trailers: Mapping[str, str]) -> GrpcStatus:
"""从 trailers 提取 gRPC 状态"""
raw = str(trailers.get("grpc-status", "")).strip()
msg = str(trailers.get("grpc-message", "")).strip()
try:
code = int(raw)
except Exception:
code = -1
return GrpcStatus(code=code, message=msg)
__all__ = [
"encode_grpc_web_payload",
"parse_grpc_web_response",
"get_grpc_status",
"GrpcStatus",
]
app/services/grok/services/image.py
+3 -271
View File
@@ -2,278 +2,10 @@
Grok Imagine WebSocket image service.
"""
import asyncio
import certifi
import json
import re
import ssl
import time
import uuid
from typing import AsyncGenerator, Dict, Optional
import aiohttp
from app.core.config import get_config
from app.core.logger import logger
from app.services.grok.utils.headers import build_sso_cookie
from app.services.reverse.utils.websocket import resolve_proxy
WS_URL = "wss://grok.com/ws/imagine/listen"
from app.services.reverse.ws_imagine import ImagineWebSocketReverse
class _BlockedError(Exception):
pass
class ImageService:
"""Grok Imagine WebSocket image service."""
def __init__(self):
self._ssl_context = ssl.create_default_context()
self._ssl_context.load_verify_locations(certifi.where())
self._url_pattern = re.compile(r"/images/([a-f0-9-]+)\.(png|jpg|jpeg)")
def _resolve_proxy(self) -> tuple[aiohttp.BaseConnector, Optional[str]]:
proxy_url = get_config("network.base_proxy_url")
return resolve_proxy(proxy_url, self._ssl_context)
def _get_ws_headers(self, token: str) -> Dict[str, str]:
cookie = build_sso_cookie(token, include_rw=True)
user_agent = get_config("security.user_agent")
return {
"Cookie": cookie,
"Origin": "https://grok.com",
"User-Agent": user_agent,
"Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
"Cache-Control": "no-cache",
"Pragma": "no-cache",
}
def _extract_image_id(self, url: str) -> Optional[str]:
match = self._url_pattern.search(url or "")
return match.group(1) if match else None
def _is_final_image(self, url: str, blob_size: int) -> bool:
return (url or "").lower().endswith(
(".jpg", ".jpeg")
) and blob_size > get_config("image.image_ws_final_min_bytes")
def _classify_image(self, url: str, blob: str) -> Optional[Dict[str, object]]:
if not url or not blob:
return None
image_id = self._extract_image_id(url) or uuid.uuid4().hex
blob_size = len(blob)
is_final = self._is_final_image(url, blob_size)
stage = (
"final"
if is_final
else (
"medium"
if blob_size > get_config("image.image_ws_medium_min_bytes")
else "preview"
)
)
return {
"type": "image",
"image_id": image_id,
"stage": stage,
"blob": blob,
"blob_size": blob_size,
"url": url,
"is_final": is_final,
}
async def stream(
self,
token: str,
prompt: str,
aspect_ratio: str = "2:3",
n: int = 1,
enable_nsfw: bool = True,
max_retries: int = None,
) -> AsyncGenerator[Dict[str, object], None]:
retries = max(1, max_retries if max_retries is not None else 1)
logger.info(
f"Image generation: prompt='{prompt[:50]}...', n={n}, ratio={aspect_ratio}, nsfw={enable_nsfw}"
)
for attempt in range(retries):
try:
yielded_any = False
async for item in self._stream_once(
token, prompt, aspect_ratio, n, enable_nsfw
):
yielded_any = True
yield item
return
except _BlockedError:
if yielded_any or attempt + 1 >= retries:
if not yielded_any:
yield {
"type": "error",
"error_code": "blocked",
"error": "blocked_no_final_image",
}
return
logger.warning(f"WebSocket blocked, retry {attempt + 1}/{retries}")
except Exception as e:
logger.error(f"WebSocket stream failed: {e}")
return
async def _stream_once(
self,
token: str,
prompt: str,
aspect_ratio: str,
n: int,
enable_nsfw: bool,
) -> AsyncGenerator[Dict[str, object], None]:
request_id = str(uuid.uuid4())
headers = self._get_ws_headers(token)
timeout = float(get_config("network.timeout"))
blocked_seconds = float(get_config("image.image_ws_blocked_seconds"))
try:
connector, proxy = self._resolve_proxy()
except Exception as e:
logger.error(f"WebSocket proxy setup failed: {e}")
return
try:
async with aiohttp.ClientSession(connector=connector) as session:
async with session.ws_connect(
WS_URL,
headers=headers,
heartbeat=20,
receive_timeout=timeout,
proxy=proxy,
) as ws:
message = {
"type": "conversation.item.create",
"timestamp": int(time.time() * 1000),
"item": {
"type": "message",
"content": [
{
"requestId": request_id,
"text": prompt,
"type": "input_text",
"properties": {
"section_count": 0,
"is_kids_mode": False,
"enable_nsfw": enable_nsfw,
"skip_upsampler": False,
"is_initial": False,
"aspect_ratio": aspect_ratio,
},
}
],
},
}
await ws.send_json(message)
logger.info(f"WebSocket request sent: {prompt[:80]}...")
images = {}
completed = 0
start_time = last_activity = time.time()
medium_received_time = None
while time.time() - start_time < timeout:
try:
ws_msg = await asyncio.wait_for(ws.receive(), timeout=5.0)
except asyncio.TimeoutError:
if (
medium_received_time
and completed == 0
and time.time() - medium_received_time
> min(10, blocked_seconds)
):
raise _BlockedError()
if completed > 0 and time.time() - last_activity > 10:
logger.info(
f"WebSocket idle timeout, collected {completed} images"
)
break
continue
if ws_msg.type == aiohttp.WSMsgType.TEXT:
last_activity = time.time()
msg = json.loads(ws_msg.data)
msg_type = msg.get("type")
if msg_type == "image":
info = self._classify_image(
msg.get("url", ""), msg.get("blob", "")
)
if not info:
continue
image_id = info["image_id"]
existing = images.get(image_id, {})
if (
info["stage"] == "medium"
and medium_received_time is None
):
medium_received_time = time.time()
if info["is_final"] and not existing.get("is_final"):
completed += 1
logger.debug(
f"Final image received: id={image_id}, size={info['blob_size']}"
)
images[image_id] = {
"is_final": info["is_final"]
or existing.get("is_final")
}
yield info
elif msg_type == "error":
logger.warning(
f"WebSocket error: {msg.get('err_code', '')} - {msg.get('err_msg', '')}"
)
yield {
"type": "error",
"error_code": msg.get("err_code", ""),
"error": msg.get("err_msg", ""),
}
return
if completed >= n:
logger.info(
f"WebSocket collected {completed} final images"
)
break
if (
medium_received_time
and completed == 0
and time.time() - medium_received_time > blocked_seconds
):
raise _BlockedError()
elif ws_msg.type in (
aiohttp.WSMsgType.CLOSED,
aiohttp.WSMsgType.ERROR,
):
logger.warning(f"WebSocket closed/error: {ws_msg.type}")
yield {
"type": "error",
"error_code": "ws_closed",
"error": f"websocket closed: {ws_msg.type}",
}
break
except aiohttp.ClientError as e:
logger.error(f"WebSocket connection error: {e}")
yield {"type": "error", "error_code": "connection_failed", "error": str(e)}
image_service = ImageService()
ImageService = ImagineWebSocketReverse
image_service = ImagineWebSocketReverse()
__all__ = ["image_service", "ImageService"]
app/services/grok/services/voice.py
+2 -1
View File
@@ -22,10 +22,11 @@ class VoiceService:
) -> Dict[str, Any]:
browser = get_config("security.browser")
async with AsyncSession(impersonate=browser) as session:
return await LivekitTokenReverse.request(
response = await LivekitTokenReverse.request(
session,
token=token,
voice=voice,
personality=personality,
speed=speed,
)
return response.json()
app/services/reverse/__init__.py
+2
View File
@@ -10,6 +10,7 @@ from .nsfw_mgmt import NsfwMgmtReverse
from .rate_limits import RateLimitsReverse
from .set_birth import SetBirthReverse
from .ws_livekit import LivekitTokenReverse, LivekitWebSocketReverse
from .ws_imagine import ImagineWebSocketReverse
from .utils.headers import build_headers
from .utils.statsig import StatsigGenerator
@@ -25,6 +26,7 @@ __all__ = [
"SetBirthReverse",
"LivekitTokenReverse",
"LivekitWebSocketReverse",
"ImagineWebSocketReverse",
"StatsigGenerator",
"build_headers",
]
app/services/reverse/utils/grpc.py
+4 -23
View File
@@ -6,7 +6,7 @@ import base64
import re
import struct
from dataclasses import dataclass
from typing import Dict, List, Mapping, Tuple
from typing import Dict, List, Mapping, Optional, Tuple
from urllib.parse import unquote
@@ -41,7 +41,7 @@ class GrpcClient:
return b"\x00" + struct.pack(">I", len(data)) + data
@staticmethod
def _maybe_decode_grpc_web_text(body: bytes, content_type: str | None) -> bytes:
def _maybe_decode_grpc_web_text(body: bytes, content_type: Optional[str]) -> bytes:
ct = (content_type or "").lower()
if "grpc-web-text" in ct:
compact = b"".join(body.split())
@@ -77,8 +77,8 @@ class GrpcClient:
def parse_response(
cls,
body: bytes,
content_type: str | None = None,
headers: Mapping[str, str] | None = None,
content_type: Optional[str] = None,
headers: Optional[Mapping[str, str]] = None,
) -> Tuple[List[bytes], Dict[str, str]]:
decoded = cls._maybe_decode_grpc_web_text(body, content_type)
@@ -128,26 +128,7 @@ class GrpcClient:
return GrpcStatus(code=code, message=msg)
def encode_grpc_web_payload(data: bytes) -> bytes:
return GrpcClient.encode_payload(data)
def parse_grpc_web_response(
body: bytes,
content_type: str | None = None,
headers: Mapping[str, str] | None = None,
) -> Tuple[List[bytes], Dict[str, str]]:
return GrpcClient.parse_response(body, content_type=content_type, headers=headers)
def get_grpc_status(trailers: Mapping[str, str]) -> GrpcStatus:
return GrpcClient.get_status(trailers)
__all__ = [
"encode_grpc_web_payload",
"parse_grpc_web_response",
"get_grpc_status",
"GrpcStatus",
"GrpcClient",
]
app/services/reverse/utils/headers.py
+43 -13
View File
@@ -10,9 +10,15 @@ from app.core.config import get_config
from app.services.reverse.utils.statsig import StatsigGenerator
def _build_sso_cookie(sso_token: str) -> str:
def build_sso_cookie(sso_token: str) -> str:
"""
Build SSO Cookie string.
Args:
sso_token: str, the SSO token.
Returns:
str: The SSO Cookie string.
"""
# Format
sso_token = sso_token[4:] if sso_token.startswith("sso=") else sso_token
@@ -28,20 +34,44 @@ def _build_sso_cookie(sso_token: str) -> str:
return cookie
def build_headers(
cookie_token: str,
content_type: Optional[str] = None,
origin: Optional[str] = None,
referer: Optional[str] = None,
) -> Dict[str, str]:
def build_ws_headers(token: Optional[str] = None, origin: Optional[str] = None, extra: Optional[Dict[str, str]] = None) -> Dict[str, str]:
"""
Build headers for WebSocket requests.
Args:
token: Optional[str], the SSO token for Cookie. Defaults to None.
origin: Optional[str], the Origin value. Defaults to "https://grok.com" if not provided.
extra: Optional[Dict[str, str]], extra headers to merge. Defaults to None.
Returns:
Dict[str, str]: The headers dictionary.
"""
headers = {
"Origin": origin or "https://grok.com",
"User-Agent": get_config("security.user_agent"),
"Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
"Cache-Control": "no-cache",
"Pragma": "no-cache",
}
if token:
headers["Cookie"] = build_sso_cookie(token)
if extra:
headers.update(extra)
return headers
def build_headers(cookie_token: str, content_type: Optional[str] = None, origin: Optional[str] = None, referer: Optional[str] = None) -> Dict[str, str]:
"""
Build headers for reverse interfaces.
Args:
cookie_token: The SSO token.
content_type: Optional Content-Type value.
origin: Optional Origin value. Defaults to "https://grok.com" if not provided.
referer: Optional Referer value. Defaults to "https://grok.com/" if not provided.
cookie_token: str, the SSO token.
content_type: Optional[str], the Content-Type value.
origin: Optional[str], the Origin value. Defaults to "https://grok.com" if not provided.
referer: Optional[str], the Referer value. Defaults to "https://grok.com/" if not provided.
Returns:
Dict[str, str]: The headers dictionary.
@@ -64,7 +94,7 @@ def build_headers(
}
# Cookie
headers["Cookie"] = _build_sso_cookie(cookie_token)
headers["Cookie"] = build_sso_cookie(cookie_token)
# Content-Type and Accept/Sec-Fetch-Dest
if content_type and content_type == "application/json":
@@ -101,4 +131,4 @@ def build_headers(
return headers
__all__ = ["build_headers"]
__all__ = ["build_headers", "build_sso_cookie", "build_ws_headers"]
app/services/reverse/utils/websocket.py
+12 -20
View File
@@ -37,14 +37,12 @@ def _normalize_socks_proxy(proxy_url: str) -> tuple[str, Optional[bool]]:
return proxy_url, rdns
def resolve_proxy(
proxy_url: str | None, ssl_context: ssl.SSLContext
) -> tuple[aiohttp.BaseConnector, Optional[str]]:
def resolve_proxy(proxy_url: Optional[str] = None, ssl_context: ssl.SSLContext = _default_ssl_context()) -> tuple[aiohttp.BaseConnector, Optional[str]]:
"""Resolve proxy connector.
Args:
proxy_url: str, the proxy URL.
ssl_context: ssl.SSLContext, the SSL context.
proxy_url: Optional[str], the proxy URL. Defaults to None.
ssl_context: ssl.SSLContext, the SSL context. Defaults to _default_ssl_context().
Returns:
tuple[aiohttp.BaseConnector, Optional[str]]: The proxy connector and the proxy URL.
@@ -73,11 +71,7 @@ def resolve_proxy(
class WebSocketConnection:
"""WebSocket connection wrapper."""
def __init__(
self,
session: aiohttp.ClientSession,
ws: aiohttp.ClientWebSocketResponse,
) -> None:
def __init__(self, session: aiohttp.ClientSession, ws: aiohttp.ClientWebSocketResponse) -> None:
self.session = session
self.ws = ws
@@ -96,22 +90,22 @@ class WebSocketConnection:
class WebSocketClient:
"""WebSocket client with proxy support."""
def __init__(self, proxy: str | None = None) -> None:
def __init__(self, proxy: Optional[str] = None) -> None:
self.proxy = proxy or get_config("network.base_proxy_url")
self._ssl_context = _default_ssl_context()
async def connect(
self,
url: str,
headers: Mapping[str, str] | None = None,
timeout: float | aiohttp.ClientTimeout | None = None,
headers: Optional[Mapping[str, str]] = None,
ws_kwargs: Optional[Mapping[str, object]] = None,
) -> WebSocketConnection:
"""Connect to the WebSocket.
Args:
url: str, the URL to connect to.
headers: Mapping[str, str], the headers to send.
timeout: float | aiohttp.ClientTimeout | None, the timeout.
headers: Optional[Mapping[str, str]], the headers to send. Defaults to None.
ws_kwargs: Optional[Mapping[str, object]], extra ws_connect kwargs. Defaults to None.
Returns:
WebSocketConnection: The WebSocket connection.
@@ -120,20 +114,18 @@ class WebSocketClient:
connector, proxy = resolve_proxy(self.proxy, self._ssl_context)
# Build client timeout
client_timeout = (
timeout
if isinstance(timeout, aiohttp.ClientTimeout)
else aiohttp.ClientTimeout(total=timeout)
)
client_timeout = aiohttp.ClientTimeout(total=get_config("network.timeout"))
# Create session
session = aiohttp.ClientSession(connector=connector, timeout=client_timeout)
try:
extra_kwargs = dict(ws_kwargs or {})
ws = await session.ws_connect(
url,
headers=headers,
proxy=proxy,
ssl=self._ssl_context,
**extra_kwargs,
)
return WebSocketConnection(session, ws)
except Exception:
app/services/reverse/ws_imagine.py
+262
View File
@@ -0,0 +1,262 @@
"""
Reverse interface: Imagine WebSocket image stream.
"""
import asyncio
import orjson
import re
import time
import uuid
from typing import AsyncGenerator, Dict, Optional
import aiohttp
from app.core.config import get_config
from app.core.logger import logger
from app.services.reverse.utils.headers import build_ws_headers
from app.services.reverse.utils.websocket import WebSocketClient
WS_IMAGINE_URL = "wss://grok.com/ws/imagine/listen"
class _BlockedError(Exception):
pass
class ImagineWebSocketReverse:
"""Imagine WebSocket reverse interface."""
def __init__(self) -> None:
self._url_pattern = re.compile(r"/images/([a-f0-9-]+)\.(png|jpg|jpeg)")
self._client = WebSocketClient()
def _extract_image_id(self, url: str) -> Optional[str]:
match = self._url_pattern.search(url or "")
return match.group(1) if match else None
def _is_final_image(self, url: str, blob_size: int, final_min_bytes: int) -> bool:
return (url or "").lower().endswith((".jpg", ".jpeg")) and blob_size > final_min_bytes
def _classify_image(self, url: str, blob: str, final_min_bytes: int, medium_min_bytes: int) -> Optional[Dict[str, object]]:
if not url or not blob:
return None
image_id = self._extract_image_id(url) or uuid.uuid4().hex
blob_size = len(blob)
is_final = self._is_final_image(url, blob_size, final_min_bytes)
stage = (
"final"
if is_final
else ("medium" if blob_size > medium_min_bytes else "preview")
)
return {
"type": "image",
"image_id": image_id,
"stage": stage,
"blob": blob,
"blob_size": blob_size,
"url": url,
"is_final": is_final,
}
def _build_request_message(self, request_id: str, prompt: str, aspect_ratio: str, enable_nsfw: bool) -> Dict[str, object]:
return {
"type": "conversation.item.create",
"timestamp": int(time.time() * 1000),
"item": {
"type": "message",
"content": [
{
"requestId": request_id,
"text": prompt,
"type": "input_text",
"properties": {
"section_count": 0,
"is_kids_mode": False,
"enable_nsfw": enable_nsfw,
"skip_upsampler": False,
"is_initial": False,
"aspect_ratio": aspect_ratio,
},
}
],
},
}
async def stream(
self,
token: str,
prompt: str,
aspect_ratio: str = "2:3",
n: int = 1,
enable_nsfw: bool = True,
max_retries: Optional[int] = None,
) -> AsyncGenerator[Dict[str, object], None]:
retries = max(1, max_retries if max_retries is not None else 1)
logger.info(
f"Image generation: prompt='{prompt[:50]}...', n={n}, ratio={aspect_ratio}, nsfw={enable_nsfw}"
)
for attempt in range(retries):
try:
yielded_any = False
async for item in self._stream_once(
token, prompt, aspect_ratio, n, enable_nsfw
):
yielded_any = True
yield item
return
except _BlockedError:
if yielded_any or attempt + 1 >= retries:
if not yielded_any:
yield {
"type": "error",
"error_code": "blocked",
"error": "blocked_no_final_image",
}
return
logger.warning(f"WebSocket blocked, retry {attempt + 1}/{retries}")
except Exception as e:
logger.error(f"WebSocket stream failed: {e}")
return
async def _stream_once(
self,
token: str,
prompt: str,
aspect_ratio: str,
n: int,
enable_nsfw: bool,
) -> AsyncGenerator[Dict[str, object], None]:
request_id = str(uuid.uuid4())
headers = build_ws_headers(token=token)
timeout = float(get_config("network.timeout"))
blocked_seconds = float(get_config("image.image_ws_blocked_seconds"))
blocked_grace = min(10.0, blocked_seconds)
final_min_bytes = int(get_config("image.image_ws_final_min_bytes"))
medium_min_bytes = int(get_config("image.image_ws_medium_min_bytes"))
try:
conn = await self._client.connect(
WS_IMAGINE_URL,
headers=headers,
ws_kwargs={
"heartbeat": 20,
"receive_timeout": timeout,
},
)
except Exception as e:
logger.error(f"WebSocket connect failed: {e}")
yield {
"type": "error",
"error_code": "connection_failed",
"error": str(e),
}
return
try:
async with conn as ws:
message = self._build_request_message(
request_id, prompt, aspect_ratio, enable_nsfw
)
await ws.send_json(message)
logger.info(f"WebSocket request sent: {prompt[:80]}...")
final_ids: set[str] = set()
completed = 0
start_time = last_activity = time.monotonic()
medium_received_time: Optional[float] = None
while time.monotonic() - start_time < timeout:
try:
ws_msg = await asyncio.wait_for(ws.receive(), timeout=5.0)
except asyncio.TimeoutError:
now = time.monotonic()
if (
medium_received_time
and completed == 0
and now - medium_received_time > blocked_grace
):
raise _BlockedError()
if completed > 0 and now - last_activity > 10:
logger.info(
f"WebSocket idle timeout, collected {completed} images"
)
break
continue
if ws_msg.type == aiohttp.WSMsgType.TEXT:
last_activity = time.monotonic()
try:
msg = orjson.loads(ws_msg.data)
except orjson.JSONDecodeError as e:
logger.warning(f"WebSocket message decode failed: {e}")
continue
msg_type = msg.get("type")
if msg_type == "image":
info = self._classify_image(
msg.get("url", ""),
msg.get("blob", ""),
final_min_bytes,
medium_min_bytes,
)
if not info:
continue
image_id = info["image_id"]
if info["stage"] == "medium" and medium_received_time is None:
medium_received_time = time.monotonic()
if info["is_final"] and image_id not in final_ids:
final_ids.add(image_id)
completed += 1
logger.debug(
f"Final image received: id={image_id}, size={info['blob_size']}"
)
yield info
elif msg_type == "error":
logger.warning(
f"WebSocket error: {msg.get('err_code', '')} - {msg.get('err_msg', '')}"
)
yield {
"type": "error",
"error_code": msg.get("err_code", ""),
"error": msg.get("err_msg", ""),
}
return
if completed >= n:
logger.info(f"WebSocket collected {completed} final images")
break
if (
medium_received_time
and completed == 0
and time.monotonic() - medium_received_time > blocked_seconds
):
raise _BlockedError()
elif ws_msg.type in (
aiohttp.WSMsgType.CLOSED,
aiohttp.WSMsgType.ERROR,
):
logger.warning(f"WebSocket closed/error: {ws_msg.type}")
yield {
"type": "error",
"error_code": "ws_closed",
"error": f"websocket closed: {ws_msg.type}",
}
break
except aiohttp.ClientError as e:
logger.error(f"WebSocket connection error: {e}")
yield {"type": "error", "error_code": "connection_failed", "error": str(e)}
__all__ = ["ImagineWebSocketReverse", "WS_IMAGINE_URL"]
app/services/reverse/ws_livekit.py
+17 -91
View File
@@ -11,7 +11,7 @@ from app.core.logger import logger
from app.core.config import get_config
from app.core.exceptions import UpstreamException
from app.services.token.service import TokenService
from app.services.reverse.utils.headers import build_headers
from app.services.reverse.utils.headers import build_headers, build_ws_headers
from app.services.reverse.utils.retry import retry_on_status
from app.services.reverse.utils.websocket import WebSocketClient, WebSocketConnection
@@ -29,7 +29,6 @@ class LivekitTokenReverse:
voice: str = "ara",
personality: str = "assistant",
speed: float = 1.0,
livekit_url: str = LIVEKIT_WS_URL,
) -> Dict[str, Any]:
"""Fetch LiveKit token.
@@ -39,7 +38,6 @@ class LivekitTokenReverse:
voice: str, the voice to use for the request.
personality: str, the personality to use for the request.
speed: float, the speed to use for the request.
livekit_url: str, the LiveKit URL to use for the request.
Returns:
Dict[str, Any]: The LiveKit token.
@@ -69,7 +67,7 @@ class LivekitTokenReverse:
}
).decode(),
"requestAgentDispatch": False,
"livekitUrl": livekit_url,
"livekitUrl": LIVEKIT_WS_URL,
"params": {"enable_markdown_transcript": "true"},
}
@@ -100,8 +98,7 @@ class LivekitTokenReverse:
return response
response = await retry_on_status(_do_request)
return response.json()
return response
except Exception as e:
# Handle upstream exception
@@ -137,108 +134,37 @@ class LivekitWebSocketReverse:
def __init__(self) -> None:
self._client = WebSocketClient()
def build_url(
self,
access_token: str,
*,
livekit_url: str = LIVEKIT_WS_URL,
auto_subscribe: bool = True,
sdk: str = "js",
version: str = "2.11.4",
protocol: int = 15,
) -> str:
"""Build LiveKit WebSocket URL.
async def connect(self, token: str) -> WebSocketConnection:
"""Connect to the LiveKit WebSocket.
Args:
access_token: str, the LiveKit access token.
livekit_url: str, the LiveKit URL to use for the request.
auto_subscribe: bool, whether to auto subscribe to the WebSocket.
sdk: str, the SDK to use for the request.
version: str, the version to use for the request.
protocol: int, the protocol to use for the request.
token: str, the SSO token.
Returns:
str: The LiveKit WebSocket URL.
WebSocketConnection: The LiveKit WebSocket connection.
"""
# Build base URL
base = livekit_url.rstrip("/")
# Format URL
base = LIVEKIT_WS_URL.rstrip("/")
if not base.endswith("/rtc"):
base = f"{base}/rtc"
# Build parameters
params = {
"access_token": access_token,
"auto_subscribe": str(int(auto_subscribe)),
"sdk": sdk,
"version": version,
"protocol": str(protocol),
"access_token": token,
"auto_subscribe": "1",
"sdk": "js",
"version": "2.11.4",
"protocol": "15",
}
return f"{base}?{urlencode(params)}"
def _build_headers(self, extra: Dict[str, str] | None = None) -> Dict[str, str]:
"""Build LiveKit WebSocket headers."""
# Build headers
headers = {
"Origin": "https://grok.com",
"User-Agent": get_config("security.user_agent"),
"Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
"Cache-Control": "no-cache",
"Pragma": "no-cache",
}
# Update headers
if extra:
headers.update(extra)
return headers
async def connect(
self,
access_token: str,
*,
livekit_url: str = LIVEKIT_WS_URL,
auto_subscribe: bool = True,
sdk: str = "js",
version: str = "2.11.4",
protocol: int = 15,
headers: Dict[str, str] | None = None,
timeout: float | None = None,
) -> WebSocketConnection:
"""Connect to the LiveKit WebSocket.
Args:
access_token: str, the LiveKit access token.
livekit_url: str, the LiveKit URL to use for the request.
auto_subscribe: bool, whether to auto subscribe to the WebSocket.
sdk: str, the SDK to use for the request.
version: str, the version to use for the request.
protocol: int, the protocol to use for the request.
headers: Dict[str, str], the headers to send.
timeout: float, the timeout to use for the request.
Returns:
WebSocketConnection: The LiveKit WebSocket connection.
"""
# Build URL
url = self.build_url(
access_token,
livekit_url=livekit_url,
auto_subscribe=auto_subscribe,
sdk=sdk,
version=version,
protocol=protocol,
)
url = f"{base}?{urlencode(params)}"
# Build WebSocket headers
ws_headers = self._build_headers(headers)
ws_headers = build_ws_headers()
# Build timeout
if timeout is None:
timeout = get_config("network.timeout")
# Connect to the LiveKit WebSocket
try:
return await self._client.connect(url, headers=ws_headers, timeout=timeout)
return await self._client.connect(url, headers=ws_headers)
except Exception as e:
logger.error(f"LivekitWebSocketReverse: Connect failed, {e}")
raise UpstreamException(