[litellm] Initial migration to LiteLLM (#1425)

* add new /api/ai/models endpoint

* add litellm dependency

* add new model ID input component

* fix custom model ID input

* update ConfigManager to accept litellm model IDs

* update Jupyternaut to use litellm

* remove LangChain partner packages from dependencies

* remove LangChain provider entry points

* pre-commit

* rename /api/ai/models => /api/ai/models/chat

* simplify model settings UI

Author: dlqqq

packages/jupyter-ai-magics/jupyter_ai_magics/tests/test_utils.py

@@ -18,6 +18,7 @@
         {"allowed_providers": [KNOWN_LM_A], "blocked_providers": None},
     ],
 )
+@pytest.skip("Update this to use LiteLLM")
 def test_get_lm_providers_not_restricted(restrictions):
     a_not_restricted = get_lm_providers(None, restrictions)
     assert KNOWN_LM_A in a_not_restricted
@@ -33,6 +34,7 @@ def test_get_lm_providers_not_restricted(restrictions):
         {"allowed_providers": [KNOWN_LM_B], "blocked_providers": None},
     ],
 )
+@pytest.skip("Update this to use LiteLLM")
 def test_get_lm_providers_restricted(restrictions):
     a_not_restricted = get_lm_providers(None, restrictions)
     assert KNOWN_LM_A not in a_not_restricted

packages/jupyter-ai-magics/pyproject.toml

@@ -30,7 +30,6 @@ dependencies = [
     "pydantic>=2.10.0,<3",
     "click>=8.1.0,<9",
     "jsonpath-ng>=1.5.3,<2",
-    "langchain-google-vertexai",
 ]
 
 [project.optional-dependencies]
