dify/api/services/model_billing_dependency_runtime.py

"""Technical adapters for fresh read-only model selection evidence."""

from collections import Counter
from collections.abc import Callable, Sequence
from typing import Literal, cast

from configs import dify_config
from core.entities.provider_entities import ProviderQuotaType, QuotaConfiguration, QuotaUnit
from core.model_invocation_routing import (
    compatibility_quota_type,
    has_compatibility_route,
    migration_routing_state,
)
from core.plugin.entities.plugin import PluginInstallationSource
from core.plugin.impl.model import PluginModelClient
from enums import DeploymentEdition
from extensions import ext_hosting_provider
from models.provider_ids import ModelProviderID
from repositories.knowledge.dataset_repository import SQLAlchemyDatasetRepository
from repositories.model_billing_dependency_repository import ModelBillingDependencyRepository
from services.credit_pool_service import CreditPoolService
from services.knowledge.model_dependencies import DatasetModelDependencyRecord
from services.model_billing_dependency_inventory import (
    DatasetModelSelectors,
    DependencyInventoryPolicy,
    ModelDependencyDocument,
    WorkflowModelSelectors,
)
from services.model_billing_dependency_provenance import (
    ModelDependencyDefault,
    ModelDependencyProviderState,
    ModelDependencyQuota,
    ModelDependencySelectionSnapshot,
    StoredDependencyProvider,
)
from services.model_billing_mapping_service import (
    additional_hosted_availability,
    hosted_embedding_source_fact,
    hosted_gemini_source_fact,
    hosted_source_variant_fact,
    hosted_tongyi_source_fact,
    provider_has_prepared_ownership,
    published_original_quotas,
    retirement_policy_for_revision,
    source_mapping_revision,
    uses_published_tokener_balance,
    uses_stable_dependency_inventory,
)
from services.model_billing_migration_inventory import saved_dependency_inventory_policy


class RevisionModelDependencyRetirements:
    def identities(self, mapping_revision: str | None) -> frozenset[tuple[str, str, str]]:
        if mapping_revision is None:
            return frozenset()
        registry = source_mapping_revision(dify_config.TOKENER_LEGACY_MODEL_MAPPING_JSON, mapping_revision)
        policy = retirement_policy_for_revision(
            dify_config.TOKENER_LEGACY_RETIRED_MODELS_JSON, mapping_revision, dict(registry.declarations)
        )
        return frozenset((provider, kind.value, model) for provider, kind, model in policy.identities)


class RevisionModelDependencyInventoryPolicies:
    def for_revision(self, mapping_revision: str | None) -> DependencyInventoryPolicy:
        return saved_dependency_inventory_policy(
            dify_config.TOKENER_LEGACY_DEPENDENCY_INVENTORIES_JSON, mapping_revision
        )


class CompositeModelDependencyDocuments:
    def __init__(self, readers: Sequence[Callable[[str], Sequence[ModelDependencyDocument]]]) -> None:
        self._readers = tuple(readers)

    def collect(self, tenant_id: str) -> Sequence[ModelDependencyDocument]:
        return tuple(document for reader in self._readers for document in reader(tenant_id))


class KnowledgeModelDependencyDocuments:
    def __init__(self, *, datasets: SQLAlchemyDatasetRepository) -> None:
        self._datasets = datasets

    def collect(self, tenant_id: str) -> Sequence[ModelDependencyDocument]:
        documents: list[ModelDependencyDocument] = []
        for item in self._datasets.collect_model_dependency_records(tenant_id):
            if isinstance(item, DatasetModelDependencyRecord):
                documents.append(
                    ModelDependencyDocument(
                        item.tenant_id,
                        "dataset",
                        item.dataset_id,
                        "configuration",
                        DatasetModelSelectors(
                            item.indexing_technique,
                            item.embedding_provider,
                            item.embedding_model,
                            item.retrieval,
                            item.external,
                            item.summary,
                            item.has_qa_documents,
                        ),
                    )
                )
            else:
                documents.append(
                    ModelDependencyDocument(
                        item.tenant_id,
                        "workflow",
                        item.pipeline_id,
                        item.workflow_id,
                        WorkflowModelSelectors(item.graph, item.features),
                    )
                )
        return tuple(documents)


