fix security update

llama_stack/distribution/ui/page/playground/chat.py

@@ -6,12 +6,12 @@
 
 import streamlit as st
 
-from llama_stack.distribution.ui.modules.api import llama_stack_api
+from llama_stack.distribution.ui.modules.api import LlamaStackApi
 
 # Sidebar configurations
 with st.sidebar:
     st.header("Configuration")
-    available_models = llama_stack_api.client.models.list()
+    available_models = LlamaStackApi().client.models.list()
     available_models = [model.identifier for model in available_models if model.model_type == "llm"]
     selected_model = st.selectbox(
         "Choose a model",
@@ -103,7 +103,7 @@ if prompt := st.chat_input("Example: What is Llama Stack?"):
         else:
             strategy = {"type": "greedy"}
 
-        response = llama_stack_api.client.inference.chat_completion(
+        response = LlamaStackApi().client.inference.chat_completion(
             messages=[
                 {"role": "system", "content": system_prompt},
                 {"role": "user", "content": prompt},

llama_stack/distribution/ui/page/playground/rag.py

@@ -10,7 +10,7 @@ import streamlit as st
 from llama_stack_client import Agent, AgentEventLogger, RAGDocument
 
 from llama_stack.apis.common.content_types import ToolCallDelta
-from llama_stack.distribution.ui.modules.api import llama_stack_api
+from llama_stack.distribution.ui.modules.api import LlamaStackApi
 from llama_stack.distribution.ui.modules.utils import data_url_from_file
 
 
@@ -57,14 +57,14 @@ def rag_chat_page():
                     for i, uploaded_file in enumerate(uploaded_files)
                 ]
 
-                providers = llama_stack_api.client.providers.list()
+                providers = LlamaStackApi().client.providers.list()
                 vector_io_provider = None
 
                 for x in providers:
                     if x.api == "vector_io":
                         vector_io_provider = x.provider_id
 
-                llama_stack_api.client.vector_dbs.register(
+                LlamaStackApi().client.vector_dbs.register(
                     vector_db_id=vector_db_name,  # Use the user-provided name
                     embedding_dimension=384,
                     embedding_model="all-MiniLM-L6-v2",
@@ -72,7 +72,7 @@ def rag_chat_page():
                 )
 
                 # insert documents using the custom vector db name
-                llama_stack_api.client.tool_runtime.rag_tool.insert(
+                LlamaStackApi().client.tool_runtime.rag_tool.insert(
                     vector_db_id=vector_db_name,  # Use the user-provided name
                     documents=documents,
                     chunk_size_in_tokens=512,
@@ -93,7 +93,7 @@ def rag_chat_page():
         )
 
         # select memory banks
-        vector_dbs = llama_stack_api.client.vector_dbs.list()
+        vector_dbs = LlamaStackApi().client.vector_dbs.list()
         vector_dbs = [vector_db.identifier for vector_db in vector_dbs]
         selected_vector_dbs = st.multiselect(
             label="Select Document Collections to use in RAG queries",
@@ -103,7 +103,7 @@ def rag_chat_page():
         )
 
         st.subheader("Inference Parameters", divider=True)
-        available_models = llama_stack_api.client.models.list()
+        available_models = LlamaStackApi().client.models.list()
         available_models = [model.identifier for model in available_models if model.model_type == "llm"]
         selected_model = st.selectbox(
             label="Choose a model",
@@ -167,7 +167,7 @@ def rag_chat_page():
     @st.cache_resource
     def create_agent():
         return Agent(
-            llama_stack_api.client,
+            LlamaStackApi().client,
             model=selected_model,
             instructions=system_prompt,
             sampling_params={
@@ -232,7 +232,7 @@ def rag_chat_page():
             st.session_state.messages.append({"role": "system", "content": system_prompt})
 
         # Query the vector DB
-        rag_response = llama_stack_api.client.tool_runtime.rag_tool.query(
+        rag_response = LlamaStackApi().client.tool_runtime.rag_tool.query(
             content=prompt, vector_db_ids=list(selected_vector_dbs)
         )
         prompt_context = rag_response.content
@@ -251,7 +251,7 @@ def rag_chat_page():
 
             # Run inference directly
             st.session_state.messages.append({"role": "user", "content": extended_prompt})
-            response = llama_stack_api.client.inference.chat_completion(
+            response = LlamaStackApi().client.inference.chat_completion(
                 messages=st.session_state.messages,
                 model_id=selected_model,
                 sampling_params={

llama_stack/distribution/ui/page/playground/tools.py

@@ -13,7 +13,7 @@ from llama_stack_client import Agent
 from llama_stack_client.lib.agents.react.agent import ReActAgent
 from llama_stack_client.lib.agents.react.tool_parser import ReActOutput
 
-from llama_stack.distribution.ui.modules.api import llama_stack_api
+from llama_stack.distribution.ui.modules.api import LlamaStackApi
 
 
 class AgentType(enum.Enum):
@@ -24,7 +24,7 @@ class AgentType(enum.Enum):
 def tool_chat_page():
     st.title("🛠 Tools")
 
-    client = llama_stack_api.client
+    client = LlamaStackApi().client
     models = client.models.list()
     model_list = [model.identifier for model in models if model.api_model_type == "llm"]
 
@@ -55,7 +55,7 @@ def tool_chat_page():
         )
 
         if "builtin::rag" in toolgroup_selection:
-            vector_dbs = llama_stack_api.client.vector_dbs.list() or []
+            vector_dbs = LlamaStackApi().client.vector_dbs.list() or []
             if not vector_dbs:
                 st.info("No vector databases available for selection.")
             vector_dbs = [vector_db.identifier for vector_db in vector_dbs]