add nvidia nim to __init__

This commit is contained in:
Ishaan Jaff 2024-06-25 08:53:06 -07:00
parent e471aca81f
commit b55c31ec3f
2 changed files with 11 additions and 9 deletions

View file

@ -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",

View file

@ -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"