diff --git a/app/services/grok/cache.py b/app/services/grok/cache.py index 5dcbfb1..d9d9220 100644 --- a/app/services/grok/cache.py +++ b/app/services/grok/cache.py @@ -1,6 +1,7 @@ """缓存服务模块""" import asyncio +import base64 from pathlib import Path from typing import Optional from curl_cffi.requests import AsyncSession @@ -19,42 +20,20 @@ class CacheService: self.cache_dir = Path(f"data/temp/{cache_type}") self.cache_dir.mkdir(parents=True, exist_ok=True) - @staticmethod - def _get_cache_filename(file_path: str) -> str: - """将文件路径转换为缓存文件名""" - filename = file_path.lstrip('/').replace('/', '-') - return filename - - def _get_cache_path(self, file_path: str) -> Path: + def _cache_path(self, file_path: str) -> Path: """获取缓存文件的完整路径""" - filename = self._get_cache_filename(file_path) + filename = file_path.lstrip('/').replace('/', '-') return self.cache_dir / filename async def download_file(self, file_path: str, auth_token: str, timeout: float = 30.0) -> Optional[Path]: - """下载并缓存文件 - - Args: - file_path: 文件路径，如 /users/xxx/generated/xxx/file.jpg - auth_token: 认证令牌 - timeout: 下载超时时间（秒） - - Returns: - 缓存文件路径，下载失败返回 None - """ - cache_path = self._get_cache_path(file_path) - + """下载并缓存文件""" + cache_path = self._cache_path(file_path) if cache_path.exists(): logger.debug(f"[{self.cache_type.upper()}Cache] 文件已缓存: {cache_path}") return cache_path - file_url = f"https://assets.grok.com{file_path}" - try: - # 构建 Cookie cf_clearance = setting.grok_config.get("cf_clearance", "") - cookie = f"{auth_token};{cf_clearance}" if cf_clearance else auth_token - - # 构建请求头 headers = { **get_dynamic_headers(pathname=file_path), "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7", @@ -64,17 +43,16 @@ class CacheService: "Sec-Fetch-User": "?1", "Upgrade-Insecure-Requests": "1", "Referer": "https://grok.com/", - "Cookie": cookie + "Cookie": f"{auth_token};{cf_clearance}" if cf_clearance else auth_token } - # 代理配置 proxy_url = setting.grok_config.get("proxy_url") proxies = {"http": proxy_url, "https": proxy_url} if proxy_url else {} async with AsyncSession() as session: - logger.debug(f"[{self.cache_type.upper()}Cache] 开始下载: {file_url}") + logger.debug(f"[{self.cache_type.upper()}Cache] 开始下载: https://assets.grok.com{file_path}") response = await session.get( - file_url, + f"https://assets.grok.com{file_path}", headers=headers, proxies=proxies, timeout=timeout, @@ -82,70 +60,44 @@ class CacheService: impersonate="chrome133a" ) response.raise_for_status() - cache_path.write_bytes(response.content) logger.debug(f"[{self.cache_type.upper()}Cache] 文件已缓存: {cache_path} ({len(response.content)} bytes)") - asyncio.create_task(self.cleanup_cache()) - return cache_path - except Exception as e: logger.error(f"[{self.cache_type.upper()}Cache] 下载文件失败: {e}") return None - def get_cached_file(self, file_path: str) -> Optional[Path]: - """获取缓存的文件路径 - - Args: - file_path: 文件路径 - - Returns: - 缓存文件路径，不存在返回 None - """ - cache_path = self._get_cache_path(file_path) + def get_cached(self, file_path: str) -> Optional[Path]: + """获取缓存的文件路径""" + cache_path = self._cache_path(file_path) return cache_path if cache_path.exists() else None async def cleanup_cache(self): """清理缓存目录，确保不超过配置的大小限制""" try: - # 获取配置的最大缓存大小 - config_key = f"{self.cache_type}_cache_max_size_mb" - max_size_mb = setting.global_config.get(config_key, 500) + max_size_mb = setting.global_config.get(f"{self.cache_type}_cache_max_size_mb", 500) max_size_bytes = max_size_mb * 1024 * 1024 - # 获取缓存大小和修改时间 - files = [] - total_size = 0 + files = [(fp, (stat := fp.stat()).st_size, stat.st_mtime) + for fp in self.cache_dir.glob("*") if fp.is_file()] + total_size = sum(size for _, size, _ in files) - for file_path in self.cache_dir.glob("*"): - if file_path.is_file(): - size = file_path.stat().st_size - mtime = file_path.stat().st_mtime - files.append((file_path, size, mtime)) - total_size += size - - # 如果总大小未超限，无需清理 if total_size <= max_size_bytes: logger.debug(f"[{self.cache_type.upper()}Cache] 缓存大小 {total_size / 1024 / 1024:.2f}MB，未超限") return logger.info(f"[{self.cache_type.upper()}Cache] 缓存大小 {total_size / 1024 / 1024:.2f}MB 超过限制 {max_size_mb}MB，开始清理") - - # 按修改时间排序 files.sort(key=lambda x: x[2]) - # 删除文件直到总大小低于限制 for file_path, size, _ in files: if total_size <= max_size_bytes: break - file_path.unlink() total_size -= size logger.debug(f"[{self.cache_type.upper()}Cache] 已删除缓存文件: {file_path}") logger.info(f"[{self.cache_type.upper()}Cache] 缓存清理完成，当前大小 {total_size / 1024 / 1024:.2f}MB") - except Exception as e: logger.error(f"[{self.cache_type.upper()}Cache] 清理缓存失败: {e}") @@ -157,27 +109,51 @@ class ImageCacheService(CacheService): super().__init__("image") async def download_image(self, image_path: str, auth_token: str) -> Optional[Path]: - """下载并缓存图片 - - Args: - image_path: 图片路径，如 /users/xxx/generated/xxx/image.jpg - auth_token: 认证令牌 - - Returns: - 缓存文件路径，下载失败返回 None` - """ + """下载并缓存图片""" return await self.download_file(image_path, auth_token, timeout=30.0) - def get_cached_image(self, image_path: str) -> Optional[Path]: - """获取缓存的图片路径 + def get_cached(self, image_path: str) -> Optional[Path]: + """获取缓存的图片路径""" + return super().get_cached(image_path) - Args: - image_path: 图片路径 + @staticmethod + def to_base64(image_path: Path) -> Optional[str]: + """将图片转换为 base64 编码""" + try: + if not image_path.exists(): + logger.error(f"[ImageCache] 图片文件不存在: {image_path}") + return None - Returns: - 缓存文件路径，不存在返回 None - """ - return self.get_cached_file(image_path) + with open(image_path, "rb") as f: + base64_data = base64.b64encode(f.read()).decode('utf-8') + + mime_type = {'.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.png': 'image/png', + '.gif': 'image/gif', '.webp': 'image/webp'}.get(image_path.suffix.lower(), 'image/jpeg') + + return f"data:{mime_type};base64,{base64_data}" + except Exception as e: + logger.error(f"[ImageCache] 图片转 base64 失败: {e}") + return None + + async def download_base64(self, image_path: str, auth_token: str) -> Optional[str]: + """下载图片并转换为 base64 编码（转换后立即删除缓存文件）""" + try: + cache_path = await self.download_file(image_path, auth_token, timeout=30.0) + if not cache_path: + return None + + base64_str = self.to_base64(cache_path) + + try: + cache_path.unlink() + logger.debug(f"[ImageCache] 已删除临时文件: {cache_path}") + except Exception as e: + logger.warning(f"[ImageCache] 删除临时文件失败: {e}") + + return base64_str + except Exception as e: + logger.error(f"[ImageCache] 下载并转换 base64 失败: {e}") + return None class VideoCacheService(CacheService): @@ -187,30 +163,14 @@ class VideoCacheService(CacheService): super().__init__("video") async def download_video(self, video_path: str, auth_token: str) -> Optional[Path]: - """下载并缓存视频 - - Args: - video_path: 视频路径，如 /users/xxx/generated/xxx/video.mp4 - auth_token: 认证令牌 - - Returns: - 缓存文件路径，下载失败返回 None - """ + """下载并缓存视频""" return await self.download_file(video_path, auth_token, timeout=60.0) - def get_cached_video(self, video_path: str) -> Optional[Path]: - """获取缓存的视频路径 - - Args: - video_path: 视频路径 - - Returns: - 缓存文件路径，不存在返回 None - """ - return self.get_cached_file(video_path) + def get_cached(self, video_path: str) -> Optional[Path]: + """获取缓存的视频路径""" + return super().get_cached(video_path) # 全局实例 image_cache_service = ImageCacheService() video_cache_service = VideoCacheService() - diff --git a/app/services/grok/processer.py b/app/services/grok/processer.py index 171cc14..a5d7b8c 100644 --- a/app/services/grok/processer.py +++ b/app/services/grok/processer.py @@ -157,18 +157,30 @@ class GrokResponseProcessor: # 提取图片数据 if images := model_response.get("generatedImageUrls"): + # 获取图片返回模式 + image_mode = setting.global_config.get("image_mode", "url") + for img in images: try: - cache_path = await image_cache_service.download_image(f"/{img}", auth_token) - if cache_path: - img_path = img.replace('/', '-') - base_url = setting.global_config.get("base_url", "") - img_url = f"{base_url}/images/{img_path}" if base_url else f"/images/{img_path}" - content += f"\n![Generated Image]({img_url})" + if image_mode == "base64": + # base64 模式：下载并转换为 base64 + base64_str = await image_cache_service.download_base64(f"/{img}", auth_token) + if base64_str: + content += f"\n![Generated Image]({base64_str})" + else: + content += f"\n![Generated Image](https://assets.grok.com/{img})" else: - content += f"\n![Generated Image](https://assets.grok.com/{img})" + # url 模式：缓存并返回链接 + cache_path = await image_cache_service.download_image(f"/{img}", auth_token) + if cache_path: + img_path = img.replace('/', '-') + base_url = setting.global_config.get("base_url", "") + img_url = f"{base_url}/images/{img_path}" if base_url else f"/images/{img_path}" + content += f"\n![Generated Image]({img_url})" + else: + content += f"\n![Generated Image](https://assets.grok.com/{img})" except Exception as e: - logger.warning(f"[Processor] 缓存图片失败: {e}") + logger.warning(f"[Processor] 处理图片失败: {e}") content += f"\n![Generated Image](https://assets.grok.com/{img})" # 返回 OpenAI 响应格式 @@ -333,19 +345,66 @@ class GrokResponseProcessor: # 提取图片数据 if is_image: if model_resp := grok_resp.get("modelResponse"): + # 获取图片返回模式 + image_mode = setting.global_config.get("image_mode", "url") + # 生成图片链接并缓存 - content = "" for img in model_resp.get("generatedImageUrls", []): try: - # 缓存图片 - await image_cache_service.download_image(f"/{img}", auth_token) - # 本地图片路径 - img_path = img.replace('/', '-') - base_url = setting.global_config.get("base_url", "") - img_url = f"{base_url}/images/{img_path}" if base_url else f"/images/{img_path}" - content += f"![Generated Image]({img_url})\n" + if image_mode == "base64": + # base64 模式：下载并转换为 base64 + base64_str = await image_cache_service.download_base64(f"/{img}", auth_token) + if base64_str: + # 分块发送 base64 数据，每 8KB 一个 chunk + markdown_prefix = "![Generated Image](data:" + markdown_suffix = ")\n" + + # 提取 data URL 的 mime 和 base64 部分 + if base64_str.startswith("data:"): + parts = base64_str.split(",", 1) + if len(parts) == 2: + mime_part = parts[0] + "," + b64_data = parts[1] + + # 发送前缀 + yield make_chunk(markdown_prefix + mime_part) + timeout_manager.mark_chunk_received() + chunk_index += 1 + + # 分块发送 base64 数据 + chunk_size = 8192 + for i in range(0, len(b64_data), chunk_size): + chunk_data = b64_data[i:i + chunk_size] + yield make_chunk(chunk_data) + timeout_manager.mark_chunk_received() + chunk_index += 1 + + # 发送后缀 + yield make_chunk(markdown_suffix) + timeout_manager.mark_chunk_received() + chunk_index += 1 + else: + yield make_chunk(f"![Generated Image]({base64_str})\n") + timeout_manager.mark_chunk_received() + chunk_index += 1 + else: + yield make_chunk(f"![Generated Image]({base64_str})\n") + timeout_manager.mark_chunk_received() + chunk_index += 1 + else: + yield make_chunk(f"![Generated Image](https://assets.grok.com/{img})\n") + timeout_manager.mark_chunk_received() + chunk_index += 1 + else: + # url 模式：缓存并返回链接 + await image_cache_service.download_image(f"/{img}", auth_token) + # 本地图片路径 + img_path = img.replace('/', '-') + base_url = setting.global_config.get("base_url", "") + img_url = f"{base_url}/images/{img_path}" if base_url else f"/images/{img_path}" + content += f"![Generated Image]({img_url})\n" except Exception as e: - logger.warning(f"[Processor] 缓存图片失败: {e}") + logger.warning(f"[Processor] 处理图片失败: {e}") content += f"![Generated Image](https://assets.grok.com/{img})\n" yield make_chunk(content.strip(), "stop") timeout_manager.mark_chunk_received() diff --git a/app/template/admin.html b/app/template/admin.html index 9bdd8d0..0b9e0c3 100644 --- a/app/template/admin.html +++ b/app/template/admin.html @@ -4,103 +4,49 @@ 管理控制台 - Grok2API +
-
- Grok2API +
+ Grok2API + by @Chenyme
- @@ -274,6 +220,18 @@