diff --git a/app/api/v1/image.py b/app/api/v1/image.py index 9d61d01..20c6545 100644 --- a/app/api/v1/image.py +++ b/app/api/v1/image.py @@ -19,7 +19,7 @@ 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.models.model import ModelService -from app.services.grok.processors.processor import ( +from app.services.grok.processors import ( ImageStreamProcessor, ImageCollectProcessor, ImageWSStreamProcessor, diff --git a/app/core/config.py b/app/core/config.py index 641b1cb..d9d3e8c 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -53,12 +53,19 @@ class Config: def __init__(self): self._config = {} self._defaults = {} + self._code_defaults = {} self._defaults_loaded = False + def register_defaults(self, defaults: Dict[str, Any]): + """注册代码中定义的默认值""" + self._code_defaults = _deep_merge(self._code_defaults, defaults) + def _ensure_defaults(self): if self._defaults_loaded: return - self._defaults = _load_defaults() + file_defaults = _load_defaults() + # 合并文件默认值和代码默认值（代码默认值优先级更低） + self._defaults = _deep_merge(self._code_defaults, file_defaults) self._defaults_loaded = True async def load(self): @@ -140,4 +147,9 @@ def get_config(key: str, default: Any = None) -> Any: return config.get(key, default) -__all__ = ["Config", "config", "get_config"] +def register_defaults(defaults: Dict[str, Any]): + """注册默认配置""" + config.register_defaults(defaults) + + +__all__ = ["Config", "config", "get_config", "register_defaults"] diff --git a/app/services/grok/defaults.py b/app/services/grok/defaults.py new file mode 100644 index 0000000..78ade9a --- /dev/null +++ b/app/services/grok/defaults.py @@ -0,0 +1,71 @@ +""" +Grok 服务默认配置 + +此文件定义所有 Grok 相关服务的默认值，会在应用启动时注册到配置系统中。 +""" + +# Grok 服务默认配置 +GROK_DEFAULTS = { + "grok": { + # 网络配置 + "browser": "chrome136", + "timeout": 120, + "user_agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36", + "base_proxy_url": "", + "asset_proxy_url": "", + "cf_clearance": "", + + # 对话配置 + "temporary": True, + "disable_memory": True, + "stream": True, + "thinking": False, + "filter_tags": ["grok:render", "xaiartifact", "xai:tool_usage_card"], + + # 重试配置 + "max_retry": 3, + "retry_status_codes": [401, 429, 403], + "retry_backoff_base": 0.5, + "retry_backoff_factor": 2.0, + "retry_backoff_max": 30.0, + "retry_budget": 90.0, + + # 超时配置 + "stream_idle_timeout": 45.0, + "video_idle_timeout": 90.0, + + # 图片配置 + "image_ws_blocked_seconds": 15, + "image_ws_final_min_bytes": 100000, + "image_ws_medium_min_bytes": 30000, + + # Statsig + "dynamic_statsig": True, + }, + + "app": { + "app_url": "", + "image_format": "url", + "video_format": "html", + }, + + "performance": { + "assets_max_concurrent": 25, + "assets_delete_batch_size": 10, + "media_max_concurrent": 50, + "usage_max_concurrent": 25, + }, + + "cache": { + "enable_auto_clean": True, + "limit_mb": 1024, + }, +} + + +def get_grok_defaults(): + """获取 Grok 默认配置""" + return GROK_DEFAULTS + + +__all__ = ["GROK_DEFAULTS", "get_grok_defaults"] diff --git a/app/services/grok/processors/__init__.py b/app/services/grok/processors/__init__.py index e69de29..04773f6 100644 --- a/app/services/grok/processors/__init__.py +++ b/app/services/grok/processors/__init__.py @@ -0,0 +1,22 @@ +""" +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 + +__all__ = [ + "BaseProcessor", + "StreamIdleTimeoutError", + "StreamProcessor", + "CollectProcessor", + "VideoStreamProcessor", + "VideoCollectProcessor", + "ImageStreamProcessor", + "ImageCollectProcessor", + "ImageWSStreamProcessor", + "ImageWSCollectProcessor", +] diff --git a/app/services/grok/processors/base.py b/app/services/grok/processors/base.py new file mode 100644 index 0000000..bfda0d2 --- /dev/null +++ b/app/services/grok/processors/base.py @@ -0,0 +1,159 @@ +""" +响应处理器基类和通用工具 +""" + +import asyncio +import time +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 + + +ASSET_URL = "https://assets.grok.com/" + +T = TypeVar("T") + + +def _is_http2_stream_error(e: Exception) -> bool: + """检查是否为 HTTP/2 流错误""" + err_str = str(e).lower() + return "http/2" in err_str or "curl: (92)" in err_str or "stream" in err_str + + +def _normalize_stream_line(line: Any) -> Optional[str]: + """规范化流式响应行，兼容 SSE data 前缀与空行""" + if line is None: + return None + if isinstance(line, (bytes, bytearray)): + text = line.decode("utf-8", errors="ignore") + else: + text = str(line) + text = text.strip() + if not text: + return None + if text.startswith("data:"): + text = text[5:].strip() + if text == "[DONE]": + return None + return text + + +def _collect_image_urls(obj: Any) -> List[str]: + """递归收集响应中的图片 URL""" + urls: List[str] = [] + seen = set() + + def add(url: str): + if not url or url in seen: + return + seen.add(url) + urls.append(url) + + def walk(value: Any): + if isinstance(value, dict): + for key, item in value.items(): + if key in {"generatedImageUrls", "imageUrls", "imageURLs"}: + if isinstance(item, list): + for url in item: + if isinstance(url, str): + add(url) + elif isinstance(item, str): + add(item) + continue + walk(item) + elif isinstance(value, list): + for item in value: + walk(item) + + walk(obj) + return urls + + +class StreamIdleTimeoutError(Exception): + """流空闲超时错误""" + + def __init__(self, idle_seconds: float): + self.idle_seconds = idle_seconds + super().__init__(f"Stream idle timeout after {idle_seconds}s") + + +async def _with_idle_timeout( + iterable: AsyncIterable[T], idle_timeout: float, model: str = "" +) -> AsyncGenerator[T, None]: + """ + 包装异步迭代器，添加空闲超时检测 + + Args: + iterable: 原始异步迭代器 + idle_timeout: 空闲超时时间(秒)，0 表示禁用 + model: 模型名称(用于日志) + """ + if idle_timeout <= 0: + async for item in iterable: + yield item + return + + iterator = iterable.__aiter__() + while True: + try: + item = await asyncio.wait_for(iterator.__anext__(), timeout=idle_timeout) + yield item + except asyncio.TimeoutError: + logger.warning( + f"Stream idle timeout after {idle_timeout}s", + extra={"model": model, "idle_timeout": idle_timeout}, + ) + raise StreamIdleTimeoutError(idle_timeout) + except StopAsyncIteration: + break + + +class BaseProcessor: + """基础处理器""" + + def __init__(self, model: str, token: str = ""): + self.model = model + self.token = token + self.created = int(time.time()) + self.app_url = get_config("app.app_url") + self._dl_service: Optional[DownloadService] = None + + def _get_dl(self) -> DownloadService: + """获取下载服务实例（复用）""" + if self._dl_service is None: + self._dl_service = DownloadService() + return self._dl_service + + async def close(self): + """释放下载服务资源""" + if self._dl_service: + await self._dl_service.close() + self._dl_service = None + + async def process_url(self, path: str, media_type: str = "image") -> str: + """处理资产 URL""" + if path.startswith("http"): + from urllib.parse import urlparse + path = urlparse(path).path + + if not path.startswith("/"): + path = f"/{path}" + + if self.app_url: + dl_service = self._get_dl() + await dl_service.download(path, self.token, media_type) + return f"{self.app_url.rstrip('/')}/v1/files/{media_type}{path}" + else: + return f"{ASSET_URL.rstrip('/')}{path}" + + +__all__ = [ + "BaseProcessor", + "StreamIdleTimeoutError", + "_with_idle_timeout", + "_normalize_stream_line", + "_collect_image_urls", + "_is_http2_stream_error", +] diff --git a/app/services/grok/processors/chat_processors.py b/app/services/grok/processors/chat_processors.py new file mode 100644 index 0000000..a7bf5d4 --- /dev/null +++ b/app/services/grok/processors/chat_processors.py @@ -0,0 +1,326 @@ +""" +聊天响应处理器 +""" + +import asyncio +import uuid +import re +from typing import Any, AsyncGenerator, AsyncIterable + +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 StreamProcessor(BaseProcessor): + """流式响应处理器""" + + def __init__(self, model: str, token: str = "", think: bool = None): + super().__init__(model, token) + self.response_id: str = None + self.fingerprint: str = "" + self.think_opened: bool = False + self.role_sent: bool = False + self.filter_tags = get_config("grok.filter_tags") + self.image_format = get_config("app.image_format") + self._tag_buffer: str = "" + self._in_filter_tag: bool = False + + if think is None: + self.show_think = get_config("grok.thinking") + else: + self.show_think = think + + def _filter_token(self, token: str) -> str: + """过滤 token 中的特殊标签（如 ...），支持跨 token 的标签过滤""" + if not self.filter_tags: + return token + + result = [] + i = 0 + while i < len(token): + char = token[i] + + if self._in_filter_tag: + self._tag_buffer += char + if char == ">": + if "/>" in self._tag_buffer: + self._in_filter_tag = False + self._tag_buffer = "" + else: + for tag in self.filter_tags: + if f"" in self._tag_buffer: + self._in_filter_tag = False + self._tag_buffer = "" + break + i += 1 + continue + + if char == "<": + remaining = token[i:] + tag_started = False + for tag in self.filter_tags: + if remaining.startswith(f"<{tag}"): + tag_started = True + break + if len(remaining) < len(tag) + 1: + for j in range(1, len(remaining) + 1): + if f"<{tag}".startswith(remaining[:j]): + tag_started = True + break + + if tag_started: + self._in_filter_tag = True + self._tag_buffer = char + i += 1 + continue + + result.append(char) + i += 1 + + return "".join(result) + + def _sse(self, content: str = "", role: str = None, finish: str = None) -> str: + """构建 SSE 响应""" + delta = {} + if role: + delta["role"] = role + delta["content"] = "" + elif content: + delta["content"] = content + + chunk = { + "id": self.response_id or f"chatcmpl-{uuid.uuid4().hex[:24]}", + "object": "chat.completion.chunk", + "created": self.created, + "model": self.model, + "system_fingerprint": self.fingerprint, + "choices": [ + {"index": 0, "delta": delta, "logprobs": None, "finish_reason": finish} + ], + } + return f"data: {orjson.dumps(chunk).decode()}\n\n" + + async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]: + """处理流式响应""" + idle_timeout = get_config("grok.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 (llm := resp.get("llmInfo")) and not self.fingerprint: + self.fingerprint = llm.get("modelHash", "") + if rid := resp.get("responseId"): + self.response_id = rid + + if not self.role_sent: + yield self._sse(role="assistant") + self.role_sent = True + + # 图像生成进度 + if img := resp.get("streamingImageGenerationResponse"): + if self.show_think: + if not self.think_opened: + yield self._sse("\n") + self.think_opened = True + idx = img.get("imageIndex", 0) + 1 + progress = img.get("progress", 0) + yield self._sse(f"正在生成第{idx}张图片中，当前进度{progress}%\n") + continue + + # modelResponse + if mr := resp.get("modelResponse"): + if self.think_opened and self.show_think: + if msg := mr.get("message"): + yield self._sse(msg + "\n") + yield self._sse("\n") + self.think_opened = False + + # 处理生成的图片 + for url in _collect_image_urls(mr): + parts = url.split("/") + img_id = parts[-2] if len(parts) >= 2 else "image" + + if self.image_format == "base64": + dl_service = self._get_dl() + base64_data = await dl_service.to_base64(url, self.token, "image") + if base64_data: + yield self._sse(f"![{img_id}]({base64_data})\n") + else: + final_url = await self.process_url(url, "image") + yield self._sse(f"![{img_id}]({final_url})\n") + else: + final_url = await self.process_url(url, "image") + yield self._sse(f"![{img_id}]({final_url})\n") + + if (meta := mr.get("metadata", {})).get("llm_info", {}).get("modelHash"): + self.fingerprint = meta["llm_info"]["modelHash"] + continue + + # 普通 token + if (token := resp.get("token")) is not None: + if token: + filtered = self._filter_token(token) + if filtered: + yield self._sse(filtered) + + if self.think_opened: + yield self._sse("\n") + yield self._sse(finish="stop") + yield "data: [DONE]\n\n" + except asyncio.CancelledError: + logger.debug("Stream cancelled by client", extra={"model": self.model}) + except StreamIdleTimeoutError as e: + raise UpstreamException( + message=f"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: {e}", extra={"model": self.model}) + raise UpstreamException( + message="Upstream connection closed unexpectedly", + status_code=502, + details={"error": str(e), "type": "http2_stream_error"}, + ) + logger.error(f"Stream request error: {e}", extra={"model": self.model}) + raise UpstreamException( + message=f"Upstream request failed: {e}", + status_code=502, + details={"error": str(e)}, + ) + except Exception as e: + logger.error(f"Stream processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__}) + raise + finally: + await self.close() + + +class CollectProcessor(BaseProcessor): + """非流式响应处理器""" + + def __init__(self, model: str, token: str = ""): + super().__init__(model, token) + self.image_format = get_config("app.image_format") + self.filter_tags = get_config("grok.filter_tags") + + def _filter_content(self, content: str) -> str: + """过滤内容中的特殊标签""" + if not content or not self.filter_tags: + return content + + result = content + for tag in self.filter_tags: + pattern = rf"<{re.escape(tag)}[^>]*>.*?|<{re.escape(tag)}[^>]*/>" + result = re.sub(pattern, "", result, flags=re.DOTALL) + + return result + + async def process(self, response: AsyncIterable[bytes]) -> dict[str, Any]: + """处理并收集完整响应""" + response_id = "" + fingerprint = "" + content = "" + idle_timeout = get_config("grok.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 (llm := resp.get("llmInfo")) and not fingerprint: + fingerprint = llm.get("modelHash", "") + + if mr := resp.get("modelResponse"): + response_id = mr.get("responseId", "") + content = mr.get("message", "") + + if urls := _collect_image_urls(mr): + content += "\n" + for url in urls: + parts = url.split("/") + img_id = parts[-2] if len(parts) >= 2 else "image" + + if self.image_format == "base64": + dl_service = self._get_dl() + base64_data = await dl_service.to_base64(url, self.token, "image") + if base64_data: + content += f"![{img_id}]({base64_data})\n" + else: + final_url = await self.process_url(url, "image") + content += f"![{img_id}]({final_url})\n" + else: + final_url = await self.process_url(url, "image") + content += f"![{img_id}]({final_url})\n" + + if (meta := mr.get("metadata", {})).get("llm_info", {}).get("modelHash"): + fingerprint = meta["llm_info"]["modelHash"] + + except asyncio.CancelledError: + logger.debug("Collect cancelled by client", extra={"model": self.model}) + except StreamIdleTimeoutError as e: + logger.warning(f"Collect idle timeout: {e}", extra={"model": self.model}) + except RequestsError as e: + if _is_http2_stream_error(e): + logger.warning(f"HTTP/2 stream error in collect: {e}", extra={"model": self.model}) + else: + logger.error(f"Collect request error: {e}", extra={"model": self.model}) + except Exception as e: + logger.error(f"Collect processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__}) + finally: + await self.close() + + content = self._filter_content(content) + + return { + "id": response_id, + "object": "chat.completion", + "created": self.created, + "model": self.model, + "system_fingerprint": fingerprint, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content, "refusal": None, "annotations": []}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "prompt_tokens_details": {"cached_tokens": 0, "text_tokens": 0, "audio_tokens": 0, "image_tokens": 0}, + "completion_tokens_details": {"text_tokens": 0, "audio_tokens": 0, "reasoning_tokens": 0}, + }, + } + + +__all__ = ["StreamProcessor", "CollectProcessor"] diff --git a/app/services/grok/processors/image_processors.py b/app/services/grok/processors/image_processors.py new file mode 100644 index 0000000..8550108 --- /dev/null +++ b/app/services/grok/processors/image_processors.py @@ -0,0 +1,211 @@ +""" +图片生成响应处理器（HTTP） +""" + +import asyncio +import random +from typing import Any, 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("grok.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 + 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) + 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("grok.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 + 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 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 new file mode 100644 index 0000000..86a6aa9 --- /dev/null +++ b/app/services/grok/processors/image_ws_processors.py @@ -0,0 +1,231 @@ +""" +图片生成响应处理器（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.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 = Path(__file__).parent.parent.parent.parent.parent / "data" / "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/processor.py b/app/services/grok/processors/processor.py deleted file mode 100644 index 49ac9d1..0000000 --- a/app/services/grok/processors/processor.py +++ /dev/null @@ -1,1262 +0,0 @@ -""" -OpenAI 响应格式处理器 -""" - -import asyncio -import base64 -import time -import uuid -import random -import re -from pathlib import Path -import orjson -from typing import Any, AsyncGenerator, Optional, AsyncIterable, List, TypeVar, Dict -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 app.services.grok.services.assets import DownloadService - - -DEFAULT_APP_URL = "" -DEFAULT_FILTER_TAGS = ["grok:render", "xaiartifact", "xai:tool_usage_card"] -DEFAULT_IMAGE_FORMAT = "url" -DEFAULT_THINKING = False -DEFAULT_STREAM_IDLE_TIMEOUT = 45.0 -DEFAULT_VIDEO_FORMAT = "html" -DEFAULT_VIDEO_IDLE_TIMEOUT = 90.0 - - -def _is_http2_stream_error(e: Exception) -> bool: - """检查是否为 HTTP/2 流错误""" - err_str = str(e).lower() - return "http/2" in err_str or "curl: (92)" in err_str or "stream" in err_str - - -def _normalize_stream_line(line: Any) -> Optional[str]: - """规范化流式响应行，兼容 SSE data 前缀与空行""" - if line is None: - return None - if isinstance(line, (bytes, bytearray)): - text = line.decode("utf-8", errors="ignore") - else: - text = str(line) - text = text.strip() - if not text: - return None - if text.startswith("data:"): - text = text[5:].strip() - if text == "[DONE]": - return None - return text - - -def _collect_image_urls(obj: Any) -> List[str]: - """递归收集响应中的图片 URL""" - urls: List[str] = [] - seen = set() - - def add(url: str): - if not url or url in seen: - return - seen.add(url) - urls.append(url) - - def walk(value: Any): - if isinstance(value, dict): - for key, item in value.items(): - if key in {"generatedImageUrls", "imageUrls", "imageURLs"}: - if isinstance(item, list): - for url in item: - if isinstance(url, str): - add(url) - elif isinstance(item, str): - add(item) - continue - walk(item) - elif isinstance(value, list): - for item in value: - walk(item) - - walk(obj) - return urls - - -T = TypeVar("T") - - -class StreamIdleTimeoutError(Exception): - """流空闲超时错误""" - - def __init__(self, idle_seconds: float): - self.idle_seconds = idle_seconds - super().__init__(f"Stream idle timeout after {idle_seconds}s") - - -async def _with_idle_timeout( - iterable: AsyncIterable[T], idle_timeout: float, model: str = "" -) -> AsyncGenerator[T, None]: - """ - 包装异步迭代器，添加空闲超时检测 - - Args: - iterable: 原始异步迭代器 - idle_timeout: 空闲超时时间(秒)，0 表示禁用 - model: 模型名称(用于日志) - - Yields: - 原始迭代器的元素 - - Raises: - StreamIdleTimeoutError: 当空闲超时时 - """ - if idle_timeout <= 0: - async for item in iterable: - yield item - return - - iterator = iterable.__aiter__() - while True: - try: - item = await asyncio.wait_for(iterator.__anext__(), timeout=idle_timeout) - yield item - except asyncio.TimeoutError: - logger.warning( - f"Stream idle timeout after {idle_timeout}s", - extra={"model": model, "idle_timeout": idle_timeout}, - ) - raise StreamIdleTimeoutError(idle_timeout) - except StopAsyncIteration: - break - - -ASSET_URL = "https://assets.grok.com/" - - -class BaseProcessor: - """基础处理器""" - - def __init__(self, model: str, token: str = ""): - self.model = model - self.token = token - self.created = int(time.time()) - self.app_url = get_config("app.app_url", DEFAULT_APP_URL) - self._dl_service: Optional[DownloadService] = None - - def _get_dl(self) -> DownloadService: - """获取下载服务实例（复用）""" - if self._dl_service is None: - self._dl_service = DownloadService() - return self._dl_service - - async def close(self): - """释放下载服务资源""" - if self._dl_service: - await self._dl_service.close() - self._dl_service = None - - async def process_url(self, path: str, media_type: str = "image") -> str: - """处理资产 URL""" - # 处理可能的绝对路径 - if path.startswith("http"): - from urllib.parse import urlparse - - path = urlparse(path).path - - if not path.startswith("/"): - path = f"/{path}" - - if self.app_url: - dl_service = self._get_dl() - await dl_service.download(path, self.token, media_type) - return f"{self.app_url.rstrip('/')}/v1/files/{media_type}{path}" - else: - return f"{ASSET_URL.rstrip('/')}{path}" - - def _sse(self, content: str = "", role: str = None, finish: str = None) -> str: - """构建 SSE 响应 (StreamProcessor 通用)""" - if not hasattr(self, "response_id"): - self.response_id = None - if not hasattr(self, "fingerprint"): - self.fingerprint = "" - - delta = {} - if role: - delta["role"] = role - delta["content"] = "" - elif content: - delta["content"] = content - - chunk = { - "id": self.response_id or f"chatcmpl-{uuid.uuid4().hex[:24]}", - "object": "chat.completion.chunk", - "created": self.created, - "model": self.model, - "system_fingerprint": self.fingerprint - if hasattr(self, "fingerprint") - else "", - "choices": [ - {"index": 0, "delta": delta, "logprobs": None, "finish_reason": finish} - ], - } - return f"data: {orjson.dumps(chunk).decode()}\n\n" - - -class StreamProcessor(BaseProcessor): - """流式响应处理器""" - - def __init__(self, model: str, token: str = "", think: bool = None): - super().__init__(model, token) - self.response_id: Optional[str] = None - self.fingerprint: str = "" - self.think_opened: bool = False - self.role_sent: bool = False - self.filter_tags = get_config("grok.filter_tags", DEFAULT_FILTER_TAGS) - self.image_format = get_config("app.image_format", DEFAULT_IMAGE_FORMAT) - # 用于过滤跨 token 的标签 - self._tag_buffer: str = "" - self._in_filter_tag: bool = False - - if think is None: - self.show_think = get_config("grok.thinking", DEFAULT_THINKING) - else: - self.show_think = think - - def _filter_token(self, token: str) -> str: - """ - 过滤 token 中的特殊标签（如 ...） - 支持跨 token 的标签过滤 - """ - if not self.filter_tags: - return token - - result = [] - i = 0 - while i < len(token): - char = token[i] - - # 如果在过滤标签内 - if self._in_filter_tag: - self._tag_buffer += char - # 检查是否到达结束标签 - if char == ">": - # 检查是否是自闭合标签 - if "/>" in self._tag_buffer: - self._in_filter_tag = False - self._tag_buffer = "" - else: - # 检查是否是结束标签 - for tag in self.filter_tags: - if f"" in self._tag_buffer: - self._in_filter_tag = False - self._tag_buffer = "" - break - # 如果不是结束标签，检查是否是开始标签结束（非自闭合） - # 继续等待结束标签 - i += 1 - continue - - # 检查是否开始一个过滤标签 - if char == "<": - # 查看后续字符 - remaining = token[i:] - tag_started = False - for tag in self.filter_tags: - if remaining.startswith(f"<{tag}"): - tag_started = True - break - # 部分匹配（可能跨 token） - if len(remaining) < len(tag) + 1: - for j in range(1, len(remaining) + 1): - if f"<{tag}".startswith(remaining[:j]): - tag_started = True - break - - if tag_started: - self._in_filter_tag = True - self._tag_buffer = char - i += 1 - continue - - result.append(char) - i += 1 - - return "".join(result) - - async def process( - self, response: AsyncIterable[bytes] - ) -> AsyncGenerator[str, None]: - """处理流式响应""" - # 获取空闲超时配置 - idle_timeout = get_config("grok.stream_idle_timeout", DEFAULT_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 (llm := resp.get("llmInfo")) and not self.fingerprint: - self.fingerprint = llm.get("modelHash", "") - if rid := resp.get("responseId"): - self.response_id = rid - - # 首次发送 role - if not self.role_sent: - yield self._sse(role="assistant") - self.role_sent = True - - # 图像生成进度 - if img := resp.get("streamingImageGenerationResponse"): - if self.show_think: - if not self.think_opened: - yield self._sse("\n") - self.think_opened = True - idx = img.get("imageIndex", 0) + 1 - progress = img.get("progress", 0) - yield self._sse( - f"正在生成第{idx}张图片中，当前进度{progress}%\n" - ) - continue - - # modelResponse - if mr := resp.get("modelResponse"): - if self.think_opened and self.show_think: - if msg := mr.get("message"): - yield self._sse(msg + "\n") - yield self._sse("\n") - self.think_opened = False - - # 处理生成的图片 - for url in _collect_image_urls(mr): - parts = url.split("/") - img_id = parts[-2] if len(parts) >= 2 else "image" - - if self.image_format == "base64": - dl_service = self._get_dl() - base64_data = await dl_service.to_base64( - url, self.token, "image" - ) - if base64_data: - yield self._sse(f"![{img_id}]({base64_data})\n") - else: - final_url = await self.process_url(url, "image") - yield self._sse(f"![{img_id}]({final_url})\n") - else: - final_url = await self.process_url(url, "image") - yield self._sse(f"![{img_id}]({final_url})\n") - - if ( - (meta := mr.get("metadata", {})) - .get("llm_info", {}) - .get("modelHash") - ): - self.fingerprint = meta["llm_info"]["modelHash"] - continue - - # 普通 token - if (token := resp.get("token")) is not None: - if token: - filtered = self._filter_token(token) - if filtered: - yield self._sse(filtered) - - if self.think_opened: - yield self._sse("\n") - yield self._sse(finish="stop") - yield "data: [DONE]\n\n" - except asyncio.CancelledError: - # 客户端断开连接，静默处理 - logger.debug("Stream cancelled by client", extra={"model": self.model}) - except StreamIdleTimeoutError as e: - # 流空闲超时 - raise UpstreamException( - message=f"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: - # HTTP/2 流错误转换为 UpstreamException - if _is_http2_stream_error(e): - logger.warning(f"HTTP/2 stream error: {e}", extra={"model": self.model}) - raise UpstreamException( - message="Upstream connection closed unexpectedly", - status_code=502, - details={"error": str(e), "type": "http2_stream_error"}, - ) - logger.error(f"Stream request error: {e}", extra={"model": self.model}) - raise UpstreamException( - message=f"Upstream request failed: {e}", - status_code=502, - details={"error": str(e)}, - ) - except Exception as e: - logger.error( - f"Stream processing error: {e}", - extra={"model": self.model, "error_type": type(e).__name__}, - ) - raise - finally: - await self.close() - - -class CollectProcessor(BaseProcessor): - """非流式响应处理器""" - - def __init__(self, model: str, token: str = ""): - super().__init__(model, token) - self.image_format = get_config("app.image_format", DEFAULT_IMAGE_FORMAT) - self.filter_tags = get_config("grok.filter_tags", DEFAULT_FILTER_TAGS) - - def _filter_content(self, content: str) -> str: - """过滤内容中的特殊标签""" - import re - - if not content or not self.filter_tags: - return content - - result = content - for tag in self.filter_tags: - # 匹配 ... 或 ，re.DOTALL 使 . 匹配换行符 - pattern = rf"<{re.escape(tag)}[^>]*>.*?|<{re.escape(tag)}[^>]*/>" - result = re.sub(pattern, "", result, flags=re.DOTALL) - - return result - - async def process(self, response: AsyncIterable[bytes]) -> dict[str, Any]: - """处理并收集完整响应""" - response_id = "" - fingerprint = "" - content = "" - # 获取空闲超时配置 - idle_timeout = get_config("grok.stream_idle_timeout", DEFAULT_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 (llm := resp.get("llmInfo")) and not fingerprint: - fingerprint = llm.get("modelHash", "") - - if mr := resp.get("modelResponse"): - response_id = mr.get("responseId", "") - content = mr.get("message", "") - - if urls := _collect_image_urls(mr): - content += "\n" - for url in urls: - parts = url.split("/") - img_id = parts[-2] if len(parts) >= 2 else "image" - - if self.image_format == "base64": - dl_service = self._get_dl() - base64_data = await dl_service.to_base64( - url, self.token, "image" - ) - if base64_data: - content += f"![{img_id}]({base64_data})\n" - else: - final_url = await self.process_url(url, "image") - content += f"![{img_id}]({final_url})\n" - else: - final_url = await self.process_url(url, "image") - content += f"![{img_id}]({final_url})\n" - - if ( - (meta := mr.get("metadata", {})) - .get("llm_info", {}) - .get("modelHash") - ): - fingerprint = meta["llm_info"]["modelHash"] - - except asyncio.CancelledError: - logger.debug("Collect cancelled by client", extra={"model": self.model}) - except StreamIdleTimeoutError as e: - logger.warning(f"Collect idle timeout: {e}", extra={"model": self.model}) - # 非流式模式下，超时后返回已收集的内容 - except RequestsError as e: - if _is_http2_stream_error(e): - logger.warning( - f"HTTP/2 stream error in collect: {e}", extra={"model": self.model} - ) - else: - logger.error(f"Collect request error: {e}", extra={"model": self.model}) - except Exception as e: - logger.error( - f"Collect processing error: {e}", - extra={"model": self.model, "error_type": type(e).__name__}, - ) - finally: - await self.close() - - # 过滤特殊标签 - content = self._filter_content(content) - - return { - "id": response_id, - "object": "chat.completion", - "created": self.created, - "model": self.model, - "system_fingerprint": fingerprint, - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": content, - "refusal": None, - "annotations": [], - }, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": 0, - "completion_tokens": 0, - "total_tokens": 0, - "prompt_tokens_details": { - "cached_tokens": 0, - "text_tokens": 0, - "audio_tokens": 0, - "image_tokens": 0, - }, - "completion_tokens_details": { - "text_tokens": 0, - "audio_tokens": 0, - "reasoning_tokens": 0, - }, - }, - } - - -class VideoStreamProcessor(BaseProcessor): - """视频流式响应处理器""" - - def __init__(self, model: str, token: str = "", think: bool = None): - super().__init__(model, token) - self.response_id: Optional[str] = None - self.think_opened: bool = False - self.role_sent: bool = False - self.video_format = str(get_config("app.video_format", DEFAULT_VIDEO_FORMAT)).lower() - - if think is None: - self.show_think = get_config("grok.thinking", DEFAULT_THINKING) - else: - self.show_think = think - - def _build_video_html(self, video_url: str, thumbnail_url: str = "") -> str: - """构建视频 HTML 标签""" - import html - - safe_video_url = html.escape(video_url) - safe_thumbnail_url = html.escape(thumbnail_url) - poster_attr = f' poster="{safe_thumbnail_url}"' if safe_thumbnail_url else "" - return f'''''' - - async def process( - self, response: AsyncIterable[bytes] - ) -> AsyncGenerator[str, None]: - """处理视频流式响应""" - # 视频生成使用更长的空闲超时 - idle_timeout = get_config("grok.video_idle_timeout", DEFAULT_VIDEO_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 rid := resp.get("responseId"): - self.response_id = rid - - # 首次发送 role - if not self.role_sent: - yield self._sse(role="assistant") - self.role_sent = True - - # 视频生成进度 - if video_resp := resp.get("streamingVideoGenerationResponse"): - progress = video_resp.get("progress", 0) - - if self.show_think: - if not self.think_opened: - yield self._sse("\n") - self.think_opened = True - yield self._sse(f"正在生成视频中，当前进度{progress}%\n") - - if progress == 100: - video_url = video_resp.get("videoUrl", "") - thumbnail_url = video_resp.get("thumbnailImageUrl", "") - - if self.think_opened and self.show_think: - yield self._sse("\n") - self.think_opened = False - - if video_url: - final_video_url = await self.process_url(video_url, "video") - final_thumbnail_url = "" - if thumbnail_url: - final_thumbnail_url = await self.process_url( - thumbnail_url, "image" - ) - - if self.video_format == "url": - yield self._sse(final_video_url) - else: - video_html = self._build_video_html( - final_video_url, final_thumbnail_url - ) - yield self._sse(video_html) - - logger.info(f"Video generated: {video_url}") - continue - - if self.think_opened: - yield self._sse("\n") - yield self._sse(finish="stop") - yield "data: [DONE]\n\n" - except asyncio.CancelledError: - logger.debug( - "Video stream cancelled by client", extra={"model": self.model} - ) - except StreamIdleTimeoutError as e: - raise UpstreamException( - message=f"Video 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 video: {e}", extra={"model": self.model} - ) - raise UpstreamException( - message="Upstream connection closed unexpectedly", - status_code=502, - details={"error": str(e), "type": "http2_stream_error"}, - ) - logger.error( - f"Video stream request error: {e}", extra={"model": self.model} - ) - raise UpstreamException( - message=f"Upstream request failed: {e}", - status_code=502, - details={"error": str(e)}, - ) - except Exception as e: - logger.error( - f"Video stream processing error: {e}", - extra={"model": self.model, "error_type": type(e).__name__}, - ) - finally: - await self.close() - - -class VideoCollectProcessor(BaseProcessor): - """视频非流式响应处理器""" - - def __init__(self, model: str, token: str = ""): - super().__init__(model, token) - self.video_format = str(get_config("app.video_format", DEFAULT_VIDEO_FORMAT)).lower() - - def _build_video_html(self, video_url: str, thumbnail_url: str = "") -> str: - poster_attr = f' poster="{thumbnail_url}"' if thumbnail_url else "" - return f'''''' - - async def process(self, response: AsyncIterable[bytes]) -> dict[str, Any]: - """处理并收集视频响应""" - response_id = "" - content = "" - # 视频生成使用更长的空闲超时 - idle_timeout = get_config("grok.video_idle_timeout", DEFAULT_VIDEO_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 video_resp := resp.get("streamingVideoGenerationResponse"): - if video_resp.get("progress") == 100: - response_id = resp.get("responseId", "") - video_url = video_resp.get("videoUrl", "") - thumbnail_url = video_resp.get("thumbnailImageUrl", "") - - if video_url: - final_video_url = await self.process_url(video_url, "video") - final_thumbnail_url = "" - if thumbnail_url: - final_thumbnail_url = await self.process_url( - thumbnail_url, "image" - ) - - if self.video_format == "url": - content = final_video_url - else: - content = self._build_video_html( - final_video_url, final_thumbnail_url - ) - logger.info(f"Video generated: {video_url}") - - except asyncio.CancelledError: - logger.debug( - "Video collect cancelled by client", extra={"model": self.model} - ) - except StreamIdleTimeoutError as e: - logger.warning( - f"Video collect idle timeout: {e}", extra={"model": self.model} - ) - except RequestsError as e: - if _is_http2_stream_error(e): - logger.warning( - f"HTTP/2 stream error in video collect: {e}", - extra={"model": self.model}, - ) - else: - logger.error( - f"Video collect request error: {e}", extra={"model": self.model} - ) - except Exception as e: - logger.error( - f"Video collect processing error: {e}", - extra={"model": self.model, "error_type": type(e).__name__}, - ) - finally: - await self.close() - - return { - "id": response_id, - "object": "chat.completion", - "created": self.created, - "model": self.model, - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": content, - "refusal": None, - }, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, - } - - -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("grok.stream_idle_timeout", DEFAULT_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 - 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) - 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("grok.stream_idle_timeout", DEFAULT_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 - 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 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 图片处理基类""" - - 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 = ( - Path(__file__).parent.parent.parent.parent.parent - / "data" - / "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: - if self.app_url: - return f"{self.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__ = [ - "StreamProcessor", - "CollectProcessor", - "VideoStreamProcessor", - "VideoCollectProcessor", - "ImageStreamProcessor", - "ImageCollectProcessor", - "ImageWSStreamProcessor", - "ImageWSCollectProcessor", -] diff --git a/app/services/grok/processors/video_processors.py b/app/services/grok/processors/video_processors.py new file mode 100644 index 0000000..df8750c --- /dev/null +++ b/app/services/grok/processors/video_processors.py @@ -0,0 +1,236 @@ +""" +视频响应处理器 +""" + +import asyncio +import uuid +from typing import Any, AsyncGenerator, AsyncIterable, 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.exceptions import UpstreamException +from .base import ( + BaseProcessor, + StreamIdleTimeoutError, + _with_idle_timeout, + _normalize_stream_line, + _is_http2_stream_error, +) + + +class VideoStreamProcessor(BaseProcessor): + """视频流式响应处理器""" + + def __init__(self, model: str, token: str = "", think: bool = None): + super().__init__(model, token) + self.response_id: Optional[str] = None + self.think_opened: bool = False + self.role_sent: bool = False + self.video_format = str(get_config("app.video_format")).lower() + + if think is None: + self.show_think = get_config("grok.thinking") + else: + self.show_think = think + + def _sse(self, content: str = "", role: str = None, finish: str = None) -> str: + """构建 SSE 响应""" + delta = {} + if role: + delta["role"] = role + delta["content"] = "" + elif content: + delta["content"] = content + + chunk = { + "id": self.response_id or f"chatcmpl-{uuid.uuid4().hex[:24]}", + "object": "chat.completion.chunk", + "created": self.created, + "model": self.model, + "choices": [ + {"index": 0, "delta": delta, "logprobs": None, "finish_reason": finish} + ], + } + return f"data: {orjson.dumps(chunk).decode()}\n\n" + + def _build_video_html(self, video_url: str, thumbnail_url: str = "") -> str: + """构建视频 HTML 标签""" + import html + safe_video_url = html.escape(video_url) + safe_thumbnail_url = html.escape(thumbnail_url) + poster_attr = f' poster="{safe_thumbnail_url}"' if safe_thumbnail_url else "" + return f'''''' + + async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]: + """处理视频流式响应""" + idle_timeout = get_config("grok.video_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 rid := resp.get("responseId"): + self.response_id = rid + + if not self.role_sent: + yield self._sse(role="assistant") + self.role_sent = True + + # 视频生成进度 + if video_resp := resp.get("streamingVideoGenerationResponse"): + progress = video_resp.get("progress", 0) + + if self.show_think: + if not self.think_opened: + yield self._sse("\n") + self.think_opened = True + yield self._sse(f"正在生成视频中，当前进度{progress}%\n") + + if progress == 100: + video_url = video_resp.get("videoUrl", "") + thumbnail_url = video_resp.get("thumbnailImageUrl", "") + + if self.think_opened and self.show_think: + yield self._sse("\n") + self.think_opened = False + + if video_url: + final_video_url = await self.process_url(video_url, "video") + final_thumbnail_url = "" + if thumbnail_url: + final_thumbnail_url = await self.process_url(thumbnail_url, "image") + + if self.video_format == "url": + yield self._sse(final_video_url) + else: + video_html = self._build_video_html(final_video_url, final_thumbnail_url) + yield self._sse(video_html) + + logger.info(f"Video generated: {video_url}") + continue + + if self.think_opened: + yield self._sse("\n") + yield self._sse(finish="stop") + yield "data: [DONE]\n\n" + except asyncio.CancelledError: + logger.debug("Video stream cancelled by client", extra={"model": self.model}) + except StreamIdleTimeoutError as e: + raise UpstreamException( + message=f"Video 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 video: {e}", extra={"model": self.model}) + raise UpstreamException( + message="Upstream connection closed unexpectedly", + status_code=502, + details={"error": str(e), "type": "http2_stream_error"}, + ) + logger.error(f"Video stream request error: {e}", extra={"model": self.model}) + raise UpstreamException( + message=f"Upstream request failed: {e}", + status_code=502, + details={"error": str(e)}, + ) + except Exception as e: + logger.error(f"Video stream processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__}) + finally: + await self.close() + + +class VideoCollectProcessor(BaseProcessor): + """视频非流式响应处理器""" + + def __init__(self, model: str, token: str = ""): + super().__init__(model, token) + self.video_format = str(get_config("app.video_format")).lower() + + def _build_video_html(self, video_url: str, thumbnail_url: str = "") -> str: + poster_attr = f' poster="{thumbnail_url}"' if thumbnail_url else "" + return f'''''' + + async def process(self, response: AsyncIterable[bytes]) -> dict[str, Any]: + """处理并收集视频响应""" + response_id = "" + content = "" + idle_timeout = get_config("grok.video_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 video_resp := resp.get("streamingVideoGenerationResponse"): + if video_resp.get("progress") == 100: + response_id = resp.get("responseId", "") + video_url = video_resp.get("videoUrl", "") + thumbnail_url = video_resp.get("thumbnailImageUrl", "") + + if video_url: + final_video_url = await self.process_url(video_url, "video") + final_thumbnail_url = "" + if thumbnail_url: + final_thumbnail_url = await self.process_url(thumbnail_url, "image") + + if self.video_format == "url": + content = final_video_url + else: + content = self._build_video_html(final_video_url, final_thumbnail_url) + logger.info(f"Video generated: {video_url}") + + except asyncio.CancelledError: + logger.debug("Video collect cancelled by client", extra={"model": self.model}) + except StreamIdleTimeoutError as e: + logger.warning(f"Video collect idle timeout: {e}", extra={"model": self.model}) + except RequestsError as e: + if _is_http2_stream_error(e): + logger.warning(f"HTTP/2 stream error in video collect: {e}", extra={"model": self.model}) + else: + logger.error(f"Video collect request error: {e}", extra={"model": self.model}) + except Exception as e: + logger.error(f"Video collect processing error: {e}", extra={"model": self.model, "error_type": type(e).__name__}) + finally: + await self.close() + + return { + "id": response_id, + "object": "chat.completion", + "created": self.created, + "model": self.model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content, "refusal": None}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, + } + + +__all__ = ["VideoStreamProcessor", "VideoCollectProcessor"] diff --git a/app/services/grok/services/assets.py b/app/services/grok/services/assets.py index 6e53634..9ac3eb7 100644 --- a/app/services/grok/services/assets.py +++ b/app/services/grok/services/assets.py @@ -4,29 +4,31 @@ Grok 文件资产服务 import asyncio import base64 -import os -import time import hashlib +import os import re -from pathlib import Path +import time from contextlib import asynccontextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple +from urllib.parse import urlparse try: import fcntl -except ImportError: # pragma: no cover - non-posix platforms +except ImportError: fcntl = None -from typing import Tuple, List, Dict, Optional, Any -from urllib.parse import urlparse import aiofiles from curl_cffi.requests import AsyncSession -from app.core.logger import logger from app.core.config import get_config from app.core.exceptions import AppException, UpstreamException, ValidationException +from app.core.logger import logger from app.services.grok.utils.headers import apply_statsig, build_sso_cookie from app.services.token.service import TokenService +# ==================== 常量 ==================== UPLOAD_API = "https://grok.com/rest/app-chat/upload-file" LIST_API = "https://grok.com/rest/assets" @@ -34,31 +36,48 @@ DELETE_API = "https://grok.com/rest/assets-metadata" DOWNLOAD_API = "https://assets.grok.com" LOCK_DIR = Path(__file__).parent.parent.parent.parent.parent / "data" / ".locks" -DEFAULT_TIMEOUT = 120 -DEFAULT_BROWSER = "chrome136" -DEFAULT_USER_AGENT = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36" -DEFAULT_MIME = "application/octet-stream" -DEFAULT_BASE_PROXY_URL = "" -DEFAULT_ASSET_PROXY_URL = "" -DEFAULT_CACHE_ENABLE_AUTO_CLEAN = True -DEFAULT_CACHE_LIMIT_MB = 1024 -DEFAULT_APP_URL = "" -DEFAULT_MAX_CONCURRENT = 25 -DEFAULT_DELETE_BATCH_SIZE = 10 -_ASSETS_SEMAPHORE = asyncio.Semaphore(DEFAULT_MAX_CONCURRENT) -_ASSETS_SEM_VALUE = DEFAULT_MAX_CONCURRENT +# 全局信号量（运行时动态初始化） +_ASSETS_SEMAPHORE = None +_ASSETS_SEM_VALUE = None + +# 常用 MIME 类型（业务数据，非配置） +MIME_TYPES = { + ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".png": "image/png", ".gif": "image/gif", + ".webp": "image/webp", ".bmp": "image/bmp", ".pdf": "application/pdf", ".txt": "text/plain", + ".md": "text/markdown", ".csv": "text/csv", ".json": "application/json", ".xml": "application/xml", + ".py": "text/x-python-script", ".js": "application/javascript", ".html": "text/html", + ".css": "text/css", ".mp4": "video/mp4", ".webm": "video/webm", +} + +IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"} +VIDEO_EXTS = {".mp4", ".mov", ".m4v", ".webm", ".avi", ".mkv"} + +# ==================== 工具函数 ==================== + +def _get_assets_semaphore() -> asyncio.Semaphore: + """获取全局并发控制信号量""" + value = max(1, int(get_config("performance.assets_max_concurrent"))) + + global _ASSETS_SEMAPHORE, _ASSETS_SEM_VALUE + if _ASSETS_SEMAPHORE is None or value != _ASSETS_SEM_VALUE: + _ASSETS_SEM_VALUE = value + _ASSETS_SEMAPHORE = asyncio.Semaphore(value) + return _ASSETS_SEMAPHORE @asynccontextmanager async def _file_lock(name: str, timeout: int = 10): + """文件锁""" if fcntl is None: yield return + LOCK_DIR.mkdir(parents=True, exist_ok=True) lock_path = LOCK_DIR / f"{name}.lock" fd = None locked = False start = time.monotonic() + try: fd = open(lock_path, "a+") while True: @@ -78,131 +97,59 @@ async def _file_lock(name: str, timeout: int = 10): fcntl.flock(fd, fcntl.LOCK_UN) except Exception: pass - try: - fd.close() - except Exception: - pass + fd.close() -MIME_TYPES = { - # 图片 - ".jpg": "image/jpeg", - ".jpeg": "image/jpeg", - ".png": "image/png", - ".gif": "image/gif", - ".webp": "image/webp", - ".bmp": "image/bmp", - # 文档 - ".pdf": "application/pdf", - ".txt": "text/plain", - ".md": "text/markdown", - ".doc": "application/msword", - ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", - ".rtf": "application/rtf", - # 表格 - ".csv": "text/csv", - ".xls": "application/vnd.ms-excel", - ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", - # 代码 - ".py": "text/x-python-script", - ".js": "application/javascript", - ".ts": "application/typescript", - ".java": "text/x-java", - ".cpp": "text/x-c++", - ".c": "text/x-c", - ".go": "text/x-go", - ".rs": "text/x-rust", - ".rb": "text/x-ruby", - ".php": "text/x-php", - ".sh": "application/x-sh", - ".html": "text/html", - ".css": "text/css", - ".sql": "application/sql", - # 数据 - ".json": "application/json", - ".xml": "application/xml", - ".yaml": "application/x-yaml", - ".yml": "application/x-yaml", - ".toml": "application/toml", - ".ini": "text/plain", - ".log": "text/plain", - ".tmp": "application/octet-stream", - # 其他 - ".graphql": "application/graphql", - ".proto": "application/x-protobuf", - ".latex": "application/x-latex", - ".wiki": "text/plain", - ".rst": "text/x-rst", -} - -IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"} -VIDEO_EXTS = {".mp4", ".mov", ".m4v", ".webm", ".avi", ".mkv"} +@dataclass +class ServiceConfig: + """服务配置""" + proxy: str + timeout: int + browser: str + user_agent: str + + @classmethod + def from_settings(cls, proxy: Optional[str] = None): + return cls( + proxy=proxy or get_config("grok.asset_proxy_url") or get_config("grok.base_proxy_url"), + timeout=get_config("grok.timeout"), + browser=get_config("grok.browser"), + user_agent=get_config("grok.user_agent"), + ) + + def get_proxies(self) -> Optional[dict]: + return {"http": self.proxy, "https": self.proxy} if self.proxy else None # ==================== 基础服务 ==================== - class BaseService: """基础服务类""" - def __init__(self, proxy: str = None): - self.proxy = ( - proxy - or get_config("grok.asset_proxy_url", DEFAULT_ASSET_PROXY_URL) - or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) - ) - self.timeout = get_config("grok.timeout", DEFAULT_TIMEOUT) + def __init__(self, proxy: Optional[str] = None): + self.config = ServiceConfig.from_settings(proxy) self._session: Optional[AsyncSession] = None - def _headers(self, token: str, referer: str = "https://grok.com/") -> dict: + def _build_headers(self, token: str, referer: str = "https://grok.com/", download: bool = False) -> dict: """构建请求头""" - user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT) - headers = { - "Accept": "*/*", - "Accept-Encoding": "gzip, deflate, br, zstd", - "Accept-Language": "zh-CN,zh;q=0.9", - "Baggage": "sentry-environment=production,sentry-release=d6add6fb0460641fd482d767a335ef72b9b6abb8,sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c", - "Cache-Control": "no-cache", - "Content-Type": "application/json", - "Origin": "https://grok.com", - "Pragma": "no-cache", - "Priority": "u=1, i", - "Referer": referer, - "Sec-Ch-Ua": '"Google Chrome";v="136", "Chromium";v="136", "Not(A:Brand";v="24"', - "Sec-Ch-Ua-Arch": "arm", - "Sec-Ch-Ua-Bitness": "64", - "Sec-Ch-Ua-Mobile": "?0", - "Sec-Ch-Ua-Model": "", - "Sec-Ch-Ua-Platform": '"macOS"', - "Sec-Fetch-Dest": "empty", - "Sec-Fetch-Mode": "cors", - "Sec-Fetch-Site": "same-origin", - "User-Agent": user_agent, - } - - apply_statsig(headers) + if download: + headers = { + "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8", + "Sec-Fetch-Dest": "document", + "Sec-Fetch-Mode": "navigate", + "Referer": referer, + } + else: + headers = { + "Accept": "*/*", + "Content-Type": "application/json", + "Origin": "https://grok.com", + "Referer": referer, + "User-Agent": self.config.user_agent, + } + apply_statsig(headers) + headers["Cookie"] = build_sso_cookie(token) - - return headers - - def _proxies(self) -> Optional[dict]: - """构建代理配置""" - return {"http": self.proxy, "https": self.proxy} if self.proxy else None - - def _dl_headers(self, token: str, file_path: str) -> dict: - """构建下载请求头""" - headers = { - "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8", - "Sec-Fetch-Dest": "document", - "Sec-Fetch-Mode": "navigate", - "Sec-Fetch-Site": "same-site", - "Sec-Fetch-User": "?1", - "Upgrade-Insecure-Requests": "1", - "Referer": "https://grok.com/", - } - - headers["Cookie"] = build_sso_cookie(token) - return headers async def _get_session(self) -> AsyncSession: @@ -218,200 +165,139 @@ class BaseService: self._session = None @staticmethod - def is_url(input_str: str) -> bool: + def is_url(s: str) -> bool: """检查是否为 URL""" try: - result = urlparse(input_str) - return all([result.scheme, result.netloc]) and result.scheme in [ - "http", - "https", - ] + r = urlparse(s) + return bool(r.scheme and r.netloc and r.scheme in ["http", "https"]) except Exception: return False @staticmethod async def fetch(url: str) -> Tuple[str, str, str]: - """ - 获取远程资源并转 Base64 - - Raises: - UpstreamException: 当获取失败时 - """ + """获取远程资源并转 Base64""" try: async with AsyncSession() as session: response = await session.get(url, timeout=10) if response.status_code >= 400: raise UpstreamException( - message=f"Failed to fetch resource: {response.status_code}", + message=f"Failed to fetch: {response.status_code}", details={"url": url, "status": response.status_code}, ) - + filename = url.split("/")[-1].split("?")[0] or "download" - content_type = response.headers.get("content-type", DEFAULT_MIME).split( - ";" - )[0] + content_type = response.headers.get("content-type", "application/octet-stream").split(";")[0] b64 = base64.b64encode(response.content).decode() - - logger.debug(f"Fetched: {url} -> {filename}") + + logger.debug(f"Fetched: {url}") return filename, b64, content_type except Exception as e: - logger.error(f"Fetch failed: {url} - {e}") if isinstance(e, AppException): - raise e - raise UpstreamException( - f"Resource fetch failed: {str(e)}", details={"url": url} - ) + raise + logger.error(f"Fetch failed: {url} - {e}") + raise UpstreamException(f"Fetch failed: {str(e)}", details={"url": url}) @staticmethod def parse_b64(data_uri: str) -> Tuple[str, str, str]: """解析 Base64 数据""" - if data_uri.startswith("data:"): - try: - header, b64 = data_uri.split(",", 1) - except ValueError: - return "file.bin", data_uri, DEFAULT_MIME - if ";base64" in header: - mime = header[5:].split(";", 1)[0] or DEFAULT_MIME - b64 = re.sub(r"\\s+", "", b64) - ext = mime.split("/")[-1] if "/" in mime else "bin" - return f"file.{ext}", b64, mime - return "file.bin", data_uri, DEFAULT_MIME + if not data_uri.startswith("data:"): + return "file.bin", data_uri, "application/octet-stream" + + try: + header, b64 = data_uri.split(",", 1) + except ValueError: + return "file.bin", data_uri, "application/octet-stream" + + if ";base64" not in header: + return "file.bin", data_uri, "application/octet-stream" + + mime = header[5:].split(";", 1)[0] or "application/octet-stream" + b64 = re.sub(r"\s+", "", b64) + ext = mime.split("/")[-1] if "/" in mime else "bin" + return f"file.{ext}", b64, mime @staticmethod def to_b64(file_path: Path, mime_type: str) -> str: - """将本地文件转为 base64 data URI""" + """文件转 base64 data URI""" try: b64_data = base64.b64encode(file_path.read_bytes()).decode() return f"data:{mime_type};base64,{b64_data}" except Exception as e: logger.error(f"File to base64 failed: {file_path} - {e}") - raise AppException( - f"Failed to read file: {file_path}", code="file_read_error" - ) + raise AppException(f"Failed to read file: {file_path}", code="file_read_error") # ==================== 上传服务 ==================== - class UploadService(BaseService): """文件上传服务""" async def upload(self, file_input: str, token: str) -> Tuple[str, str]: """ 上传文件到 Grok - + Returns: (file_id, file_uri) - - Raises: - ValidationException: 输入无效 - UpstreamException: 上传失败 """ - value = get_config("performance.assets_max_concurrent", DEFAULT_MAX_CONCURRENT) - try: - value = int(value) - except Exception: - value = DEFAULT_MAX_CONCURRENT - value = max(1, value) - global _ASSETS_SEMAPHORE, _ASSETS_SEM_VALUE - if value != _ASSETS_SEM_VALUE: - _ASSETS_SEM_VALUE = value - _ASSETS_SEMAPHORE = asyncio.Semaphore(value) - async with _ASSETS_SEMAPHORE: - try: - # 处理输入 - if self.is_url(file_input): - filename, b64, mime = await self.fetch(file_input) - else: - filename, b64, mime = self.parse_b64(file_input) + async with _get_assets_semaphore(): + # 处理输入 + if self.is_url(file_input): + filename, b64, mime = await self.fetch(file_input) + else: + filename, b64, mime = self.parse_b64(file_input) + + if not b64: + raise ValidationException("Invalid file input: empty content") - if not b64: - raise ValidationException("Invalid file input: empty content") + # 执行上传 + session = await self._get_session() + response = await session.post( + UPLOAD_API, + headers=self._build_headers(token), + json={"fileName": filename, "fileMimeType": mime, "content": b64}, + impersonate=self.config.browser, + timeout=self.config.timeout, + proxies=self.config.get_proxies(), + ) - # 构建请求 - headers = self._headers(token) - payload = { - "fileName": filename, - "fileMimeType": mime, - "content": b64, - } + # 处理响应 + if response.status_code == 200: + result = response.json() + file_id = result.get("fileMetadataId", "") + file_uri = result.get("fileUri", "") + logger.info(f"Upload success: {filename} -> {file_id}") + return file_id, file_uri - # 执行上传 - session = await self._get_session() - response = await session.post( - UPLOAD_API, - headers=headers, - json=payload, - impersonate=get_config("grok.browser", DEFAULT_BROWSER), - timeout=self.timeout, - proxies=self._proxies(), - ) - - if response.status_code == 200: - result = response.json() - file_id = result.get("fileMetadataId", "") - file_uri = result.get("fileUri", "") - logger.info( - f"Upload success: {filename} -> {file_id}", - extra={"file_id": file_id}, - ) - return file_id, file_uri - - # 401/403 表示 token 失效，记录失败并标记冷却 - if response.status_code in (401, 403): - logger.warning( - f"Upload token invalid: {response.status_code}, recording failure", - extra={"response": response.text[:200]}, - ) - try: - await TokenService.record_fail( - token, - status_code=response.status_code, - reason="upload_auth_failed", - ) - except Exception as mark_err: - logger.error(f"Failed to record token failure: {mark_err}") - - raise UpstreamException( - message=f"Upload authentication failed: {response.status_code}", - details={ - "status": response.status_code, - "response": response.text[:200], - "token_invalidated": True, - }, - ) - - logger.error( - f"Upload failed: {filename} - {response.status_code}", - extra={"response": response.text[:200]}, - ) + # 认证失败 + if response.status_code in (401, 403): + logger.warning(f"Upload auth failed: {response.status_code}") + try: + await TokenService.record_fail(token, response.status_code, "upload_auth_failed") + except Exception as e: + logger.error(f"Failed to record token failure: {e}") + raise UpstreamException( - message=f"Upload failed with status {response.status_code}", - details={ - "status": response.status_code, - "response": response.text[:200], - }, + message=f"Upload authentication failed: {response.status_code}", + details={"status": response.status_code, "token_invalidated": True}, ) - except Exception as e: - logger.error(f"Upload error: {e}") - if isinstance(e, AppException): - raise e - raise UpstreamException(f"Upload process error: {str(e)}") + # 其他错误 + logger.error(f"Upload failed: {filename} - {response.status_code}") + raise UpstreamException( + message=f"Upload failed: {response.status_code}", + details={"status": response.status_code}, + ) # ==================== 列表服务 ==================== - class ListService(BaseService): """文件列表查询服务""" async def iter_assets(self, token: str): - """ - 分页迭代资产列表 - """ - headers = self._headers(token, referer="https://grok.com/files") - base_params = { + """分页迭代资产列表""" + headers = self._build_headers(token, referer="https://grok.com/files") + params = { "pageSize": 50, "orderBy": "ORDER_BY_LAST_USE_TIME", "source": "SOURCE_ANY", @@ -422,29 +308,27 @@ class ListService(BaseService): async with AsyncSession() as session: while True: - params = dict(base_params) if page_token: if page_token in seen_tokens: - logger.warning( - "List pagination stopped due to repeated page token" - ) + logger.warning("Pagination stopped: repeated page token") break seen_tokens.add(page_token) params["pageToken"] = page_token + else: + params.pop("pageToken", None) response = await session.get( LIST_API, headers=headers, params=params, - impersonate=get_config("grok.browser", DEFAULT_BROWSER), - timeout=self.timeout, - proxies=self._proxies(), + impersonate=self.config.browser, + timeout=self.config.timeout, + proxies=self.config.get_proxies(), ) if response.status_code != 200: - logger.error(f"List failed: {response.status_code}") raise UpstreamException( - message=f"List assets failed: {response.status_code}", + message=f"List failed: {response.status_code}", details={"status": response.status_code}, ) @@ -457,161 +341,107 @@ class ListService(BaseService): break async def list(self, token: str) -> List[Dict]: - """ - 查询文件列表 - - Raises: - UpstreamException: 查询失败 - """ - try: - assets: List[Dict] = [] - async for page_assets in self.iter_assets(token): - assets.extend(page_assets) - - logger.info(f"List success: {len(assets)} files") - return assets - - except Exception as e: - logger.error(f"List error: {e}") - if isinstance(e, AppException): - raise e - raise UpstreamException(f"List assets error: {str(e)}") + """查询文件列表""" + assets = [] + async for page_assets in self.iter_assets(token): + assets.extend(page_assets) + logger.info(f"List success: {len(assets)} files") + return assets async def count(self, token: str) -> int: - """ - 统计资产数量（不保留明细） - """ - try: - total = 0 - async for page_assets in self.iter_assets(token): - total += len(page_assets) - return total - except Exception as e: - logger.error(f"List count error: {e}") - if isinstance(e, AppException): - raise e - raise UpstreamException(f"List assets error: {str(e)}") + """统计资产数量""" + total = 0 + async for page_assets in self.iter_assets(token): + total += len(page_assets) + return total # ==================== 删除服务 ==================== - class DeleteService(BaseService): """文件删除服务""" async def delete(self, token: str, asset_id: str) -> bool: - """ - 删除单个文件 + """删除单个文件""" + async with _get_assets_semaphore(): + session = await self._get_session() + response = await session.delete( + f"{DELETE_API}/{asset_id}", + headers=self._build_headers(token, referer="https://grok.com/files"), + impersonate=self.config.browser, + timeout=self.config.timeout, + proxies=self.config.get_proxies(), + ) - Raises: - UpstreamException: 删除失败 - """ - value = get_config("performance.assets_max_concurrent", DEFAULT_MAX_CONCURRENT) - try: - value = int(value) - except Exception: - value = DEFAULT_MAX_CONCURRENT - value = max(1, value) - global _ASSETS_SEMAPHORE, _ASSETS_SEM_VALUE - if value != _ASSETS_SEM_VALUE: - _ASSETS_SEM_VALUE = value - _ASSETS_SEMAPHORE = asyncio.Semaphore(value) - async with _ASSETS_SEMAPHORE: - try: - headers = self._headers(token, referer="https://grok.com/files") - url = f"{DELETE_API}/{asset_id}" + if response.status_code == 200: + logger.debug(f"Deleted: {asset_id}") + return True - session = await self._get_session() - response = await session.delete( - url, - headers=headers, - impersonate=get_config("grok.browser", DEFAULT_BROWSER), - timeout=self.timeout, - proxies=self._proxies(), - ) - - if response.status_code == 200: - logger.debug(f"Delete success: {asset_id}") - return True - - logger.error(f"Delete failed: {asset_id} - {response.status_code}") - #: Note: Returning False or raising Exception? - #: Assuming caller handles Exception for stricter control, or False for loose control. - #: Given "optimization" and "standardization", raising exceptions is better for API feedback. - raise UpstreamException( - message=f"Delete failed: {asset_id}", - details={"status": response.status_code}, - ) - - except Exception as e: - logger.error(f"Delete error: {asset_id} - {e}") - if isinstance(e, AppException): - raise e - raise UpstreamException(f"Delete error: {str(e)}") + logger.error(f"Delete failed: {asset_id} - {response.status_code}") + raise UpstreamException( + message=f"Delete failed: {asset_id}", + details={"status": response.status_code}, + ) async def delete_all(self, token: str) -> Dict[str, int]: - """ - 删除所有文件 - """ - total = 0 - success = 0 - failed = 0 - list_service = ListService(self.proxy) + """删除所有文件""" + total = success = failed = 0 + list_service = ListService(self.config.proxy) + try: async for assets in list_service.iter_assets(token): if not assets: continue + total += len(assets) - - # 批量并发删除 - async def _delete_one(asset: Dict, index: int) -> bool: - await asyncio.sleep(0.01 * index) - asset_id = asset.get("assetId", "") - if asset_id: - try: - return await self.delete(token, asset_id) - except Exception: - return False - return False - - batch_size = get_config( - "performance.assets_delete_batch_size", DEFAULT_DELETE_BATCH_SIZE - ) - try: - batch_size = int(batch_size) - except Exception: - batch_size = DEFAULT_DELETE_BATCH_SIZE - batch_size = max(1, batch_size) - for i in range(0, len(assets), batch_size): - batch = assets[i : i + batch_size] - results = await asyncio.gather( - *[_delete_one(asset, idx) for idx, asset in enumerate(batch)] - ) - success += sum(results) - failed += len(batch) - sum(results) + batch_result = await self._delete_batch(token, assets) + success += batch_result["success"] + failed += batch_result["failed"] if total == 0: logger.info("No assets to delete") return {"total": 0, "success": 0, "failed": 0, "skipped": True} - except Exception as e: - logger.error(f"Delete all failed during list: {e}") - return {"total": total, "success": success, "failed": failed} finally: await list_service.close() logger.info(f"Delete all: total={total}, success={success}, failed={failed}") return {"total": total, "success": success, "failed": failed} + async def _delete_batch(self, token: str, assets: List[Dict]) -> Dict[str, int]: + """批量删除""" + batch_size = max(1, int(get_config("performance.assets_delete_batch_size"))) + success = failed = 0 + + for i in range(0, len(assets), batch_size): + batch = assets[i:i + batch_size] + results = await asyncio.gather( + *[self._delete_one(token, asset, idx) for idx, asset in enumerate(batch)], + return_exceptions=True + ) + success += sum(1 for r in results if r is True) + failed += sum(1 for r in results if r is not True) + + return {"success": success, "failed": failed} + + async def _delete_one(self, token: str, asset: Dict, index: int) -> bool: + """删除单个资产（带延迟）""" + await asyncio.sleep(0.01 * index) + asset_id = asset.get("assetId", "") + if not asset_id: + return False + try: + return await self.delete(token, asset_id) + except Exception: + return False + # ==================== 下载服务 ==================== - class DownloadService(BaseService): """文件下载服务""" - def __init__(self, proxy: str = None): + def __init__(self, proxy: Optional[str] = None): super().__init__(proxy) - # 创建缓存目录 self.base_dir = Path(__file__).parent.parent.parent.parent.parent / "data" / "tmp" self.image_dir = self.base_dir / "image" self.video_dir = self.base_dir / "video" @@ -625,214 +455,159 @@ class DownloadService(BaseService): filename = file_path.lstrip("/").replace("/", "-") return cache_dir / filename - async def download( - self, file_path: str, token: str, media_type: str = "image" - ) -> Tuple[Optional[Path], str]: - """ - 下载文件到本地 + def _get_mime(self, cache_path: Path, response=None) -> str: + """获取 MIME 类型""" + if response: + return response.headers.get("content-type", "application/octet-stream").split(";")[0] + return MIME_TYPES.get(cache_path.suffix.lower(), "application/octet-stream") - Raises: - UpstreamException: 下载失败 - """ - value = get_config("performance.assets_max_concurrent", DEFAULT_MAX_CONCURRENT) - try: - value = int(value) - except Exception: - value = DEFAULT_MAX_CONCURRENT - value = max(1, value) - global _ASSETS_SEMAPHORE, _ASSETS_SEM_VALUE - if value != _ASSETS_SEM_VALUE: - _ASSETS_SEM_VALUE = value - _ASSETS_SEMAPHORE = asyncio.Semaphore(value) - async with _ASSETS_SEMAPHORE: - try: - cache_path = self._cache_path(file_path, media_type) + async def download(self, file_path: str, token: str, media_type: str = "image") -> Tuple[Optional[Path], str]: + """下载文件到本地""" + async with _get_assets_semaphore(): + cache_path = self._cache_path(file_path, media_type) - # 如果已缓存 + # 检查缓存 + if cache_path.exists(): + logger.debug(f"Cache hit: {cache_path}") + return cache_path, self._get_mime(cache_path) + + # 文件锁防止并发下载 + lock_name = f"dl_{media_type}_{hashlib.sha1(str(cache_path).encode()).hexdigest()[:16]}" + async with _file_lock(lock_name, timeout=10): + # 双重检查 if cache_path.exists(): - logger.debug(f"Cache hit: {cache_path}") - mime_type = MIME_TYPES.get(cache_path.suffix.lower(), DEFAULT_MIME) - return cache_path, mime_type + return cache_path, self._get_mime(cache_path) - lock_name = f"download_{media_type}_{hashlib.sha1(str(cache_path).encode('utf-8')).hexdigest()[:16]}" - async with _file_lock(lock_name, timeout=10): - # Double-check after lock - if cache_path.exists(): - logger.debug(f"Cache hit after lock: {cache_path}") - mime_type = MIME_TYPES.get( - cache_path.suffix.lower(), DEFAULT_MIME - ) - return cache_path, mime_type + # 执行下载 + mime = await self._download_file(file_path, token, cache_path) + logger.info(f"Downloaded: {file_path}") - # 下载文件 - if not file_path.startswith("/"): - file_path = f"/{file_path}" + # 异步检查缓存限制 + asyncio.create_task(self.check_limit()) - url = f"{DOWNLOAD_API}{file_path}" - headers = self._dl_headers(token, file_path) + return cache_path, mime - session = await self._get_session() - response = await session.get( - url, - headers=headers, - proxies=self._proxies(), - timeout=self.timeout, - allow_redirects=True, - impersonate=get_config("grok.browser", DEFAULT_BROWSER), - stream=True, - ) + async def _download_file(self, file_path: str, token: str, cache_path: Path) -> str: + """执行下载""" + if not file_path.startswith("/"): + file_path = f"/{file_path}" - if response.status_code != 200: - raise UpstreamException( - message=f"Download failed: {response.status_code}", - details={"path": file_path, "status": response.status_code}, - ) + url = f"{DOWNLOAD_API}{file_path}" + headers = self._build_headers(token, download=True) - # 保存文件（分块写入，避免大文件占用内存） - tmp_path = cache_path.with_suffix(cache_path.suffix + ".tmp") - try: - async with aiofiles.open(tmp_path, "wb") as f: - if hasattr(response, "aiter_content"): - async for chunk in response.aiter_content(): - if chunk: - await f.write(chunk) - elif hasattr(response, "aiter_bytes"): - async for chunk in response.aiter_bytes(): - if chunk: - await f.write(chunk) - elif hasattr(response, "aiter_raw"): - async for chunk in response.aiter_raw(): - if chunk: - await f.write(chunk) - else: - await f.write(response.content) - os.replace(tmp_path, cache_path) - finally: - if tmp_path.exists() and not cache_path.exists(): - try: - tmp_path.unlink() - except Exception: - pass - mime_type = response.headers.get( - "content-type", DEFAULT_MIME - ).split(";")[0] + session = await self._get_session() + response = await session.get( + url, + headers=headers, + proxies=self.config.get_proxies(), + timeout=self.config.timeout, + allow_redirects=True, + impersonate=self.config.browser, + stream=True, + ) - logger.info(f"Download success: {file_path}") + if response.status_code != 200: + raise UpstreamException( + message=f"Download failed: {response.status_code}", + details={"path": file_path, "status": response.status_code}, + ) - # 检查缓存限制 - asyncio.create_task(self.check_limit()) - - return cache_path, mime_type - - except Exception as e: - logger.error(f"Download failed: {file_path} - {e}") - if isinstance(e, AppException): - raise e - raise UpstreamException(f"Download error: {str(e)}") - - async def to_base64( - self, file_path: str, token: str, media_type: str = "image" - ) -> str: - """ - 下载文件并转为 base64 - """ + # 保存文件 + tmp_path = cache_path.with_suffix(cache_path.suffix + ".tmp") try: - cache_path, mime_type = await self.download(file_path, token, media_type) - - if not cache_path or not cache_path.exists(): - raise AppException("File download returned invalid path") - - # 使用基础服务的工具方法转换 - data_uri = self.to_b64(cache_path, mime_type) - - # 删除临时文件 - if data_uri: + async with aiofiles.open(tmp_path, "wb") as f: + # 尝试流式写入 + if hasattr(response, "aiter_content"): + async for chunk in response.aiter_content(): + if chunk: + await f.write(chunk) + else: + await f.write(response.content) + os.replace(tmp_path, cache_path) + finally: + if tmp_path.exists() and not cache_path.exists(): try: - cache_path.unlink() - except Exception as e: - logger.warning(f"Delete temp file failed: {e}") + tmp_path.unlink() + except Exception: + pass - return data_uri + return self._get_mime(cache_path, response) - except Exception as e: - logger.error(f"To base64 failed: {file_path} - {e}") - if isinstance(e, AppException): - raise e - raise AppException(f"Base64 conversion failed: {str(e)}") + async def to_base64(self, file_path: str, token: str, media_type: str = "image") -> str: + """下载并转 base64""" + cache_path, mime = await self.download(file_path, token, media_type) + if not cache_path or not cache_path.exists(): + raise AppException("Download failed: invalid path") + + data_uri = self.to_b64(cache_path, mime) + + # 删除临时文件 + if data_uri: + try: + cache_path.unlink() + except Exception: + pass + + return data_uri def get_stats(self, media_type: str = "image") -> Dict[str, Any]: """获取缓存统计""" cache_dir = self.image_dir if media_type == "image" else self.video_dir if not cache_dir.exists(): return {"count": 0, "size_mb": 0.0} - - files = [f for f in cache_dir.glob("*") if f.is_file()] - allowed_exts = IMAGE_EXTS if media_type == "image" else VIDEO_EXTS - files = [f for f in files if f.suffix.lower() in allowed_exts] + + allowed = IMAGE_EXTS if media_type == "image" else VIDEO_EXTS + files = [f for f in cache_dir.glob("*") if f.is_file() and f.suffix.lower() in allowed] total_size = sum(f.stat().st_size for f in files) - return {"count": len(files), "size_mb": round(total_size / 1024 / 1024, 2)} - def list_files( - self, media_type: str = "image", page: int = 1, page_size: int = 1000 - ) -> Dict[str, Any]: - """列出本地缓存文件""" + def list_files(self, media_type: str = "image", page: int = 1, page_size: int = 1000) -> Dict[str, Any]: + """列出缓存文件""" cache_dir = self.image_dir if media_type == "image" else self.video_dir if not cache_dir.exists(): return {"total": 0, "page": page, "page_size": page_size, "items": []} + + allowed = IMAGE_EXTS if media_type == "image" else VIDEO_EXTS + files = [f for f in cache_dir.glob("*") if f.is_file() and f.suffix.lower() in allowed] - files = [f for f in cache_dir.glob("*") if f.is_file()] - allowed_exts = IMAGE_EXTS if media_type == "image" else VIDEO_EXTS - files = [f for f in files if f.suffix.lower() in allowed_exts] + # 构建文件列表 items = [] for f in files: try: stat = f.stat() - items.append( - { - "name": f.name, - "size_bytes": stat.st_size, - "mtime_ms": int(stat.st_mtime * 1000), - } - ) + items.append({ + "name": f.name, + "size_bytes": stat.st_size, + "mtime_ms": int(stat.st_mtime * 1000), + }) except Exception: continue items.sort(key=lambda x: x["mtime_ms"], reverse=True) + + # 分页 total = len(items) start = max(0, (page - 1) * page_size) - end = start + page_size - paged = items[start:end] + paged = items[start:start + page_size] - if media_type == "image": - for item in paged: - item["view_url"] = f"/v1/files/image/{item['name']}" - else: - preview_map = {} - if self.image_dir.exists(): - for img in self.image_dir.glob("*"): - if img.is_file() and img.suffix.lower() in IMAGE_EXTS: - preview_map.setdefault(img.stem, img.name) - for item in paged: - item["view_url"] = f"/v1/files/video/{item['name']}" - preview_name = preview_map.get(Path(item["name"]).stem) - if preview_name: - item["preview_url"] = f"/v1/files/image/{preview_name}" + # 添加 URL + for item in paged: + item["view_url"] = f"/v1/files/{media_type}/{item['name']}" return {"total": total, "page": page, "page_size": page_size, "items": paged} def delete_file(self, media_type: str, name: str) -> Dict[str, Any]: - """删除单个缓存文件""" + """删除缓存文件""" cache_dir = self.image_dir if media_type == "image" else self.video_dir - safe_name = name.replace("/", "-") - file_path = cache_dir / safe_name - if not file_path.exists(): - return {"deleted": False} - try: - file_path.unlink() - return {"deleted": True} - except Exception: - return {"deleted": False} + file_path = cache_dir / name.replace("/", "-") + + if file_path.exists(): + try: + file_path.unlink() + return {"deleted": True} + except Exception: + pass + return {"deleted": False} def clear(self, media_type: str = "image") -> Dict[str, Any]: """清空缓存""" @@ -841,61 +616,41 @@ class DownloadService(BaseService): return {"count": 0, "size_mb": 0.0} files = list(cache_dir.glob("*")) - total_size = sum(f.stat().st_size for f in files) + total_size = sum(f.stat().st_size for f in files if f.is_file()) count = 0 for f in files: - try: - f.unlink() - count += 1 - except Exception as e: - logger.error(f"Failed to delete {f}: {e}") + if f.is_file(): + try: + f.unlink() + count += 1 + except Exception: + pass return {"count": count, "size_mb": round(total_size / 1024 / 1024, 2)} async def check_limit(self): - """检查并清理缓存限制""" - if self._cleanup_running: + """检查并清理缓存""" + if self._cleanup_running or not get_config("cache.enable_auto_clean"): return + self._cleanup_running = True try: async with _file_lock("cache_cleanup", timeout=5): - if not get_config( - "cache.enable_auto_clean", DEFAULT_CACHE_ENABLE_AUTO_CLEAN - ): - return - - limit_mb = get_config("cache.limit_mb", DEFAULT_CACHE_LIMIT_MB) - - # 统计总大小 - total_size = 0 - all_files = [] - - for d in [self.image_dir, self.video_dir]: - if d.exists(): - for f in d.glob("*"): - try: - stat = f.stat() - total_size += stat.st_size - all_files.append((f, stat.st_mtime, stat.st_size)) - except Exception: - pass - + limit_mb = get_config("cache.limit_mb") + all_files, total_size = self._collect_files() current_mb = total_size / 1024 / 1024 + if current_mb <= limit_mb: return - # 需要清理 - logger.info( - f"Cache limit exceeded ({current_mb:.2f}MB > {limit_mb}MB), cleaning up..." - ) - - # 按时间排序 - all_files.sort(key=lambda x: x[1]) - + # 清理到 80% + logger.info(f"Cache limit exceeded ({current_mb:.2f}MB > {limit_mb}MB), cleaning...") + all_files.sort(key=lambda x: x[1]) # 按时间排序 + deleted_count = 0 deleted_size = 0 - target_mb = limit_mb * 0.8 # 清理到 80% + target_mb = limit_mb * 0.8 for f, _, size in all_files: try: @@ -903,39 +658,32 @@ class DownloadService(BaseService): deleted_count += 1 deleted_size += size total_size -= size - if (total_size / 1024 / 1024) <= target_mb: break - except Exception as e: - logger.error(f"Cleanup failed for {f}: {e}") + except Exception: + pass - logger.info( - f"Cache cleanup: deleted {deleted_count} files ({deleted_size / 1024 / 1024:.2f}MB)" - ) + logger.info(f"Cache cleanup: {deleted_count} files ({deleted_size / 1024 / 1024:.2f}MB)") finally: self._cleanup_running = False - def get_public_url(self, file_path: str) -> str: - """ - 获取文件的公共访问 URL + def _collect_files(self) -> Tuple[List[Tuple[Path, float, int]], int]: + """收集所有缓存文件""" + total_size = 0 + all_files = [] - 如果配置了 app_url，则返回自托管 URL，否则返回 Grok 原始 URL - """ - app_url = get_config("app.app_url", DEFAULT_APP_URL) - if not app_url: - return f"{DOWNLOAD_API}{file_path if file_path.startswith('/') else '/' + file_path}" + for d in [self.image_dir, self.video_dir]: + if d.exists(): + for f in d.glob("*"): + if f.is_file(): + try: + stat = f.stat() + total_size += stat.st_size + all_files.append((f, stat.st_mtime, stat.st_size)) + except Exception: + pass - if not file_path.startswith("/"): - file_path = f"/{file_path}" - - # 自动添加 /v1/files 前缀 - return f"{app_url.rstrip('/')}/v1/files{file_path}" + return all_files, total_size -__all__ = [ - "BaseService", - "UploadService", - "ListService", - "DeleteService", - "DownloadService", -] +__all__ = ["BaseService", "UploadService", "ListService", "DeleteService", "DownloadService"] diff --git a/app/services/grok/services/chat.py b/app/services/grok/services/chat.py index f0c1b3b..ee7de10 100644 --- a/app/services/grok/services/chat.py +++ b/app/services/grok/services/chat.py @@ -10,39 +10,22 @@ from curl_cffi.requests import AsyncSession from app.core.logger import logger from app.core.config import get_config -from app.core.exceptions import ( - AppException, - UpstreamException, - ValidationException, - ErrorType, -) +from app.core.exceptions import AppException, UpstreamException, ValidationException, ErrorType from app.services.grok.models.model import ModelService from app.services.grok.services.assets import UploadService -from app.services.grok.processors.processor import StreamProcessor, CollectProcessor +from app.services.grok.processors import StreamProcessor, CollectProcessor from app.services.grok.utils.retry import retry_on_status from app.services.grok.utils.headers import apply_statsig, build_sso_cookie +from app.services.grok.utils.stream import wrap_stream_with_usage from app.services.token import get_token_manager, EffortType CHAT_API = "https://grok.com/rest/app-chat/conversations/new" -DEFAULT_TIMEOUT = 120 -DEFAULT_BROWSER = "chrome136" -DEFAULT_USER_AGENT = ( - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " - "AppleWebKit/537.36 (KHTML, like Gecko) " - "Chrome/136.0.0.0 Safari/537.36" -) -DEFAULT_TEMPORARY = True -DEFAULT_DISABLE_MEMORY = True -DEFAULT_STREAM = True -DEFAULT_BASE_PROXY_URL = "" - @dataclass class ChatRequest: """聊天请求数据""" - model: str messages: List[Dict[str, Any]] stream: bool = None @@ -52,85 +35,47 @@ class ChatRequest: class MessageExtractor: """消息内容提取器""" - # 需要上传的类型 - UPLOAD_TYPES = {"image_url", "input_audio", "file"} - # 视频模式不支持的类型 - VIDEO_UNSUPPORTED = {"input_audio", "file"} - @staticmethod - def extract( - messages: List[Dict[str, Any]], is_video: bool = False - ) -> tuple[str, List[tuple[str, str]]]: - """ - 从 OpenAI 消息格式提取内容 - - Args: - messages: OpenAI 格式消息列表 - is_video: 是否为视频模型 - - Returns: - (text, attachments): 拼接后的文本和需要上传的附件列表 - - Raises: - ValueError: 视频模型遇到不支持的内容类型 - """ + def extract(messages: List[Dict[str, Any]], is_video: bool = False) -> tuple[str, List[tuple[str, str]]]: + """从 OpenAI 消息格式提取内容，返回 (text, attachments)""" texts = [] - attachments = [] # 需要上传的附件 (URL 或 base64) - - # 先抽取每条消息的文本，保留角色信息用于合并 - extracted: List[Dict[str, str]] = [] + attachments = [] + extracted = [] for msg in messages: role = msg.get("role", "") content = msg.get("content", "") parts = [] - # 简单字符串内容 if isinstance(content, str): if content.strip(): parts.append(content) - - # 列表格式内容 elif isinstance(content, list): for item in content: item_type = item.get("type", "") - # 文本类型 if item_type == "text": - text = item.get("text", "") - if text.strip(): + if text := item.get("text", "").strip(): parts.append(text) - # 图片类型 elif item_type == "image_url": image_data = item.get("image_url", {}) - url = ( - image_data.get("url", "") - if isinstance(image_data, dict) - else str(image_data) - ) + url = image_data.get("url", "") if isinstance(image_data, dict) else str(image_data) if url: attachments.append(("image", url)) - # 音频类型 elif item_type == "input_audio": if is_video: raise ValueError("视频模型不支持 input_audio 类型") audio_data = item.get("input_audio", {}) - data = ( - audio_data.get("data", "") - if isinstance(audio_data, dict) - else str(audio_data) - ) + data = audio_data.get("data", "") if isinstance(audio_data, dict) else str(audio_data) if data: attachments.append(("audio", data)) - # 文件类型 elif item_type == "file": if is_video: raise ValueError("视频模型不支持 file 类型") file_data = item.get("file", {}) - # file 可能是 URL 或 base64 url = file_data.get("url", "") or file_data.get("data", "") if isinstance(file_data, str): url = file_data @@ -140,30 +85,15 @@ class MessageExtractor: if parts: extracted.append({"role": role, "text": "\n".join(parts)}) - # 合并文本 - last_user_index = None - for i in range(len(extracted) - 1, -1, -1): - if extracted[i]["role"] == "user": - last_user_index = i - break + # 找到最后一条 user 消息 + last_user_index = next((i for i in range(len(extracted) - 1, -1, -1) if extracted[i]["role"] == "user"), None) for i, item in enumerate(extracted): role = item["role"] or "user" text = item["text"] - if i == last_user_index: - texts.append(text) - else: - texts.append(f"{role}: {text}") + texts.append(text if i == last_user_index else f"{role}: {text}") - # 换行拼接文本 - message = "\n\n".join(texts) - return message, attachments - - @staticmethod - def extract_text_only(messages: List[Dict[str, Any]]) -> str: - """仅提取文本内容""" - text, _ = MessageExtractor.extract(messages, is_video=True) - return text + return "\n\n".join(texts), attachments class ChatRequestBuilder: @@ -172,7 +102,7 @@ class ChatRequestBuilder: @staticmethod def build_headers(token: str) -> Dict[str, str]: """构造请求头""" - user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT) + user_agent = get_config("grok.user_agent") headers = { "Accept": "*/*", "Accept-Encoding": "gzip, deflate, br, zstd", @@ -202,33 +132,17 @@ class ChatRequestBuilder: return headers @staticmethod - def build_payload( - message: str, - model: str, - mode: str = None, - file_attachments: List[str] = None, - image_attachments: List[str] = None, - ) -> Dict[str, Any]: - """ - 构造请求体 - - Args: - message: 消息文本 - model: 模型名称 - mode: 模型模式 - file_attachments: 文件附件 ID 列表 - image_attachments: 图片附件 ID 列表（合并到 fileAttachments） - """ - temporary = get_config("grok.temporary", DEFAULT_TEMPORARY) - - merged_attachments: List[str] = [] + def build_payload(message: str, model: str, mode: str = None, + file_attachments: List[str] = None, image_attachments: List[str] = None) -> Dict[str, Any]: + """构造请求体""" + merged_attachments = [] if file_attachments: merged_attachments.extend(file_attachments) if image_attachments: merged_attachments.extend(image_attachments) - payload: Dict[str, Any] = { - "temporary": temporary, + payload = { + "temporary": get_config("grok.temporary"), "modelName": model, "message": message, "fileAttachments": merged_attachments, @@ -236,23 +150,17 @@ class ChatRequestBuilder: "disableSearch": False, "enableImageGeneration": True, "returnImageBytes": False, - "returnRawGrokInXaiRequest": False, "enableImageStreaming": True, "imageGenerationCount": 2, "forceConcise": False, "toolOverrides": {}, "enableSideBySide": True, "sendFinalMetadata": True, - "isReasoning": False, - "disableTextFollowUps": False, "responseMetadata": { "modelConfigOverride": {"modelMap": {}}, "requestModelDetails": {"modelId": model}, }, - "disableMemory": get_config("grok.disable_memory", DEFAULT_DISABLE_MEMORY), - "forceSideBySide": False, - "isAsyncChat": False, - "disableSelfHarmShortCircuit": False, + "disableMemory": get_config("grok.disable_memory"), "deviceEnvInfo": { "darkModeEnabled": False, "devicePixelRatio": 2, @@ -269,70 +177,31 @@ class ChatRequestBuilder: return payload -# ==================== Grok 服务 ==================== - - class GrokChatService: """Grok API 调用服务""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) + self.proxy = proxy or get_config("grok.base_proxy_url") - async def chat( - self, - token: str, - message: str, - model: str = "grok-3", - mode: str = None, - stream: bool = None, - file_attachments: List[str] = None, - image_attachments: List[str] = None, - raw_payload: Dict[str, Any] = None, - ): - """ - 发送聊天请求 - - Args: - token: 认证 Token - message: 消息文本 - model: Grok 模型名称 - mode: 模型模式 - stream: 是否流式 - file_attachments: 文件附件 ID 列表 - image_attachments: 图片附件 URL 列表 - - Raises: - UpstreamException: 当 Grok API 返回错误且重试耗尽时 - """ + async def chat(self, token: str, message: str, model: str = "grok-3", mode: str = None, + stream: bool = None, file_attachments: List[str] = None, + image_attachments: List[str] = None, raw_payload: Dict[str, Any] = None): + """发送聊天请求""" if stream is None: - stream = get_config("grok.stream", DEFAULT_STREAM) + stream = get_config("grok.stream") headers = ChatRequestBuilder.build_headers(token) - if raw_payload is not None: - payload = raw_payload - else: - payload = ChatRequestBuilder.build_payload( - message, - model, - mode, - file_attachments, - image_attachments, - ) + payload = raw_payload if raw_payload is not None else ChatRequestBuilder.build_payload( + message, model, mode, file_attachments, image_attachments + ) proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None - timeout = get_config("grok.timeout", DEFAULT_TIMEOUT) - - # 状态码提取器 - def extract_status(e: Exception) -> int | None: - if isinstance(e, UpstreamException) and e.details: - return e.details.get("status") - return None + timeout = get_config("grok.timeout") logger.debug(f"Payload: {payload}") - # 建立连接函数 + # 建立连接 async def establish_connection(): - """建立连接并返回 response 对象""" - browser = get_config("grok.browser", DEFAULT_BROWSER) + browser = get_config("grok.browser") session = AsyncSession(impersonate=browser) try: response = await session.post( @@ -345,63 +214,44 @@ class GrokChatService: ) if response.status_code != 200: + content = "" try: content = await response.text() - except Exception: - content = "" - - logger.error( - f"Chat failed: {response.status_code}", - extra={ - "status": response.status_code, - "token": token[:10] + "...", - }, - ) - body_for_log = content if content else "" - logger.debug( - "Grok API error response " - f"(status={response.status_code}): {body_for_log}" - ) - # 关闭 session 并抛出异常 - try: - await session.close() except Exception: pass + + logger.error(f"Chat failed: {response.status_code}", + extra={"status": response.status_code, "token": token[:10] + "..."}) + + await session.close() raise UpstreamException( message=f"Grok API request failed: {response.status_code}", - details={ - "status": response.status_code, - "body": content, - }, + details={"status": response.status_code, "body": content}, ) - # 返回 session 和 response return session, response except UpstreamException: - # 已经处理过的异常，直接抛出 raise except Exception as e: - # 其他异常，关闭 session 并包装 logger.error(f"Chat request error: {e}") - try: - await session.close() - except Exception: - pass + await session.close() raise UpstreamException( message=f"Chat connection failed: {str(e)}", details={"error": str(e)}, ) - # 建立连接 + # 重试机制 + def extract_status(e: Exception) -> int | None: + if isinstance(e, UpstreamException) and e.details: + return e.details.get("status") + return None + session = None response = None try: - session, response = await retry_on_status( - establish_connection, extract_status=extract_status - ) + session, response = await retry_on_status(establish_connection, extract_status=extract_status) except Exception as e: - # 记录失败 status_code = extract_status(e) if status_code: token_mgr = await get_token_manager() @@ -436,124 +286,47 @@ class GrokChatService: # 提取消息和附件 try: - message, attachments = MessageExtractor.extract( - request.messages, is_video=is_video - ) + message, attachments = MessageExtractor.extract(request.messages, is_video=is_video) except ValueError as e: raise ValidationException(str(e)) - # 处理附件上传 + # 上传附件 file_ids = [] - if attachments: upload_service = UploadService() try: - for attach_type, attach_data in attachments: - # 获取 ID + for _, attach_data in attachments: file_id, _ = await upload_service.upload(attach_data, token) - - if attach_type == "image": - file_ids.append(file_id) - logger.debug(f"Image uploaded: {file_id}") - else: - file_ids.append(file_id) - logger.debug(f"File uploaded: {file_id}") + file_ids.append(file_id) finally: await upload_service.close() - stream = ( - request.stream - if request.stream is not None - else get_config("grok.stream", DEFAULT_STREAM) - ) + stream = request.stream if request.stream is not None else get_config("grok.stream") response = await self.chat( - token, - message, - grok_model, - mode, - stream, - file_attachments=file_ids, - image_attachments=[], + token, message, grok_model, mode, stream, + file_attachments=file_ids, image_attachments=[] ) return response, stream, request.model -# ==================== Chat 业务服务 ==================== - - class ChatService: """Chat 业务服务""" @staticmethod - async def _wrap_stream(stream: AsyncGenerator, token_mgr, token: str, model: str): - """ - 包装流式响应，在完成时记录使用 - - Args: - stream: 原始 AsyncGenerator - token_mgr: TokenManager 实例 - token: Token 字符串 - model: 模型名称 - """ - success = False - try: - async for chunk in stream: - yield chunk - success = True - finally: - # 只在成功完成时记录使用，失败/异常时不扣费 - if success: - try: - model_info = ModelService.get(model) - effort = ( - EffortType.HIGH - if (model_info and model_info.cost.value == "high") - else EffortType.LOW - ) - await token_mgr.consume(token, effort) - logger.debug( - f"Stream completed, recorded usage for token {token[:10]}... (effort={effort.value})" - ) - except Exception as e: - logger.warning(f"Failed to record stream usage: {e}") - - @staticmethod - async def completions( - model: str, - messages: List[Dict[str, Any]], - stream: bool = None, - thinking: str = None, - ): - """ - Chat Completions 入口 - - Args: - model: 模型名称 - messages: 消息列表 - stream: 是否流式 - thinking: 思考模式 - - Returns: - AsyncGenerator 或 dict - """ + async def completions(model: str, messages: List[Dict[str, Any]], + stream: bool = None, thinking: str = None): + """Chat Completions 入口""" # 获取 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(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", - ) + token_mgr = await get_token_manager() + await token_mgr.reload_if_stale() + + token = None + for pool_name in ModelService.pool_candidates_for_model(model): + token = token_mgr.get_token(pool_name) + if token: + break if not token: raise AppException( @@ -564,61 +337,30 @@ class ChatService: ) # 解析参数 - think = None - if thinking == "enabled": - think = True - elif thinking == "disabled": - think = False - - is_stream = stream if stream is not None else get_config("grok.stream", DEFAULT_STREAM) + think = {"enabled": True, "disabled": False}.get(thinking) + is_stream = stream if stream is not None else get_config("grok.stream") # 构造请求 - chat_request = ChatRequest( - model=model, messages=messages, stream=is_stream, think=think - ) + chat_request = ChatRequest(model=model, messages=messages, stream=is_stream, think=think) # 请求 Grok service = GrokChatService() - try: - response, _, model_name = await service.chat_openai(token, chat_request) - except AppException: - raise - except Exception as e: - logger.error(f"Chat service error: {e}") - raise UpstreamException( - message=f"Service processing failed: {str(e)}", - details={"error": str(e)}, - ) + response, _, model_name = await service.chat_openai(token, chat_request) # 处理响应 if is_stream: processor = StreamProcessor(model_name, token, think) - return ChatService._wrap_stream( - processor.process(response), token_mgr, token, model - ) - else: - result = await CollectProcessor(model_name, token).process(response) - # 非流式：处理完成后立即记录使用 - try: - model_info = ModelService.get(model) - effort = ( - EffortType.HIGH - if (model_info and model_info.cost.value == "high") - else EffortType.LOW - ) - await token_mgr.consume(token, effort) - logger.debug( - f"Collect completed, recorded usage for token {token[:10]}... (effort={effort.value})" - ) - except Exception as e: - logger.warning(f"Failed to record collect usage: {e}") - return result + return wrap_stream_with_usage(processor.process(response), token_mgr, token, model) + + # 非流式 + result = await CollectProcessor(model_name, token).process(response) + try: + model_info = ModelService.get(model) + 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 record usage: {e}") + return result -__all__ = [ - "GrokChatService", - "ChatRequest", - "ChatRequestBuilder", - "MessageExtractor", - "ChatService", -] +__all__ = ["GrokChatService", "ChatRequest", "ChatRequestBuilder", "MessageExtractor", "ChatService"] diff --git a/app/services/grok/services/image.py b/app/services/grok/services/image.py index 7ca9ade..06b1dd7 100644 --- a/app/services/grok/services/image.py +++ b/app/services/grok/services/image.py @@ -18,25 +18,12 @@ from app.core.config import get_config from app.core.logger import logger from app.services.grok.utils.headers import build_sso_cookie - WS_URL = "wss://grok.com/ws/imagine/listen" -DEFAULT_TIMEOUT = 120 -DEFAULT_BLOCKED_SECONDS = 15 -DEFAULT_FINAL_MIN_BYTES = 100000 -DEFAULT_MEDIUM_MIN_BYTES = 30000 -DEFAULT_USER_AGENT = ( - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " - "AppleWebKit/537.36 (KHTML, like Gecko) " - "Chrome/136.0.0.0 Safari/537.36" -) -DEFAULT_BASE_PROXY_URL = "" - class _BlockedError(Exception): pass - class ImageService: """Grok Imagine WebSocket image service.""" @@ -44,25 +31,22 @@ class ImageService: self._ssl_context = ssl.create_default_context() self._url_pattern = re.compile(r"/images/([a-f0-9-]+)\.(png|jpg|jpeg)") - def _get_ws_url(self) -> str: - return WS_URL - def _resolve_proxy(self) -> tuple[aiohttp.BaseConnector, Optional[str]]: - proxy_url = get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) + proxy_url = get_config("grok.base_proxy_url") if not proxy_url: return aiohttp.TCPConnector(ssl=self._ssl_context), None scheme = urlparse(proxy_url).scheme.lower() if scheme.startswith("socks"): - logger.info(f"Grok Imagine WebSocket using SOCKS proxy: {proxy_url}") + logger.info(f"Using SOCKS proxy: {proxy_url}") return ProxyConnector.from_url(proxy_url, ssl=self._ssl_context), None - logger.info(f"Grok Imagine WebSocket using HTTP proxy: {proxy_url}") + logger.info(f"Using HTTP proxy: {proxy_url}") return aiohttp.TCPConnector(ssl=self._ssl_context), proxy_url def _get_ws_headers(self, token: str) -> Dict[str, str]: cookie = build_sso_cookie(token, include_rw=True) - user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT) + user_agent = get_config("grok.user_agent") return { "Cookie": cookie, "Origin": "https://grok.com", @@ -74,15 +58,10 @@ class ImageService: def _extract_image_id(self, url: str) -> Optional[str]: match = self._url_pattern.search(url or "") - if match: - return match.group(1) - return None + return match.group(1) if match else None def _is_final_image(self, url: str, blob_size: int) -> bool: - min_bytes = int( - get_config("grok.image_ws_final_min_bytes", DEFAULT_FINAL_MIN_BYTES) - ) - return (url or "").lower().endswith((".jpg", ".jpeg")) and blob_size > min_bytes + return (url or "").lower().endswith((".jpg", ".jpeg")) and blob_size > get_config("grok.image_ws_final_min_bytes") def _classify_image(self, url: str, blob: str) -> Optional[Dict[str, object]]: if not url or not blob: @@ -91,12 +70,8 @@ class ImageService: image_id = self._extract_image_id(url) or uuid.uuid4().hex blob_size = len(blob) is_final = self._is_final_image(url, blob_size) - if is_final: - stage = "final" - elif blob_size > DEFAULT_MEDIUM_MIN_BYTES: - stage = "medium" - else: - stage = "preview" + + stage = "final" if is_final else ("medium" if blob_size > get_config("grok.image_ws_medium_min_bytes") else "preview") return { "type": "image", @@ -117,37 +92,23 @@ class ImageService: enable_nsfw: bool = True, max_retries: int = None, ) -> AsyncGenerator[Dict[str, object], None]: - if max_retries is None: - retries = 1 - else: - retries = max_retries - retries = max(1, retries) + retries = max(1, max_retries if max_retries is not None else 1) for attempt in range(retries): try: yielded_any = False - async for item in self._stream_once( - token, prompt, aspect_ratio, n, enable_nsfw - ): + 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: + if yielded_any or attempt + 1 >= retries: + if not yielded_any: + yield {"type": "error", "error_code": "blocked", "error": "blocked_no_final_image"} return - if attempt + 1 < retries: - logger.warning( - f"Grok Imagine WebSocket blocked, retry {attempt + 1}/{retries}" - ) - continue - 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"Grok Imagine WebSocket stream failed: {e}") + logger.error(f"WebSocket stream failed: {e}") return async def _stream_once( @@ -158,25 +119,21 @@ class ImageService: n: int, enable_nsfw: bool, ) -> AsyncGenerator[Dict[str, object], None]: - ws_url = self._get_ws_url() - request_id = str(uuid.uuid4()) headers = self._get_ws_headers(token) - timeout = float(get_config("grok.timeout", DEFAULT_TIMEOUT)) - blocked_seconds = float( - get_config("grok.image_ws_blocked_seconds", DEFAULT_BLOCKED_SECONDS) - ) + timeout = float(get_config("grok.timeout")) + blocked_seconds = float(get_config("grok.image_ws_blocked_seconds")) try: connector, proxy = self._resolve_proxy() except Exception as e: - logger.error(f"Grok Imagine WebSocket proxy setup failed: {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, + WS_URL, headers=headers, heartbeat=20, receive_timeout=timeout, @@ -206,27 +163,21 @@ class ImageService: } await ws.send_json(message) - logger.info(f"Grok Imagine WebSocket request sent: {prompt[:80]}...") + logger.info(f"WebSocket request sent: {prompt[:80]}...") - images: Dict[str, Dict[str, object]] = {} + images = {} completed = 0 - start_time = time.time() - last_activity = time.time() + 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: - if time.time() - medium_received_time > min( - 10, blocked_seconds - ): - raise _BlockedError() + 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"Grok Imagine WebSocket idle timeout, collected {completed} images" - ) + logger.info(f"WebSocket idle timeout, collected {completed} images") break continue @@ -236,73 +187,43 @@ class ImageService: msg_type = msg.get("type") if msg_type == "image": - info = self._classify_image( - msg.get("url", ""), msg.get("blob", "") - ) + 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 - images[image_id] = { - "is_final": info["is_final"] or existing.get("is_final") - } - + images[image_id] = {"is_final": info["is_final"] or existing.get("is_final")} yield info elif msg_type == "error": - error_code = msg.get("err_code", "") - error_msg = msg.get("err_msg", "") - logger.warning( - f"Grok Imagine WebSocket error: {error_code} - {error_msg}" - ) - yield { - "type": "error", - "error_code": error_code, - "error": error_msg, - } + 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"Grok Imagine WebSocket collected {completed} final images" - ) + logger.info(f"WebSocket collected {completed} final images") break - if medium_received_time and completed == 0: - if time.time() - medium_received_time > blocked_seconds: - raise _BlockedError() + 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"Grok Imagine WebSocket closed/error: {ws_msg.type}" - ) - yield { - "type": "error", - "error_code": "ws_closed", - "error": f"websocket closed: {ws_msg.type}", - } + 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"Grok Imagine WebSocket connection error: {e}") - yield { - "type": "error", - "error_code": "connection_failed", - "error": str(e), - } - + logger.error(f"WebSocket connection error: {e}") + yield {"type": "error", "error_code": "connection_failed", "error": str(e)} image_service = ImageService() - __all__ = ["image_service", "ImageService"] diff --git a/app/services/grok/services/media.py b/app/services/grok/services/media.py index 0780957..8004028 100644 --- a/app/services/grok/services/media.py +++ b/app/services/grok/services/media.py @@ -18,39 +18,39 @@ from app.core.exceptions import ( ) from app.services.grok.models.model import ModelService from app.services.token import get_token_manager, EffortType -from app.services.grok.processors.processor import VideoStreamProcessor, VideoCollectProcessor +from app.services.grok.processors import VideoStreamProcessor, VideoCollectProcessor from app.services.grok.utils.headers import apply_statsig, build_sso_cookie - +from app.services.grok.utils.stream import wrap_stream_with_usage CREATE_POST_API = "https://grok.com/rest/media/post/create" CHAT_API = "https://grok.com/rest/app-chat/conversations/new" -DEFAULT_BROWSER = "chrome136" -DEFAULT_USER_AGENT = ( - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " - "AppleWebKit/537.36 (KHTML, like Gecko) " - "Chrome/136.0.0.0 Safari/537.36" -) -DEFAULT_TIMEOUT = 300 -DEFAULT_MAX_CONCURRENT = 50 -DEFAULT_BASE_PROXY_URL = "" -DEFAULT_STREAM = True -_MEDIA_SEMAPHORE = asyncio.Semaphore(DEFAULT_MAX_CONCURRENT) -_MEDIA_SEM_VALUE = DEFAULT_MAX_CONCURRENT +_MEDIA_SEMAPHORE = None +_MEDIA_SEM_VALUE = 0 + + +def _get_semaphore() -> asyncio.Semaphore: + """获取或更新信号量""" + global _MEDIA_SEMAPHORE, _MEDIA_SEM_VALUE + value = max(1, int(get_config("performance.media_max_concurrent"))) + if value != _MEDIA_SEM_VALUE: + _MEDIA_SEM_VALUE = value + _MEDIA_SEMAPHORE = asyncio.Semaphore(value) + return _MEDIA_SEMAPHORE class VideoService: """视频生成服务""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) - self.timeout = get_config("grok.timeout", DEFAULT_TIMEOUT) + self.proxy = proxy or get_config("grok.base_proxy_url") + self.timeout = get_config("grok.timeout") def _build_headers( self, token: str, referer: str = "https://grok.com/imagine" ) -> dict: """构建请求头""" - user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT) + user_agent = get_config("grok.user_agent") headers = { "Accept": "*/*", "Accept-Encoding": "gzip, deflate, br, zstd", @@ -84,27 +84,14 @@ class VideoService: return {"http": self.proxy, "https": self.proxy} if self.proxy else None async def create_post( - self, - token: str, - prompt: str, + self, token: str, prompt: str, media_type: str = "MEDIA_POST_TYPE_VIDEO", - media_url: str = None, + media_url: str = None ) -> str: - """ - 创建媒体帖子 - - Args: - token: 认证 Token - prompt: 提示词（视频生成用） - media_type: 媒体类型 (MEDIA_POST_TYPE_VIDEO 或 MEDIA_POST_TYPE_IMAGE) - media_url: 媒体 URL（图片模式用） - - Returns: - post ID - """ + """创建媒体帖子，返回 post ID""" try: headers = self._build_headers(token) - + # 根据类型构建不同的载荷 if media_type == "MEDIA_POST_TYPE_IMAGE" and media_url: payload = {"mediaType": media_type, "mediaUrl": media_url} @@ -112,352 +99,160 @@ class VideoService: payload = {"mediaType": media_type, "prompt": prompt} async with AsyncSession() as session: - browser = get_config("grok.browser", DEFAULT_BROWSER) response = await session.post( CREATE_POST_API, headers=headers, json=payload, - impersonate=browser, + impersonate=get_config("grok.browser"), timeout=30, proxies=self._build_proxies(), ) if response.status_code != 200: logger.error(f"Create post failed: {response.status_code}") - raise UpstreamException( - f"Failed to create post: {response.status_code}" - ) - - data = response.json() - post_id = data.get("post", {}).get("id", "") + raise UpstreamException(f"Failed to create post: {response.status_code}") + post_id = response.json().get("post", {}).get("id", "") if not post_id: raise UpstreamException("No post ID in response") logger.info(f"Media post created: {post_id} (type={media_type})") return post_id + except AppException: + raise except Exception as e: logger.error(f"Create post error: {e}") - if isinstance(e, AppException): - raise e raise UpstreamException(f"Create post error: {str(e)}") async def create_image_post(self, token: str, image_url: str) -> str: - """ - 创建图片帖子 - - Args: - token: 认证 Token - image_url: 完整的图片 URL (https://assets.grok.com/...) - - Returns: - post ID - """ - return await self.create_post( - token, prompt="", media_type="MEDIA_POST_TYPE_IMAGE", media_url=image_url - ) + """创建图片帖子，返回 post ID""" + return await self.create_post(token, prompt="", media_type="MEDIA_POST_TYPE_IMAGE", media_url=image_url) def _build_payload( - self, - prompt: str, - post_id: str, + self, prompt: str, post_id: str, aspect_ratio: str = "3:2", video_length: int = 6, resolution_name: str = "480p", - preset: str = "normal", + preset: str = "normal" ) -> dict: """构建视频生成载荷""" - mode_flag = "--mode=custom" - if preset == "fun": - mode_flag = "--mode=extremely-crazy" - elif preset == "normal": - mode_flag = "--mode=normal" - elif preset == "spicy": - mode_flag = "--mode=extremely-spicy-or-crazy" - - full_prompt = f"{prompt} {mode_flag}" - video_gen_config = { - "parentPostId": post_id, - "aspectRatio": aspect_ratio, - "videoLength": video_length, - "resolutionName": resolution_name, - } + mode_map = {"fun": "--mode=extremely-crazy", "normal": "--mode=normal", "spicy": "--mode=extremely-spicy-or-crazy"} + mode_flag = mode_map.get(preset, "--mode=custom") + return { "temporary": True, "modelName": "grok-3", - "message": full_prompt, + "message": f"{prompt} {mode_flag}", "toolOverrides": {"videoGen": True}, "enableSideBySide": True, "responseMetadata": { "experiments": [], "modelConfigOverride": { "modelMap": { - "videoGenModelConfig": video_gen_config + "videoGenModelConfig": { + "parentPostId": post_id, + "aspectRatio": aspect_ratio, + "videoLength": video_length, + "resolutionName": resolution_name, + } } }, }, } + async def _generate_internal( + self, token: str, post_id: str, prompt: str, + aspect_ratio: str, video_length: int, + resolution_name: str, preset: str + ) -> AsyncGenerator[bytes, None]: + """内部生成逻辑""" + session = None + try: + headers = self._build_headers(token) + payload = self._build_payload(prompt, post_id, aspect_ratio, video_length, resolution_name, preset) + + session = AsyncSession(impersonate=get_config("grok.browser")) + response = await session.post( + CHAT_API, + headers=headers, + data=orjson.dumps(payload), + timeout=self.timeout, + stream=True, + proxies=self._build_proxies(), + ) + + if response.status_code != 200: + logger.error(f"Video generation failed: {response.status_code}") + raise UpstreamException( + message=f"Video generation failed: {response.status_code}", + details={"status": response.status_code}, + ) + + async def stream_response(): + try: + async for line in response.aiter_lines(): + yield line + finally: + await session.close() + + return stream_response() + + except Exception as e: + if session: + try: + await session.close() + except Exception: + pass + logger.error(f"Video generation error: {e}") + if isinstance(e, AppException): + raise + raise UpstreamException(f"Video generation error: {str(e)}") + async def generate( - self, - token: str, - prompt: str, + self, token: str, prompt: str, aspect_ratio: str = "3:2", video_length: int = 6, resolution_name: str = "480p", - preset: str = "normal", + preset: str = "normal" ) -> AsyncGenerator[bytes, None]: - """ - 生成视频 - - Args: - token: 认证 Token - prompt: 视频描述 - aspect_ratio: 宽高比 - video_length: 视频时长 - resolution_name: 分辨率 - preset: 预设 - - Returns: - AsyncGenerator，流式传输 - - Raises: - UpstreamException: 连接失败时 - """ - value = get_config("performance.media_max_concurrent", DEFAULT_MAX_CONCURRENT) - try: - value = int(value) - except Exception: - value = DEFAULT_MAX_CONCURRENT - value = max(1, value) - global _MEDIA_SEMAPHORE, _MEDIA_SEM_VALUE - if value != _MEDIA_SEM_VALUE: - _MEDIA_SEM_VALUE = value - _MEDIA_SEMAPHORE = asyncio.Semaphore(value) - async with _MEDIA_SEMAPHORE: - session = None - try: - # Step 1: 创建帖子 - post_id = await self.create_post(token, prompt) - - # Step 2: 建立连接 - headers = self._build_headers(token) - payload = self._build_payload( - prompt, post_id, aspect_ratio, video_length, resolution_name, preset - ) - - browser = get_config("grok.browser", DEFAULT_BROWSER) - session = AsyncSession(impersonate=browser) - response = await session.post( - CHAT_API, - headers=headers, - data=orjson.dumps(payload), - timeout=self.timeout, - stream=True, - proxies=self._build_proxies(), - ) - - if response.status_code != 200: - logger.error(f"Video generation failed: {response.status_code}") - try: - await session.close() - except Exception: - pass - raise UpstreamException( - message=f"Video generation failed: {response.status_code}", - details={"status": response.status_code}, - ) - - # Step 3: 流式传输 - async def stream_response(): - try: - async for line in response.aiter_lines(): - yield line - finally: - if session: - await session.close() - - return stream_response() - - except Exception as e: - if session: - try: - await session.close() - except Exception: - pass - logger.error(f"Video generation error: {e}") - if isinstance(e, AppException): - raise e - raise UpstreamException(f"Video generation error: {str(e)}") + """生成视频""" + async with _get_semaphore(): + post_id = await self.create_post(token, prompt) + return await self._generate_internal(token, post_id, prompt, aspect_ratio, video_length, resolution_name, preset) async def generate_from_image( - self, - token: str, - prompt: str, - image_url: str, + self, token: str, prompt: str, image_url: str, aspect_ratio: str = "3:2", video_length: int = 6, resolution: str = "480p", - preset: str = "normal", + preset: str = "normal" ) -> AsyncGenerator[bytes, None]: - """ - 从图片生成视频 - - Args: - token: 认证 Token - prompt: 视频描述 - image_url: 图片 URL - aspect_ratio: 宽高比 - video_length: 视频时长 - resolution: 分辨率 - preset: 预设 - - Returns: - AsyncGenerator，流式传输 - """ - value = get_config("performance.media_max_concurrent", DEFAULT_MAX_CONCURRENT) - try: - value = int(value) - except Exception: - value = DEFAULT_MAX_CONCURRENT - value = max(1, value) - global _MEDIA_SEMAPHORE, _MEDIA_SEM_VALUE - if value != _MEDIA_SEM_VALUE: - _MEDIA_SEM_VALUE = value - _MEDIA_SEMAPHORE = asyncio.Semaphore(value) - async with _MEDIA_SEMAPHORE: - session = None - try: - # Step 1: 创建帖子 - post_id = await self.create_image_post(token, image_url) - - # Step 2: 建立连接 - headers = self._build_headers(token) - payload = self._build_payload( - prompt, post_id, aspect_ratio, video_length, resolution, preset - ) - - browser = get_config("grok.browser", DEFAULT_BROWSER) - session = AsyncSession(impersonate=browser) - response = await session.post( - CHAT_API, - headers=headers, - data=orjson.dumps(payload), - timeout=self.timeout, - stream=True, - proxies=self._build_proxies(), - ) - - if response.status_code != 200: - logger.error(f"Video from image failed: {response.status_code}") - try: - await session.close() - except Exception: - pass - raise UpstreamException( - message=f"Video from image failed: {response.status_code}", - details={"status": response.status_code}, - ) - - # Step 3: 流式传输 - async def stream_response(): - try: - async for line in response.aiter_lines(): - yield line - finally: - if session: - await session.close() - - return stream_response() - - except Exception as e: - if session: - try: - await session.close() - except Exception: - pass - logger.error(f"Video from image error: {e}") - if isinstance(e, AppException): - raise e - raise UpstreamException(f"Video from image error: {str(e)}") - - @staticmethod - async def _wrap_stream(stream: AsyncGenerator, token_mgr, token: str, model: str): - """ - 包装流式响应，在完成时记录使用 - - Args: - stream: 原始 AsyncGenerator - token_mgr: TokenManager 实例 - token: Token 字符串 - model: 模型名称 - """ - success = False - try: - async for chunk in stream: - yield chunk - success = True - finally: - # 只在成功完成时记录使用，失败/异常时不扣费 - if success: - try: - model_info = ModelService.get(model) - effort = ( - EffortType.HIGH - if (model_info and model_info.cost.value == "high") - else EffortType.LOW - ) - await token_mgr.consume(token, effort) - logger.debug( - f"Video stream completed, recorded usage for token {token[:10]}... (effort={effort.value})" - ) - except Exception as e: - logger.warning(f"Failed to record video stream usage: {e}") + """从图片生成视频""" + async with _get_semaphore(): + post_id = await self.create_image_post(token, image_url) + return await self._generate_internal(token, post_id, prompt, aspect_ratio, video_length, resolution, preset) @staticmethod async def completions( - model: str, - messages: list, + model: str, messages: list, stream: bool = None, thinking: str = None, aspect_ratio: str = "3:2", video_length: int = 6, resolution: str = "480p", - preset: str = "normal", + preset: str = "normal" ): - """ - 视频生成入口 - - Args: - model: 模型名称 - messages: 消息列表 - stream: 是否流式 - thinking: 思考模式 - aspect_ratio: 宽高比 - video_length: 视频时长 - resolution: 分辨率 - preset: 预设模式 - - Returns: - AsyncGenerator (流式) 或 dict (非流式) - """ + """视频生成入口""" # 获取 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(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", - ) + token_mgr = await get_token_manager() + await token_mgr.reload_if_stale() + + token = None + for pool_name in ModelService.pool_candidates_for_model(model): + token = token_mgr.get_token(pool_name) + if token: + break if not token: raise AppException( @@ -467,14 +262,8 @@ class VideoService: status_code=429, ) - # 解析参数 - think = None - if thinking == "enabled": - think = True - elif thinking == "disabled": - think = False - - is_stream = stream if stream is not None else get_config("grok.stream", DEFAULT_STREAM) + think = {"enabled": True, "disabled": False}.get(thinking) + is_stream = stream if stream is not None else get_config("grok.stream") # 提取内容 from app.services.grok.services.chat import MessageExtractor @@ -492,27 +281,18 @@ class VideoService: try: for attach_type, attach_data in attachments: if attach_type == "image": - # 上传图片 _, file_uri = await upload_service.upload(attach_data, token) image_url = f"https://assets.grok.com/{file_uri}" logger.info(f"Image uploaded for video: {image_url}") - break # 视频模型只使用第一张图片 + break finally: await upload_service.close() # 生成视频 service = VideoService() - - # 图片转视频 if image_url: response = await service.generate_from_image( - token, - prompt, - image_url, - aspect_ratio, - video_length, - resolution, - preset, + token, prompt, image_url, aspect_ratio, video_length, resolution, preset ) else: response = await service.generate( @@ -522,26 +302,16 @@ class VideoService: # 处理响应 if is_stream: processor = VideoStreamProcessor(model, token, think) - return VideoService._wrap_stream( - processor.process(response), token_mgr, token, model - ) - else: - result = await VideoCollectProcessor(model, token).process(response) - # 非流式：处理完成后立即记录使用 - try: - model_info = ModelService.get(model) - effort = ( - EffortType.HIGH - if (model_info and model_info.cost.value == "high") - else EffortType.LOW - ) - await token_mgr.consume(token, effort) - logger.debug( - f"Video completed, recorded usage for token {token[:10]}... (effort={effort.value})" - ) - except Exception as e: - logger.warning(f"Failed to record video usage: {e}") - return result - + return wrap_stream_with_usage(processor.process(response), token_mgr, token, model) + + result = await VideoCollectProcessor(model, token).process(response) + try: + model_info = ModelService.get(model) + effort = EffortType.HIGH if (model_info and model_info.cost.value == "high") else EffortType.LOW + await token_mgr.consume(token, effort) + logger.debug(f"Video completed, recorded usage (effort={effort.value})") + except Exception as e: + logger.warning(f"Failed to record video usage: {e}") + return result __all__ = ["VideoService"] diff --git a/app/services/grok/services/nsfw.py b/app/services/grok/services/nsfw.py index 88411be..169ed4d 100644 --- a/app/services/grok/services/nsfw.py +++ b/app/services/grok/services/nsfw.py @@ -20,16 +20,9 @@ from app.services.grok.protocols.grpc_web import ( ) from app.services.grok.utils.headers import build_sso_cookie - NSFW_API = "https://grok.com/auth_mgmt.AuthManagement/UpdateUserFeatureControls" BIRTH_DATE_API = "https://grok.com/rest/auth/set-birth-date" -DEFAULT_BROWSER = "chrome136" -DEFAULT_USER_AGENT = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36" -DEFAULT_TIMEOUT = 30 -DEFAULT_BASE_PROXY_URL = "" - - @dataclass class NSFWResult: """NSFW 操作结果""" @@ -40,20 +33,22 @@ class NSFWResult: grpc_message: Optional[str] = None error: Optional[str] = None - class NSFWService: """NSFW 模式服务""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) - self.timeout = float(get_config("grok.timeout", DEFAULT_TIMEOUT)) + self.proxy = proxy or get_config("grok.base_proxy_url") + self.timeout = float(get_config("grok.timeout")) + + def _build_proxies(self) -> Optional[dict]: + """构建代理配置""" + return {"http": self.proxy, "https": self.proxy} if self.proxy else None @staticmethod def _random_birth_date() -> str: - """生成随机出生日期""" + """生成随机出生日期（20-40岁）""" today = datetime.date.today() - age = random.randint(20, 40) - birth_year = today.year - age + birth_year = today.year - random.randint(20, 40) birth_month = random.randint(1, 12) birth_day = random.randint(1, 28) hour = random.randint(0, 23) @@ -65,7 +60,7 @@ class NSFWService: def _build_headers(self, token: str) -> dict: """构造 gRPC-Web 请求头""" cookie = build_sso_cookie(token, include_rw=True) - user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT) + user_agent = get_config("grok.user_agent") return { "accept": "*/*", "content-type": "application/grpc-web+proto", @@ -80,7 +75,7 @@ class NSFWService: def _build_birth_headers(self, token: str) -> dict: """构造设置出生日期请求头""" cookie = build_sso_cookie(token, include_rw=True) - user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT) + user_agent = get_config("grok.user_agent") return { "accept": "*/*", "content-type": "application/json", @@ -102,18 +97,20 @@ class NSFWService: protobuf = b"\x0a\x02\x10\x01\x12" + bytes([len(inner)]) + inner return encode_grpc_web_payload(protobuf) - async def _set_birth_date(self, session: AsyncSession, token: str) -> tuple[bool, int, Optional[str]]: + async def _set_birth_date( + self, session: AsyncSession, token: str + ) -> tuple[bool, int, Optional[str]]: """设置出生日期""" headers = self._build_birth_headers(token) payload = {"birthDate": self._random_birth_date()} - proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None + try: response = await session.post( BIRTH_DATE_API, json=payload, headers=headers, timeout=self.timeout, - proxies=proxies, + proxies=self._build_proxies(), ) if response.status_code in (200, 204): return True, response.status_code, None @@ -125,31 +122,27 @@ class NSFWService: """为单个 token 开启 NSFW 模式""" headers = self._build_headers(token) payload = self._build_payload() - logger.debug( - "NSFW payload: len={} hex={}", - len(payload), - payload.hex(), - ) - proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None + logger.debug(f"NSFW payload: len={len(payload)} hex={payload.hex()}") try: - browser = get_config("grok.browser", DEFAULT_BROWSER) + browser = get_config("grok.browser") async with AsyncSession(impersonate=browser) as session: - ok, birth_status, birth_err = await self._set_birth_date( - session, token - ) + # 先设置出生日期 + ok, birth_status, birth_err = await self._set_birth_date(session, token) if not ok: return NSFWResult( success=False, http_status=birth_status, error=f"Set birth date failed: {birth_err}", ) + + # 开启 NSFW response = await session.post( NSFW_API, data=payload, headers=headers, timeout=self.timeout, - proxies=proxies, + proxies=self._build_proxies(), ) if response.status_code != 200: @@ -160,18 +153,15 @@ class NSFWService: ) # 解析 gRPC-Web 响应 - content_type = response.headers.get("content-type") _, trailers = parse_grpc_web_response( - response.content, content_type=content_type + response.content, + content_type=response.headers.get("content-type") ) grpc_status = get_grpc_status(trailers) logger.debug( - "NSFW response: http={} grpc={} msg={} trailers={}", - response.status_code, - grpc_status.code, - grpc_status.message, - trailers, + f"NSFW response: http={response.status_code} grpc={grpc_status.code} " + f"msg={grpc_status.message} trailers={trailers}" ) # HTTP 200 且无 grpc-status（空响应）或 grpc-status=0 都算成功 @@ -188,5 +178,4 @@ class NSFWService: logger.error(f"NSFW enable failed: {e}") return NSFWResult(success=False, http_status=0, error=str(e)[:100]) - __all__ = ["NSFWService", "NSFWResult"] diff --git a/app/services/grok/services/usage.py b/app/services/grok/services/usage.py index dc2ff52..8c386f4 100644 --- a/app/services/grok/services/usage.py +++ b/app/services/grok/services/usage.py @@ -13,32 +13,22 @@ from app.core.exceptions import UpstreamException from app.services.grok.utils.headers import apply_statsig, build_sso_cookie from app.services.grok.utils.retry import retry_on_status - LIMITS_API = "https://grok.com/rest/rate-limits" -DEFAULT_BROWSER = "chrome136" -DEFAULT_USER_AGENT = ( - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " - "AppleWebKit/537.36 (KHTML, like Gecko) " - "Chrome/136.0.0.0 Safari/537.36" -) -DEFAULT_TIMEOUT = 10 -DEFAULT_MAX_CONCURRENT = 25 -DEFAULT_BASE_PROXY_URL = "" -_USAGE_SEMAPHORE = asyncio.Semaphore(DEFAULT_MAX_CONCURRENT) -_USAGE_SEM_VALUE = DEFAULT_MAX_CONCURRENT +_USAGE_SEMAPHORE = asyncio.Semaphore(25) +_USAGE_SEM_VALUE = 25 class UsageService: """用量查询服务""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) - self.timeout = get_config("grok.timeout", DEFAULT_TIMEOUT) + self.proxy = proxy or get_config("grok.base_proxy_url") + self.timeout = get_config("grok.timeout") def _build_headers(self, token: str) -> dict: """构建请求头""" - user_agent = get_config("grok.user_agent", DEFAULT_USER_AGENT) + user_agent = get_config("grok.user_agent") headers = { "Accept": "*/*", "Accept-Encoding": "gzip, deflate, br, zstd", @@ -85,11 +75,11 @@ class UsageService: Raises: UpstreamException: 当获取失败且重试耗尽时 """ - value = get_config("performance.usage_max_concurrent", DEFAULT_MAX_CONCURRENT) + value = get_config("performance.usage_max_concurrent") try: value = int(value) except Exception: - value = DEFAULT_MAX_CONCURRENT + value = 25 value = max(1, value) global _USAGE_SEMAPHORE, _USAGE_SEM_VALUE if value != _USAGE_SEM_VALUE: @@ -107,7 +97,7 @@ class UsageService: try: headers = self._build_headers(token) payload = {"requestKind": "DEFAULT", "modelName": model_name} - browser = get_config("grok.browser", DEFAULT_BROWSER) + browser = get_config("grok.browser") async with AsyncSession() as session: response = await session.post( @@ -152,5 +142,4 @@ class UsageService: # 最后一次失败已经被记录 raise - __all__ = ["UsageService"] diff --git a/app/services/grok/services/voice.py b/app/services/grok/services/voice.py index 17a4de7..b887d31 100644 --- a/app/services/grok/services/voice.py +++ b/app/services/grok/services/voice.py @@ -12,18 +12,13 @@ from app.core.config import get_config from app.core.exceptions import UpstreamException from app.services.grok.utils.headers import apply_statsig, build_sso_cookie - LIVEKIT_TOKEN_API = "https://grok.com/rest/livekit/tokens" -DEFAULT_TIMEOUT = 30 -DEFAULT_BROWSER = "chrome136" -DEFAULT_BASE_PROXY_URL = "" - class VoiceService: """Voice Mode Service (LiveKit)""" def __init__(self, proxy: str = None): - self.proxy = proxy or get_config("grok.base_proxy_url", DEFAULT_BASE_PROXY_URL) + self.proxy = proxy or get_config("grok.base_proxy_url") async def get_token( self, @@ -46,8 +41,8 @@ class VoiceService: proxies = {"http": self.proxy, "https": self.proxy} if self.proxy else None try: - browser = get_config("grok.browser", DEFAULT_BROWSER) - timeout = get_config("grok.timeout", DEFAULT_TIMEOUT) + browser = get_config("grok.browser") + timeout = get_config("grok.timeout") async with AsyncSession(impersonate=browser) as session: response = await session.post( LIVEKIT_TOKEN_API, diff --git a/app/services/grok/utils/headers.py b/app/services/grok/utils/headers.py index 44c94ae..146e704 100644 --- a/app/services/grok/utils/headers.py +++ b/app/services/grok/utils/headers.py @@ -11,16 +11,13 @@ from app.core.config import get_config from app.services.grok.utils.statsig import StatsigService -DEFAULT_CF_CLEARANCE = "" - - def _normalize_token(token: str) -> str: return token[4:] if token.startswith("sso=") else token def build_sso_cookie(token: str, include_rw: bool = False) -> str: token = _normalize_token(token) - cf = get_config("grok.cf_clearance", DEFAULT_CF_CLEARANCE) + cf = get_config("grok.cf_clearance") cookie = f"sso={token}" if include_rw: cookie = f"{cookie}; sso-rw={token}" diff --git a/app/services/grok/utils/retry.py b/app/services/grok/utils/retry.py index 9a69e87..2cd52c5 100644 --- a/app/services/grok/utils/retry.py +++ b/app/services/grok/utils/retry.py @@ -18,38 +18,22 @@ from app.core.config import get_config from app.core.exceptions import UpstreamException -DEFAULT_MAX_RETRY = 3 -DEFAULT_RETRY_STATUS_CODES = [401, 429, 403] -DEFAULT_RETRY_BACKOFF_BASE = 0.5 -DEFAULT_RETRY_BACKOFF_FACTOR = 2.0 -DEFAULT_RETRY_BACKOFF_MAX = 30.0 -DEFAULT_RETRY_BUDGET = 90.0 - - class RetryContext: """重试上下文""" def __init__(self): self.attempt = 0 - self.max_retry = int(get_config("grok.max_retry", DEFAULT_MAX_RETRY)) - self.retry_codes = get_config( - "grok.retry_status_codes", DEFAULT_RETRY_STATUS_CODES - ) + self.max_retry = int(get_config("grok.max_retry")) + self.retry_codes = get_config("grok.retry_status_codes") self.last_error = None self.last_status = None self.total_delay = 0.0 - self.retry_budget = float(get_config("grok.retry_budget", DEFAULT_RETRY_BUDGET)) + self.retry_budget = float(get_config("grok.retry_budget")) # 退避参数 - self.backoff_base = float( - get_config("grok.retry_backoff_base", DEFAULT_RETRY_BACKOFF_BASE) - ) - self.backoff_factor = float( - get_config("grok.retry_backoff_factor", DEFAULT_RETRY_BACKOFF_FACTOR) - ) - self.backoff_max = float( - get_config("grok.retry_backoff_max", DEFAULT_RETRY_BACKOFF_MAX) - ) + self.backoff_base = float(get_config("grok.retry_backoff_base")) + self.backoff_factor = float(get_config("grok.retry_backoff_factor")) + self.backoff_max = float(get_config("grok.retry_backoff_max")) # decorrelated jitter 状态 self._last_delay = self.backoff_base diff --git a/app/services/grok/utils/statsig.py b/app/services/grok/utils/statsig.py index d5813b6..bf78741 100644 --- a/app/services/grok/utils/statsig.py +++ b/app/services/grok/utils/statsig.py @@ -9,9 +9,6 @@ import string from app.core.config import get_config -DEFAULT_DYNAMIC_STATSIG = True - - class StatsigService: """Statsig ID 生成服务""" @@ -33,8 +30,7 @@ class StatsigService: Returns: Base64 编码的 ID """ - # 读取配置 - dynamic = get_config("grok.dynamic_statsig", DEFAULT_DYNAMIC_STATSIG) + dynamic = get_config("grok.dynamic_statsig") if not dynamic: return "ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk=" diff --git a/app/services/grok/utils/stream.py b/app/services/grok/utils/stream.py new file mode 100644 index 0000000..fb107b7 --- /dev/null +++ b/app/services/grok/utils/stream.py @@ -0,0 +1,43 @@ +""" +流式响应通用工具 +""" + +from typing import AsyncGenerator + +from app.core.logger import logger +from app.services.grok.models.model import ModelService +from app.services.token import EffortType + + +async def wrap_stream_with_usage( + stream: AsyncGenerator, + token_mgr, + token: str, + model: str +) -> AsyncGenerator: + """ + 包装流式响应，在完成时记录使用 + + Args: + stream: 原始 AsyncGenerator + token_mgr: TokenManager 实例 + token: Token 字符串 + model: 模型名称 + """ + success = False + try: + async for chunk in stream: + yield chunk + success = True + finally: + if success: + try: + model_info = ModelService.get(model) + effort = EffortType.HIGH if (model_info and model_info.cost.value == "high") else EffortType.LOW + await token_mgr.consume(token, effort) + logger.debug(f"Stream completed, recorded usage for token {token[:10]}... (effort={effort.value})") + except Exception as e: + logger.warning(f"Failed to record stream usage: {e}") + + +__all__ = ["wrap_stream_with_usage"] diff --git a/data/config.toml b/data/config.toml index 3db391c..22296cf 100644 --- a/data/config.toml +++ b/data/config.toml @@ -1,47 +1,36 @@ [grok] -temporary = true -stream = true -thinking = true -dynamic_statsig = true -filter_tags = ["xaiartifact","xai:tool_usage_card","grok:render"] +browser = "chrome136" timeout = 120 +user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36" base_proxy_url = "" asset_proxy_url = "" cf_clearance = "" -browser = "chrome136" -user_agent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36" +temporary = true disable_memory = true +stream = true +thinking = true +filter_tags = ["xaiartifact","xai:tool_usage_card","grok:render"] max_retry = 3 retry_status_codes = [401,429,403] -image_ws = true -image_ws_blocked_seconds = 15 -image_ws_final_min_bytes = 100000 -image_ws_nsfw = true retry_backoff_base = 0.5 retry_backoff_factor = 2.0 retry_backoff_max = 30.0 retry_budget = 90.0 stream_idle_timeout = 120.0 video_idle_timeout = 90.0 +image_ws_blocked_seconds = 15 +image_ws_final_min_bytes = 100000 +image_ws_medium_min_bytes = 30000 +dynamic_statsig = true +image_ws = true +image_ws_nsfw = true [app] app_url = "http://127.0.0.1:8000" -app_key = "grok2api" -api_key = "" image_format = "url" video_format = "html" - -[token] -auto_refresh = true -refresh_interval_hours = 8 -super_refresh_interval_hours = 2 -fail_threshold = 5 -save_delay_ms = 200 -reload_interval_sec = 30 - -[cache] -enable_auto_clean = true -limit_mb = 1024 +app_key = "grok2api" +api_key = "" [performance] assets_max_concurrent = 25 @@ -55,3 +44,15 @@ usage_max_tokens = 1000 nsfw_max_concurrent = 10 nsfw_batch_size = 50 nsfw_max_tokens = 1000 + +[cache] +enable_auto_clean = true +limit_mb = 1024 + +[token] +auto_refresh = true +refresh_interval_hours = 8 +super_refresh_interval_hours = 2 +fail_threshold = 5 +save_delay_ms = 200 +reload_interval_sec = 30 diff --git a/main.py b/main.py index bd81ff2..97da6ac 100644 --- a/main.py +++ b/main.py @@ -41,17 +41,21 @@ setup_logging( @asynccontextmanager async def lifespan(app: FastAPI): """应用生命周期管理""" - # 1. 加载配置 - from app.core.config import config - + # 1. 注册服务默认配置 + from app.core.config import config, register_defaults + from app.services.grok.defaults import get_grok_defaults + + register_defaults(get_grok_defaults()) + + # 2. 加载配置 await config.load() - # 2. 启动服务显示 + # 3. 启动服务显示 logger.info("Starting Grok2API...") logger.info(f"Platform: {platform.system()} {platform.release()}") logger.info(f"Python: {sys.version.split()[0]}") - # 3. 启动 Token 刷新调度器 + # 4. 启动 Token 刷新调度器 refresh_enabled = get_config("token.auto_refresh", True) if refresh_enabled: basic_interval = get_config("token.refresh_interval_hours", 8) diff --git a/tests/test_model.py b/tests/test_model.py index e58f263..b4e96e7 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -18,7 +18,7 @@ load_dotenv(ROOT / ".env") from app.core.config import config from app.core.logger import setup_logging from app.services.grok.services.chat import GrokChatService -from app.services.grok.processors.processor import CollectProcessor +from app.services.grok.processors import CollectProcessor from app.services.grok.services.usage import UsageService