From 84d5f35a48c7ff28e1372958c203f0a0247e2385 Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Date: Tue, 19 Nov 2024 00:22:24 -0800
Subject: [PATCH] Update the model alias for llama guard models in ollama

---
 .../remote/inference/ollama/ollama.py          | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/llama_stack/providers/remote/inference/ollama/ollama.py b/llama_stack/providers/remote/inference/ollama/ollama.py
index 1c5d26a84..f53ed4e14 100644
--- a/llama_stack/providers/remote/inference/ollama/ollama.py
+++ b/llama_stack/providers/remote/inference/ollama/ollama.py
@@ -73,14 +73,6 @@ model_aliases = [
         "llama3.2:3b",
         CoreModelId.llama3_2_3b_instruct.value,
     ),
-    build_model_alias_with_just_provider_model_id(
-        "llama-guard3:8b",
-        CoreModelId.llama_guard_3_8b.value,
-    ),
-    build_model_alias_with_just_provider_model_id(
-        "llama-guard3:1b",
-        CoreModelId.llama_guard_3_1b.value,
-    ),
     build_model_alias(
         "llama3.2-vision:11b-instruct-fp16",
         CoreModelId.llama3_2_11b_vision_instruct.value,
@@ -89,6 +81,16 @@ model_aliases = [
         "llama3.2-vision",
         CoreModelId.llama3_2_11b_vision_instruct.value,
     ),
+    # The Llama Guard models don't have their full fp16 versions
+    # so we are going to alias their default version to the canonical SKU
+    build_model_alias(
+        "llama-guard3:8b",
+        CoreModelId.llama_guard_3_8b.value,
+    ),
+    build_model_alias(
+        "llama-guard3:1b",
+        CoreModelId.llama_guard_3_1b.value,
+    ),
 ]