diff --git a/docs/my-website/docs/providers/nvidia_nim.md b/docs/my-website/docs/providers/nvidia_nim.md new file mode 100644 index 00000000000..7f895aa337c --- /dev/null +++ b/docs/my-website/docs/providers/nvidia_nim.md @@ -0,0 +1,103 @@ +# Nvidia NIM +https://docs.api.nvidia.com/nim/reference/ + +:::tip + +**We support ALL Nvidia NIM models, just set `model=nvidia_nim/` as a prefix when sending litellm requests** + +::: + +## API Key +```python +# env variable +os.environ['NVIDIA_NIM_API_KEY'] +``` + +## Sample Usage +```python +from litellm import completion +import os + +os.environ['NVIDIA_NIM_API_KEY'] = "" +response = completion( + model="nvidia_nim/meta/llama3-70b-instruct", + messages=[ + { + "role": "user", + "content": "What's the weather like in Boston today in Fahrenheit?", + } + ], + temperature=0.2, # optional + top_p=0.9, # optional + frequency_penalty=0.1, # optional + presence_penalty=0.1, # optional + max_tokens=10, # optional + stop=["\n\n"], # optional +) +print(response) +``` + +## Sample Usage - Streaming +```python +from litellm import completion +import os + +os.environ['NVIDIA_NIM_API_KEY'] = "" +response = completion( + model="nvidia_nim/meta/llama3-70b-instruct", + messages=[ + { + "role": "user", + "content": "What's the weather like in Boston today in Fahrenheit?", + } + ], + stream=True, + temperature=0.2, # optional + top_p=0.9, # optional + frequency_penalty=0.1, # optional + presence_penalty=0.1, # optional + max_tokens=10, # optional + stop=["\n\n"], # optional +) + +for chunk in response: + print(chunk) +``` + + +## Supported Models - 💥 ALL Nvidia NIM Models Supported! +We support ALL `nvidia_nim` models, just set `nvidia_nim/` as a prefix when sending completion requests + +| Model Name | Function Call | +|------------|---------------| +| nvidia/nemotron-4-340b-reward | `completion(model="nvidia_nim/nvidia/nemotron-4-340b-reward", messages)` | +| 01-ai/yi-large | `completion(model="nvidia_nim/01-ai/yi-large", messages)` | +| aisingapore/sea-lion-7b-instruct | `completion(model="nvidia_nim/aisingapore/sea-lion-7b-instruct", messages)` | +| databricks/dbrx-instruct | `completion(model="nvidia_nim/databricks/dbrx-instruct", messages)` | +| google/gemma-7b | `completion(model="nvidia_nim/google/gemma-7b", messages)` | +| google/gemma-2b | `completion(model="nvidia_nim/google/gemma-2b", messages)` | +| google/codegemma-1.1-7b | `completion(model="nvidia_nim/google/codegemma-1.1-7b", messages)` | +| google/codegemma-7b | `completion(model="nvidia_nim/google/codegemma-7b", messages)` | +| google/recurrentgemma-2b | `completion(model="nvidia_nim/google/recurrentgemma-2b", messages)` | +| ibm/granite-34b-code-instruct | `completion(model="nvidia_nim/ibm/granite-34b-code-instruct", messages)` | +| ibm/granite-8b-code-instruct | `completion(model="nvidia_nim/ibm/granite-8b-code-instruct", messages)` | +| mediatek/breeze-7b-instruct | `completion(model="nvidia_nim/mediatek/breeze-7b-instruct", messages)` | +| meta/codellama-70b | `completion(model="nvidia_nim/meta/codellama-70b", messages)` | +| meta/llama2-70b | `completion(model="nvidia_nim/meta/llama2-70b", messages)` | +| meta/llama3-8b | `completion(model="nvidia_nim/meta/llama3-8b", messages)` | +| meta/llama3-70b | `completion(model="nvidia_nim/meta/llama3-70b", messages)` | +| microsoft/phi-3-medium-4k-instruct | `completion(model="nvidia_nim/microsoft/phi-3-medium-4k-instruct", messages)` | +| microsoft/phi-3-mini-128k-instruct | `completion(model="nvidia_nim/microsoft/phi-3-mini-128k-instruct", messages)` | +| microsoft/phi-3-mini-4k-instruct | `completion(model="nvidia_nim/microsoft/phi-3-mini-4k-instruct", messages)` | +| microsoft/phi-3-small-128k-instruct | `completion(model="nvidia_nim/microsoft/phi-3-small-128k-instruct", messages)` | +| microsoft/phi-3-small-8k-instruct | `completion(model="nvidia_nim/microsoft/phi-3-small-8k-instruct", messages)` | +| mistralai/codestral-22b-instruct-v0.1 | `completion(model="nvidia_nim/mistralai/codestral-22b-instruct-v0.1", messages)` | +| mistralai/mistral-7b-instruct | `completion(model="nvidia_nim/mistralai/mistral-7b-instruct", messages)` | +| mistralai/mistral-7b-instruct-v0.3 | `completion(model="nvidia_nim/mistralai/mistral-7b-instruct-v0.3", messages)` | +| mistralai/mixtral-8x7b-instruct | `completion(model="nvidia_nim/mistralai/mixtral-8x7b-instruct", messages)` | +| mistralai/mixtral-8x22b-instruct | `completion(model="nvidia_nim/mistralai/mixtral-8x22b-instruct", messages)` | +| mistralai/mistral-large | `completion(model="nvidia_nim/mistralai/mistral-large", messages)` | +| nvidia/nemotron-4-340b-instruct | `completion(model="nvidia_nim/nvidia/nemotron-4-340b-instruct", messages)` | +| seallms/seallm-7b-v2.5 | `completion(model="nvidia_nim/seallms/seallm-7b-v2.5", messages)` | +| snowflake/arctic | `completion(model="nvidia_nim/snowflake/arctic", messages)` | +| upstage/solar-10.7b-instruct | `completion(model="nvidia_nim/upstage/solar-10.7b-instruct", messages)` | \ No newline at end of file diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 2673933f4ca..9835a260b34 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -146,13 +146,14 @@ const sidebars = { "providers/databricks", "providers/watsonx", "providers/predibase", - "providers/clarifai", + "providers/nvidia_nim", "providers/triton-inference-server", "providers/ollama", "providers/perplexity", "providers/groq", "providers/deepseek", - "providers/fireworks_ai", + "providers/fireworks_ai", + "providers/clarifai", "providers/vllm", "providers/xinference", "providers/cloudflare_workers", diff --git a/litellm/__init__.py b/litellm/__init__.py index f07ce880928..08ee84aaad5 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -401,6 +401,7 @@ openai_compatible_endpoints: List = [ "codestral.mistral.ai/v1/chat/completions", "codestral.mistral.ai/v1/fim/completions", "api.groq.com/openai/v1", + "https://integrate.api.nvidia.com/v1", "api.deepseek.com/v1", "api.together.xyz/v1", "inference.friendli.ai/v1", @@ -411,6 +412,7 @@ openai_compatible_providers: List = [ "anyscale", "mistral", "groq", + "nvidia_nim", "codestral", "deepseek", "deepinfra", @@ -640,6 +642,7 @@ provider_list: List = [ "anyscale", "mistral", "groq", + "nvidia_nim", "codestral", "text-completion-codestral", "deepseek", @@ -813,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 00000000000..ebcc84c13eb --- /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/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py index e359d36f4c2..a97d6812c8c 100644 --- a/litellm/llms/prompt_templates/factory.py +++ b/litellm/llms/prompt_templates/factory.py @@ -135,7 +135,7 @@ def convert_to_ollama_image(openai_image_url: str): def ollama_pt( - model, messages + model, messages ): # https://github.com/ollama/ollama/blob/af4cf55884ac54b9e637cd71dadfe9b7a5685877/docs/modelfile.md#template if "instruct" in model: prompt = custom_prompt( @@ -185,19 +185,18 @@ def ollama_pt( function_name: str = call["function"]["name"] arguments = json.loads(call["function"]["arguments"]) - tool_calls.append({ - "id": call_id, - "type": "function", - "function": { - "name": function_name, - "arguments": arguments + tool_calls.append( + { + "id": call_id, + "type": "function", + "function": {"name": function_name, "arguments": arguments}, } - }) + ) prompt += f"### Assistant:\nTool Calls: {json.dumps(tool_calls, indent=2)}\n\n" elif "tool_call_id" in message: - prompt += f"### User:\n{message["content"]}\n\n" + prompt += f"### User:\n{message['content']}\n\n" elif content: prompt += f"### {role.capitalize()}:\n{content}\n\n" diff --git a/litellm/main.py b/litellm/main.py index 307659c8a2f..8c531643b8e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -348,6 +348,7 @@ async def acompletion( or custom_llm_provider == "deepinfra" or custom_llm_provider == "perplexity" or custom_llm_provider == "groq" + or custom_llm_provider == "nvidia_nim" or custom_llm_provider == "codestral" or custom_llm_provider == "text-completion-codestral" or custom_llm_provider == "deepseek" @@ -1171,6 +1172,7 @@ def completion( or custom_llm_provider == "deepinfra" or custom_llm_provider == "perplexity" or custom_llm_provider == "groq" + or custom_llm_provider == "nvidia_nim" or custom_llm_provider == "codestral" or custom_llm_provider == "deepseek" or custom_llm_provider == "anyscale" @@ -2932,6 +2934,7 @@ async def aembedding(*args, **kwargs) -> EmbeddingResponse: or custom_llm_provider == "deepinfra" or custom_llm_provider == "perplexity" or custom_llm_provider == "groq" + or custom_llm_provider == "nvidia_nim" or custom_llm_provider == "deepseek" or custom_llm_provider == "fireworks_ai" or custom_llm_provider == "ollama" @@ -3507,6 +3510,7 @@ async def atext_completion( or custom_llm_provider == "deepinfra" or custom_llm_provider == "perplexity" or custom_llm_provider == "groq" + or custom_llm_provider == "nvidia_nim" or custom_llm_provider == "text-completion-codestral" or custom_llm_provider == "deepseek" or custom_llm_provider == "fireworks_ai" diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 830b3acd380..30ae1d0ab13 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -23,7 +23,7 @@ from litellm import RateLimitError, Timeout, completion, completion_cost, embedd from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.prompt_templates.factory import anthropic_messages_pt -# litellm.num_retries=3 +# litellm.num_retries = 3 litellm.cache = None litellm.success_callback = [] user_message = "Write a short poem about the sky" @@ -3470,6 +3470,28 @@ def test_completion_deep_infra_mistral(): # test_completion_deep_infra_mistral() +def test_completion_nvidia_nim(): + model_name = "nvidia_nim/databricks/dbrx-instruct" + try: + response = completion( + model=model_name, + messages=[ + { + "role": "user", + "content": "What's the weather like in Boston today in Fahrenheit?", + } + ], + ) + # Add any assertions here to check the response + print(response) + assert response.choices[0].message.content is not None + assert len(response.choices[0].message.content) > 0 + except litellm.exceptions.Timeout as e: + pass + except Exception as e: + pytest.fail(f"Error occurred: {e}") + + # Gemini tests @pytest.mark.parametrize( "model", diff --git a/litellm/utils.py b/litellm/utils.py index 1bc8bf771f4..7709e882105 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"