From 6e829f8f516f7edf6f5775638e8ac10bc77379df Mon Sep 17 00:00:00 2001
From: Chantal D Gama Rose <cdgamarose@nvidia.com>
Date: Thu, 19 Dec 2024 11:45:07 -0800
Subject: [PATCH] Fixed imports for inference

---
 llama_stack/providers/remote/inference/nvidia/openai_utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/llama_stack/providers/remote/inference/nvidia/openai_utils.py b/llama_stack/providers/remote/inference/nvidia/openai_utils.py
index ba8ff0fa4..ffca32c44 100644
--- a/llama_stack/providers/remote/inference/nvidia/openai_utils.py
+++ b/llama_stack/providers/remote/inference/nvidia/openai_utils.py
@@ -10,9 +10,7 @@ from typing import Any, AsyncGenerator, Dict, Generator, List, Optional
 
 from llama_models.llama3.api.datatypes import (
     BuiltinTool,
-    CompletionMessage,
     StopReason,
-    TokenLogProbs,
     ToolCall,
     ToolDefinition,
 )
@@ -42,12 +40,14 @@ from llama_stack.apis.inference import (
     ChatCompletionResponseEvent,
     ChatCompletionResponseEventType,
     ChatCompletionResponseStreamChunk,
+    CompletionMessage,
     CompletionRequest,
     CompletionResponse,
     CompletionResponseStreamChunk,
     JsonSchemaResponseFormat,
     Message,
     SystemMessage,
+    TokenLogProbs,
     ToolCallDelta,
     ToolCallParseStatus,
     ToolResponseMessage,