"""模型缓存服务 - 基于 Redis 的跨进程模型信息缓存。 本模块将数据库中的 model_providers 表数据序列化到 Redis， 供 API 和 Worker 等多进程同步读取，避免在同步函数中查询异步数据库。 v2 模型 spec 格式: provider_id:model_id（冒号分隔） v1 模型 spec 格式: provider/model_name（斜杠分隔） """ from __future__ import annotations import json import os import time from dataclasses import dataclass, field from typing import Any from yuxi.utils.logging_config import logger REDIS_CACHE_KEY = "yuxi:model_cache:v2" _DEFAULT_REDIS_URL = "redis://redis:6379/0" _CACHE_TTL_SECONDS = 5 @dataclass(frozen=True) class ModelInfo: """不可变的模型信息，供运行时使用。""" provider_id: str model_id: str model_type: str # chat / embedding / rerank display_name: str # 运行时配置 api_key: str base_url: str provider_type: str # openai / anthropic / gemini / ollama / openrouter / lmstudio # 可选配置 headers: dict[str, str] = field(default_factory=dict) extra: dict[str, Any] = field(default_factory=dict) # Embedding 专属 dimension: int | None = None batch_size: int = 40 @property def spec(self) -> str: """完整的 spec 标识，格式: provider_id:model_id""" return f"{self.provider_id}:{self.model_id}" def to_dict(self) -> dict: return { "provider_id": self.provider_id, "model_id": self.model_id, "model_type": self.model_type, "display_name": self.display_name, "api_key": self.api_key, "base_url": self.base_url, "provider_type": self.provider_type, "headers": self.headers, "extra": self.extra, "dimension": self.dimension, "batch_size": self.batch_size, } @classmethod def from_dict(cls, data: dict) -> ModelInfo: return cls( provider_id=data["provider_id"], model_id=data["model_id"], model_type=data["model_type"], display_name=data["display_name"], api_key=data["api_key"], base_url=data["base_url"], provider_type=data["provider_type"], headers=data.get("headers", {}), extra=data.get("extra", {}), dimension=data.get("dimension"), batch_size=data.get("batch_size", 40), ) class ModelCache: """基于 Redis 的模型缓存，所有写入均走 Redis，保证跨进程一致。 查询接口使用本地内存缓存（带 TTL），避免热路径反复反序列化 JSON。 """ def __init__(self) -> None: self._redis = None self._local_cache: dict[str, ModelInfo] | None = None self._local_cache_at: float = 0.0 def _get_redis(self): """懒加载同步 Redis 客户端，异常时重置以便下次重试。""" if self._redis is None: try: import redis redis_url = os.getenv("REDIS_URL", _DEFAULT_REDIS_URL) self._redis = redis.from_url(redis_url, decode_responses=True) except Exception as e: logger.warning(f"Redis client unavailable: {e}") return self._redis def _reset_redis(self) -> None: """重置 Redis 连接，下次调用时重新初始化。""" self._redis = None def _load_cache(self) -> dict[str, ModelInfo]: """从 Redis 同步读取完整缓存，带本地 TTL。""" now = time.monotonic() if self._local_cache is not None and (now - self._local_cache_at) < _CACHE_TTL_SECONDS: return self._local_cache r = self._get_redis() if r is None: return {} try: raw = r.get(REDIS_CACHE_KEY) if not raw: self._local_cache = {} self._local_cache_at = now return {} items = json.loads(raw) cache = {spec: ModelInfo.from_dict(data) for spec, data in items.items()} self._local_cache = cache self._local_cache_at = now return cache except Exception as e: logger.warning(f"Failed to load model cache from Redis: {e}") self._reset_redis() return {} def _invalidate_local(self) -> None: """使本地缓存失效，下次读取将重新从 Redis 加载。""" self._local_cache = None self._local_cache_at = 0.0 def get_model_info(self, spec: str) -> ModelInfo | None: """根据 spec 获取模型信息（同步）。""" cache = self._load_cache() return cache.get(spec) def is_v2_spec(self, spec: str) -> bool: """判断 spec 是否为 v2 格式（存在于缓存中）。""" return self.get_model_info(spec) is not None def get_all_specs(self, model_type: str | None = None) -> list[ModelInfo]: """获取所有缓存的模型信息，可按类型过滤。""" cache = self._load_cache() if model_type is None: return list(cache.values()) return [info for info in cache.values() if info.model_type == model_type] def get_specs_grouped_by_provider(self, model_type: str = "chat") -> dict[str, list[ModelInfo]]: """按 provider 分组获取模型列表（供前端使用）。""" cache = self._load_cache() grouped: dict[str, list[ModelInfo]] = {} for info in cache.values(): if info.model_type != model_type: continue grouped.setdefault(info.provider_id, []).append(info) return grouped def rebuild(self, providers: list[Any]) -> None: """从数据库 provider 列表重建缓存并写入 Redis。 所有修改操作（启动初始化、CRUD）都通过此方法写入， 保证 Redis 中的数据始终是最新的。 """ from yuxi.services.model_provider_service import _resolve_api_key new_cache: dict[str, ModelInfo] = {} for provider in providers: if not provider.is_enabled: continue api_key = _resolve_api_key(provider) for model in provider.enabled_models or []: model_type = model.get("type", "chat") base_url = self._get_base_url_for_type(provider, model_type) info = ModelInfo( provider_id=provider.provider_id, model_id=model["id"], model_type=model_type, display_name=model.get("display_name", model["id"]), api_key=api_key or "", base_url=base_url, provider_type=provider.provider_type, headers=dict(provider.headers_json or {}), extra=dict(provider.extra_json or {}), dimension=model.get("dimension"), batch_size=model.get("batch_size", 40), ) new_cache[info.spec] = info self._save_cache(new_cache) self._invalidate_local() logger.info(f"Model cache rebuilt: {len(new_cache)} models → Redis") def _save_cache(self, cache: dict[str, ModelInfo]) -> None: """将完整缓存写入 Redis。""" r = self._get_redis() if r is None: logger.warning("Redis unavailable, cache not saved") return try: data = {spec: info.to_dict() for spec, info in cache.items()} r.set(REDIS_CACHE_KEY, json.dumps(data, ensure_ascii=False)) except Exception as e: logger.error(f"Failed to save model cache to Redis: {e}") self._reset_redis() @staticmethod def _get_base_url_for_type(provider: Any, model_type: str) -> str: """根据模型类型获取对应的 base_url。""" if model_type == "embedding" and provider.embedding_base_url: return provider.embedding_base_url if model_type == "rerank" and provider.rerank_base_url: return provider.rerank_base_url return provider.base_url # 全局单例 model_cache = ModelCache()