221 lines
7.3 KiB
Python
221 lines
7.3 KiB
Python
"""模型缓存服务 - 基于 Redis 的跨进程模型信息缓存。
|
||
|
||
本模块将数据库中的 model_providers 表数据序列化到 Redis，
|
||
供 API 和 Worker 等多进程同步读取，避免在同步函数中查询异步数据库。
|
||
|
||
模型 spec 格式: provider_id:model_id（冒号分隔）。model_id 允许包含斜杠。
|
||
"""
|
||
|
||
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"
|
||
_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 / 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:
|
||
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，保证跨进程一致。"""
|
||
|
||
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):
|
||
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:
|
||
self._redis = None
|
||
|
||
def _load_cache(self) -> dict[str, ModelInfo]:
|
||
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:
|
||
self._local_cache = None
|
||
self._local_cache_at = 0.0
|
||
|
||
def get_model_info(self, spec: str) -> ModelInfo | None:
|
||
cache = self._load_cache()
|
||
return cache.get(spec)
|
||
|
||
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]]:
|
||
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:
|
||
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:
|
||
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:
|
||
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()
|
||
|
||
|
||
def resolve_model_spec(spec: str) -> ModelInfo:
|
||
"""根据 spec 返回 ModelInfo。"""
|
||
if not spec:
|
||
raise ValueError("model spec 不能为空")
|
||
|
||
info = model_cache.get_model_info(spec)
|
||
if info:
|
||
return info
|
||
|
||
all_specs = model_cache.get_all_specs()
|
||
available = [item.spec for item in all_specs[:10]]
|
||
raise ValueError(f"未找到模型: '{spec}'。可用模型 ({len(all_specs)}): {available}")
|