Spaces:
Running
Running
| # -*- coding: utf-8 -*- | |
| """Provider-neutral model routing primitives for SmartPaper. | |
| Legacy tasks continue to route by coarse capability. A5-R3-R4B adds a | |
| separate, trusted product-policy seam: academic callers provide an explicit | |
| logical_node, while the frozen TaskPolicySnapshot supplies provider/model and | |
| execution settings. Product routing never falls back to legacy capability | |
| routes. | |
| """ | |
| from __future__ import annotations | |
| from copy import deepcopy | |
| from dataclasses import dataclass | |
| from typing import Any, Callable, Iterable | |
| class ModelGatewayError(RuntimeError): | |
| pass | |
| class ModelProviderUnavailableError(ModelGatewayError): | |
| pass | |
| class ModelAuthenticationError(ModelGatewayError): | |
| pass | |
| class ModelBalanceError(ModelGatewayError): | |
| pass | |
| class ModelRequestError(ModelGatewayError): | |
| pass | |
| class ModelRoute: | |
| provider: str | |
| provider_role: str | |
| fallback_providers: tuple[str, ...] = () | |
| class TrustedProductRoute: | |
| provider: str | |
| model_id: str | |
| thinking_enabled: bool | |
| reasoning_effort: str | None | |
| max_output_tokens: int | |
| retry_budget: int | |
| class ProviderNeutralModelGateway: | |
| """Route legacy capabilities or trusted product logical nodes to adapters.""" | |
| def __init__(self, config: dict[str, Any] | None = None): | |
| cfg = dict(config or {}) | |
| self.default_provider = str(cfg.get("default_provider") or "deepseek").strip().lower() | |
| raw_routes = dict(cfg.get("capability_routes") or {}) | |
| self._routes: dict[str, ModelRoute] = {} | |
| for capability, row in raw_routes.items(): | |
| cap = str(capability or "").strip().lower() | |
| if not cap: | |
| continue | |
| if isinstance(row, str): | |
| provider = str(row or self.default_provider).strip().lower() | |
| provider_role = cap | |
| fallbacks = () | |
| else: | |
| item = dict(row or {}) | |
| provider = str(item.get("provider") or self.default_provider).strip().lower() | |
| provider_role = str(item.get("provider_role") or cap).strip().lower() | |
| raw_fallbacks = item.get("fallback_providers") or () | |
| if isinstance(raw_fallbacks, str): | |
| raw_fallbacks = [part.strip() for part in raw_fallbacks.split(",") if part.strip()] | |
| fallbacks = tuple( | |
| str(name or "").strip().lower() | |
| for name in raw_fallbacks | |
| if str(name or "").strip() | |
| ) | |
| self._routes[cap] = ModelRoute( | |
| provider=provider, | |
| provider_role=provider_role, | |
| fallback_providers=fallbacks, | |
| ) | |
| self._call_adapters: dict[str, Callable[..., str]] = {} | |
| self._stream_adapters: dict[str, Callable[..., Iterable[str]]] = {} | |
| self._policy_snapshot: dict[str, Any] = {} | |
| self._trusted_product_policy_snapshot: dict[str, Any] = {} | |
| def register_provider( | |
| self, | |
| provider: str, | |
| *, | |
| call_adapter: Callable[..., str], | |
| stream_adapter: Callable[..., Iterable[str]], | |
| ) -> None: | |
| name = str(provider or "").strip().lower() | |
| if not name: | |
| raise ValueError("provider name is required") | |
| self._call_adapters[name] = call_adapter | |
| self._stream_adapters[name] = stream_adapter | |
| # Legacy config/request policy. It remains a separate trust domain. | |
| def set_policy_snapshot(self, snapshot: dict[str, Any] | None) -> None: | |
| self._policy_snapshot = deepcopy(dict(snapshot or {})) | |
| def policy_snapshot(self) -> dict[str, Any]: | |
| return deepcopy(self._policy_snapshot) | |
| def set_trusted_product_policy_snapshot(self, snapshot: dict[str, Any] | None) -> None: | |
| """Attach a Worker/D1-verified TaskPolicySnapshot for product execution. | |
| The runner validates task identity and snapshot SHA before calling this | |
| setter. The gateway still validates route shape and fails closed before | |
| invoking any provider adapter. | |
| """ | |
| self._trusted_product_policy_snapshot = deepcopy(dict(snapshot or {})) | |
| def trusted_product_policy_snapshot(self) -> dict[str, Any]: | |
| return deepcopy(self._trusted_product_policy_snapshot) | |
| def _route(self, capability: str) -> ModelRoute: | |
| cap = str(capability or "reasoning").strip().lower() or "reasoning" | |
| policy_routes = dict(self._policy_snapshot.get("model_routes") or {}) | |
| row = policy_routes.get(cap) | |
| if isinstance(row, str): | |
| return ModelRoute(provider=row.strip().lower(), provider_role=cap) | |
| if isinstance(row, dict): | |
| raw_fallbacks = row.get("fallback_providers") or () | |
| if isinstance(raw_fallbacks, str): | |
| raw_fallbacks = [part.strip() for part in raw_fallbacks.split(",") if part.strip()] | |
| return ModelRoute( | |
| provider=str(row.get("provider") or self.default_provider).strip().lower(), | |
| provider_role=str(row.get("provider_role") or cap).strip().lower(), | |
| fallback_providers=tuple( | |
| str(name or "").strip().lower() | |
| for name in raw_fallbacks | |
| if str(name or "").strip() | |
| ), | |
| ) | |
| return self._routes.get(cap) or ModelRoute(self.default_provider, cap) | |
| def _provider_sequence(route: ModelRoute) -> tuple[str, ...]: | |
| return tuple(dict.fromkeys( | |
| name for name in (route.provider, *route.fallback_providers) if name | |
| )) | |
| def _trusted_product_context(self) -> tuple[dict[str, Any], dict[str, Any]] | None: | |
| snapshot = self._trusted_product_policy_snapshot | |
| if not snapshot: | |
| return None | |
| resolved = snapshot.get("resolved_policy") | |
| if not isinstance(resolved, dict): | |
| raise ModelRequestError("trusted_product_policy_invalid") | |
| model = resolved.get("model") | |
| if not isinstance(model, dict): | |
| raise ModelRequestError("trusted_product_model_policy_missing") | |
| routes = model.get("routes") | |
| if not isinstance(routes, dict) or not routes: | |
| raise ModelRequestError("trusted_product_model_routes_missing") | |
| return resolved, routes | |
| def _trusted_product_route(self, logical_node: str | None) -> tuple[TrustedProductRoute, str]: | |
| context = self._trusted_product_context() | |
| if context is None: | |
| raise ModelRequestError("trusted_product_policy_missing") | |
| resolved, routes = context | |
| node = str(logical_node or "").strip() | |
| if not node: | |
| raise ModelRequestError("trusted_product_logical_node_required") | |
| row = routes.get(node) | |
| if not isinstance(row, dict): | |
| raise ModelRequestError(f"trusted_product_route_missing:{node}") | |
| provider = str(row.get("provider") or "").strip().lower() | |
| model_id = str(row.get("model_id_or_alias") or "").strip() | |
| if not provider or not model_id: | |
| raise ModelRequestError(f"trusted_product_route_invalid:{node}") | |
| thinking = row.get("thinking_enabled") | |
| if not isinstance(thinking, bool): | |
| raise ModelRequestError(f"trusted_product_route_thinking_invalid:{node}") | |
| effort = row.get("reasoning_effort") | |
| if thinking: | |
| if str(effort or "").strip().lower() != "high": | |
| raise ModelRequestError(f"trusted_product_route_effort_invalid:{node}") | |
| effort = "high" | |
| else: | |
| if effort not in (None, ""): | |
| raise ModelRequestError(f"trusted_product_route_effort_invalid:{node}") | |
| effort = None | |
| try: | |
| max_output_tokens = int(row.get("max_output_tokens")) | |
| retry_budget = int(row.get("retry_budget")) | |
| except (TypeError, ValueError): | |
| raise ModelRequestError(f"trusted_product_route_budget_invalid:{node}") | |
| if max_output_tokens <= 0 or retry_budget < 0: | |
| raise ModelRequestError(f"trusted_product_route_budget_invalid:{node}") | |
| if row.get("escalation_model_id_or_alias") not in (None, ""): | |
| raise ModelRequestError(f"trusted_product_route_escalation_forbidden:{node}") | |
| if row.get("escalation_condition_code") not in (None, ""): | |
| raise ModelRequestError(f"trusted_product_route_escalation_forbidden:{node}") | |
| model_policy_version = str( | |
| resolved.get("model_policy_version") | |
| or snapshot_model_policy_version(self._trusted_product_policy_snapshot) | |
| or "" | |
| ).strip() | |
| return TrustedProductRoute( | |
| provider=provider, | |
| model_id=model_id, | |
| thinking_enabled=thinking, | |
| reasoning_effort=effort, | |
| max_output_tokens=max_output_tokens, | |
| retry_budget=retry_budget, | |
| ), model_policy_version | |
| def _effective_max_tokens(requested: Any, ceiling: int) -> int: | |
| if requested is None: | |
| return int(ceiling) | |
| try: | |
| value = int(requested) | |
| except (TypeError, ValueError): | |
| raise ModelRequestError("trusted_product_max_tokens_invalid") | |
| if value <= 0: | |
| raise ModelRequestError("trusted_product_max_tokens_invalid") | |
| return min(value, int(ceiling)) | |
| def _product_payload(self, *, capability: str, logical_node: str | None, kwargs: dict[str, Any]): | |
| route, model_policy_version = self._trusted_product_route(logical_node) | |
| payload = dict(kwargs) | |
| payload["model_role"] = str(capability or "reasoning").strip().lower() or "reasoning" | |
| payload["model_name"] = route.model_id | |
| payload["thinking"] = route.thinking_enabled | |
| payload["reasoning_effort"] = route.reasoning_effort | |
| payload["allow_fallback"] = False | |
| payload["max_tokens"] = self._effective_max_tokens( | |
| payload.get("max_tokens"), route.max_output_tokens | |
| ) | |
| payload["retry_budget"] = route.retry_budget | |
| payload["logical_node"] = str(logical_node or "").strip() | |
| payload["model_policy_version"] = model_policy_version | |
| payload["route_enforced"] = True | |
| return route.provider, payload | |
| def call(self, *, capability: str, logical_node: str | None = None, **kwargs) -> str: | |
| if self._trusted_product_context() is not None: | |
| provider, payload = self._product_payload( | |
| capability=capability, logical_node=logical_node, kwargs=kwargs | |
| ) | |
| adapter = self._call_adapters.get(provider) | |
| if not callable(adapter): | |
| raise ModelProviderUnavailableError( | |
| f"model_provider_unavailable:{provider}:{logical_node or capability}" | |
| ) | |
| return adapter(**payload) | |
| route = self._route(capability) | |
| missing = [] | |
| for provider in self._provider_sequence(route): | |
| adapter = self._call_adapters.get(provider) | |
| if not callable(adapter): | |
| missing.append(provider) | |
| continue | |
| payload = dict(kwargs) | |
| payload["model_role"] = route.provider_role | |
| return adapter(**payload) | |
| raise ModelProviderUnavailableError( | |
| f"model_provider_unavailable:{','.join(missing) or route.provider}:{capability}" | |
| ) | |
| def stream(self, *, capability: str, logical_node: str | None = None, **kwargs): | |
| if self._trusted_product_context() is not None: | |
| provider, payload = self._product_payload( | |
| capability=capability, logical_node=logical_node, kwargs=kwargs | |
| ) | |
| adapter = self._stream_adapters.get(provider) | |
| if not callable(adapter): | |
| raise ModelProviderUnavailableError( | |
| f"model_provider_unavailable:{provider}:{logical_node or capability}" | |
| ) | |
| yield from adapter(**payload) | |
| return | |
| route = self._route(capability) | |
| missing = [] | |
| for provider in self._provider_sequence(route): | |
| adapter = self._stream_adapters.get(provider) | |
| if not callable(adapter): | |
| missing.append(provider) | |
| continue | |
| payload = dict(kwargs) | |
| payload["model_role"] = route.provider_role | |
| yield from adapter(**payload) | |
| return | |
| raise ModelProviderUnavailableError( | |
| f"model_provider_unavailable:{','.join(missing) or route.provider}:{capability}" | |
| ) | |
| def snapshot_model_policy_version(snapshot: dict[str, Any] | None) -> str: | |
| data = dict(snapshot or {}) | |
| return str(data.get("model_policy_version") or "").strip() | |
| def resolve_plan_policy_snapshot( | |
| app_config: dict[str, Any], request: dict[str, Any] | None = None | |
| ) -> dict[str, Any]: | |
| """Freeze a config-driven legacy plan policy without touching payment authority.""" | |
| req = dict(request or {}) | |
| policies = dict(app_config.get("plan_policies") or {}) | |
| entitlement = req.get("entitlement_snapshot") if isinstance(req.get("entitlement_snapshot"), dict) else {} | |
| requested = str( | |
| req.get("trusted_plan_id") or entitlement.get("plan_id") or "" | |
| ).strip().lower() | |
| selected = requested if requested in policies else "legacy_current" | |
| policy = deepcopy(dict(policies.get(selected) or policies.get("legacy_current") or {})) | |
| policy["plan_id"] = selected | |
| policy["policy_version"] = str( | |
| policy.get("policy_version") or app_config.get("plan_policy_version") or "plan-policy-v1" | |
| ) | |
| return policy | |