Chatbot / src /ui /display_results.py
pratikcsv's picture
works
b967fd3
import streamlit as st
from langchain_core.messages import AIMessage, HumanMessage
class DisplayResults:
"""Class to display results in the Streamlit UI."""
def __init__(self, user_message, use_case, graph):
"""
Initialize the DisplayResults with user message, use case, and graph.
:param user_message: The message input by the user.
:param use_case: The selected use case for the chatbot.
:param graph: The state graph to be processed.
"""
self.user_message = user_message
self.use_case = use_case
self.graph = graph
def display_results(self):
"""
Display the results in the Streamlit UI.
This method processes the user message through the graph and displays the response.
"""
use_case = self.use_case
user_message = self.user_message
graph = self.graph
if use_case == 'Basic Chatbot':
for event in graph.stream({'messages':("user",user_message)}):
print(event.values())
for value in event.values():
print(value['messages'])
with st.chat_message("user"):
st.write(user_message)
with st.chat_message("assistant"):
st.write(value["messages"].content)