diff --git a/app/api/v1/admin.py b/app/api/v1/admin.py index 8c05b05..d88b98d 100644 --- a/app/api/v1/admin.py +++ b/app/api/v1/admin.py @@ -1121,13 +1121,13 @@ async def admin_cache_page(): @router.get("/api/v1/admin/cache", dependencies=[Depends(verify_api_key)]) async def get_cache_stats_api(request: Request): """获取缓存统计""" - from app.services.grok.utils.download import DownloadService + from app.services.grok.utils.cache import CacheService from app.services.token.manager import get_token_manager try: - dl_service = DownloadService() - image_stats = dl_service.get_stats("image") - video_stats = dl_service.get_stats("video") + cache_service = CacheService() + image_stats = cache_service.get_stats("image") + video_stats = cache_service.get_stats("video") mgr = await get_token_manager() pools = mgr.pools @@ -1308,7 +1308,7 @@ async def get_cache_stats_api(request: Request): ) async def load_online_cache_api_async(data: dict): """在线资产统计（异步批量 + SSE 进度）""" - from app.services.grok.utils.download import DownloadService + from app.services.grok.utils.cache import CacheService from app.services.token.manager import get_token_manager mgr = await get_token_manager() @@ -1361,9 +1361,9 @@ async def load_online_cache_api_async(data: dict): async def _run(): try: - dl_service = DownloadService() - image_stats = dl_service.get_stats("image") - video_stats = dl_service.get_stats("video") + cache_service = CacheService() + image_stats = cache_service.get_stats("image") + video_stats = cache_service.get_stats("video") async def _on_item(item: str, res: dict): ok = bool(res.get("data", {}).get("ok")) @@ -1430,13 +1430,13 @@ async def load_online_cache_api_async(data: dict): @router.post("/api/v1/admin/cache/clear", dependencies=[Depends(verify_api_key)]) async def clear_local_cache_api(data: dict): """清理本地缓存""" - from app.services.grok.utils.download import DownloadService + from app.services.grok.utils.cache import CacheService cache_type = data.get("type", "image") try: - dl_service = DownloadService() - result = dl_service.clear(cache_type) + cache_service = CacheService() + result = cache_service.clear(cache_type) return {"status": "success", "result": result} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @@ -1450,13 +1450,13 @@ async def list_local_cache_api( page_size: int = 1000, ): """列出本地缓存文件""" - from app.services.grok.utils.download import DownloadService + from app.services.grok.utils.cache import CacheService try: if type_: cache_type = type_ - dl_service = DownloadService() - result = dl_service.list_files(cache_type, page, page_size) + cache_service = CacheService() + result = cache_service.list_files(cache_type, page, page_size) return {"status": "success", **result} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @@ -1465,15 +1465,15 @@ async def list_local_cache_api( @router.post("/api/v1/admin/cache/item/delete", dependencies=[Depends(verify_api_key)]) async def delete_local_cache_item_api(data: dict): """删除单个本地缓存文件""" - from app.services.grok.utils.download import DownloadService + from app.services.grok.utils.cache import CacheService cache_type = data.get("type", "image") name = data.get("name") if not name: raise HTTPException(status_code=400, detail="Missing file name") try: - dl_service = DownloadService() - result = dl_service.delete_file(cache_type, name) + cache_service = CacheService() + result = cache_service.delete_file(cache_type, name) return {"status": "success", "result": result} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) diff --git a/app/services/grok/batch_services/assets.py b/app/services/grok/batch_services/assets.py index 3ad213e..c22045c 100644 --- a/app/services/grok/batch_services/assets.py +++ b/app/services/grok/batch_services/assets.py @@ -2,12 +2,155 @@ Batch assets service. """ -from typing import Callable, Awaitable, Dict, Any, Optional +import asyncio +from typing import Callable, Awaitable, Dict, Any, Optional, List -from app.services.grok.services.assets import ListService, DeleteService +from curl_cffi.requests import AsyncSession + +from app.core.config import get_config +from app.core.logger import logger +from app.services.reverse import AssetsListReverse, AssetsDeleteReverse +from app.services.grok.utils.locks import _get_assets_semaphore from app.services.grok.utils.batch import run_in_batches +class BaseAssetsService: + """Base assets service.""" + + def __init__(self): + self._session: Optional[AsyncSession] = None + + async def _get_session(self) -> AsyncSession: + if self._session is None: + self._session = AsyncSession() + return self._session + + async def close(self): + if self._session: + await self._session.close() + self._session = None + + +class ListService(BaseAssetsService): + """Assets list service.""" + + async def iter_assets(self, token: str): + params = { + "pageSize": 50, + "orderBy": "ORDER_BY_LAST_USE_TIME", + "source": "SOURCE_ANY", + "isLatest": "true", + } + page_token = None + seen_tokens = set() + + async with AsyncSession() as session: + while True: + if page_token: + if page_token in seen_tokens: + logger.warning("Pagination stopped: repeated page token") + break + seen_tokens.add(page_token) + params["pageToken"] = page_token + else: + params.pop("pageToken", None) + + response = await AssetsListReverse.request( + session, + token, + params, + ) + + result = response.json() + page_assets = result.get("assets", []) + yield page_assets + + page_token = result.get("nextPageToken") + if not page_token: + break + + async def list(self, token: str) -> List[Dict]: + 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: + total = 0 + async for page_assets in self.iter_assets(token): + total += len(page_assets) + logger.debug(f"Asset count: {total}") + return total + + +class DeleteService(BaseAssetsService): + """Assets delete service.""" + + async def delete(self, token: str, asset_id: str) -> bool: + async with _get_assets_semaphore(): + session = await self._get_session() + await AssetsDeleteReverse.request( + session, + token, + asset_id, + ) + + logger.debug(f"Deleted: {asset_id}") + return True + + async def delete_all(self, token: str) -> Dict[str, int]: + total = success = failed = 0 + list_service = ListService() + + try: + async for assets in list_service.iter_assets(token): + if not assets: + continue + + total += len(assets) + 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} + 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 BatchAssetsService: """Batch assets orchestration.""" diff --git a/app/services/grok/batch_services/nsfw.py b/app/services/grok/batch_services/nsfw.py index 3c8ae06..b701bc3 100644 --- a/app/services/grok/batch_services/nsfw.py +++ b/app/services/grok/batch_services/nsfw.py @@ -2,12 +2,69 @@ Batch NSFW service. """ +from dataclasses import dataclass from typing import Callable, Awaitable, Dict, Any, Optional -from app.services.grok.services.nsfw import NSFWService +from curl_cffi.requests import AsyncSession + +from app.core.logger import logger +from app.core.config import get_config +from app.core.exceptions import UpstreamException +from app.services.reverse import NsfwMgmtReverse, SetBirthReverse +from app.services.reverse.utils.grpc import GrpcStatus from app.services.grok.utils.batch import run_in_batches +@dataclass +class NSFWResult: + """NSFW 操作结果""" + + success: bool + http_status: int + grpc_status: Optional[int] = None + grpc_message: Optional[str] = None + error: Optional[str] = None + + +class NSFWService: + """NSFW 模式服务""" + + async def enable(self, token: str) -> NSFWResult: + """为单个 token 开启 NSFW 模式""" + try: + browser = get_config("security.browser") + async with AsyncSession(impersonate=browser) as session: + # 先设置出生日期 + try: + await SetBirthReverse.request(session, token) + except UpstreamException as e: + status = None + if e.details and "status" in e.details: + status = e.details["status"] + else: + status = getattr(e, "status_code", None) + return NSFWResult( + success=False, + http_status=status or 0, + error=f"Set birth date failed: {str(e)}", + ) + + # 开启 NSFW + grpc_status: GrpcStatus = await NsfwMgmtReverse.request(session, token) + success = grpc_status.code in (-1, 0) + + return NSFWResult( + success=success, + http_status=200, + grpc_status=grpc_status.code, + grpc_message=grpc_status.message or None, + ) + + except Exception as e: + logger.error(f"NSFW enable failed: {e}") + return NSFWResult(success=False, http_status=0, error=str(e)[:100]) + + class BatchNSFWService: """Batch NSFW orchestration.""" @@ -45,4 +102,4 @@ class BatchNSFWService: ) -__all__ = ["BatchNSFWService"] +__all__ = ["BatchNSFWService", "NSFWService", "NSFWResult"] diff --git a/app/services/grok/batch_services/usage.py b/app/services/grok/batch_services/usage.py index 29e1b4d..edac6f4 100644 --- a/app/services/grok/batch_services/usage.py +++ b/app/services/grok/batch_services/usage.py @@ -2,10 +2,61 @@ Batch usage service. """ +import asyncio from typing import Callable, Awaitable, Dict, Any, Optional +from curl_cffi.requests import AsyncSession + +from app.core.logger import logger +from app.core.config import get_config +from app.services.reverse import RateLimitsReverse from app.services.grok.utils.batch import run_in_batches +_USAGE_SEMAPHORE = asyncio.Semaphore(25) +_USAGE_SEM_VALUE = 25 + + +class UsageService: + """用量查询服务""" + + async def get(self, token: str) -> Dict: + """ + 获取速率限制信息 + + Args: + token: 认证 Token + + Returns: + 响应数据 + + Raises: + UpstreamException: 当获取失败且重试耗尽时 + """ + value = get_config("performance.usage_max_concurrent") + try: + value = int(value) + except Exception: + value = 25 + value = max(1, value) + global _USAGE_SEMAPHORE, _USAGE_SEM_VALUE + if value != _USAGE_SEM_VALUE: + _USAGE_SEM_VALUE = value + _USAGE_SEMAPHORE = asyncio.Semaphore(value) + async with _USAGE_SEMAPHORE: + try: + async with AsyncSession() as session: + response = await RateLimitsReverse.request(session, token) + data = response.json() + remaining = data.get("remainingTokens", 0) + logger.info( + f"Usage sync success: remaining={remaining}, token={token[:10]}..." + ) + return data + + except Exception: + # 最后一次失败已经被记录 + raise + class BatchUsageService: """Batch usage orchestration.""" @@ -33,4 +84,4 @@ class BatchUsageService: ) -__all__ = ["BatchUsageService"] +__all__ = ["BatchUsageService", "UsageService"] diff --git a/app/services/grok/processors/base.py b/app/services/grok/processors/base.py index f1c208e..1d2d5c9 100644 --- a/app/services/grok/processors/base.py +++ b/app/services/grok/processors/base.py @@ -144,7 +144,7 @@ class BaseProcessor: if self.app_url: dl_service = self._get_dl() - await dl_service.download(path, self.token, media_type) + await dl_service.download_file(path, self.token, media_type) return f"{self.app_url.rstrip('/')}/v1/files/{media_type}{path}" else: return f"{ASSET_URL.rstrip('/')}{path}" diff --git a/app/services/grok/processors/chat.py b/app/services/grok/processors/chat.py index f55f853..1e04437 100644 --- a/app/services/grok/processors/chat.py +++ b/app/services/grok/processors/chat.py @@ -168,7 +168,7 @@ class StreamProcessor(BaseProcessor): if self.image_format == "base64": try: dl_service = self._get_dl() - base64_data = await dl_service.to_base64( + base64_data = await dl_service.parse_b64( url, self.token, "image" ) if base64_data: @@ -296,7 +296,7 @@ class CollectProcessor(BaseProcessor): if self.image_format == "base64": try: dl_service = self._get_dl() - base64_data = await dl_service.to_base64( + base64_data = await dl_service.parse_b64( url, self.token, "image" ) if base64_data: diff --git a/app/services/grok/processors/image.py b/app/services/grok/processors/image.py index a53b94a..f8f1f29 100644 --- a/app/services/grok/processors/image.py +++ b/app/services/grok/processors/image.py @@ -99,7 +99,7 @@ class ImageStreamProcessor(BaseProcessor): continue try: dl_service = self._get_dl() - base64_data = await dl_service.to_base64( + base64_data = await dl_service.parse_b64( url, self.token, "image" ) if base64_data: @@ -212,7 +212,7 @@ class ImageCollectProcessor(BaseProcessor): continue try: dl_service = self._get_dl() - base64_data = await dl_service.to_base64( + base64_data = await dl_service.parse_b64( url, self.token, "image" ) if base64_data: diff --git a/app/services/grok/services/assets.py b/app/services/grok/services/assets.py deleted file mode 100644 index 12d31d6..0000000 --- a/app/services/grok/services/assets.py +++ /dev/null @@ -1,666 +0,0 @@ -""" -Grok 文件资产服务 -""" - -import asyncio -import base64 -import hashlib -import os -import re -import time -from contextlib import asynccontextmanager -from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple -from urllib.parse import urlparse - -try: - import fcntl -except ImportError: - fcntl = None - -import aiofiles -from curl_cffi.requests import AsyncSession - -from app.core.config import get_config -from app.core.exceptions import AppException, UpstreamException, ValidationException -from app.core.logger import logger -from app.core.storage import DATA_DIR -from app.services.reverse import ( - AssetsDeleteReverse, - AssetsDownloadReverse, - AssetsListReverse, - AssetsUploadReverse, -) - -# ==================== 常量 ==================== - -LOCK_DIR = DATA_DIR / ".locks" - -# 全局信号量（运行时动态初始化） -_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: - try: - fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) - locked = True - break - except BlockingIOError: - if time.monotonic() - start >= timeout: - break - await asyncio.sleep(0.05) - yield - finally: - if fd: - if locked: - try: - fcntl.flock(fd, fcntl.LOCK_UN) - except Exception: - pass - fd.close() - - -# ==================== 基础服务 ==================== - - -class BaseService: - """基础服务类""" - - def __init__(self): - self._session: Optional[AsyncSession] = None - - async def _get_session(self) -> AsyncSession: - """获取复用 Session""" - if self._session is None: - self._session = AsyncSession() - return self._session - - async def close(self): - """关闭 Session""" - if self._session: - await self._session.close() - self._session = None - - @staticmethod - def is_url(s: str) -> bool: - """检查是否为 URL""" - try: - 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""" - 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: {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", "application/octet-stream" - ).split(";")[0] - b64 = base64.b64encode(response.content).decode() - - logger.debug(f"Fetched: {url}") - return filename, b64, content_type - except Exception as e: - if isinstance(e, AppException): - 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 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""" - try: - if not file_path.exists(): - logger.warning(f"File not found for base64 conversion: {file_path}") - raise AppException( - f"File not found: {file_path}", code="file_not_found" - ) - - if not file_path.is_file(): - logger.warning(f"Path is not a file: {file_path}") - raise AppException( - f"Invalid file path: {file_path}", code="invalid_file_path" - ) - - b64_data = base64.b64encode(file_path.read_bytes()).decode() - return f"data:{mime_type};base64,{b64_data}" - except AppException: - raise - 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" - ) - - -# ==================== 上传服务 ==================== - - -class UploadService(BaseService): - """文件上传服务""" - - async def upload(self, file_input: str, token: str) -> Tuple[str, str]: - """ - 上传文件到 Grok - - Returns: - (file_id, file_uri) - """ - 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) - - logger.debug( - f"Upload prepare: filename={filename}, type={mime}, size={len(b64)}" - ) - - if not b64: - raise ValidationException("Invalid file input: empty content") - - # 执行上传 - session = await self._get_session() - response = await AssetsUploadReverse.request( - session, - token, - filename, - mime, - b64, - ) - - 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 - - -# ==================== 列表服务 ==================== - - -class ListService(BaseService): - """文件列表查询服务""" - - async def iter_assets(self, token: str): - """分页迭代资产列表""" - params = { - "pageSize": 50, - "orderBy": "ORDER_BY_LAST_USE_TIME", - "source": "SOURCE_ANY", - "isLatest": "true", - } - page_token = None - seen_tokens = set() - - async with AsyncSession() as session: - while True: - if page_token: - if page_token in seen_tokens: - logger.warning("Pagination stopped: repeated page token") - break - seen_tokens.add(page_token) - params["pageToken"] = page_token - else: - params.pop("pageToken", None) - - response = await AssetsListReverse.request( - session, - token, - params, - ) - - result = response.json() - page_assets = result.get("assets", []) - yield page_assets - - page_token = result.get("nextPageToken") - if not page_token: - break - - async def list(self, token: str) -> List[Dict]: - """查询文件列表""" - 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: - """统计资产数量""" - total = 0 - async for page_assets in self.iter_assets(token): - total += len(page_assets) - logger.debug(f"Asset count: {total}") - 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 AssetsDeleteReverse.request( - session, - token, - asset_id, - ) - - logger.debug(f"Deleted: {asset_id}") - return True - - async def delete_all(self, token: str) -> Dict[str, int]: - """删除所有文件""" - total = success = failed = 0 - list_service = ListService() - - try: - async for assets in list_service.iter_assets(token): - if not assets: - continue - - total += len(assets) - 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} - 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): - super().__init__() - self.base_dir = DATA_DIR / "tmp" - self.image_dir = self.base_dir / "image" - self.video_dir = self.base_dir / "video" - self.image_dir.mkdir(parents=True, exist_ok=True) - self.video_dir.mkdir(parents=True, exist_ok=True) - self._cleanup_running = False - - def _cache_path(self, file_path: str, media_type: str) -> Path: - """获取缓存路径""" - cache_dir = self.image_dir if media_type == "image" else self.video_dir - filename = file_path.lstrip("/").replace("/", "-") - return cache_dir / filename - - 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") - - 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(): - return cache_path, self._get_mime(cache_path) - - # 执行下载 - mime = await self._download_file(file_path, token, cache_path) - logger.info(f"Downloaded: {file_path}") - - # 异步检查缓存限制 - asyncio.create_task(self.check_limit()) - - return cache_path, mime - - async def _download_file(self, file_path: str, token: str, cache_path: Path) -> str: - """执行下载""" - if not file_path.startswith("/"): - file_path = f"/{file_path}" - - session = await self._get_session() - response = await AssetsDownloadReverse.request(session, token, file_path) - - # 保存文件 - 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) - 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 - - return self._get_mime(cache_path, response) - - async def to_base64( - self, file_path: str, token: str, media_type: str = "image" - ) -> str: - """下载并转 base64""" - try: - cache_path, mime = await self.download(file_path, token, media_type) - if not cache_path or not cache_path.exists(): - logger.warning(f"Download failed for {file_path}: invalid path") - raise AppException( - "Download failed: invalid path", code="download_failed" - ) - - data_uri = self.to_b64(cache_path, mime) - - # 删除临时文件 - if data_uri: - try: - cache_path.unlink() - except Exception as e: - logger.debug(f"Failed to cleanup temp file {cache_path}: {e}") - - return data_uri - except Exception as e: - logger.error(f"Failed to convert {file_path} to base64: {e}") - raise - - 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} - - 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]: - """列出缓存文件""" - 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 - ] - - # 构建文件列表 - 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), - } - ) - except Exception: - continue - - items.sort(key=lambda x: x["mtime_ms"], reverse=True) - - # 分页 - total = len(items) - start = max(0, (page - 1) * page_size) - paged = items[start : start + page_size] - - # 添加 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 - 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]: - """清空缓存""" - 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 = list(cache_dir.glob("*")) - total_size = sum(f.stat().st_size for f in files if f.is_file()) - count = 0 - - for f in files: - 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 or not get_config("cache.enable_auto_clean"): - return - - self._cleanup_running = True - try: - async with _file_lock("cache_cleanup", timeout=5): - 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 - - # 清理到 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 - - for f, _, size in all_files: - try: - f.unlink() - deleted_count += 1 - deleted_size += size - total_size -= size - if (total_size / 1024 / 1024) <= target_mb: - break - except Exception: - pass - - logger.info( - f"Cache cleanup: {deleted_count} files ({deleted_size / 1024 / 1024:.2f}MB)" - ) - finally: - self._cleanup_running = False - - def _collect_files(self) -> Tuple[List[Tuple[Path, float, int]], int]: - """收集所有缓存文件""" - total_size = 0 - all_files = [] - - 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 - - return all_files, total_size - - -__all__ = [ - "BaseService", - "UploadService", - "ListService", - "DeleteService", - "DownloadService", -] diff --git a/app/services/grok/services/chat.py b/app/services/grok/services/chat.py index fc8d285..aa43923 100644 --- a/app/services/grok/services/chat.py +++ b/app/services/grok/services/chat.py @@ -16,7 +16,7 @@ from app.core.exceptions import ( UpstreamException, ) from app.services.grok.models.model import ModelService -from app.services.grok.services.assets import UploadService +from app.services.grok.utils.upload import UploadService from app.services.grok.processors import StreamProcessor, CollectProcessor from app.services.reverse import AppChatReverse from app.services.grok.utils.stream import wrap_stream_with_usage @@ -209,7 +209,7 @@ class GrokChatService: upload_service = UploadService() try: for attach_type, attach_data in attachments: - file_id, _ = await upload_service.upload(attach_data, token) + file_id, _ = await upload_service.upload_file(attach_data, token) file_ids.append(file_id) logger.debug( f"Attachment uploaded: type={attach_type}, file_id={file_id}" diff --git a/app/services/grok/services/image_edit.py b/app/services/grok/services/image_edit.py index ca460bc..0eb8177 100644 --- a/app/services/grok/services/image_edit.py +++ b/app/services/grok/services/image_edit.py @@ -11,7 +11,7 @@ from typing import AsyncGenerator, List, Union, Any from app.core.exceptions import AppException, ErrorType from app.core.logger import logger from app.services.grok.processors import ImageCollectProcessor, ImageStreamProcessor -from app.services.grok.services.assets import UploadService +from app.services.grok.utils.upload import UploadService from app.services.grok.services.chat import GrokChatService from app.services.grok.services.video import VideoService from app.services.grok.utils.stream import wrap_stream_with_usage @@ -98,7 +98,7 @@ class ImageEditService: upload_service = UploadService() try: for image in images: - _, file_uri = await upload_service.upload(image, token) + _, file_uri = await upload_service.upload_file(image, token) if file_uri: if file_uri.startswith("http"): image_urls.append(file_uri) diff --git a/app/services/grok/services/nsfw.py b/app/services/grok/services/nsfw.py deleted file mode 100644 index a919d54..0000000 --- a/app/services/grok/services/nsfw.py +++ /dev/null @@ -1,68 +0,0 @@ -""" -NSFW (Unhinged) 模式服务 - -使用 gRPC-Web 协议开启账号的 NSFW 功能。 -""" - -from dataclasses import dataclass -from typing import Optional - -from curl_cffi.requests import AsyncSession - -from app.core.logger import logger -from app.core.config import get_config -from app.core.exceptions import UpstreamException -from app.services.reverse import NsfwMgmtReverse, SetBirthReverse -from app.services.reverse.utils.grpc import GrpcStatus - -@dataclass -class NSFWResult: - """NSFW 操作结果""" - - success: bool - http_status: int - grpc_status: Optional[int] = None - grpc_message: Optional[str] = None - error: Optional[str] = None - - -class NSFWService: - """NSFW 模式服务""" - - async def enable(self, token: str) -> NSFWResult: - """为单个 token 开启 NSFW 模式""" - try: - browser = get_config("security.browser") - async with AsyncSession(impersonate=browser) as session: - # 先设置出生日期 - try: - await SetBirthReverse.request(session, token) - except UpstreamException as e: - status = None - if e.details and "status" in e.details: - status = e.details["status"] - else: - status = getattr(e, "status_code", None) - return NSFWResult( - success=False, - http_status=status or 0, - error=f"Set birth date failed: {str(e)}", - ) - - # 开启 NSFW - grpc_status: GrpcStatus = await NsfwMgmtReverse.request(session, token) - success = grpc_status.code in (-1, 0) - - return NSFWResult( - success=success, - http_status=200, - grpc_status=grpc_status.code, - grpc_message=grpc_status.message or None, - ) - - except Exception as e: - 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 deleted file mode 100644 index e273433..0000000 --- a/app/services/grok/services/usage.py +++ /dev/null @@ -1,60 +0,0 @@ -""" -Grok 用量服务 -""" - -import asyncio -from typing import Dict - -from curl_cffi.requests import AsyncSession - -from app.core.logger import logger -from app.core.config import get_config -from app.services.reverse import RateLimitsReverse - -_USAGE_SEMAPHORE = asyncio.Semaphore(25) -_USAGE_SEM_VALUE = 25 - - -class UsageService: - """用量查询服务""" - - async def get(self, token: str) -> Dict: - """ - 获取速率限制信息 - - Args: - token: 认证 Token - - Returns: - 响应数据 - - Raises: - UpstreamException: 当获取失败且重试耗尽时 - """ - value = get_config("performance.usage_max_concurrent") - try: - value = int(value) - except Exception: - value = 25 - value = max(1, value) - global _USAGE_SEMAPHORE, _USAGE_SEM_VALUE - if value != _USAGE_SEM_VALUE: - _USAGE_SEM_VALUE = value - _USAGE_SEMAPHORE = asyncio.Semaphore(value) - async with _USAGE_SEMAPHORE: - try: - async with AsyncSession() as session: - response = await RateLimitsReverse.request(session, token) - data = response.json() - remaining = data.get("remainingTokens", 0) - logger.info( - f"Usage sync success: remaining={remaining}, token={token[:10]}..." - ) - return data - - except Exception: - # 最后一次失败已经被记录 - raise - - -__all__ = ["UsageService"] diff --git a/app/services/grok/services/video.py b/app/services/grok/services/video.py index bb53e53..cd65c8a 100644 --- a/app/services/grok/services/video.py +++ b/app/services/grok/services/video.py @@ -260,7 +260,7 @@ class VideoService: # Extract content. from app.services.grok.services.chat import MessageExtractor - from app.services.grok.services.assets import UploadService + from app.services.grok.utils.upload import UploadService try: prompt, attachments = MessageExtractor.extract(messages, is_video=True) @@ -274,7 +274,7 @@ class VideoService: try: for attach_type, attach_data in attachments: if attach_type == "image": - _, file_uri = await upload_service.upload(attach_data, token) + _, file_uri = await upload_service.upload_file(attach_data, token) image_url = f"https://assets.grok.com/{file_uri}" logger.info(f"Image uploaded for video: {image_url}") break diff --git a/app/services/grok/utils/cache.py b/app/services/grok/utils/cache.py new file mode 100644 index 0000000..a728df1 --- /dev/null +++ b/app/services/grok/utils/cache.py @@ -0,0 +1,110 @@ +""" +Local cache utilities. +""" + +from typing import Any, Dict + +from app.core.storage import DATA_DIR + +IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"} +VIDEO_EXTS = {".mp4", ".mov", ".m4v", ".webm", ".avi", ".mkv"} + + +class CacheService: + """Local cache service.""" + + def __init__(self): + base_dir = DATA_DIR / "tmp" + self.image_dir = base_dir / "image" + self.video_dir = base_dir / "video" + self.image_dir.mkdir(parents=True, exist_ok=True) + self.video_dir.mkdir(parents=True, exist_ok=True) + + def _cache_dir(self, media_type: str): + return self.image_dir if media_type == "image" else self.video_dir + + def _allowed_exts(self, media_type: str): + return IMAGE_EXTS if media_type == "image" else VIDEO_EXTS + + def get_stats(self, media_type: str = "image") -> Dict[str, Any]: + cache_dir = self._cache_dir(media_type) + if not cache_dir.exists(): + return {"count": 0, "size_mb": 0.0} + + allowed = self._allowed_exts(media_type) + 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]: + cache_dir = self._cache_dir(media_type) + if not cache_dir.exists(): + return {"total": 0, "page": page, "page_size": page_size, "items": []} + + allowed = self._allowed_exts(media_type) + files = [ + f for f in cache_dir.glob("*") if f.is_file() and f.suffix.lower() in allowed + ] + + 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), + } + ) + except Exception: + continue + + items.sort(key=lambda x: x["mtime_ms"], reverse=True) + + total = len(items) + start = max(0, (page - 1) * page_size) + paged = items[start : start + page_size] + + 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._cache_dir(media_type) + 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]: + cache_dir = self._cache_dir(media_type) + if not cache_dir.exists(): + return {"count": 0, "size_mb": 0.0} + + files = list(cache_dir.glob("*")) + total_size = sum(f.stat().st_size for f in files if f.is_file()) + count = 0 + + for f in files: + if f.is_file(): + try: + f.unlink() + count += 1 + except Exception: + pass + + return {"count": count, "size_mb": round(total_size / 1024 / 1024, 2)} + + +__all__ = ["CacheService"] diff --git a/app/services/grok/utils/download.py b/app/services/grok/utils/download.py index b3b930f..9179610 100644 --- a/app/services/grok/utils/download.py +++ b/app/services/grok/utils/download.py @@ -1,5 +1,233 @@ -"""Download service (compat wrapper).""" +""" +Download service. +""" + +import asyncio +import base64 +import hashlib +import os +from pathlib import Path +from typing import List, Optional, Tuple +from urllib.parse import urlparse + +import aiofiles +from curl_cffi.requests import AsyncSession + +from app.core.logger import logger +from app.core.storage import DATA_DIR +from app.core.config import get_config +from app.core.exceptions import AppException +from app.services.reverse import AssetsDownloadReverse +from app.services.grok.utils.locks import _get_assets_semaphore, _file_lock + + +class DownloadService: + """Assets download service.""" + + def __init__(self): + self._session: Optional[AsyncSession] = None + base_dir = DATA_DIR / "tmp" + self.image_dir = base_dir / "image" + self.video_dir = base_dir / "video" + self.image_dir.mkdir(parents=True, exist_ok=True) + self.video_dir.mkdir(parents=True, exist_ok=True) + self._cleanup_running = False + + async def create(self) -> AsyncSession: + """Create or reuse a session.""" + if self._session is None: + self._session = AsyncSession() + return self._session + + async def close(self): + """Close the session.""" + if self._session: + await self._session.close() + self._session = None + + @staticmethod + def _is_url(value: str) -> bool: + """Check if the value is a URL.""" + try: + parsed = urlparse(value) + return bool(parsed.scheme and parsed.netloc and parsed.scheme in ["http", "https"]) + except Exception: + return False + + async def parse_b64(self, file_path: str, token: str, media_type: str = "image") -> str: + """Download and return data URI.""" + try: + cache_path, mime = await self.download_file(file_path, token, media_type) + if not cache_path or not cache_path.exists(): + logger.warning(f"Download failed for {file_path}: invalid path") + raise AppException( + "Download failed: invalid path", code="download_failed" + ) + + data_uri = await self.format_b64(cache_path, mime) + + if data_uri: + try: + cache_path.unlink() + except Exception as e: + logger.debug(f"Failed to cleanup temp file {cache_path}: {e}") + + return data_uri + except Exception as e: + logger.error(f"Failed to convert {file_path} to base64: {e}") + raise + + @staticmethod + async def format_b64(file_path: Path, mime_type: str) -> str: + """Format local file to data URI.""" + try: + if not file_path.exists(): + logger.warning(f"File not found for base64 conversion: {file_path}") + raise AppException( + f"File not found: {file_path}", code="file_not_found" + ) + + if not file_path.is_file(): + logger.warning(f"Path is not a file: {file_path}") + raise AppException( + f"Invalid file path: {file_path}", code="invalid_file_path" + ) + + async with aiofiles.open(file_path, "rb") as f: + data = await f.read() + b64_data = base64.b64encode(data).decode() + return f"data:{mime_type};base64,{b64_data}" + except AppException: + raise + 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" + ) + + def check_format(self, file_path: str) -> str: + """Normalize file path for download.""" + if not isinstance(file_path, str) or not file_path.strip(): + raise AppException("Invalid file path", code="invalid_file_path") + if self._is_url(file_path): + file_path = urlparse(file_path).path or "" + if not file_path.startswith("/"): + file_path = f"/{file_path}" + return file_path + + async def download_file(self, file_path: str, token: str, media_type: str = "image") -> Tuple[Optional[Path], str]: + """Download asset to local cache. + + Args: + file_path: str, the path of the file to download. + token: str, the SSO token. + media_type: str, the media type of the file. + + Returns: + Tuple[Optional[Path], str]: The path of the downloaded file and the MIME type. + """ + async with _get_assets_semaphore(): + file_path = self.check_format(file_path) + cache_dir = self.image_dir if media_type == "image" else self.video_dir + filename = file_path.lstrip("/").replace("/", "-") + cache_path = cache_dir / filename + + lock_name = ( + f"dl_{media_type}_{hashlib.sha1(str(cache_path).encode()).hexdigest()[:16]}" + ) + async with _file_lock(lock_name, timeout=10): + session = await self.create() + response = await AssetsDownloadReverse.request( + session, token, file_path + ) + + 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) + 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 = response.headers.get( + "content-type", "application/octet-stream" + ).split(";")[0] + logger.info(f"Downloaded: {file_path}") + + asyncio.create_task(self.check_limit()) + + return cache_path, mime + + async def check_limit(self): + """Check cache limit and cleanup. + + Args: + self: DownloadService, the download service instance. + + Returns: + None + """ + 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): + limit_mb = get_config("cache.limit_mb") + total_size = 0 + all_files: List[Tuple[Path, float, int]] = [] + + 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 + 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..." + ) + all_files.sort(key=lambda x: x[1]) + + deleted_count = 0 + deleted_size = 0 + target_mb = limit_mb * 0.8 + + for f, _, size in all_files: + try: + f.unlink() + deleted_count += 1 + deleted_size += size + total_size -= size + if (total_size / 1024 / 1024) <= target_mb: + break + except Exception: + pass + + logger.info( + f"Cache cleanup: {deleted_count} files ({deleted_size / 1024 / 1024:.2f}MB)" + ) + finally: + self._cleanup_running = False -from app.services.grok.services.assets import DownloadService __all__ = ["DownloadService"] diff --git a/app/services/grok/utils/locks.py b/app/services/grok/utils/locks.py new file mode 100644 index 0000000..654ae9e --- /dev/null +++ b/app/services/grok/utils/locks.py @@ -0,0 +1,71 @@ +""" +Shared locking helpers for assets operations. +""" + +import asyncio +import time +from contextlib import asynccontextmanager +from pathlib import Path + +from app.core.config import get_config +from app.core.storage import DATA_DIR + +try: + import fcntl +except ImportError: + fcntl = None + + +LOCK_DIR = DATA_DIR / ".locks" + +_ASSETS_SEMAPHORE = None +_ASSETS_SEM_VALUE = None + + +def _get_assets_semaphore() -> asyncio.Semaphore: + """Return global semaphore for assets operations.""" + 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): + """File lock guard.""" + if fcntl is None: + yield + return + + LOCK_DIR.mkdir(parents=True, exist_ok=True) + lock_path = Path(LOCK_DIR) / f"{name}.lock" + fd = None + locked = False + start = time.monotonic() + + try: + fd = open(lock_path, "a+") + while True: + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + locked = True + break + except BlockingIOError: + if time.monotonic() - start >= timeout: + break + await asyncio.sleep(0.05) + yield + finally: + if fd: + if locked: + try: + fcntl.flock(fd, fcntl.LOCK_UN) + except Exception: + pass + fd.close() + + +__all__ = ["_get_assets_semaphore", "_file_lock"] diff --git a/app/services/grok/utils/upload.py b/app/services/grok/utils/upload.py new file mode 100644 index 0000000..85302ac --- /dev/null +++ b/app/services/grok/utils/upload.py @@ -0,0 +1,137 @@ +""" +Upload service. +""" + +import base64 +import re +from typing import Optional, Tuple +from urllib.parse import urlparse + +from curl_cffi.requests import AsyncSession + +from app.core.exceptions import AppException, UpstreamException, ValidationException +from app.core.logger import logger +from app.services.reverse import AssetsUploadReverse +from app.services.grok.utils.locks import _get_assets_semaphore + + +class UploadService: + """Assets upload service.""" + + def __init__(self): + self._session: Optional[AsyncSession] = None + + async def create(self) -> AsyncSession: + """Create or reuse a session.""" + if self._session is None: + self._session = AsyncSession() + return self._session + + async def close(self): + """Close the session.""" + if self._session: + await self._session.close() + self._session = None + + @staticmethod + def _is_url(value: str) -> bool: + """Check if the value is a URL.""" + try: + parsed = urlparse(value) + return bool(parsed.scheme and parsed.netloc and parsed.scheme in ["http", "https"]) + except Exception: + return False + + @staticmethod + async def parse_b64(url: str) -> Tuple[str, str, str]: + """Fetch URL content and return (filename, base64, mime).""" + 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: {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", "application/octet-stream" + ).split(";")[0] + b64 = base64.b64encode(response.content).decode() + + logger.debug(f"Fetched: {url}") + return filename, b64, content_type + except Exception as e: + if isinstance(e, AppException): + raise + logger.error(f"Fetch failed: {url} - {e}") + raise UpstreamException(f"Fetch failed: {str(e)}", details={"url": url}) + + @staticmethod + def format_b64(data_uri: str) -> Tuple[str, str, str]: + """Format data URI to (filename, base64, 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 + + async def check_format(self, file_input: str) -> Tuple[str, str, str]: + """Check file input format and return (filename, base64, mime).""" + if not isinstance(file_input, str) or not file_input.strip(): + raise ValidationException("Invalid file input: empty content") + + if self._is_url(file_input): + return await self.parse_b64(file_input) + + return self.format_b64(file_input) + + async def upload_file(self, file_input: str, token: str) -> Tuple[str, str]: + """ + Upload file to Grok. + + Args: + file_input: str, the file input. + token: str, the SSO token. + + Returns: + Tuple[str, str]: The file ID and URI. + """ + async with _get_assets_semaphore(): + filename, b64, mime = await self.check_format(file_input) + + logger.debug( + f"Upload prepare: filename={filename}, type={mime}, size={len(b64)}" + ) + + if not b64: + raise ValidationException("Invalid file input: empty content") + + session = await self.create() + response = await AssetsUploadReverse.request( + session, + token, + filename, + mime, + b64, + ) + + 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 + + +__all__ = ["UploadService"] diff --git a/app/services/token/manager.py b/app/services/token/manager.py index f4b547e..4f056d8 100644 --- a/app/services/token/manager.py +++ b/app/services/token/manager.py @@ -14,9 +14,10 @@ from app.services.token.models import ( BASIC__DEFAULT_QUOTA, SUPER_DEFAULT_QUOTA, ) -from app.core.storage import get_storage +from app.core.storage import get_storage, LocalStorage from app.core.config import get_config from app.services.token.pool import TokenPool +from app.services.grok.batch_services.usage import UsageService DEFAULT_REFRESH_BATCH_SIZE = 10 @@ -70,8 +71,6 @@ class TokenManager: # 如果后端返回 None 或空数据，尝试从本地 data/token.json 初始化后端 if not data: - from app.core.storage import LocalStorage - local_storage = LocalStorage() local_data = await local_storage.load_tokens() if local_data: @@ -363,8 +362,6 @@ class TokenManager: # 尝试 API 同步 try: - from app.services.grok.services.usage import UsageService - usage_service = UsageService() result = await usage_service.get(token_str) @@ -634,8 +631,6 @@ class TokenManager: Returns: {"checked": int, "refreshed": int, "recovered": int, "expired": int} """ - from app.services.grok.services.usage import UsageService - # 收集需要刷新的 token to_refresh: List[TokenInfo] = [] for pool in self.pools.values(): @@ -676,7 +671,7 @@ class TokenManager: # 重试逻辑：最多 2 次重试 for retry in range(3): # 0, 1, 2 try: - result = await usage_service.get(token_str, model_name="grok-3") + result = await usage_service.get(token_str) if result and "remainingTokens" in result: new_quota = result["remainingTokens"]