Retire ZenMux subscription and keep BYOK as terminal fallback

2026-11-01 23:25:00 -08:00
parent be589ff5e0
commit c42456be22
13 changed files with 338 additions and 13 deletions
@@ -0,0 +1,118 @@
+"""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