Update app.py

app.py

@@ -1,114 +1,82 @@
 import streamlit as st
 import requests
 import json
-from langchain.chains import ConversationChain
-from langchain.memory import ConversationBufferMemory
-from langchain.llms.base import LLM
-from typing import Optional
 
 # Constants
 SPACE_URL = "https://abanm-dubs.hf.space/api/generate"
+API_KEY = "s3cr3t_k3y"  # Replace with your actual API key
 
-# Streamlit Configurations
+# Streamlit Configurations (must come first)
 st.set_page_config(page_title="🐶 DUBSChat", layout="centered")
 
 # Sidebar Configuration
 with st.sidebar:
-    if "dubs_key" not in st.session_state:
-        st.session_state["dubs_key"] = ""
-    st.session_state["dubs_key"] = st.text_input("Enter Dubs Key", type="password", key="chatbot_api_key")
+    dubs_key = st.text_input("Enter Dubs Key", key="chatbot_api_key", type="password")
+    st.markdown("Dubs Recall")
+    st.markdown("""Coming Soon!""")
 
-    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, index=session_keys.index(st.experimental_get_query_params().get("session", ["Default"])[0]))
-    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_set_query_params(session=new_session_name)
-
-# Ensure Dubs Key is Provided
-if not st.session_state["dubs_key"]:
-    st.warning("Please enter a valid Dubs Key in the sidebar to proceed.")
-    st.stop()
-
-# Define Custom LLM for the fine-tuned model
-class CustomLLM(LLM):
+def call_api(message):
     """
-    Custom LLM to interact with the fine-tuned model hosted on SPACE_URL.
+    Calls the backend API with the user's input using streaming.
+    
+    Args:
+        message (str): User's message.
+    
+    Yields:
+        str: Streamed response chunks from the backend API.
     """
-
-    def _call(self, prompt: str, stop: Optional[list] = None) -> str:
-        dubs_key = st.session_state.get("dubs_key")
-        if not dubs_key:
-            return "Error: Missing Dubs Key. Please provide a valid key in the sidebar."
-
-        headers = {
-            "Authorization": f"Bearer {dubs_key}",
-            "Content-Type": "application/json",
-        }
-        payload = {"prompt": prompt}
-
-        try:
-            response = requests.post(SPACE_URL, json=payload, headers=headers)
-            response.raise_for_status()
-            result = response.json()
-            return result.get("response", "No response from the model.")
-
-        except requests.exceptions.HTTPError as e:
-            if response.status_code == 401:
-                return "Error: Unauthorized. Please check your Dubs Key."
-            return f"HTTP Error: {response.status_code} - {response.reason}"
-        except requests.exceptions.RequestException as e:
-            return f"Error: {str(e)}"
-
-    @property
-    def _identifying_params(self):
-        """Return unique parameters for this LLM."""
-        return {"model": "custom_model"}
-
-    @property
-    def _llm_type(self):
-        """Return type of LLM."""
-        return "custom_llm"
-
-# LangChain Configuration
-if "conversation" not in st.session_state:
-    st.session_state["conversation"] = ConversationChain(
-        llm=CustomLLM(),
-        memory=st.session_state["sessions"][selected_session],
-    )
+    headers = {
+        "Authorization": f"Bearer {dubs_key}",
+        "Content-Type": "application/json",
+    }
+    payload = {"prompt": message}  # Match the API's expected input key
+
+    try:
+        # Send request with streaming enabled
+        response = requests.post(SPACE_URL, json=payload, headers=headers, stream=True)
+        response.raise_for_status()
+
+        for chunk in response.iter_lines(decode_unicode=True):
+            if chunk.strip():  # Ensure the chunk is not empty
+                try:
+                    # Parse JSON chunk
+                    data = json.loads(chunk)
+                    if "response" in data:
+                        yield data["response"]
+                except json.JSONDecodeError:
+                    yield "Error decoding response chunk."
+    except requests.exceptions.Timeout:
+        yield "Error: The API call timed out."
+    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 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)
+# 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"])
 
 # User Input
 if prompt := st.chat_input():
-    # Append the user's message to the chat memory
-    memory.chat_memory.add_user_message(prompt)
+    # Append the user's message to the chat history
+    st.session_state["messages"].append({"role": "user", "content": prompt})
     st.chat_message("user").write(prompt)
 
     with st.spinner("Translating Dubs language (Woof Woof!)"):
         assistant_message_placeholder = st.chat_message("assistant").empty()
-        full_response = st.session_state["conversation"].run(
-            f"System:The following is a friendly conversation between a user and an AI. "
-            f"The AI is talkative and provides lots of specific details from its context. "
-            f"If the AI does not know the answer to a question, it truthfully says it does not know.\n\n"
-            f"Current conversation:\nUser: {prompt}\nAssistant:"
-        )
-        assistant_message_placeholder.write(full_response)
-
-        # Append the assistant's response to the chat memory
-        memory.chat_memory.add_ai_message(full_response)
+        full_response = ""  # To accumulate the entire response
+        for response_chunk in call_api(prompt):
+            full_response += response_chunk
+            assistant_message_placeholder.write(full_response)  # Update progressively
 
-# Update the LangChain conversation to use the selected session's memory
-st.session_state["conversation"].memory = memory
+        # Append the final assistant's response to the chat history
+        st.session_state["messages"].append({"role": "assistant", "content": full_response})