updated playground rag page to use session id for persistent conversation

Signed-off-by: Michael Clifford <mcliffor@redhat.com>
This commit is contained in:
Michael Clifford 2025-04-03 15:51:31 -04:00
parent c52ccc4bbd
commit 84fd317783

View file

@ -4,6 +4,8 @@
# This source code is licensed under the terms described in the LICENSE file in # This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree. # the root directory of this source tree.
import uuid
import streamlit as st import streamlit as st
from llama_stack_client import Agent, AgentEventLogger, RAGDocument from llama_stack_client import Agent, AgentEventLogger, RAGDocument
@ -102,8 +104,8 @@ def rag_chat_page():
# Add clear chat button to sidebar # Add clear chat button to sidebar
if st.button("Clear Chat", use_container_width=True): if st.button("Clear Chat", use_container_width=True):
st.session_state.messages = [] st.session_state.clear()
st.rerun() st.cache_resource.clear()
# Chat Interface # Chat Interface
if "messages" not in st.session_state: if "messages" not in st.session_state:
@ -123,23 +125,31 @@ def rag_chat_page():
else: else:
strategy = {"type": "greedy"} strategy = {"type": "greedy"}
agent = Agent( @st.cache_resource
llama_stack_api.client, def create_agent():
model=selected_model, return Agent(
instructions=system_prompt, llama_stack_api.client,
sampling_params={ model=selected_model,
"strategy": strategy, instructions=system_prompt,
}, sampling_params={
tools=[ "strategy": strategy,
dict( },
name="builtin::rag/knowledge_search", tools=[
args={ dict(
"vector_db_ids": list(selected_vector_dbs), name="builtin::rag/knowledge_search",
}, args={
) "vector_db_ids": list(selected_vector_dbs),
], },
) )
session_id = agent.create_session("rag-session") ],
)
agent = create_agent()
if "agent_session_id" not in st.session_state:
st.session_state["agent_session_id"] = agent.create_session(session_name=f"rag_demo_{uuid.uuid4()}")
session_id = st.session_state["agent_session_id"]
# Chat input # Chat input
if prompt := st.chat_input("Ask a question about your documents"): if prompt := st.chat_input("Ask a question about your documents"):