From b55c31ec3fd5127e8595a3933d8eef0b87e4e48f Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 25 Jun 2024 08:53:06 -0700
Subject: [PATCH] add nvidia nim to __init__

---
 litellm/__init__.py                      |  3 +++
 litellm/llms/prompt_templates/factory.py | 17 ++++++++---------
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/litellm/__init__.py b/litellm/__init__.py
index f07ce880928..d23247d5313 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",
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"