Update app.py

app.py

@@ -1,25 +1,37 @@
 import streamlit as st
 import requests
 import json
+from langchain.chains import ConversationChain
 from langchain.memory import ConversationBufferMemory
+from langchain.chat_models import ChatOpenAI
 
 # Constants
 SPACE_URL = "https://abanm-dubs.hf.space/api/generate"
+API_KEY = "s3cr3t_k3y"  # Replace with your actual API key
 
-# Streamlit Configurations (must come first)
+# Streamlit Configurations
 st.set_page_config(page_title="🐶 DUBSChat", layout="centered")
 
-# Initialize memory for context retention
-if "memory" not in st.session_state:
-    st.session_state["memory"] = ConversationBufferMemory()
-
 # Sidebar Configuration
 with st.sidebar:
     dubs_key = st.text_input("Enter Dubs Key", key="chatbot_api_key", type="password")
-    st.markdown("Dubs Recall")
-    if st.button("Show Conversation History"):
-        if "memory" in st.session_state:
-            st.text_area("Conversation History", st.session_state["memory"].buffer)
+    st.markdown("### Chat Sessions")
+    if "sessions" not in st.session_state:
+        st.session_state["sessions"] = {"Default": ConversationBufferMemory()}
+    session_keys = list(st.session_state["sessions"].keys())
+    selected_session = st.selectbox("Select Chat Session", session_keys)
+    new_session_name = st.text_input("New Session Name")
+    if st.button("Create New Session"):
+        if new_session_name and new_session_name not in st.session_state["sessions"]:
+            st.session_state["sessions"][new_session_name] = ConversationBufferMemory()
+            st.experimental_rerun()
+
+# LangChain Configuration
+if "conversation" not in st.session_state:
+    st.session_state["conversation"] = ConversationChain(
+        llm=ChatOpenAI(temperature=0),
+        memory=st.session_state["sessions"][selected_session],
+    )
 
 def call_api(message):
     """
@@ -35,10 +47,7 @@ def call_api(message):
         "Authorization": f"Bearer {dubs_key}",
         "Content-Type": "application/json",
     }
-    # Add context to the prompt
-    context = st.session_state["memory"].buffer
-    prompt = f"{context}\nUser: {message}\nAI:"
-    payload = {"prompt": prompt}  # Match the API's expected input key
+    payload = {"prompt": message}  # Match the API's expected input key
 
     try:
         # Send request with streaming enabled
@@ -59,25 +68,20 @@ def call_api(message):
     except requests.exceptions.RequestException as e:
         yield f"Error: {str(e)}"
 
-# Chat History Initialization
-if "messages" not in st.session_state:
-    st.session_state["messages"] = [{"role": "assistant", "content": "Hello! How can I assist you today?"}]
-
 # Main Chat UI
 st.title("🐶 DUBSChat")
 st.markdown("Empowering you with a Sustainable AI")
 
-# Chat Display
-for message in st.session_state["messages"]:
-    if message["role"] == "user":
-        st.chat_message("user").write(message["content"])
-    elif message["role"] == "assistant":
-        st.chat_message("assistant").write(message["content"])
+# Chat Display for the selected session
+memory = st.session_state["sessions"][selected_session]
+for message in memory.chat_memory.messages:
+    role = "user" if message.type == "human" else "assistant"
+    st.chat_message(role).write(message.content)
 
 # User Input
 if prompt := st.chat_input():
-    # Append the user's message to the chat history
-    st.session_state["messages"].append({"role": "user", "content": prompt})
+    # Append the user's message to the chat memory
+    memory.chat_memory.add_user_message(prompt)
     st.chat_message("user").write(prompt)
 
     with st.spinner("Translating Dubs language (Woof Woof!)"):
@@ -87,8 +91,8 @@ if prompt := st.chat_input():
             full_response += response_chunk
             assistant_message_placeholder.write(full_response)  # Update progressively
 
-        # Append the final assistant's response to the chat history
-        st.session_state["messages"].append({"role": "assistant", "content": full_response})
+        # Append the assistant's response to the chat memory
+        memory.chat_memory.add_ai_message(full_response)
 
-        # Save user input and bot response to memory for context retention
-        st.session_state["memory"].save_context({"User": prompt}, {"AI": full_response})
+# Update the LangChain conversation to use the selected session's memory
+st.session_state["conversation"].memory = memory