rag-chatbot/streamlit_frontend_database.py at main · codeantik/rag-chatbot · GitHub

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
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
# ==============================
# 📦 Imports
# ==============================
import streamlit as st
from langgraph_database_backend import chatbot, retrieve_all_threads
from langchain_core.messages import HumanMessage
import uuid


# ==============================
# ⚙️ Utility Functions
# ==============================
def get_thread_id():
    return str(uuid.uuid4())


def init_session():
    ss = st.session_state
    if 'current_thread_id' not in ss:
        ss.current_thread_id = get_thread_id()

    if 'message_history' not in ss:
        ss.message_history = []

    if 'thread_history' not in ss:
        ss.thread_history = retrieve_all_threads()

    if 'chat_started' not in ss:
        ss.chat_started = False


def reset_chat():
    ss = st.session_state
    thread_id = get_thread_id()

    ss.current_thread_id = thread_id
    ss.message_history = []

    if thread_id not in ss.thread_history:
        ss.thread_history.append(thread_id)


def resume_thread(thread_id):
    ss = st.session_state
    ss.current_thread_id = thread_id

    messages = load_conversation(thread_id)

    ss.message_history = [
        {
            "role": "user" if isinstance(m, HumanMessage) else "assistant",
            "content": m.content
        }
        for m in messages
    ]


def load_conversation(thread_id):
    state = chatbot.get_state(
        config={'configurable': {'thread_id': thread_id}}
    )
    return state.values.get('messages', [])


# ==============================
# 🧠 Initialize Session
# ==============================
init_session()
ss = st.session_state


# ==============================
# 🏠 Welcome Screen
# ==============================
if not ss.chat_started:
    st.markdown(
        """
        <div style="display:flex; flex-direction:column; justify-content:center; align-items:center; height:60vh;">
            <h1>LangGraph Chatbot</h1>
            <p>Welcome to the LangGraph Chatbot!</p>
        </div>
        """,
        unsafe_allow_html=True
    )


# ==============================
# 💬 Sidebar
# ==============================
with st.sidebar:
    st.title("Create New Chat")

    if st.button("New Chat"):
        reset_chat()

    st.title("Conversation History")

    for thread_id in reversed(ss.thread_history):
        messages = load_conversation(thread_id)
        label = messages[0].content if messages else "New Conversation"

        if st.button(label, key=thread_id):
            resume_thread(thread_id)


# ==============================
# 💬 Chat History
# ==============================
for msg in ss.message_history:
    with st.chat_message(msg["role"]):
        st.text(msg["content"])


# ==============================
# ⌨️ User Input
# ==============================
user_input = st.chat_input("Type here")


# ==============================
# 🚀 Chat Handling
# ==============================
if user_input:
    ss.chat_started = True

    # User message
    ss.message_history.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.text(user_input)

    # CONFIG = {'configurable': {'thread_id': ss.current_thread_id}}
    CONFIG = {
        'configurable': {
            'thread_id': ss.current_thread_id
        },
        'metadata': {
            'thread_id': ss.current_thread_id,
        },
        'run_name': "Chat Run"
    }

    # AI response (stream)
    with st.chat_message("assistant"):
        stream = chatbot.stream(
            {"messages": [HumanMessage(content=user_input)]},
            config=CONFIG,
            stream_mode="messages"
        )

        ai_response = st.write_stream(
            chunk.content for chunk, _ in stream
        )

    ss.message_history.append({"role": "assistant", "content": ai_response})