fix security update
All checks were successful
Build and Push playground container / build-playground (push) Successful in 1m57s
Build and Push container / build (push) Successful in 4m33s

This commit is contained in:
Angel Nunez Mencias 2025-06-03 20:07:06 +02:00
parent 8943b283e9
commit b174effe05
Signed by: angel.nunez
SSH key fingerprint: SHA256:z1nFAg1v1AfbhEHrgBetByUJUwziv2R2f4VyN75opcg
14 changed files with 45 additions and 41 deletions

View file

@ -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={