feat(api): cache workflow provider configurations (#37980)

This commit is contained in:
林玮 (Jade Lin)
2026-06-26 08:04:37 +00:00
committed by GitHub
parent 3f2ef24755
commit 1dbda1463e
5 changed files with 1004 additions and 97 deletions
+114 -15
View File
@@ -22,7 +22,10 @@ from core.entities.provider_entities import (
SystemConfigurationStatus,
)
from core.helper import encrypter
from core.helper.model_provider_cache import ProviderCredentialsCache, ProviderCredentialsCacheType
from core.helper.model_provider_cache import (
ProviderCredentialsCache,
ProviderCredentialsCacheType,
)
from core.plugin.impl.model_runtime_factory import create_model_type_instance, create_plugin_model_assembly
from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType
from graphon.model_runtime.entities.provider_entities import (
@@ -473,6 +476,39 @@ class ProviderConfiguration(BaseModel):
provider_names.append(model_provider_id.provider_name)
return provider_names
def _invalidate_provider_configuration_cache(
self,
*,
provider_models: bool = False,
preferred_model_providers: bool = False,
provider_model_settings: bool = False,
provider_model_credentials: bool = False,
provider_credentials: bool = False,
provider_load_balancing_configs: bool = False,
) -> None:
"""Invalidate tenant-scoped provider snapshots after committing configuration writes."""
from core.provider_manager import ProviderConfigurationCacheSource, ProviderManager
sources: list[ProviderConfigurationCacheSource] = []
if provider_models:
sources.append(ProviderConfigurationCacheSource.PROVIDER_MODELS)
if preferred_model_providers:
sources.append(ProviderConfigurationCacheSource.PREFERRED_MODEL_PROVIDERS)
if provider_model_settings:
sources.append(ProviderConfigurationCacheSource.PROVIDER_MODEL_SETTINGS)
if provider_model_credentials:
sources.append(ProviderConfigurationCacheSource.PROVIDER_MODEL_CREDENTIALS)
if provider_credentials:
sources.append(ProviderConfigurationCacheSource.PROVIDER_CREDENTIALS)
if provider_load_balancing_configs:
sources.append(ProviderConfigurationCacheSource.PROVIDER_LOAD_BALANCING_CONFIGS)
if not sources:
logger.warning("No provider configuration cache source selected for invalidation")
return
ProviderManager.invalidate_configurations_cache(self.tenant_id, sources=sources)
def create_provider_credential(self, credentials: dict[str, Any], credential_name: str | None):
"""
Add custom provider credentials.
@@ -489,6 +525,7 @@ class ProviderConfiguration(BaseModel):
credentials = self.validate_provider_credentials(credentials=credentials)
preferred_model_providers_changed = False
with Session(db.engine) as session:
provider_record = self._get_provider_record(session)
try:
@@ -518,7 +555,9 @@ class ProviderConfiguration(BaseModel):
)
provider_model_credentials_cache.delete()
self.switch_preferred_provider_type(provider_type=ProviderType.CUSTOM, session=session)
preferred_model_providers_changed = self.switch_preferred_provider_type(
provider_type=ProviderType.CUSTOM, session=session
)
else:
provider_record.is_valid = True
@@ -533,12 +572,18 @@ class ProviderConfiguration(BaseModel):
)
provider_model_credentials_cache.delete()
self.switch_preferred_provider_type(provider_type=ProviderType.CUSTOM, session=session)
preferred_model_providers_changed = self.switch_preferred_provider_type(
provider_type=ProviderType.CUSTOM, session=session
)
session.commit()
except Exception:
session.rollback()
raise
self._invalidate_provider_configuration_cache(
preferred_model_providers=preferred_model_providers_changed,
provider_credentials=True,
)
def update_provider_credential(
self,
@@ -562,6 +607,7 @@ class ProviderConfiguration(BaseModel):
credentials = self.validate_provider_credentials(credentials=credentials, credential_id=credential_id)
load_balancing_configs_changed = False
with Session(db.engine) as session:
provider_record = self._get_provider_record(session)
stmt = select(ProviderCredential).where(
@@ -588,7 +634,7 @@ class ProviderConfiguration(BaseModel):
)
provider_model_credentials_cache.delete()
self._update_load_balancing_configs_with_credential(
load_balancing_configs_changed = self._update_load_balancing_configs_with_credential(
credential_id=credential_id,
credential_record=credential_record,
credential_source=CredentialSourceType.PROVIDER,
@@ -597,6 +643,10 @@ class ProviderConfiguration(BaseModel):
except Exception:
session.rollback()
raise
self._invalidate_provider_configuration_cache(
provider_credentials=True,
provider_load_balancing_configs=load_balancing_configs_changed,
)
def _update_load_balancing_configs_with_credential(
self,
@@ -604,7 +654,7 @@ class ProviderConfiguration(BaseModel):
credential_record: ProviderCredential | ProviderModelCredential,
credential_source: str,
session: Session,
):
) -> bool:
"""
Update load balancing configurations that reference the given credential_id.
@@ -625,7 +675,7 @@ class ProviderConfiguration(BaseModel):
load_balancing_configs = session.execute(stmt).scalars().all()
if not load_balancing_configs:
return
return False
# Update each load balancing config with the new credentials
for lb_config in load_balancing_configs:
@@ -643,6 +693,7 @@ class ProviderConfiguration(BaseModel):
lb_credentials_cache.delete()
session.commit()
return True
def delete_provider_credential(self, credential_id: str):
"""
@@ -651,6 +702,8 @@ class ProviderConfiguration(BaseModel):
:param credential_id: credential id
:return:
"""
preferred_model_providers_changed = False
load_balancing_configs_changed = False
with Session(db.engine) as session:
stmt = select(ProviderCredential).where(
ProviderCredential.id == credential_id,
@@ -671,6 +724,7 @@ class ProviderConfiguration(BaseModel):
LoadBalancingModelConfig.credential_source_type == CredentialSourceType.PROVIDER,
)
lb_configs_using_credential = session.execute(lb_stmt).scalars().all()
load_balancing_configs_changed = bool(lb_configs_using_credential)
try:
for lb_config in lb_configs_using_credential:
lb_credentials_cache = ProviderCredentialsCache(
@@ -703,7 +757,9 @@ class ProviderConfiguration(BaseModel):
cache_type=ProviderCredentialsCacheType.PROVIDER,
)
provider_model_credentials_cache.delete()
self.switch_preferred_provider_type(provider_type=ProviderType.SYSTEM, session=session)
preferred_model_providers_changed = self.switch_preferred_provider_type(
provider_type=ProviderType.SYSTEM, session=session
)
elif provider_record and provider_record.credential_id == credential_id:
provider_record.credential_id = None
provider_record.updated_at = naive_utc_now()
@@ -714,12 +770,19 @@ class ProviderConfiguration(BaseModel):
cache_type=ProviderCredentialsCacheType.PROVIDER,
)
provider_model_credentials_cache.delete()
self.switch_preferred_provider_type(provider_type=ProviderType.SYSTEM, session=session)
preferred_model_providers_changed = self.switch_preferred_provider_type(
provider_type=ProviderType.SYSTEM, session=session
)
session.commit()
except Exception:
session.rollback()
raise
self._invalidate_provider_configuration_cache(
preferred_model_providers=preferred_model_providers_changed,
provider_credentials=True,
provider_load_balancing_configs=load_balancing_configs_changed,
)
def switch_active_provider_credential(self, credential_id: str):
"""
@@ -728,6 +791,7 @@ class ProviderConfiguration(BaseModel):
:param credential_id: credential id
:return:
"""
preferred_model_providers_changed = False
with Session(db.engine) as session:
stmt = select(ProviderCredential).where(
ProviderCredential.id == credential_id,
@@ -753,10 +817,14 @@ class ProviderConfiguration(BaseModel):
cache_type=ProviderCredentialsCacheType.PROVIDER,
)
provider_model_credentials_cache.delete()
self.switch_preferred_provider_type(ProviderType.CUSTOM, session=session)
preferred_model_providers_changed = self.switch_preferred_provider_type(
ProviderType.CUSTOM, session=session
)
except Exception:
session.rollback()
raise
if preferred_model_providers_changed:
self._invalidate_provider_configuration_cache(preferred_model_providers=True)
def _get_custom_model_record(
self,
@@ -1017,6 +1085,10 @@ class ProviderConfiguration(BaseModel):
except Exception:
session.rollback()
raise
self._invalidate_provider_configuration_cache(
provider_models=True,
provider_model_credentials=True,
)
def update_custom_model_credential(
self,
@@ -1053,6 +1125,7 @@ class ProviderConfiguration(BaseModel):
credential_id=credential_id,
)
load_balancing_configs_changed = False
with Session(db.engine) as session:
provider_model_record = self._get_custom_model_record(model_type=model_type, model=model, session=session)
@@ -1082,7 +1155,7 @@ class ProviderConfiguration(BaseModel):
)
provider_model_credentials_cache.delete()
self._update_load_balancing_configs_with_credential(
load_balancing_configs_changed = self._update_load_balancing_configs_with_credential(
credential_id=credential_id,
credential_record=credential_record,
credential_source=CredentialSourceType.CUSTOM_MODEL,
@@ -1091,6 +1164,11 @@ class ProviderConfiguration(BaseModel):
except Exception:
session.rollback()
raise
self._invalidate_provider_configuration_cache(
provider_models=True,
provider_model_credentials=True,
provider_load_balancing_configs=load_balancing_configs_changed,
)
def delete_custom_model_credential(self, model_type: ModelType, model: str, credential_id: str):
"""
@@ -1099,6 +1177,7 @@ class ProviderConfiguration(BaseModel):
:param credential_id: credential id
:return:
"""
load_balancing_configs_changed = False
with Session(db.engine) as session:
stmt = select(ProviderModelCredential).where(
ProviderModelCredential.id == credential_id,
@@ -1118,6 +1197,7 @@ class ProviderConfiguration(BaseModel):
LoadBalancingModelConfig.credential_source_type == CredentialSourceType.CUSTOM_MODEL,
)
lb_configs_using_credential = session.execute(lb_stmt).scalars().all()
load_balancing_configs_changed = bool(lb_configs_using_credential)
try:
for lb_config in lb_configs_using_credential:
@@ -1161,6 +1241,11 @@ class ProviderConfiguration(BaseModel):
except Exception:
session.rollback()
raise
self._invalidate_provider_configuration_cache(
provider_models=True,
provider_model_credentials=True,
provider_load_balancing_configs=load_balancing_configs_changed,
)
def add_model_credential_to_model(self, model_type: ModelType, model: str, credential_id: str):
"""
@@ -1213,6 +1298,7 @@ class ProviderConfiguration(BaseModel):
session.add(provider_model_record)
session.commit()
self._invalidate_provider_configuration_cache(provider_models=True)
def switch_custom_model_credential(self, model_type: ModelType, model: str, credential_id: str):
"""
@@ -1251,6 +1337,7 @@ class ProviderConfiguration(BaseModel):
cache_type=ProviderCredentialsCacheType.MODEL,
)
provider_model_credentials_cache.delete()
self._invalidate_provider_configuration_cache(provider_models=True)
def delete_custom_model(self, model_type: ModelType, model: str):
"""
@@ -1259,6 +1346,7 @@ class ProviderConfiguration(BaseModel):
:param model: model name
:return:
"""
provider_models_changed = False
with Session(db.engine) as session:
# get provider model
provider_model_record = self._get_custom_model_record(model_type=model_type, model=model, session=session)
@@ -1267,6 +1355,7 @@ class ProviderConfiguration(BaseModel):
if provider_model_record:
session.delete(provider_model_record)
session.commit()
provider_models_changed = True
provider_model_credentials_cache = ProviderCredentialsCache(
tenant_id=self.tenant_id,
@@ -1275,6 +1364,8 @@ class ProviderConfiguration(BaseModel):
)
provider_model_credentials_cache.delete()
if provider_models_changed:
self._invalidate_provider_configuration_cache(provider_models=True)
def _get_provider_model_setting(
self, model_type: ModelType, model: str, session: Session
@@ -1314,6 +1405,7 @@ class ProviderConfiguration(BaseModel):
)
session.add(model_setting)
session.commit()
self._invalidate_provider_configuration_cache(provider_model_settings=True)
return model_setting
@@ -1340,6 +1432,7 @@ class ProviderConfiguration(BaseModel):
)
session.add(model_setting)
session.commit()
self._invalidate_provider_configuration_cache(provider_model_settings=True)
return model_setting
@@ -1392,6 +1485,7 @@ class ProviderConfiguration(BaseModel):
)
session.add(model_setting)
session.commit()
self._invalidate_provider_configuration_cache(provider_model_settings=True)
return model_setting
@@ -1419,6 +1513,7 @@ class ProviderConfiguration(BaseModel):
)
session.add(model_setting)
session.commit()
self._invalidate_provider_configuration_cache(provider_model_settings=True)
return model_setting
@@ -1454,19 +1549,19 @@ class ProviderConfiguration(BaseModel):
credentials=credentials or {},
)
def switch_preferred_provider_type(self, provider_type: ProviderType, session: Session | None = None):
def switch_preferred_provider_type(self, provider_type: ProviderType, session: Session | None = None) -> bool:
"""
Switch preferred provider type.
:param provider_type:
:return:
"""
if provider_type == self.preferred_provider_type:
return
return False
if provider_type == ProviderType.SYSTEM and not self.system_configuration.enabled:
return
return False
def _switch(s: Session):
def _switch(s: Session) -> bool:
stmt = select(TenantPreferredModelProvider).where(
TenantPreferredModelProvider.tenant_id == self.tenant_id,
TenantPreferredModelProvider.provider_name.in_(self._get_provider_names()),
@@ -1483,12 +1578,16 @@ class ProviderConfiguration(BaseModel):
)
s.add(preferred_model_provider)
s.commit()
return True
if session:
return _switch(session)
else:
with Session(db.engine) as session:
return _switch(session)
changed = _switch(session)
if changed:
self._invalidate_provider_configuration_cache(preferred_model_providers=True)
return changed
def extract_secret_variables(self, credential_form_schemas: list[CredentialFormSchema]) -> list[str]:
"""
+620 -55
View File
@@ -1,10 +1,14 @@
from __future__ import annotations
import contextlib
import json
import logging
from collections import defaultdict
from collections.abc import Sequence
from collections.abc import Callable, Sequence
from dataclasses import asdict, dataclass
from enum import StrEnum
from json import JSONDecodeError
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, Protocol, Self
from pydantic import TypeAdapter
from sqlalchemy import select
@@ -41,6 +45,7 @@ from graphon.model_runtime.entities.provider_entities import (
ProviderEntity,
)
from graphon.model_runtime.model_providers.model_provider_factory import ModelProviderFactory
from models.enums import CredentialSourceType
from models.provider import (
LoadBalancingModelConfig,
Provider,
@@ -59,7 +64,494 @@ if TYPE_CHECKING:
from graphon.model_runtime.protocols.runtime import ModelRuntime
from models.account import Account
logger = logging.getLogger(__name__)
_credentials_adapter: TypeAdapter[dict[str, Any]] = TypeAdapter(dict[str, Any])
_PROVIDER_CONFIGURATION_CACHE_TTL_SECONDS = 300
_PROVIDER_CONFIGURATION_CACHE_VERSION_TTL_SECONDS = 360
_PROVIDER_CONFIGURATION_CACHE_VERSION_KEY = "provider_configurations:tenant:{tenant_id}:source:{source}:version"
_PROVIDER_CONFIGURATION_CACHE_SOURCE_KEY = "provider_configurations:tenant:{tenant_id}:source:{source}:v:{version}"
class ProviderConfigurationCacheSource(StrEnum):
PROVIDER_MODELS = "provider_models"
PREFERRED_MODEL_PROVIDERS = "preferred_model_providers"
PROVIDER_MODEL_SETTINGS = "provider_model_settings"
PROVIDER_MODEL_CREDENTIALS = "provider_model_credentials"
PROVIDER_CREDENTIALS = "provider_credentials"
PROVIDER_LOAD_BALANCING_CONFIGS = "provider_load_balancing_configs"
_PROVIDER_CONFIGURATION_SOURCES = tuple(ProviderConfigurationCacheSource)
class _CacheEntry(Protocol):
@classmethod
def from_cache_row(cls, row: dict[str, Any]) -> Self: ...
def to_cache_row(self) -> dict[str, Any]: ...
@dataclass(frozen=True, slots=True)
class _ProviderConfigurationCacheSourceSpec[T: _CacheEntry]:
name: ProviderConfigurationCacheSource
entry_cls: type[T]
load_records: Callable[[str], list[T]]
@dataclass(frozen=True, slots=True)
class _ProviderModelCacheEntry:
id: str
provider_name: str
model_name: str
model_type: ModelType
credential_id: str | None
credential_name: str | None
encrypted_config: str | None
@classmethod
def from_record(cls, record: ProviderModel) -> _ProviderModelCacheEntry:
credential = record.__dict__.get("credential")
return cls(
id=record.id,
provider_name=record.provider_name,
model_name=record.model_name,
model_type=record.model_type,
credential_id=record.credential_id,
credential_name=credential.credential_name if credential else None,
encrypted_config=credential.encrypted_config if credential else None,
)
@classmethod
def from_cache_row(cls, row: dict[str, Any]) -> _ProviderModelCacheEntry:
return cls(
id=row["id"],
provider_name=row["provider_name"],
model_name=row["model_name"],
model_type=ModelType(row["model_type"]),
credential_id=row.get("credential_id"),
credential_name=row.get("credential_name"),
encrypted_config=row.get("encrypted_config"),
)
def to_cache_row(self) -> dict[str, Any]:
row = asdict(self)
row["model_type"] = self.model_type.value
return row
@dataclass(frozen=True, slots=True)
class _TenantPreferredModelProviderCacheEntry:
provider_name: str
preferred_provider_type: ProviderType
@classmethod
def from_record(cls, record: TenantPreferredModelProvider) -> _TenantPreferredModelProviderCacheEntry:
return cls(
provider_name=record.provider_name,
preferred_provider_type=record.preferred_provider_type,
)
@classmethod
def from_cache_row(cls, row: dict[str, Any]) -> _TenantPreferredModelProviderCacheEntry:
return cls(
provider_name=row["provider_name"],
preferred_provider_type=ProviderType(row["preferred_provider_type"]),
)
def to_cache_row(self) -> dict[str, Any]:
return {
"provider_name": self.provider_name,
"preferred_provider_type": self.preferred_provider_type.value,
}
@dataclass(frozen=True, slots=True)
class _ProviderModelSettingCacheEntry:
provider_name: str
model_name: str
model_type: ModelType
enabled: bool
load_balancing_enabled: bool
@classmethod
def from_record(cls, record: ProviderModelSetting) -> _ProviderModelSettingCacheEntry:
return cls(
provider_name=record.provider_name,
model_name=record.model_name,
model_type=record.model_type,
enabled=record.enabled,
load_balancing_enabled=record.load_balancing_enabled,
)
@classmethod
def from_cache_row(cls, row: dict[str, Any]) -> _ProviderModelSettingCacheEntry:
return cls(
provider_name=row["provider_name"],
model_name=row["model_name"],
model_type=ModelType(row["model_type"]),
enabled=row["enabled"],
load_balancing_enabled=row["load_balancing_enabled"],
)
def to_cache_row(self) -> dict[str, Any]:
return {
"provider_name": self.provider_name,
"model_name": self.model_name,
"model_type": self.model_type.value,
"enabled": self.enabled,
"load_balancing_enabled": self.load_balancing_enabled,
}
@dataclass(frozen=True, slots=True)
class _ProviderModelCredentialCacheEntry:
id: str
provider_name: str
model_name: str
model_type: ModelType
credential_name: str
@classmethod
def from_record(cls, record: ProviderModelCredential) -> _ProviderModelCredentialCacheEntry:
return cls(
id=record.id,
provider_name=record.provider_name,
model_name=record.model_name,
model_type=record.model_type,
credential_name=record.credential_name,
)
@classmethod
def from_cache_row(cls, row: dict[str, Any]) -> _ProviderModelCredentialCacheEntry:
return cls(
id=row["id"],
provider_name=row["provider_name"],
model_name=row["model_name"],
model_type=ModelType(row["model_type"]),
credential_name=row["credential_name"],
)
def to_cache_row(self) -> dict[str, Any]:
return {
"id": self.id,
"provider_name": self.provider_name,
"model_name": self.model_name,
"model_type": self.model_type.value,
"credential_name": self.credential_name,
}
@dataclass(frozen=True, slots=True)
class _ProviderCredentialCacheEntry:
id: str
provider_name: str
credential_name: str
@classmethod
def from_record(cls, record: ProviderCredential) -> _ProviderCredentialCacheEntry:
return cls(
id=record.id,
provider_name=record.provider_name,
credential_name=record.credential_name,
)
@classmethod
def from_cache_row(cls, row: dict[str, Any]) -> _ProviderCredentialCacheEntry:
return cls(
id=row["id"],
provider_name=row["provider_name"],
credential_name=row["credential_name"],
)
def to_cache_row(self) -> dict[str, Any]:
return asdict(self)
@dataclass(frozen=True, slots=True)
class _LoadBalancingModelConfigCacheEntry:
id: str
tenant_id: str
provider_name: str
model_name: str
model_type: ModelType
name: str
encrypted_config: str | None
credential_id: str | None
credential_source_type: CredentialSourceType | None
enabled: bool
@classmethod
def from_record(cls, record: LoadBalancingModelConfig) -> _LoadBalancingModelConfigCacheEntry:
return cls(
id=record.id,
tenant_id=record.tenant_id,
provider_name=record.provider_name,
model_name=record.model_name,
model_type=record.model_type,
name=record.name,
encrypted_config=record.encrypted_config,
credential_id=record.credential_id,
credential_source_type=record.credential_source_type,
enabled=record.enabled,
)
@classmethod
def from_cache_row(cls, row: dict[str, Any]) -> _LoadBalancingModelConfigCacheEntry:
return cls(
id=row["id"],
tenant_id=row["tenant_id"],
provider_name=row["provider_name"],
model_name=row["model_name"],
model_type=ModelType(row["model_type"]),
name=row["name"],
encrypted_config=row.get("encrypted_config"),
credential_id=row.get("credential_id"),
credential_source_type=CredentialSourceType(row["credential_source_type"])
if row.get("credential_source_type")
else None,
enabled=row["enabled"],
)
def to_cache_row(self) -> dict[str, Any]:
return {
"id": self.id,
"tenant_id": self.tenant_id,
"provider_name": self.provider_name,
"model_name": self.model_name,
"model_type": self.model_type.value,
"name": self.name,
"encrypted_config": self.encrypted_config,
"credential_id": self.credential_id,
"credential_source_type": self.credential_source_type.value if self.credential_source_type else None,
"enabled": self.enabled,
}
class _ProviderConfigurationSourceCache:
"""Redis-backed cache for tenant provider DB cache entries.
The assembled ``ProviderConfigurations`` object is intentionally not cached
here because it carries request-scoped runtime bindings. Cache only the DB
rows that are stable enough to reuse across processes, then let each
``ProviderManager`` assemble and bind fresh runtime-aware entities.
"""
@classmethod
def get_records[T: _CacheEntry](
cls,
*,
tenant_id: str,
source: ProviderConfigurationCacheSource,
entry_cls: type[T],
) -> tuple[list[T] | None, str | None]:
version: str | None = None
try:
version = cls._get_version(tenant_id=tenant_id, source=source)
cache_key = cls._source_key(tenant_id=tenant_id, source=source, version=version)
cached_records = redis_client.get(cache_key)
if cached_records is None:
return None, version
cached_text = cached_records.decode("utf-8") if isinstance(cached_records, bytes) else cached_records
rows = json.loads(cached_text)
if not isinstance(rows, list):
return None, version
return [entry_cls.from_cache_row(row) for row in rows if isinstance(row, dict)], version
except Exception:
logger.warning("Failed to read provider configuration source cache", exc_info=True)
return None, version
@classmethod
def set_records(
cls,
*,
tenant_id: str,
source: ProviderConfigurationCacheSource,
records: Sequence[_CacheEntry],
expected_version: str | None = None,
) -> None:
try:
version = cls._get_version(tenant_id=tenant_id, source=source)
if expected_version is not None and version != expected_version:
return
cache_key = cls._source_key(tenant_id=tenant_id, source=source, version=version)
rows = [record.to_cache_row() for record in records]
redis_client.setex(cache_key, _PROVIDER_CONFIGURATION_CACHE_TTL_SECONDS, json.dumps(rows))
except Exception:
logger.warning("Failed to write provider configuration source cache", exc_info=True)
@classmethod
def invalidate_tenant(
cls,
tenant_id: str,
sources: Sequence[ProviderConfigurationCacheSource] | None = None,
) -> None:
try:
if sources is None:
sources = _PROVIDER_CONFIGURATION_SOURCES
for source in sources:
version_key = _PROVIDER_CONFIGURATION_CACHE_VERSION_KEY.format(tenant_id=tenant_id, source=source.value)
redis_client.incr(version_key)
redis_client.expire(version_key, _PROVIDER_CONFIGURATION_CACHE_VERSION_TTL_SECONDS)
except Exception:
logger.warning("Failed to invalidate provider configuration source cache", exc_info=True)
@classmethod
def _get_version(cls, *, tenant_id: str, source: ProviderConfigurationCacheSource) -> str:
version_key = _PROVIDER_CONFIGURATION_CACHE_VERSION_KEY.format(tenant_id=tenant_id, source=source.value)
version = redis_client.get(version_key)
if version is None:
redis_client.set(version_key, "0", ex=_PROVIDER_CONFIGURATION_CACHE_VERSION_TTL_SECONDS)
return "0"
redis_client.expire(version_key, _PROVIDER_CONFIGURATION_CACHE_VERSION_TTL_SECONDS)
return version.decode("utf-8") if isinstance(version, bytes) else str(version)
@staticmethod
def _source_key(*, tenant_id: str, source: ProviderConfigurationCacheSource, version: str) -> str:
return _PROVIDER_CONFIGURATION_CACHE_SOURCE_KEY.format(
tenant_id=tenant_id,
source=source.value,
version=version,
)
def _get_cached_or_load_records[T: _CacheEntry](
*,
tenant_id: str,
cache_source: _ProviderConfigurationCacheSourceSpec[T],
) -> list[T]:
cached_records, cache_version = _ProviderConfigurationSourceCache.get_records(
tenant_id=tenant_id,
source=cache_source.name,
entry_cls=cache_source.entry_cls,
)
if cached_records is not None:
return cached_records
records = cache_source.load_records(tenant_id)
_ProviderConfigurationSourceCache.set_records(
tenant_id=tenant_id,
source=cache_source.name,
records=records,
expected_version=cache_version,
)
return records
def _attach_active_credentials(
*,
session: Any,
records: Sequence[Provider | ProviderModel],
credential_model_cls: type[ProviderCredential | ProviderModelCredential],
) -> None:
credential_ids = [record.credential_id for record in records if getattr(record, "credential_id", None)]
if not credential_ids:
return
credentials = session.scalars(select(credential_model_cls).where(credential_model_cls.id.in_(credential_ids))).all()
credential_by_id = {credential.id: credential for credential in credentials}
for record in records:
if getattr(record, "credential_id", None):
record.__dict__["credential"] = credential_by_id.get(record.credential_id)
def _load_provider_model_cache_entries(tenant_id: str) -> list[_ProviderModelCacheEntry]:
with session_factory.create_session() as session:
stmt = select(ProviderModel).where(ProviderModel.tenant_id == tenant_id, ProviderModel.is_valid == True)
provider_models = list(session.scalars(stmt))
_attach_active_credentials(
session=session,
records=provider_models,
credential_model_cls=ProviderModelCredential,
)
return [_ProviderModelCacheEntry.from_record(provider_model) for provider_model in provider_models]
def _load_preferred_model_provider_cache_entries(tenant_id: str) -> list[_TenantPreferredModelProviderCacheEntry]:
with session_factory.create_session() as session:
stmt = select(TenantPreferredModelProvider).where(TenantPreferredModelProvider.tenant_id == tenant_id)
return [
_TenantPreferredModelProviderCacheEntry.from_record(preferred_model_provider)
for preferred_model_provider in session.scalars(stmt)
]
def _load_provider_model_setting_cache_entries(tenant_id: str) -> list[_ProviderModelSettingCacheEntry]:
with session_factory.create_session() as session:
stmt = select(ProviderModelSetting).where(ProviderModelSetting.tenant_id == tenant_id)
return [
_ProviderModelSettingCacheEntry.from_record(provider_model_setting)
for provider_model_setting in session.scalars(stmt)
]
def _load_provider_model_credential_cache_entries(tenant_id: str) -> list[_ProviderModelCredentialCacheEntry]:
with session_factory.create_session() as session:
stmt = (
select(ProviderModelCredential)
.where(ProviderModelCredential.tenant_id == tenant_id)
.order_by(ProviderModelCredential.created_at.desc())
)
return [
_ProviderModelCredentialCacheEntry.from_record(provider_model_credential)
for provider_model_credential in session.scalars(stmt)
]
def _load_provider_credential_cache_entries(tenant_id: str) -> list[_ProviderCredentialCacheEntry]:
with session_factory.create_session() as session:
stmt = (
select(ProviderCredential)
.where(ProviderCredential.tenant_id == tenant_id)
.order_by(ProviderCredential.created_at.desc())
)
return [
_ProviderCredentialCacheEntry.from_record(provider_credential)
for provider_credential in session.scalars(stmt)
]
def _load_provider_load_balancing_config_cache_entries(tenant_id: str) -> list[_LoadBalancingModelConfigCacheEntry]:
with session_factory.create_session() as session:
stmt = select(LoadBalancingModelConfig).where(LoadBalancingModelConfig.tenant_id == tenant_id)
return [
_LoadBalancingModelConfigCacheEntry.from_record(load_balancing_model_config)
for load_balancing_model_config in session.scalars(stmt)
]
_PROVIDER_MODELS_CACHE_SOURCE = _ProviderConfigurationCacheSourceSpec(
name=ProviderConfigurationCacheSource.PROVIDER_MODELS,
entry_cls=_ProviderModelCacheEntry,
load_records=_load_provider_model_cache_entries,
)
_PREFERRED_MODEL_PROVIDERS_CACHE_SOURCE = _ProviderConfigurationCacheSourceSpec(
name=ProviderConfigurationCacheSource.PREFERRED_MODEL_PROVIDERS,
entry_cls=_TenantPreferredModelProviderCacheEntry,
load_records=_load_preferred_model_provider_cache_entries,
)
_PROVIDER_MODEL_SETTINGS_CACHE_SOURCE = _ProviderConfigurationCacheSourceSpec(
name=ProviderConfigurationCacheSource.PROVIDER_MODEL_SETTINGS,
entry_cls=_ProviderModelSettingCacheEntry,
load_records=_load_provider_model_setting_cache_entries,
)
_PROVIDER_MODEL_CREDENTIALS_CACHE_SOURCE = _ProviderConfigurationCacheSourceSpec(
name=ProviderConfigurationCacheSource.PROVIDER_MODEL_CREDENTIALS,
entry_cls=_ProviderModelCredentialCacheEntry,
load_records=_load_provider_model_credential_cache_entries,
)
_PROVIDER_CREDENTIALS_CACHE_SOURCE = _ProviderConfigurationCacheSourceSpec(
name=ProviderConfigurationCacheSource.PROVIDER_CREDENTIALS,
entry_cls=_ProviderCredentialCacheEntry,
load_records=_load_provider_credential_cache_entries,
)
_PROVIDER_LOAD_BALANCING_CONFIGS_CACHE_SOURCE = _ProviderConfigurationCacheSourceSpec(
name=ProviderConfigurationCacheSource.PROVIDER_LOAD_BALANCING_CONFIGS,
entry_cls=_LoadBalancingModelConfigCacheEntry,
load_records=_load_provider_load_balancing_config_cache_entries,
)
class ProviderManager:
@@ -98,6 +590,14 @@ class ProviderManager:
self._configurations_cache.pop(tenant_id, None)
@staticmethod
def invalidate_configurations_cache(
tenant_id: str,
sources: Sequence[ProviderConfigurationCacheSource] | None = None,
) -> None:
"""Invalidate cross-process provider configuration source cache for a tenant."""
_ProviderConfigurationSourceCache.invalidate_tenant(tenant_id, sources=sources)
def get_configurations(self, tenant_id: str) -> ProviderConfigurations:
"""
Get model provider configurations.
@@ -192,6 +692,9 @@ class ProviderManager:
# Get All provider model credentials
provider_name_to_provider_model_credentials_dict = self._get_all_provider_model_credentials(tenant_id)
# Get All provider credentials
provider_name_to_provider_credentials_dict = self._get_all_provider_credentials(tenant_id)
provider_configurations = ProviderConfigurations(tenant_id=tenant_id)
# Construct ProviderConfiguration objects for each provider
@@ -224,7 +727,12 @@ class ProviderManager:
# Convert to custom configuration
custom_configuration = self._to_custom_configuration(
tenant_id, provider_entity, provider_records, provider_model_records, provider_model_credentials
tenant_id,
provider_entity,
provider_records,
provider_model_records,
provider_model_credentials,
provider_name_to_provider_credentials_dict,
)
# Convert to system configuration
@@ -448,84 +956,115 @@ class ProviderManager:
provider_name_to_provider_records_dict = defaultdict(list)
with session_factory.create_session() as session:
stmt = select(Provider).where(Provider.tenant_id == tenant_id, Provider.is_valid == True)
providers = session.scalars(stmt)
providers = list(session.scalars(stmt))
_attach_active_credentials(
session=session,
records=providers,
credential_model_cls=ProviderCredential,
)
for provider in providers:
# Use provider name with prefix after the data migration
provider_name_to_provider_records_dict[str(ModelProviderID(provider.provider_name))].append(provider)
return provider_name_to_provider_records_dict
@staticmethod
def _get_all_provider_models(tenant_id: str) -> dict[str, list[ProviderModel]]:
def _get_all_provider_models(tenant_id: str) -> dict[str, list[_ProviderModelCacheEntry]]:
"""
Get all provider model records of the workspace.
:param tenant_id: workspace id
:return:
"""
provider_models = _get_cached_or_load_records(
tenant_id=tenant_id,
cache_source=_PROVIDER_MODELS_CACHE_SOURCE,
)
provider_name_to_provider_model_records_dict = defaultdict(list)
with session_factory.create_session() as session:
stmt = select(ProviderModel).where(ProviderModel.tenant_id == tenant_id, ProviderModel.is_valid == True)
provider_models = session.scalars(stmt)
for provider_model in provider_models:
provider_name_to_provider_model_records_dict[provider_model.provider_name].append(provider_model)
for provider_model in provider_models:
provider_name_to_provider_model_records_dict[provider_model.provider_name].append(provider_model)
return provider_name_to_provider_model_records_dict
@staticmethod
def _get_all_preferred_model_providers(tenant_id: str) -> dict[str, TenantPreferredModelProvider]:
def _get_all_preferred_model_providers(tenant_id: str) -> dict[str, _TenantPreferredModelProviderCacheEntry]:
"""
Get All preferred provider types of the workspace.
:param tenant_id: workspace id
:return:
"""
provider_name_to_preferred_provider_type_records_dict = {}
with session_factory.create_session() as session:
stmt = select(TenantPreferredModelProvider).where(TenantPreferredModelProvider.tenant_id == tenant_id)
preferred_provider_types = session.scalars(stmt)
provider_name_to_preferred_provider_type_records_dict = {
preferred_provider_type.provider_name: preferred_provider_type
for preferred_provider_type in preferred_provider_types
}
return provider_name_to_preferred_provider_type_records_dict
preferred_provider_types = _get_cached_or_load_records(
tenant_id=tenant_id,
cache_source=_PREFERRED_MODEL_PROVIDERS_CACHE_SOURCE,
)
return {
preferred_provider_type.provider_name: preferred_provider_type
for preferred_provider_type in preferred_provider_types
}
@staticmethod
def _get_all_provider_model_settings(tenant_id: str) -> dict[str, list[ProviderModelSetting]]:
def _get_all_provider_model_settings(tenant_id: str) -> dict[str, list[_ProviderModelSettingCacheEntry]]:
"""
Get All provider model settings of the workspace.
:param tenant_id: workspace id
:return:
"""
provider_model_settings = _get_cached_or_load_records(
tenant_id=tenant_id,
cache_source=_PROVIDER_MODEL_SETTINGS_CACHE_SOURCE,
)
provider_name_to_provider_model_settings_dict = defaultdict(list)
with session_factory.create_session() as session:
stmt = select(ProviderModelSetting).where(ProviderModelSetting.tenant_id == tenant_id)
provider_model_settings = session.scalars(stmt)
for provider_model_setting in provider_model_settings:
provider_name_to_provider_model_settings_dict[provider_model_setting.provider_name].append(
provider_model_setting
)
for provider_model_setting in provider_model_settings:
provider_name_to_provider_model_settings_dict[provider_model_setting.provider_name].append(
provider_model_setting
)
return provider_name_to_provider_model_settings_dict
@staticmethod
def _get_all_provider_model_credentials(tenant_id: str) -> dict[str, list[ProviderModelCredential]]:
def _get_all_provider_model_credentials(tenant_id: str) -> dict[str, list[_ProviderModelCredentialCacheEntry]]:
"""
Get All provider model credentials of the workspace.
:param tenant_id: workspace id
:return:
"""
provider_model_credentials = _get_cached_or_load_records(
tenant_id=tenant_id,
cache_source=_PROVIDER_MODEL_CREDENTIALS_CACHE_SOURCE,
)
provider_name_to_provider_model_credentials_dict = defaultdict(list)
with session_factory.create_session() as session:
stmt = select(ProviderModelCredential).where(ProviderModelCredential.tenant_id == tenant_id)
provider_model_credentials = session.scalars(stmt)
for provider_model_credential in provider_model_credentials:
provider_name_to_provider_model_credentials_dict[provider_model_credential.provider_name].append(
provider_model_credential
)
for provider_model_credential in provider_model_credentials:
provider_name_to_provider_model_credentials_dict[provider_model_credential.provider_name].append(
provider_model_credential
)
return provider_name_to_provider_model_credentials_dict
@staticmethod
def _get_all_provider_load_balancing_configs(tenant_id: str) -> dict[str, list[LoadBalancingModelConfig]]:
def _get_all_provider_credentials(tenant_id: str) -> dict[str, list[_ProviderCredentialCacheEntry]]:
"""
Get All provider credentials of the workspace.
:param tenant_id: workspace id
:return:
"""
provider_credentials = _get_cached_or_load_records(
tenant_id=tenant_id,
cache_source=_PROVIDER_CREDENTIALS_CACHE_SOURCE,
)
provider_name_to_provider_credentials_dict = defaultdict(list)
for provider_credential in provider_credentials:
provider_name_to_provider_credentials_dict[provider_credential.provider_name].append(provider_credential)
return provider_name_to_provider_credentials_dict
@staticmethod
def _get_all_provider_load_balancing_configs(
tenant_id: str,
) -> dict[str, list[_LoadBalancingModelConfigCacheEntry]]:
"""
Get All provider load balancing configs of the workspace.
@@ -546,14 +1085,16 @@ class ProviderManager:
if not model_load_balancing_enabled:
return {}
provider_load_balancing_configs = _get_cached_or_load_records(
tenant_id=tenant_id,
cache_source=_PROVIDER_LOAD_BALANCING_CONFIGS_CACHE_SOURCE,
)
provider_name_to_provider_load_balancing_model_configs_dict = defaultdict(list)
with session_factory.create_session() as session:
stmt = select(LoadBalancingModelConfig).where(LoadBalancingModelConfig.tenant_id == tenant_id)
provider_load_balancing_configs = session.scalars(stmt)
for provider_load_balancing_config in provider_load_balancing_configs:
provider_name_to_provider_load_balancing_model_configs_dict[
provider_load_balancing_config.provider_name
].append(provider_load_balancing_config)
for provider_load_balancing_config in provider_load_balancing_configs:
provider_name_to_provider_load_balancing_model_configs_dict[
provider_load_balancing_config.provider_name
].append(provider_load_balancing_config)
return provider_name_to_provider_load_balancing_model_configs_dict
@@ -722,8 +1263,9 @@ class ProviderManager:
tenant_id: str,
provider_entity: ProviderEntity,
provider_records: list[Provider],
provider_model_records: list[ProviderModel],
provider_model_credentials: list[ProviderModelCredential],
provider_model_records: list[_ProviderModelCacheEntry],
provider_model_credentials: list[_ProviderModelCredentialCacheEntry],
provider_credentials_by_name: dict[str, list[_ProviderCredentialCacheEntry]],
) -> CustomConfiguration:
"""
Convert to custom configuration.
@@ -736,7 +1278,10 @@ class ProviderManager:
"""
# Get custom provider configuration
custom_provider_configuration = self._get_custom_provider_configuration(
tenant_id, provider_entity, provider_records
tenant_id,
provider_entity,
provider_records,
provider_credentials_by_name,
)
# Get custom models which have not been added to the model list yet
@@ -758,7 +1303,11 @@ class ProviderManager:
)
def _get_custom_provider_configuration(
self, tenant_id: str, provider_entity: ProviderEntity, provider_records: list[Provider]
self,
tenant_id: str,
provider_entity: ProviderEntity,
provider_records: list[Provider],
provider_credentials_by_name: dict[str, list[_ProviderCredentialCacheEntry]],
) -> CustomProviderConfiguration | None:
"""Get custom provider configuration."""
# Find custom provider record (non-system)
@@ -790,13 +1339,29 @@ class ProviderManager:
credentials=provider_credentials,
current_credential_name=custom_provider_record.credential_name,
current_credential_id=custom_provider_record.credential_id,
available_credentials=self.get_provider_available_credentials(
tenant_id, custom_provider_record.provider_name
available_credentials=self._get_provider_available_credentials_from_records(
custom_provider_record.provider_name,
provider_credentials_by_name,
),
)
@staticmethod
def _get_provider_available_credentials_from_records(
provider_name: str,
provider_credentials_by_name: dict[str, list[_ProviderCredentialCacheEntry]],
) -> list[CredentialConfiguration]:
available_credentials: list[CredentialConfiguration] = []
for candidate_provider_name in ProviderManager._get_provider_names(provider_name):
available_credentials.extend(
CredentialConfiguration(credential_id=credential.id, credential_name=credential.credential_name)
for credential in provider_credentials_by_name.get(candidate_provider_name, [])
)
return available_credentials
def _get_can_added_models(
self, provider_model_records: list[ProviderModel], all_model_credentials: Sequence[ProviderModelCredential]
self,
provider_model_records: list[_ProviderModelCacheEntry],
all_model_credentials: Sequence[_ProviderModelCredentialCacheEntry],
) -> list[dict]:
"""Get the custom models and credentials from enterprise version which haven't add to the model list"""
existing_model_set = {(record.model_name, record.model_type) for record in provider_model_records}
@@ -829,9 +1394,9 @@ class ProviderManager:
self,
tenant_id: str,
provider_entity: ProviderEntity,
provider_model_records: list[ProviderModel],
provider_model_records: list[_ProviderModelCacheEntry],
can_added_models: list[dict],
all_model_credentials: Sequence[ProviderModelCredential],
all_model_credentials: Sequence[_ProviderModelCredentialCacheEntry],
) -> list[CustomModelConfiguration]:
"""Get custom model configurations."""
# Get model credential secret variables
@@ -1151,8 +1716,8 @@ class ProviderManager:
def _to_model_settings(
self,
provider_entity: ProviderEntity,
provider_model_settings: list[ProviderModelSetting] | None = None,
load_balancing_model_configs: list[LoadBalancingModelConfig] | None = None,
provider_model_settings: list[_ProviderModelSettingCacheEntry] | None = None,
load_balancing_model_configs: list[_LoadBalancingModelConfigCacheEntry] | None = None,
) -> list[ModelSettings]:
"""
Convert to model settings.