"""缓存服务模块 - 提供图片和视频的下载、缓存和清理功能""" import asyncio import base64 from pathlib import Path from typing import Optional, Tuple from curl_cffi.requests import AsyncSession from app.core.config import setting from app.core.logger import logger from app.services.grok.statsig import get_dynamic_headers # 常量定义 MIME_TYPES = { '.jpg': 'image/jpeg', '.jpeg': 'image/jpeg', '.png': 'image/png', '.gif': 'image/gif', '.webp': 'image/webp', '.bmp': 'image/bmp', } DEFAULT_MIME = 'image/jpeg' ASSETS_URL = "https://assets.grok.com" class CacheService: """缓存服务基类""" def __init__(self, cache_type: str, timeout: float = 30.0): self.cache_type = cache_type self.cache_dir = Path(f"data/temp/{cache_type}") self.cache_dir.mkdir(parents=True, exist_ok=True) self.timeout = timeout self._cleanup_lock = asyncio.Lock() def _get_path(self, file_path: str) -> Path: """转换文件路径为缓存路径""" return self.cache_dir / file_path.lstrip('/').replace('/', '-') def _log(self, level: str, msg: str): """统一日志输出""" getattr(logger, level)(f"[{self.cache_type.upper()}Cache] {msg}") def _build_headers(self, file_path: str, auth_token: str) -> dict: """构建请求头""" cf = setting.grok_config.get("cf_clearance", "") return { **get_dynamic_headers(pathname=file_path), "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/", "Cookie": f"{auth_token};{cf}" if cf else auth_token } async def download(self, file_path: str, auth_token: str, timeout: Optional[float] = None) -> Optional[Path]: """下载并缓存文件""" cache_path = self._get_path(file_path) if cache_path.exists(): self._log("debug", "文件已缓存") return cache_path # 外层重试：可配置状态码（401/429等） retry_codes = setting.grok_config.get("retry_status_codes", [401, 429]) MAX_OUTER_RETRY = 3 for outer_retry in range(MAX_OUTER_RETRY + 1): # +1 确保实际重试3次 try: # 内层重试：403代理池重试（cache使用缓存代理，不支持代理池，所以403只重试一次） max_403_retries = 5 retry_403_count = 0 while retry_403_count <= max_403_retries: proxy = await setting.get_proxy_async("cache") proxies = {"http": proxy, "https": proxy} if proxy else {} if proxy and outer_retry == 0 and retry_403_count == 0: self._log("debug", f"使用代理: {proxy.split('@')[-1] if '@' in proxy else proxy}") async with AsyncSession() as session: url = f"{ASSETS_URL}{file_path}" if outer_retry == 0 and retry_403_count == 0: self._log("debug", f"下载: {url}") response = await session.get( url, headers=self._build_headers(file_path, auth_token), proxies=proxies, timeout=timeout or self.timeout, allow_redirects=True, impersonate="chrome133a" ) # 检查403错误 - 内层重试(cache不使用代理池，所以直接失败) if response.status_code == 403: retry_403_count += 1 if retry_403_count <= max_403_retries: self._log("warning", f"遇到403错误，正在重试 ({retry_403_count}/{max_403_retries})...") await asyncio.sleep(0.5) continue self._log("error", f"403错误，已重试{retry_403_count-1}次，放弃") return None # 检查可配置状态码错误 - 外层重试 if response.status_code in retry_codes: if outer_retry < MAX_OUTER_RETRY: delay = (outer_retry + 1) * 0.1 # 渐进延迟：0.1s, 0.2s, 0.3s self._log("warning", f"遇到{response.status_code}错误，外层重试 ({outer_retry+1}/{MAX_OUTER_RETRY})，等待{delay}s...") await asyncio.sleep(delay) break # 跳出内层循环，进入外层重试 else: self._log("error", f"{response.status_code}错误，已重试{outer_retry}次，放弃") return None response.raise_for_status() await asyncio.to_thread(cache_path.write_bytes, response.content) if outer_retry > 0 or retry_403_count > 0: self._log("info", f"重试成功！") else: self._log("debug", "缓存成功") # 异步清理（带错误处理） asyncio.create_task(self._safe_cleanup()) return cache_path except Exception as e: if outer_retry < MAX_OUTER_RETRY - 1: self._log("warning", f"下载异常: {e}，外层重试 ({outer_retry+1}/{MAX_OUTER_RETRY})...") await asyncio.sleep(0.5) continue self._log("error", f"下载失败: {e}（已重试{outer_retry}次）") return None return None def get_cached(self, file_path: str) -> Optional[Path]: """获取已缓存的文件""" path = self._get_path(file_path) return path if path.exists() else None async def _safe_cleanup(self): """安全清理（捕获异常）""" try: await self.cleanup() except Exception as e: self._log("error", f"后台清理失败: {e}") async def cleanup(self): """清理超限缓存""" if self._cleanup_lock.locked(): return async with self._cleanup_lock: try: max_mb = setting.global_config.get(f"{self.cache_type}_cache_max_size_mb", 500) max_bytes = max_mb * 1024 * 1024 # 获取文件信息 (path, size, mtime) files = [(f, (s := f.stat()).st_size, s.st_mtime) for f in self.cache_dir.glob("*") if f.is_file()] total = sum(size for _, size, _ in files) if total <= max_bytes: return self._log("info", f"清理缓存 {total/1024/1024:.1f}MB -> {max_mb}MB") # 删除最旧的文件 for path, size, _ in sorted(files, key=lambda x: x[2]): if total <= max_bytes: break await asyncio.to_thread(path.unlink) total -= size self._log("info", f"清理完成: {total/1024/1024:.1f}MB") except Exception as e: self._log("error", f"清理失败: {e}") class ImageCache(CacheService): """图片缓存服务""" def __init__(self): super().__init__("image", timeout=30.0) async def download_image(self, path: str, token: str) -> Optional[Path]: """下载图片""" return await self.download(path, token) @staticmethod def to_base64(image_path: Path) -> Optional[str]: """图片转base64""" try: if not image_path.exists(): logger.error(f"[ImageCache] 文件不存在: {image_path}") return None data = base64.b64encode(image_path.read_bytes()).decode() mime = MIME_TYPES.get(image_path.suffix.lower(), DEFAULT_MIME) return f"data:{mime};base64,{data}" except Exception as e: logger.error(f"[ImageCache] 转换失败: {e}") return None async def download_base64(self, path: str, token: str) -> Optional[str]: """下载并转为base64（自动删除临时文件）""" try: cache_path = await self.download(path, token) if not cache_path: return None result = self.to_base64(cache_path) # 清理临时文件 try: cache_path.unlink() except Exception as e: logger.warning(f"[ImageCache] 删除临时文件失败: {e}") return result except Exception as e: logger.error(f"[ImageCache] 下载base64失败: {e}") return None class VideoCache(CacheService): """视频缓存服务""" def __init__(self): super().__init__("video", timeout=60.0) async def download_video(self, path: str, token: str) -> Optional[Path]: """下载视频""" return await self.download(path, token) # 全局实例 image_cache_service = ImageCache() video_cache_service = VideoCache()