From 07829514d11df6543d85c38e92ba477e1691cff8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 25 Jun 2024 09:13:08 -0700 Subject: [PATCH] feat - add param mapping for nvidia nim --- litellm/__init__.py | 1 + litellm/llms/nvidia_nim.py | 79 ++++++++++++++++++++++++++++++++++++++ litellm/utils.py | 23 +++++++++++ 3 files changed, 103 insertions(+) create mode 100644 litellm/llms/nvidia_nim.py diff --git a/litellm/__init__.py b/litellm/__init__.py index d23247d53..08ee84aaa 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -816,6 +816,7 @@ from .llms.openai import ( DeepInfraConfig, AzureAIStudioConfig, ) +from .llms.nvidia_nim import NvidiaNimConfig from .llms.text_completion_codestral import MistralTextCompletionConfig from .llms.azure import ( AzureOpenAIConfig, diff --git a/litellm/llms/nvidia_nim.py b/litellm/llms/nvidia_nim.py new file mode 100644 index 000000000..ebcc84c13 --- /dev/null +++ b/litellm/llms/nvidia_nim.py @@ -0,0 +1,79 @@ +""" +Nvidia NIM endpoint: https://docs.api.nvidia.com/nim/reference/databricks-dbrx-instruct-infer + +This is OpenAI compatible + +This file only contains param mapping logic + +API calling is done using the OpenAI SDK with an api_base +""" + +import types +from typing import Optional, Union + + +class NvidiaNimConfig: + """ + Reference: https://docs.api.nvidia.com/nim/reference/databricks-dbrx-instruct-infer + + The class `NvidiaNimConfig` provides configuration for the Nvidia NIM's Chat Completions API interface. Below are the parameters: + """ + + temperature: Optional[int] = None + top_p: Optional[int] = None + frequency_penalty: Optional[int] = None + presence_penalty: Optional[int] = None + max_tokens: Optional[int] = None + stop: Optional[Union[str, list]] = None + + def __init__( + self, + temperature: Optional[int] = None, + top_p: Optional[int] = None, + frequency_penalty: Optional[int] = None, + presence_penalty: Optional[int] = None, + max_tokens: Optional[int] = None, + stop: Optional[Union[str, list]] = None, + ) -> None: + locals_ = locals().copy() + for key, value in locals_.items(): + if key != "self" and value is not None: + setattr(self.__class__, key, value) + + @classmethod + def get_config(cls): + return { + k: v + for k, v in cls.__dict__.items() + if not k.startswith("__") + and not isinstance( + v, + ( + types.FunctionType, + types.BuiltinFunctionType, + classmethod, + staticmethod, + ), + ) + and v is not None + } + + def get_supported_openai_params(self): + return [ + "stream", + "temperature", + "top_p", + "frequency_penalty", + "presence_penalty", + "max_tokens", + "stop", + ] + + def map_openai_params( + self, non_default_params: dict, optional_params: dict + ) -> dict: + supported_openai_params = self.get_supported_openai_params() + for param, value in non_default_params.items(): + if param in supported_openai_params: + optional_params[param] = value + return optional_params diff --git a/litellm/utils.py b/litellm/utils.py index 1bc8bf771..7709e8821 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2410,6 +2410,7 @@ def get_optional_params( and custom_llm_provider != "anyscale" and custom_llm_provider != "together_ai" and custom_llm_provider != "groq" + and custom_llm_provider != "nvidia_nim" and custom_llm_provider != "deepseek" and custom_llm_provider != "codestral" and custom_llm_provider != "mistral" @@ -3060,6 +3061,14 @@ def get_optional_params( optional_params = litellm.DatabricksConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params ) + elif custom_llm_provider == "nvidia_nim": + supported_params = get_supported_openai_params( + model=model, custom_llm_provider=custom_llm_provider + ) + _check_valid_arg(supported_params=supported_params) + optional_params = litellm.NvidiaNimConfig().map_openai_params( + non_default_params=non_default_params, optional_params=optional_params + ) elif custom_llm_provider == "groq": supported_params = get_supported_openai_params( model=model, custom_llm_provider=custom_llm_provider @@ -3626,6 +3635,8 @@ def get_supported_openai_params( return litellm.OllamaChatConfig().get_supported_openai_params() elif custom_llm_provider == "anthropic": return litellm.AnthropicConfig().get_supported_openai_params() + elif custom_llm_provider == "nvidia_nim": + return litellm.NvidiaNimConfig().get_supported_openai_params() elif custom_llm_provider == "groq": return [ "temperature", @@ -3986,6 +3997,10 @@ def get_llm_provider( # groq is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.groq.com/openai/v1 api_base = "https://api.groq.com/openai/v1" dynamic_api_key = get_secret("GROQ_API_KEY") + elif custom_llm_provider == "nvidia_nim": + # nvidia_nim is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.endpoints.anyscale.com/v1 + api_base = "https://integrate.api.nvidia.com/v1" + dynamic_api_key = get_secret("NVIDIA_NIM_API_KEY") elif custom_llm_provider == "codestral": # codestral is openai compatible, we just need to set this to custom_openai and have the api_base be https://codestral.mistral.ai/v1 api_base = "https://codestral.mistral.ai/v1" @@ -4087,6 +4102,9 @@ def get_llm_provider( elif endpoint == "api.groq.com/openai/v1": custom_llm_provider = "groq" dynamic_api_key = get_secret("GROQ_API_KEY") + elif endpoint == "https://integrate.api.nvidia.com/v1": + custom_llm_provider = "nvidia_nim" + dynamic_api_key = get_secret("NVIDIA_NIM_API_KEY") elif endpoint == "https://codestral.mistral.ai/v1": custom_llm_provider = "codestral" dynamic_api_key = get_secret("CODESTRAL_API_KEY") @@ -4900,6 +4918,11 @@ def validate_environment(model: Optional[str] = None) -> dict: keys_in_environment = True else: missing_keys.append("GROQ_API_KEY") + elif custom_llm_provider == "nvidia_nim": + if "NVIDIA_NIM_API_KEY" in os.environ: + keys_in_environment = True + else: + missing_keys.append("NVIDIA_NIM_API_KEY") elif ( custom_llm_provider == "codestral" or custom_llm_provider == "text-completion-codestral"