119 lines
5.4 KiB
Python
119 lines
5.4 KiB
Python
"""Subscription retirement: existing aliases stay usable, BYOK is terminal only."""
|
|
from copy import deepcopy
|
|
|
|
|
|
def subscription_disabled(catalog):
|
|
return catalog.get('zenmux_subscription_enabled') is False
|
|
|
|
|
|
def without_subscription(catalog, env):
|
|
if not subscription_disabled(catalog):
|
|
return env
|
|
return {key: value for key, value in env.items()
|
|
if key != 'ZENMUX_SUBSCRIPTION_API_KEY'}
|
|
|
|
|
|
def apply_provider_policy(config, catalog, env):
|
|
if not subscription_disabled(catalog):
|
|
return config
|
|
routes = config['model_list']
|
|
if any(r['litellm_params'].get('api_key') ==
|
|
'os.environ/ZENMUX_SUBSCRIPTION_API_KEY' for r in routes):
|
|
raise ValueError('Retired subscription route remains')
|
|
indexed = {r['model_name']: r for r in routes}
|
|
groups = {k: list(v) for row in config['router_settings']['fallbacks']
|
|
for k, v in row.items()}
|
|
|
|
def add(name, params):
|
|
if name in indexed:
|
|
return indexed[name]
|
|
route = {'model_name': name, 'litellm_params': deepcopy(params)}
|
|
indexed[name] = route
|
|
routes.append(route)
|
|
return route
|
|
|
|
def fallback(source, target):
|
|
if source == target or source not in indexed or target not in indexed:
|
|
return
|
|
targets = groups.setdefault(source, [])
|
|
if target not in targets:
|
|
targets.append(target)
|
|
|
|
def payg_params(model, anthropic=False):
|
|
return {'model': ('anthropic/' if anthropic else 'openai/') + model,
|
|
'api_key': 'os.environ/ZENMUX_PAYG_API_KEY',
|
|
'api_base': 'https://zenmux.ai/api/anthropic' if anthropic
|
|
else 'https://zenmux.ai/api/v1',
|
|
'max_retries': 0}
|
|
|
|
# Explicit BYOK aliases remain direct; automatic/legacy aliases prefer OpenRouter.
|
|
for family in ('zenmux', 'zenmux_open_models'):
|
|
for alias, model in catalog.get(family, {}).items():
|
|
byok = 'zenmux-payg/' + alias
|
|
if env.get('ZENMUX_PAYG_API_KEY'):
|
|
add(byok, payg_params(model, anthropic=family == 'zenmux'))
|
|
primary = None
|
|
if env.get('OPENROUTER_API_KEY'):
|
|
primary = {'model': 'openrouter/' + model,
|
|
'api_key': 'os.environ/OPENROUTER_API_KEY', 'max_retries': 0}
|
|
elif byok in indexed:
|
|
primary = indexed[byok]['litellm_params']
|
|
if primary is None:
|
|
continue
|
|
names = [alias, 'auto/' + alias, 'zenmux/' + alias]
|
|
if env.get('OPENROUTER_API_KEY'):
|
|
names += ['openrouter/' + alias]
|
|
if family == 'zenmux_open_models':
|
|
names += [alias + '-openrouter']
|
|
for name in names:
|
|
route = add(name, primary)
|
|
if route['litellm_params'].get('api_key') != 'os.environ/ZENMUX_PAYG_API_KEY':
|
|
fallback(name, byok)
|
|
|
|
# Former Gemini subscription aliases follow the existing Google primary.
|
|
for alias, model in catalog.get('zenmux_chat', {}).items():
|
|
byok = 'zenmux-payg/' + alias.replace('-zenmux', '-latest')
|
|
if env.get('ZENMUX_PAYG_API_KEY'):
|
|
add(byok, payg_params(model))
|
|
google = indexed.get('aistudio/gemini-flash-latest') or indexed.get('vertex/gemini-flash-latest')
|
|
primary = google or indexed.get(byok)
|
|
if primary:
|
|
for name in [alias, 'zenmux/' + alias.replace('-zenmux', '-latest')]:
|
|
route = add(name, primary['litellm_params'])
|
|
if route['litellm_params'].get('api_key') != 'os.environ/ZENMUX_PAYG_API_KEY':
|
|
fallback(name, byok)
|
|
|
|
# Public ASR compatibility aliases use their same-model direct provider first.
|
|
asr = {'asr-gpt': ('openai/gpt-transcribe', 'openai/gpt-transcribe'),
|
|
'asr-qwen': ('bailian-cn/qwen-audio-3.1-asr-flash', 'qwen/qwen-audio-3.1-asr-flash'),
|
|
'asr-doubao': (None, 'bytedance/doubao-seed-asr-2.0'),
|
|
'asr-mimo': (None, 'xiaomi/mimo-v2.5-asr')}
|
|
for alias, (direct, model) in asr.items():
|
|
# Do not introduce new ASR models to the LAN instance.
|
|
if not config.get('_public_gateway'):
|
|
break
|
|
byok = 'zenmux-payg/' + model.split('/', 1)[1]
|
|
if env.get('ZENMUX_PAYG_API_KEY'):
|
|
add(byok, payg_params(model))
|
|
primary = indexed.get(direct) or indexed.get(byok)
|
|
if not primary:
|
|
continue
|
|
if alias in indexed:
|
|
indexed[alias]['litellm_params'] = deepcopy(primary['litellm_params'])
|
|
else:
|
|
add(alias, primary['litellm_params'])
|
|
# Old explicit ZenMux names are compatibility aliases, no subscription calls.
|
|
add('zenmux/' + model.split('/', 1)[1], primary['litellm_params'])
|
|
if primary['litellm_params'].get('api_key') != 'os.environ/ZENMUX_PAYG_API_KEY':
|
|
fallback(alias, byok)
|
|
fallback('zenmux/' + model.split('/', 1)[1], byok)
|
|
|
|
config.pop('_public_gateway', None)
|
|
config['router_settings']['fallbacks'] = [{k: v} for k, v in groups.items() if v]
|
|
for source, targets in groups.items():
|
|
if source not in indexed or any(t not in indexed for t in targets):
|
|
raise ValueError('Dangling fallback: ' + source)
|
|
if indexed[source]['litellm_params'].get('api_key') == 'os.environ/ZENMUX_PAYG_API_KEY':
|
|
raise ValueError('BYOK must be terminal: ' + source)
|
|
return config
|