@@ -45,64 +44,10 @@ dev = [
 test = ["coverage", "pytest", "pytest-asyncio", "pytest-cov"]
 
 all = [
-    "ai21",
-    "gpt4all",
-    "huggingface_hub",
-    "ipywidgets",
-    "langchain_anthropic",
-    "langchain_aws",
-    "langchain_cohere",
-    # Pin cohere to <5.16 to prevent langchain_cohere from breaking as it uses ChatResponse removed in cohere 5.16.0
-    # TODO: remove cohere pin when langchain_cohere is updated to work with cohere >=5.16
-    "cohere<5.16",
-    "langchain_google_genai",
-    "langchain_mistralai",
-    "langchain_nvidia_ai_endpoints",
-    "langchain_openai",
-    "langchain_ollama",
-    "pillow",
+    # Required for using Amazon Bedrock
     "boto3",
-    "qianfan",
-    "together",
-    "langchain-google-vertexai",
 ]
 
-[project.entry-points."jupyter_ai.model_providers"]
-ai21 = "jupyter_ai_magics:AI21Provider"
-anthropic-chat = "jupyter_ai_magics.partner_providers.anthropic:ChatAnthropicProvider"
-cohere = "jupyter_ai_magics.partner_providers.cohere:CohereProvider"
-gpt4all = "jupyter_ai_magics:GPT4AllProvider"
-huggingface_hub = "jupyter_ai_magics:HfHubProvider"
-ollama = "jupyter_ai_magics.partner_providers.ollama:OllamaProvider"
-openai = "jupyter_ai_magics.partner_providers.openai:OpenAIProvider"
-openai-chat = "jupyter_ai_magics.partner_providers.openai:ChatOpenAIProvider"
-openai-chat-custom = "jupyter_ai_magics.partner_providers.openai:ChatOpenAICustomProvider"
-azure-chat-openai = "jupyter_ai_magics.partner_providers.openai:AzureChatOpenAIProvider"
-sagemaker-endpoint = "jupyter_ai_magics.partner_providers.aws:SmEndpointProvider"
-amazon-bedrock = "jupyter_ai_magics.partner_providers.aws:BedrockProvider"
-amazon-bedrock-chat = "jupyter_ai_magics.partner_providers.aws:BedrockChatProvider"
-amazon-bedrock-custom = "jupyter_ai_magics.partner_providers.aws:BedrockCustomProvider"
-qianfan = "jupyter_ai_magics:QianfanProvider"
-nvidia-chat = "jupyter_ai_magics.partner_providers.nvidia:ChatNVIDIAProvider"
-together-ai = "jupyter_ai_magics:TogetherAIProvider"
-gemini = "jupyter_ai_magics.partner_providers.gemini:GeminiProvider"
-mistralai = "jupyter_ai_magics.partner_providers.mistralai:MistralAIProvider"
-openrouter = "jupyter_ai_magics.partner_providers.openrouter:OpenRouterProvider"
-vertexai = "jupyter_ai_magics.partner_providers.vertexai:VertexAIProvider"
-
-[project.entry-points."jupyter_ai.embeddings_model_providers"]
-azure = "jupyter_ai_magics.partner_providers.openai:AzureOpenAIEmbeddingsProvider"
-bedrock = "jupyter_ai_magics.partner_providers.aws:BedrockEmbeddingsProvider"
-cohere = "jupyter_ai_magics.partner_providers.cohere:CohereEmbeddingsProvider"
-mistralai = "jupyter_ai_magics.partner_providers.mistralai:MistralAIEmbeddingsProvider"
-gpt4all = "jupyter_ai_magics:GPT4AllEmbeddingsProvider"
-huggingface_hub = "jupyter_ai_magics:HfHubEmbeddingsProvider"
-ollama = "jupyter_ai_magics.partner_providers.ollama:OllamaEmbeddingsProvider"
-openai = "jupyter_ai_magics.partner_providers.openai:OpenAIEmbeddingsProvider"
-openai-custom = "jupyter_ai_magics.partner_providers.openai:OpenAIEmbeddingsCustomProvider"
-qianfan = "jupyter_ai_magics:QianfanEmbeddingsEndpointProvider"
-vertexai = "jupyter_ai_magics.partner_providers.vertexai:VertexAIEmbeddingsProvider"
-
 [tool.hatch.version]
 source = "nodejs"
 

packages/jupyter-ai/jupyter_ai/config_manager.py

@@ -3,7 +3,7 @@
 import os
 import time
 from copy import deepcopy
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
 from deepmerge import always_merger
 from jupyter_ai_magics.utils import (
@@ -175,69 +175,10 @@ def _process_existing_config(self):
         with open(self.config_path, encoding="utf-8") as f:
             existing_config = json.loads(f.read())
             config = JaiConfig(**existing_config)
-            validated_config = self._validate_model_ids(config)
 
         # re-write to the file to validate the config and apply any
         # updates to the config file immediately
-        self._write_config(validated_config)
-
-    def _validate_model_ids(self, config):
-        lm_provider_keys = ["model_provider_id", "completions_model_provider_id"]
-        em_provider_keys = ["embeddings_provider_id"]
-        clm_provider_keys = ["completions_model_provider_id"]
-
-        # if the currently selected language or embedding model are
-        # forbidden, set them to `None` and log a warning.
-        for lm_key in lm_provider_keys:
-            lm_id = getattr(config, lm_key)
-            if lm_id is not None and not self._validate_model(lm_id, raise_exc=False):
-                self.log.warning(
-                    f"Language model {lm_id} is forbidden by current allow/blocklists. Setting to None."
-                )
-                setattr(config, lm_key, None)
-        for em_key in em_provider_keys:
-            em_id = getattr(config, em_key)
-            if em_id is not None and not self._validate_model(em_id, raise_exc=False):
-                self.log.warning(
-                    f"Embedding model {em_id} is forbidden by current allow/blocklists. Setting to None."
-                )
-                setattr(config, em_key, None)
-        for clm_key in clm_provider_keys:
-            clm_id = getattr(config, clm_key)
-            if clm_id is not None and not self._validate_model(clm_id, raise_exc=False):
-                self.log.warning(
-                    f"Completion model {clm_id} is forbidden by current allow/blocklists. Setting to None."
-                )
-                setattr(config, clm_key, None)
-
-        # if the currently selected language or embedding model ids are
-        # not associated with models, set them to `None` and log a warning.
-        for lm_key in lm_provider_keys:
-            lm_id = getattr(config, lm_key)
-            if lm_id is not None and not get_lm_provider(lm_id, self._lm_providers)[1]:
-                self.log.warning(
-                    f"No language model is associated with '{lm_id}'. Setting to None."
-                )
-                setattr(config, lm_key, None)
-        for em_key in em_provider_keys:
-            em_id = getattr(config, em_key)
-            if em_id is not None and not get_em_provider(em_id, self._em_providers)[1]:
-                self.log.warning(
-                    f"No embedding model is associated with '{em_id}'. Setting to None."
-                )
-                setattr(config, em_key, None)
-        for clm_key in clm_provider_keys:
-            clm_id = getattr(config, clm_key)
-            if (
-                clm_id is not None
-                and not get_lm_provider(clm_id, self._lm_providers)[1]
-            ):
-                self.log.warning(
-                    f"No completion model is associated with '{clm_id}'. Setting to None."
-                )
-                setattr(config, clm_key, None)
-
-        return config
+        self._write_config(config)
 
     def _read_config(self) -> JaiConfig:
         """
@@ -268,78 +209,79 @@ def _validate_config(self, config: JaiConfig):
         user has specified authentication for all configured models that require
         it.
         """
+        # TODO: re-implement this w/ liteLLM
         # validate language model config
-        if config.model_provider_id:
-            _, lm_provider = get_lm_provider(
-                config.model_provider_id, self._lm_providers
-            )
+        # if config.model_provider_id:
+        #     _, lm_provider = get_lm_provider(
+        #         config.model_provider_id, self._lm_providers
+        #     )
 
-            # verify model is declared by some provider
-            if not lm_provider:
-                raise ValueError(
-                    f"No language model is associated with '{config.model_provider_id}'."
-                )
+        #     # verify model is declared by some provider
+        #     if not lm_provider:
+        #         raise ValueError(
+        #             f"No language model is associated with '{config.model_provider_id}'."
+        #         )
 
-            # verify model is not blocked
-            self._validate_model(config.model_provider_id)
+        #     # verify model is not blocked
+        #     self._validate_model(config.model_provider_id)
 
-            # verify model is authenticated
-            _validate_provider_authn(config, lm_provider)
+        #     # verify model is authenticated
+        #     _validate_provider_authn(config, lm_provider)
 
-            # verify fields exist for this model if needed
-            if lm_provider.fields and config.model_provider_id not in config.fields:
-                config.fields[config.model_provider_id] = {}
+        #     # verify fields exist for this model if needed
+        #     if lm_provider.fields and config.model_provider_id not in config.fields:
+        #         config.fields[config.model_provider_id] = {}
 
         # validate completions model config
-        if config.completions_model_provider_id:
-            _, completions_provider = get_lm_provider(
-                config.completions_model_provider_id, self._lm_providers
-            )
-
-            # verify model is declared by some provider
-            if not completions_provider:
-                raise ValueError(
-                    f"No language model is associated with '{config.completions_model_provider_id}'."
-                )
-
-            # verify model is not blocked
-            self._validate_model(config.completions_model_provider_id)
-
-            # verify model is authenticated
-            _validate_provider_authn(config, completions_provider)
-
-            # verify completions fields exist for this model if needed
-            if (
-                completions_provider.fields
-                and config.completions_model_provider_id
-                not in config.completions_fields
-            ):
-                config.completions_fields[config.completions_model_provider_id] = {}
-
-        # validate embedding model config
-        if config.embeddings_provider_id:
-            _, em_provider = get_em_provider(
-                config.embeddings_provider_id, self._em_providers
-            )
-
-            # verify model is declared by some provider
-            if not em_provider:
-                raise ValueError(
-                    f"No embedding model is associated with '{config.embeddings_provider_id}'."
-                )
-
-            # verify model is not blocked
-            self._validate_model(config.embeddings_provider_id)
-
-            # verify model is authenticated
-            _validate_provider_authn(config, em_provider)
-
-            # verify embedding fields exist for this model if needed
-            if (
-                em_provider.fields
-                and config.embeddings_provider_id not in config.embeddings_fields
-            ):
-                config.embeddings_fields[config.embeddings_provider_id] = {}
+        # if config.completions_model_provider_id:
+        #     _, completions_provider = get_lm_provider(
+        #         config.completions_model_provider_id, self._lm_providers
+        #     )
+
+        #     # verify model is declared by some provider
+        #     if not completions_provider:
+        #         raise ValueError(
+        #             f"No language model is associated with '{config.completions_model_provider_id}'."
+        #         )
+
+        #     # verify model is not blocked
+        #     self._validate_model(config.completions_model_provider_id)
+
+        #     # verify model is authenticated
+        #     _validate_provider_authn(config, completions_provider)
+
+        #     # verify completions fields exist for this model if needed
+        #     if (
+        #         completions_provider.fields
+        #         and config.completions_model_provider_id
+        #         not in config.completions_fields
+        #     ):
+        #         config.completions_fields[config.completions_model_provider_id] = {}
+
+        # # validate embedding model config
+        # if config.embeddings_provider_id:
+        #     _, em_provider = get_em_provider(
+        #         config.embeddings_provider_id, self._em_providers
+        #     )
+
+        #     # verify model is declared by some provider
+        #     if not em_provider:
+        #         raise ValueError(
+        #             f"No embedding model is associated with '{config.embeddings_provider_id}'."
+        #         )
+
+        #     # verify model is not blocked
+        #     self._validate_model(config.embeddings_provider_id)
+
+        #     # verify model is authenticated
+        #     _validate_provider_authn(config, em_provider)
+
+        #     # verify embedding fields exist for this model if needed
+        #     if (
+        #         em_provider.fields
+        #         and config.embeddings_provider_id not in config.embeddings_fields
+        #     ):
+        #         config.embeddings_fields[config.embeddings_provider_id] = {}
 
     def _validate_model(self, model_id: str, raise_exc=True):
         """
@@ -449,23 +391,30 @@ def get_config(self):
         )
 
     @property
-    def lm_gid(self):
+    def chat_model(self) -> str | None:
+        """
+        Returns the model ID of the chat model from AI settings, if any.
+        """
         config = self._read_config()
         return config.model_provider_id
 
     @property
-    def em_gid(self):
-        config = self._read_config()
-        return config.embeddings_provider_id
+    def chat_model_params(self) -> dict[str, Any]:
+        return self._provider_params("model_provider_id", self._lm_providers)
 
     @property
-    def lm_provider(self):
-        return self._get_provider("model_provider_id", self._lm_providers)
+    def embedding_model(self) -> str | None:
+        """
+        Returns the model ID of the embedding model from AI settings, if any.
+        """
+        config = self._read_config()
+        return config.embeddings_provider_id
 
     @property
-    def em_provider(self):
-        return self._get_provider("embeddings_provider_id", self._em_providers)
+    def embedding_model_params(self) -> dict[str, Any]:
+        return self._provider_params("embeddings_provider_id", self._em_providers)
 
+    # TODO: use LiteLLM in completions
     @property
     def completions_lm_provider(self):
         return self._get_provider("completions_model_provider_id", self._lm_providers)
@@ -479,14 +428,7 @@ def _get_provider(self, key, listing):
         _, Provider = get_lm_provider(gid, listing)
         return Provider
 
-    @property
-    def lm_provider_params(self):
-        return self._provider_params("model_provider_id", self._lm_providers)
-
-    @property
-    def em_provider_params(self):
-        return self._provider_params("embeddings_provider_id", self._em_providers)
-
+    # TODO: use LiteLLM in completions
     @property
     def completions_lm_provider_params(self):
         return self._provider_params(