diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py b/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py index 176e30ada..a00d4877c 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/__init__.py @@ -22,7 +22,6 @@ BedrockChatProvider, BedrockProvider, ChatAnthropicProvider, - ChatNVIDIAProvider, ChatOpenAIProvider, CohereProvider, GPT4AllProvider, diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/nvidia.py b/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/nvidia.py new file mode 100644 index 000000000..26137eb9f --- /dev/null +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/nvidia.py @@ -0,0 +1,23 @@ +from jupyter_ai_magics.providers import BaseProvider, EnvAuthStrategy +from langchain_nvidia_ai_endpoints import ChatNVIDIA + + +class ChatNVIDIAProvider(BaseProvider, ChatNVIDIA): + id = "nvidia-chat" + name = "NVIDIA" + models = [ + "playground_llama2_70b", + "playground_nemotron_steerlm_8b", + "playground_mistral_7b", + "playground_nv_llama2_rlhf_70b", + "playground_llama2_13b", + "playground_steerlm_llama_70b", + "playground_llama2_code_13b", + "playground_yi_34b", + "playground_mixtral_8x7b", + "playground_neva_22b", + "playground_llama2_code_34b", + ] + model_id_key = "model" + auth_strategy = EnvAuthStrategy(name="NVIDIA_API_KEY") + pypi_package_deps = ["langchain_nvidia_ai_endpoints"] diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py index 9053a2752..850e24a99 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py @@ -38,7 +38,6 @@ OpenAI, SagemakerEndpoint, ) -from langchain_nvidia_ai_endpoints import ChatNVIDIA # this is necessary because `langchain.pydantic_v1.main` does not include # `ModelMetaclass`, as it is not listed in `__all__` by the `pydantic.main` @@ -859,23 +858,3 @@ class QianfanProvider(BaseProvider, QianfanChatEndpoint): model_id_key = "model_name" pypi_package_deps = ["qianfan"] auth_strategy = MultiEnvAuthStrategy(names=["QIANFAN_AK", "QIANFAN_SK"]) - - -class ChatNVIDIAProvider(BaseProvider, ChatNVIDIA): - id = "nvidia-chat" - name = "NVIDIA" - models = [ - "playground_llama2_70b", - "playground_nemotron_steerlm_8b", - "playground_mistral_7b", - "playground_nv_llama2_rlhf_70b", - "playground_llama2_13b", - "playground_steerlm_llama_70b", - "playground_llama2_code_13b", - "playground_yi_34b", - "playground_mixtral_8x7b", - "playground_neva_22b", - "playground_llama2_code_34b", - ] - model_id_key = "model" - auth_strategy = EnvAuthStrategy(name="NVIDIA_API_KEY") diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/utils.py b/packages/jupyter-ai-magics/jupyter_ai_magics/utils.py index 0441d707c..983bbf2d5 100644 --- a/packages/jupyter-ai-magics/jupyter_ai_magics/utils.py +++ b/packages/jupyter-ai-magics/jupyter_ai_magics/utils.py @@ -26,16 +26,22 @@ def get_lm_providers( restrictions = {"allowed_providers": None, "blocked_providers": None} providers = {} eps = entry_points() - model_provider_eps = eps.select(group="jupyter_ai.model_providers") - for model_provider_ep in model_provider_eps: + provider_ep_group = eps.select(group="jupyter_ai.model_providers") + for provider_ep in provider_ep_group: try: - provider = model_provider_ep.load() + provider = provider_ep.load() + except ImportError as e: + log.warning( + f"Unable to load model provider `{provider_ep.name}`. Please install the `{e.name}` package." + ) + continue except Exception as e: log.error( - f"Unable to load model provider class from entry point `{model_provider_ep.name}`: %s.", - e, + f"Unable to load model provider `{provider_ep.name}`. Printing full exception below." ) + log.exception(e) continue + if not is_provider_allowed(provider.id, restrictions): log.info(f"Skipping blocked provider `{provider.id}`.") continue diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml index 138b79ee3..377776fba 100644 --- a/packages/jupyter-ai-magics/pyproject.toml +++ b/packages/jupyter-ai-magics/pyproject.toml @@ -62,7 +62,7 @@ amazon-bedrock = "jupyter_ai_magics:BedrockProvider" anthropic-chat = "jupyter_ai_magics:ChatAnthropicProvider" amazon-bedrock-chat = "jupyter_ai_magics:BedrockChatProvider" qianfan = "jupyter_ai_magics:QianfanProvider" -nvidia-chat = "jupyter_ai_magics:ChatNVIDIAProvider" +nvidia-chat = "jupyter_ai_magics.partner_providers.nvidia:ChatNVIDIAProvider" [project.entry-points."jupyter_ai.embeddings_model_providers"] bedrock = "jupyter_ai_magics:BedrockEmbeddingsProvider"