Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,96 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+from chat_langraph import system, workflow, HumanMessage, AIMessage, get_all_chat_ids, ToolMessage
+import uuid
+from chat_title_giver import give_me_the_title
+import re
+import os
+st.title("College chatbot")
+st.set_page_config(layout='wide')
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+
+def set_title(messages):
+    if messages:
+        title = give_me_the_title(messages)
+        st.session_state.chat_dict[st.session_state.current_chat_id] = title
+
+
+def set_config():
+    return {'configurable': {'thread_id': st.session_state.current_chat_id}}
+
+
+def load_session_state():
+    if "chats" not in st.session_state:
+        st.session_state.chats = get_all_chat_ids()
+    if "current_chat_id" not in st.session_state:
+        if len(st.session_state.chats) > 0:
+            st.session_state.current_chat_id = st.session_state.chats[-1]
+        else:
+            new_id = str(uuid.uuid4())
+            st.session_state.chats.append(new_id)
+            st.session_state.current_chat_id = new_id
+    if "chat_dict" not in st.session_state:
+        st.session_state.chat_dict = {}
+
+
+def render_sidebar():
+    with st.sidebar:
+        st.title("Chats")
+        if st.button("➕ New Chat"):
+            new_id = str(uuid.uuid4())
+            st.session_state.chats.append(new_id)
+            st.session_state.current_chat_id = new_id
+            config = {'configurable': {'thread_id': new_id}}
+            workflow.update_state(config, {"messages": [system]})
+            st.session_state.chat_dict[new_id] = "New Chat"
+            st.experimental_rerun()
+        for chat_id in st.session_state.chats:
+            if st.button(st.session_state.chat_dict.get(chat_id, "New Chat"), key=chat_id):
+                st.session_state.current_chat_id = chat_id
+
+
+def loadchats():
+    if "current_chat_id" not in st.session_state:
+        return []
+    config = {'configurable': {'thread_id': st.session_state.current_chat_id}}
+    state = workflow.get_state(config)
+    messages = state.values.get("messages", [])
+    for message in messages:
+        if isinstance(message, HumanMessage):
+            if message.content and message.content.strip():
+                with st.chat_message("human"):
+                    st.write(message.content)
+        elif isinstance(message, AIMessage):
+            if message.content and message.content.strip():
+                with st.chat_message("assistant"):
+                    st.write(message.content)
+        elif isinstance(message, ToolMessage):
+            with st.chat_message("assistant"):
+                st.info(f"Used a tool")
+    return messages
+
+
+load_session_state()
+render_sidebar()
+
+if "current_chat_id" in st.session_state:
+    loadchats()
+    user_input = st.chat_input("Your message: ")
+    if user_input is not None:
+        with st.chat_message("human"):
+            st.write(user_input)
+
+        with st.chat_message("assistant"):
+            response_placeholder = st.empty()
+            full_response = ""
+            for messages, metadata in workflow.stream(
+                {"messages": [system, HumanMessage(user_input)]},
+                config={'configurable': {'thread_id': st.session_state.current_chat_id}},
+                stream_mode="messages"
+            ):
+                if isinstance(messages, AIMessage) and messages.content.strip():
+                    full_response += messages.content
+                if isinstance(messages, ToolMessage) and messages.content.strip():
+                    st.info(f"Using appropriate tool")
+            response_placeholder.markdown(full_response)
+
+        st.experimental_rerun()