def _provider(value: str, *, plugin_id: str | None = None) -> str:
    return str(ModelProviderID(value if value.count("/") == 2 or plugin_id is None else plugin_id + "/" + value))


class SavedModelSelectionRuntimeQuery:
    def __init__(self, *, records: ModelBillingDependencyRepository) -> None:
        self._records = records

    def selection_snapshot(self, tenant_id: str) -> ModelDependencySelectionSnapshot:
        stored = self._records.stored_selection_snapshot(tenant_id)
        if stored.billing_source not in {None, "tokener"}:
            raise ValueError("unknown_model_billing_source")
        state = migration_routing_state(tenant_id)
        compatibility = has_compatibility_route(state)
        published = uses_published_tokener_balance(state)
        stable_inventory = uses_stable_dependency_inventory(
            dify_config.TOKENER_LEGACY_DEPENDENCY_INVENTORIES_JSON, state
        )
        cloud = dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD
        client = PluginModelClient()
        declarations = client.fetch_model_providers(tenant_id)
        bindings = client.fetch_model_provider_bindings(tenant_id)
        counts = Counter(_provider(item.provider, plugin_id=item.plugin_id) for item in bindings)
        fresh = {_provider(item.provider, plugin_id=item.plugin_id): item for item in bindings}
        declaration_counts = Counter(_provider(item.provider, plugin_id=item.plugin_id) for item in declarations)
        pools = {kind: (limit, used) for kind, limit, used in stored.pools}
        if cloud and not compatibility:
            # This is a balance-only read, outside every SQL read session.
            for kind in ("trial", "paid"):
                balance = CreditPoolService.get_pool(tenant_id, kind)
                if balance is not None:
                    pools[kind] = (balance.quota_limit, balance.quota_used)
        raw_records: dict[str, list[StoredDependencyProvider]] = {}
        raw_preferences: dict[str, list[str]] = {}
        for item in stored.providers:
            raw_records.setdefault(item.provider, []).append(item)
        for preference in stored.preferences:
            raw_preferences.setdefault(preference.provider, []).append(preference.preferred)
        providers: list[ModelDependencyProviderState] = []
        for declaration in declarations:
            name = _provider(declaration.provider, plugin_id=declaration.plugin_id)
            provider_id = ModelProviderID(name)
            # Exact canonical entries override a bare legacy alias, matching
            # ProviderManager. Never merge conflicting Provider quotas/prefs.
            records = raw_records.get(
                name, raw_records.get(provider_id.provider_name, []) if provider_id.is_langgenius() else []
            )
            preferences = raw_preferences.get(
                name, raw_preferences.get(provider_id.provider_name, []) if provider_id.is_langgenius() else []
            )
            custom = any(item.provider_type != "system" for item in records) or any(
                _provider(item) == name for item in stored.custom_model_providers
            )
            binding = fresh.get(name)
            trustworthy = bool(
                counts[name] == 1
                and declaration_counts[name] == 1
                and binding is not None
                and binding.runtime_type in {"local", "serverless"}
                and binding.source != PluginInstallationSource.Remote
                and binding.plugin_unique_identifier == declaration.plugin_unique_identifier
                and declaration.tenant_id == tenant_id
                and len(preferences) <= 1
            )
            hosting = ext_hosting_provider.hosting_configuration.provider_map.get(name)
            enabled = bool(
                trustworthy
                and binding
                and binding.verified
                # Management declarations may omit this optional enrichment;
                # the fresh exact-UID installation binding owns source/signature
                # authority. An explicit conflicting declaration is not ignored.
                and binding.source == PluginInstallationSource.Marketplace
                and declaration.installation_source in {None, PluginInstallationSource.Marketplace}
                and hosting
                and hosting.enabled
                and (stored.billing_source is None or compatibility)
            )
            quotas: list[QuotaConfiguration] = []
            current = None
            if enabled and hosting:
                by_pool = {item.quota_type: item for item in records if item.provider_type == "system"}
                for quota in hosting.quotas:
                    row = by_pool.get(quota.quota_type.value)
                    # Normal initialization creates missing paid/trial rows at
                    # zero. Model that fact without performing the write.
                    limit, used = (row.quota_limit, row.quota_used) if row else (0, 0)
                    if cloud and not compatibility and quota.quota_type.value in pools:
                        limit, used = pools[quota.quota_type.value]
                    if not isinstance(limit, int) or not isinstance(used, int):
                        trustworthy = False
                        continue
                    quotas.append(
                        QuotaConfiguration(
                            quota_type=quota.quota_type,
                            quota_unit=hosting.quota_unit or QuotaUnit.TOKENS,
                            quota_limit=limit,
                            quota_used=used,
                            is_valid=limit == -1 or limit > used,
                            restrict_models=quota.restrict_models,
                        )
                    )
                if compatibility and state is not None:
                    additional = not provider_has_prepared_ownership(name, state)
                    if published or additional:
                        preserve = None
                        if published:
                            quotas, preserve = published_original_quotas(name, quotas, state)
                        available = additional_hosted_availability(
                            provider=name,
                            hosting=hosting,
                            installed_models=declaration.declaration.models,
                            original_quotas=quotas,
                            raw_registry=dify_config.TOKENER_LEGACY_MODEL_MAPPING_JSON,
                            revision=state.get("model_mapping_version", ""),
                            published_balance=published,
                            preserve_independent_pool=preserve,
                            preserve_independent_units=stable_inventory,
                            source_uid=declaration.plugin_unique_identifier,
                            source_fact=hosted_source_variant_fact(
                                provider=name,
                                source_uid=declaration.plugin_unique_identifier,
                                models=declaration.declaration.models,
                                credentials=hosting.credentials,
                            ),
                            embedding_source_fact=hosted_embedding_source_fact(
                                provider=name,
                                source_uid=declaration.plugin_unique_identifier,
                                models=declaration.declaration.models,
                                credentials=hosting.credentials,
                            ),
                            gemini_source_fact=hosted_gemini_source_fact(
                                provider=name,
                                source_uid=declaration.plugin_unique_identifier,
                                models=declaration.declaration.models,
                                credentials=hosting.credentials,
                            ),
                            tongyi_source_fact=hosted_tongyi_source_fact(
                                provider=name,
                                source_uid=declaration.plugin_unique_identifier,
                                models=declaration.declaration.models,
                                credentials=hosting.credentials,
                            ),
                            allow_source_variants=stable_inventory,
                        )
                        if available is not None:
                            quotas, current = available.quotas, available.current_quota_type
                        else:
                            enabled = False
                    else:
                        current = compatibility_quota_type(name, quotas, state)
                elif quotas:
                    current = next(
                        (
                            kind
                            for kind in (ProviderQuotaType.PAID, ProviderQuotaType.FREE, ProviderQuotaType.TRIAL)
                            if any(item.quota_type == kind and item.is_valid for item in quotas)
                        ),
                        quotas[-1].quota_type,
                    )
            preferred = preferences[0] if preferences else None
            if preferred not in {None, "system", "custom"}:
                trustworthy = False
            providers.append(
                ModelDependencyProviderState(
                    name,
                    True,
                    enabled,
                    custom,
                    cast(Literal["system", "custom"] | None, preferred),
                    tuple(
                        ModelDependencyQuota(
                            cast(Literal["paid", "free", "trial"], item.quota_type.value),
                            item.quota_unit.value,
                            item.quota_limit,
                            item.quota_used,
                            tuple((model.model_type.value, model.model) for model in item.restrict_models),
                        )
                        for item in quotas
                    ),
                    current.value if current is not None else None,
                    cloud,
                    compatibility,
                    stored.billing_source == "tokener",
                    trustworthy,
                    len(preferences) <= 1 and preferred in {None, "system", "custom"},
                    binding.plugin_unique_identifier if binding is not None else None,
                    stable_inventory,
                )
            )
        defaults = tuple(
            ModelDependencyDefault(item.model_type, _provider(item.provider), item.model) for item in stored.defaults
        )
        return ModelDependencySelectionSnapshot(tuple(providers), defaults)