phoenix-oss/llama-stack
1
0
Fork 0
forked from phoenix-oss/llama-stack-mirror
Code Pull requests Releases Packages 2 Activity Actions

feat: Hide tool output under an expander in Playground UI (#2003)

Browse source 
# What does this PR do?
Now, tool outputs and retrieved chunks from the vector DB (i.e.,
everything except for the actual model reply) are hidden under an
expander form when presented to the user.

# Test Plan
Navigate to the RAG page in the Playground UI.
This commit is contained in:
Ilya Kolchinsky 2025-04-23 15:32:12 +02:00 • committed by GitHub
parent d6e88e0bc6
commit d39462d073
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 16 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

25
llama_stack/distribution/ui/page/playground/rag.py
View file

@ -24,6 +24,13 @@ def rag_chat_page():
def should_disable_input(): def should_disable_input():
return "displayed_messages" in st.session_state and len(st.session_state.displayed_messages) > 0 return "displayed_messages" in st.session_state and len(st.session_state.displayed_messages) > 0
def log_message(message):
with st.chat_message(message["role"]):
if "tool_output" in message and message["tool_output"]:
with st.expander(label="Tool Output", expanded=False, icon="🛠"):
st.write(message["tool_output"])
st.markdown(message["content"])
with st.sidebar: with st.sidebar:
# File/Directory Upload Section # File/Directory Upload Section
st.subheader("Upload Documents", divider=True) st.subheader("Upload Documents", divider=True)
@ -146,8 +153,7 @@ def rag_chat_page():
# Display chat history # Display chat history
for message in st.session_state.displayed_messages: for message in st.session_state.displayed_messages:
with st.chat_message(message["role"]): log_message(message)
st.markdown(message["content"])
if temperature > 0.0: if temperature > 0.0:
strategy = { strategy = {
@ -201,7 +207,7 @@ def rag_chat_page():
# Display assistant response # Display assistant response
with st.chat_message("assistant"): with st.chat_message("assistant"):
retrieval_message_placeholder = st.empty() retrieval_message_placeholder = st.expander(label="Tool Output", expanded=False, icon="🛠")
message_placeholder = st.empty() message_placeholder = st.empty()
full_response = "" full_response = ""
retrieval_response = "" retrieval_response = ""
@ -209,14 +215,16 @@ def rag_chat_page():
log.print() log.print()
if log.role == "tool_execution": if log.role == "tool_execution":
retrieval_response += log.content.replace("====", "").strip() retrieval_response += log.content.replace("====", "").strip()
retrieval_message_placeholder.info(retrieval_response) retrieval_message_placeholder.write(retrieval_response)
else: else:
full_response += log.content full_response += log.content
message_placeholder.markdown(full_response + "▌") message_placeholder.markdown(full_response + "▌")
message_placeholder.markdown(full_response) message_placeholder.markdown(full_response)
st.session_state.messages.append({"role": "assistant", "content": full_response}) st.session_state.messages.append({"role": "assistant", "content": full_response})
st.session_state.displayed_messages.append({"role": "assistant", "content": full_response}) st.session_state.displayed_messages.append(
{"role": "assistant", "content": full_response, "tool_output": retrieval_response}
)
def direct_process_prompt(prompt): def direct_process_prompt(prompt):
# Add the system prompt in the beginning of the conversation # Add the system prompt in the beginning of the conversation
@ -230,15 +238,14 @@ def rag_chat_page():
prompt_context = rag_response.content prompt_context = rag_response.content
with st.chat_message("assistant"): with st.chat_message("assistant"):
with st.expander(label="Retrieval Output", expanded=False):
st.write(prompt_context)
retrieval_message_placeholder = st.empty() retrieval_message_placeholder = st.empty()
message_placeholder = st.empty() message_placeholder = st.empty()
full_response = "" full_response = ""
retrieval_response = "" retrieval_response = ""
# Display the retrieved content
retrieval_response += str(prompt_context)
retrieval_message_placeholder.info(retrieval_response)
# Construct the extended prompt # Construct the extended prompt
extended_prompt = f"Please answer the following query using the context below.\n\nCONTEXT:\n{prompt_context}\n\nQUERY:\n{prompt}" extended_prompt = f"Please answer the following query using the context below.\n\nCONTEXT:\n{prompt_context}\n\nQUERY:\n{prompt}"