CoolFace
Apppublic

chwellofficial/nt360Slides

sourceHugging Faceupdated 5mo agoView on Hugging Face
0likes
model_availability.py154 linesDownload Raw Back to utils
1from constants.supported_ollama_models import SUPPORTED_OLLAMA_MODELS2from constants.llm import OPENAI_URL3from enums.image_provider import ImageProvider4from enums.llm_provider import LLMProvider5from utils.available_models import (6    list_available_anthropic_models,7    list_available_google_models,8    list_available_openai_compatible_models,9)10from utils.get_env import (11    get_anthropic_api_key_env,12    get_anthropic_model_env,13    get_can_change_keys_env,14    get_google_model_env,15    get_openai_api_key_env,16    get_openai_model_env,17    get_pixabay_api_key_env,18    get_pexels_api_key_env,19    get_comfyui_url_env,20    get_comfyui_workflow_env,21)22from utils.get_env import get_google_api_key_env23from utils.get_env import get_ollama_model_env24from utils.get_env import get_custom_llm_api_key_env25from utils.get_env import get_custom_llm_url_env26from utils.get_env import get_custom_model_env27from utils.llm_provider import (28    get_llm_provider,29    is_custom_llm_selected,30    is_ollama_selected,31)32from utils.ollama import pull_ollama_model33from utils.image_provider import (34    get_selected_image_provider,35    is_image_generation_disabled,36)37 38 39async def check_llm_and_image_provider_api_or_model_availability():40    can_change_keys = get_can_change_keys_env() != "false"41    if not can_change_keys:42        if get_llm_provider() == LLMProvider.OPENAI:43            openai_api_key = get_openai_api_key_env()44            if not openai_api_key:45                raise Exception("OPENAI_API_KEY must be provided")46            openai_model = get_openai_model_env()47            if openai_model:48                available_models = await list_available_openai_compatible_models(49                    OPENAI_URL, openai_api_key50                )51                if openai_model not in available_models:52                    print("-" * 50)53                    print("Available models: ", available_models)54                    raise Exception(f"Model {openai_model} is not available")55 56        elif get_llm_provider() == LLMProvider.GOOGLE:57            google_api_key = get_google_api_key_env()58            if not google_api_key:59                raise Exception("GOOGLE_API_KEY must be provided")60            google_model = get_google_model_env()61            if google_model:62                available_models = await list_available_google_models(google_api_key)63                if google_model not in available_models:64                    print("-" * 50)65                    print("Available models: ", available_models)66                    raise Exception(f"Model {google_model} is not available")67 68        elif get_llm_provider() == LLMProvider.ANTHROPIC:69            anthropic_api_key = get_anthropic_api_key_env()70            if not anthropic_api_key:71                raise Exception("ANTHROPIC_API_KEY must be provided")72            anthropic_model = get_anthropic_model_env()73            if anthropic_model:74                available_models = await list_available_anthropic_models(75                    anthropic_api_key76                )77                if anthropic_model not in available_models:78                    print("-" * 50)79                    print("Available models: ", available_models)80                    raise Exception(f"Model {anthropic_model} is not available")81 82        elif is_ollama_selected():83            ollama_model = get_ollama_model_env()84            if not ollama_model:85                raise Exception("OLLAMA_MODEL must be provided")86 87            if ollama_model not in SUPPORTED_OLLAMA_MODELS:88                raise Exception(f"Model {ollama_model} is not supported")89 90            print("-" * 50)91            print("Pulling model: ", ollama_model)92            async for event in pull_ollama_model(ollama_model):93                print(event)94            print("Pulled model: ", ollama_model)95            print("-" * 50)96 97        elif is_custom_llm_selected():98            custom_model = get_custom_model_env()99            custom_llm_url = get_custom_llm_url_env()100            if not custom_model:101                raise Exception("CUSTOM_MODEL must be provided")102            if not custom_llm_url:103                raise Exception("CUSTOM_LLM_URL must be provided")104            available_models = await list_available_openai_compatible_models(105                custom_llm_url, get_custom_llm_api_key_env() or "null"106            )107            print("-" * 50)108            print("Available models: ", available_models)109            if custom_model not in available_models:110                raise Exception(f"Model {custom_model} is not available")111 112        # Skip image provider and API key checks if image generation is disabled113        if is_image_generation_disabled():114            return115 116        # Check for Image Provider and API keys117        selected_image_provider = get_selected_image_provider()118        if not selected_image_provider:119            raise Exception("IMAGE_PROVIDER must be provided")120 121        if selected_image_provider == ImageProvider.PEXELS:122            pexels_api_key = get_pexels_api_key_env()123            if not pexels_api_key:124                raise Exception("PEXELS_API_KEY must be provided")125 126        elif selected_image_provider == ImageProvider.PIXABAY:127            pixabay_api_key = get_pixabay_api_key_env()128            if not pixabay_api_key:129                raise Exception("PIXABAY_API_KEY must be provided")130 131        elif (132            selected_image_provider == ImageProvider.GEMINI_FLASH133            or selected_image_provider == ImageProvider.NANOBANANA_PRO134        ):135            google_api_key = get_google_api_key_env()136            if not google_api_key:137                raise Exception("GOOGLE_API_KEY must be provided")138 139        elif (140            selected_image_provider == ImageProvider.DALLE3141            or selected_image_provider == ImageProvider.GPT_IMAGE_1_5142        ):143            openai_api_key = get_openai_api_key_env()144            if not openai_api_key:145                raise Exception("OPENAI_API_KEY must be provided")146 147        elif selected_image_provider == ImageProvider.COMFYUI:148            comfyui_url = get_comfyui_url_env()149            if not comfyui_url:150                raise Exception("COMFYUI_URL must be provided")151            workflow_json = get_comfyui_workflow_env()152            if not workflow_json:153                raise Exception("COMFYUI_WORKFLOW must be provided")154