mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
Merge pull request #4403 from BerriAI/litellm_add_nvidia_nim
[Feat-New Provider] Add Nvidia NIM
This commit is contained in:
Ishaan Jaff
GitHub
commit
6f51da4e78
8 changed files with 247 additions and 12 deletions
103
docs/my-website/docs/providers/nvidia_nim.md
Normal file
103
docs/my-website/docs/providers/nvidia_nim.md
Normal file
|
|
@ -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/<any-model-on-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)` |
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
79
litellm/llms/nvidia_nim.py
Normal file
79
litellm/llms/nvidia_nim.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue