model selection playground fix

meta-llama · Dec 18, 2024 · af8f1b3 · af8f1b3
1 parent 3700022
commit af8f1b3
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/llama_stack/distribution/ui/page/playground/chat.py b/llama_stack/distribution/ui/page/playground/chat.py
@@ -11,7 +11,11 @@
 with st.sidebar:
     st.header("Configuration")
     available_models = llama_stack_api.client.models.list()
-    available_models = [model.identifier for model in available_models]
+    available_models = [
+        model.identifier
+        for model in available_models
+        if model.identifier.startswith("meta-llama")
+    ]
     selected_model = st.selectbox(
         "Choose a model",
         available_models,

diff --git a/llama_stack/distribution/ui/page/playground/rag.py b/llama_stack/distribution/ui/page/playground/rag.py
@@ -74,7 +74,11 @@ def rag_chat_page():
         ]
 
         available_models = llama_stack_api.client.models.list()
-        available_models = [model.identifier for model in available_models]
+        available_models = [
+            model.identifier
+            for model in available_models
+            if model.identifier.startswith("meta-llama")
+        ]
         selected_model = st.selectbox(
             "Choose a model",
             available_models,
@@ -116,8 +120,6 @@ def rag_chat_page():
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
 
-    selected_model = llama_stack_api.client.models.list()[0].identifier
-
     agent_config = AgentConfig(
         model=selected_model,
         instructions=system_prompt,