diff --git a/llama_stack/providers/tests/inference/test_vision_inference.py b/llama_stack/providers/tests/inference/test_vision_inference.py
index d58164676..1bdee051f 100644
--- a/llama_stack/providers/tests/inference/test_vision_inference.py
+++ b/llama_stack/providers/tests/inference/test_vision_inference.py
@@ -8,11 +8,16 @@ from pathlib import Path
 
 import pytest
 
-
-from llama_models.llama3.api.datatypes import *  # noqa: F403
-from llama_stack.apis.inference import *  # noqa: F403
 from llama_stack.apis.common.content_types import ImageContentItem, TextContentItem, URL
 
+from llama_stack.apis.inference import (
+    ChatCompletionResponse,
+    ChatCompletionResponseEventType,
+    ChatCompletionResponseStreamChunk,
+    SamplingParams,
+    UserMessage,
+)
+
 from .utils import group_chunks
 
 THIS_DIR = Path(__file__).parent