Spaces:
Sleeping
Sleeping
File size: 2,069 Bytes
4531c67 c6dbe7a 4531c67 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
import streamlit as st
from agents import investigator, router
# from agents_openai_fc import investigator
st.title('Sonic Cyber Assistant')
st.caption("🚀 A cyber assistant chatbot for threat hunting")
if "messages" not in st.session_state:
st.session_state["messages"] = [{"role": "assistant", "content": "How can I help you?"}]
for msg in st.session_state.messages:
st.chat_message(msg["role"]).write(msg["content"])
if prompt := st.chat_input():
st.session_state.messages.append({"role": "user", "content": prompt})
st.chat_message("user").write(prompt)
response = router.invoke(prompt)
msg = response["output"]
print("Title : ", router.generate_title(prompt))
# print(response['intermediate_steps'])
# If there's an existing assistant message, update it with the new response
if st.session_state.messages[-1]["role"] == "assistant":
st.session_state.messages[-1]["content"] = msg
else:
st.session_state.messages.append({"role": "assistant", "content": msg})
st.chat_message("assistant").write(msg)
# Add a button to regenerate response
if st.button("Regenerate Response"):
st.session_state.messages.pop() # Remove the latest assistant message
prompt = st.session_state.messages[-1]["content"] # Retrieve user's last input
with st.spinner("Searching..."):
response = router.invoke(prompt)
msg = response["output"]
# If there's an existing assistant message, update it with the new response
if st.session_state.messages[-1]["role"] == "assistant":
st.session_state.messages[-1]["content"] = msg
else:
st.session_state.messages.append({"role": "assistant", "content": msg})
st.chat_message("assistant").write(msg)
if st.button("Clear Chat"):
st.session_state.messages = [{"role": "assistant", "content": "How can I help you?"}]
router.memory.clear()
for msg in st.session_state.messages:
st.chat_message(msg["role"]).write(msg["content"])
st.success("Chat cleared!")
|