From da9f0b786932f7c6995f7c65781da44ec7a25605 Mon Sep 17 00:00:00 2001
From: Rashmi Pawar <168514198+raspawar@users.noreply.github.com>
Date: Fri, 21 Feb 2025 21:39:17 +0530
Subject: [PATCH] test(client-sdk): Update embedding test types to use latest
 imports (#1203)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

# What does this PR do?
- Updates ImageContentItemImageURL import
- fixes `embedding_dimensions` metadata param

## Test Plan
- Ran pytest locally, verified embedding tests pass with new types

![Screenshot 2025-02-21 at 6 54
27 PM](https://github.com/user-attachments/assets/f80e3785-04c3-415e-9276-88aa8136bf00)

cc: @dglogo @sumitb
---
 llama_stack/providers/remote/inference/nvidia/models.py | 2 +-
 llama_stack/templates/nvidia/run.yaml                   | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/llama_stack/providers/remote/inference/nvidia/models.py b/llama_stack/providers/remote/inference/nvidia/models.py
index fa9944be1..4305f4c6f 100644
--- a/llama_stack/providers/remote/inference/nvidia/models.py
+++ b/llama_stack/providers/remote/inference/nvidia/models.py
@@ -52,7 +52,7 @@ _MODEL_ENTRIES = [
         provider_model_id="baai/bge-m3",
         model_type=ModelType.embedding,
         metadata={
-            "embedding_dimensions": 1024,
+            "embedding_dimension": 1024,
             "context_length": 8192,
         },
     ),
diff --git a/llama_stack/templates/nvidia/run.yaml b/llama_stack/templates/nvidia/run.yaml
index 891fd112a..4c38ec24e 100644
--- a/llama_stack/templates/nvidia/run.yaml
+++ b/llama_stack/templates/nvidia/run.yaml
@@ -136,7 +136,7 @@ models:
   provider_model_id: meta/llama-3.2-90b-vision-instruct
   model_type: llm
 - metadata:
-    embedding_dimensions: 1024
+    embedding_dimension: 1024
     context_length: 8192
   model_id: baai/bge-m3
   provider_id: nvidia