Underground-Digital/Workflow-Engine

1import logging2import os3 4from core.model_runtime.entities.model_entities import ModelType5from core.model_runtime.entities.provider_entities import ProviderConfig, ProviderEntity, SimpleProviderEntity6from core.model_runtime.model_providers.model_provider_factory import ModelProviderExtension, ModelProviderFactory7 8logger = logging.getLogger(__name__)9 10 11def test_get_providers():12    factory = ModelProviderFactory()13    providers = factory.get_providers()14 15    for provider in providers:16        logger.debug(provider)17 18    assert len(providers) >= 119    assert isinstance(providers[0], ProviderEntity)20 21 22def test_get_models():23    factory = ModelProviderFactory()24    providers = factory.get_models(25        model_type=ModelType.LLM,26        provider_configs=[27            ProviderConfig(provider="openai", credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")})28        ],29    )30 31    logger.debug(providers)32 33    assert len(providers) >= 134    assert isinstance(providers[0], SimpleProviderEntity)35 36    # all provider models type equals to ModelType.LLM37    for provider in providers:38        for provider_model in provider.models:39            assert provider_model.model_type == ModelType.LLM40 41    providers = factory.get_models(42        provider="openai",43        provider_configs=[44            ProviderConfig(provider="openai", credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")})45        ],46    )47 48    assert len(providers) == 149    assert isinstance(providers[0], SimpleProviderEntity)50    assert providers[0].provider == "openai"51 52 53def test_provider_credentials_validate():54    factory = ModelProviderFactory()55    factory.provider_credentials_validate(56        provider="openai", credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}57    )58 59 60def test__get_model_provider_map():61    factory = ModelProviderFactory()62    model_providers = factory._get_model_provider_map()63 64    for name, model_provider in model_providers.items():65        logger.debug(name)66        logger.debug(model_provider.provider_instance)67 68    assert len(model_providers) >= 169    assert isinstance(model_providers["openai"], ModelProviderExtension)70