Upload app.py

app.py

@@ -1,13 +1,15 @@
+import logging
+from concurrent.futures import ThreadPoolExecutor
+
 import openai
 import streamlit as st
 from langchain_core.messages import AIMessage, ChatMessage, HumanMessage
-from langchain_core.tracers.context import collect_runs
-from langsmith import Client
-from streamlit_feedback import streamlit_feedback
 
 from rag.runnable import get_runnable
 from utils.error_message_template import ERROR_MESSAGE
 
+logging.basicConfig(level=logging.ERROR)
+
 # Streamlit page configuration
 st.set_page_config(
     page_title="ELLA AI Assistant",
@@ -22,24 +24,33 @@ with open("styles/styles.css") as css:
 
 
 # Get runnable and memory
-@st.cache_resource(show_spinner=False)
-def get_runnable_and_memory():
+def initialize_runnable_and_memory():
     try:
-        return get_runnable(model="gpt-4-turbo", temperature=0)
+        return get_runnable(model="gpt-4o", temperature=0)
     except Exception:
         st.warning(ERROR_MESSAGE, icon="🙁")
         st.stop()
 
 
-chain, memory = get_runnable_and_memory()
+# Get the ThreadPoolExecutor
+if "executor" not in st.session_state:
+    st.session_state.executor = ThreadPoolExecutor(max_workers=4)
 
+executor = st.session_state.executor
+
+# Submit initialization task if not already done
+if "initialization_future" not in st.session_state:
+    st.session_state["initialization_future"] = executor.submit(
+        initialize_runnable_and_memory
+    )
 
-# Set up session state variables
-# Clean memory (important! to clean the memory at the end of each session)
-if "history" not in st.session_state:
-    st.session_state["history"] = []
-    memory.clear()
+# Check if initialization is complete
+future = st.session_state["initialization_future"]
+if future.done() and "runnable" not in st.session_state:
+    st.session_state["runnable"], st.session_state["memory"] = future.result()
+    st.session_state["memory"].clear()
 
+# Other session state variables
 if "messages" not in st.session_state:
     st.session_state["messages"] = []
 
@@ -85,14 +96,14 @@ user_input = st.chat_input(
     "Ask ELLA...", disabled=st.session_state["disable_chat_input"]
 )
 
-if user_input:
+if user_input and user_input.strip():
     st.session_state["messages"].append(ChatMessage(role="user", content=user_input))
     prompt = f"{user_input}\nLocation preference: {st.session_state.selected_location}."
 
 else:
     prompt = None
 
-# Display previous messages
+# Render chat messages
 user_avatar = "images/user.png"
 ai_avatar = "images/tall-tree-logo.png"
 for msg in st.session_state["messages"]:
@@ -100,26 +111,19 @@ for msg in st.session_state["messages"]:
     with st.chat_message(msg.role, avatar=avatar):
         st.markdown(msg.content)
 
-# Chat interface
-if prompt:
-    # Add all previous messages to memory
-    for human, ai in st.session_state["history"]:
-        memory.chat_memory.add_user_message(HumanMessage(content=human))
-        memory.chat_memory.add_ai_message(AIMessage(content=ai))
-
-    # render the assistant's response
+# Chat interface (we have to wait for the runnable initialization to complete)
+if "runnable" in st.session_state and prompt:
+    # Render the assistant's response
     with st.chat_message("assistant", avatar=ai_avatar):
         message_placeholder = st.empty()
-
         try:
-            partial_message = ""
-            # Collect runs for feedback using Langsmith.
-            with st.spinner(" "), collect_runs() as cb:
-                for chunk in chain.stream({"message": prompt}):
-                    partial_message += chunk
-                    message_placeholder.markdown(partial_message + "|")
-                st.session_state.run_id = cb.traced_runs[0].id
-            message_placeholder.markdown(partial_message)
+            response = ""
+            with st.spinner(" "):
+                for chunk in st.session_state["runnable"].stream({"message": prompt}):
+                    response += chunk
+                    message_placeholder.markdown(response + "|")
+
+            message_placeholder.markdown(response)
         except openai.BadRequestError:
             st.warning(ERROR_MESSAGE, icon="🙁")
             st.stop()
@@ -127,60 +131,17 @@ if prompt:
             st.warning(ERROR_MESSAGE, icon="🙁")
             st.stop()
 
-        # Add the full response to the history
-        st.session_state["history"].append((prompt, partial_message))
-
-        # Add AI message to memory after the response is generated
-        memory.chat_memory.add_ai_message(AIMessage(content=partial_message))
-
-        # Add the full response to the message history
+        # Add response to the message history
         st.session_state["messages"].append(
-            ChatMessage(role="assistant", content=partial_message)
+            ChatMessage(role="assistant", content=response)
         )
 
-
-# Feedback system using streamlit-feedback and Langsmith
-
-# Langsmith client for the feedback system
-ls_client = Client()
-
-# Feedback option
-feedback_option = "thumbs"
-
-if st.session_state.get("run_id"):
-    run_id = st.session_state.run_id
-    feedback = streamlit_feedback(
-        feedback_type=feedback_option,
-        optional_text_label="[Optional] Please provide an explanation",
-        key=f"feedback_{run_id}",
-    )
-    score_mappings = {
-        "thumbs": {"👍": 1, "👎": 0},
-        "faces": {"😀": 1, "🙂": 0.75, "😐": 0.5, "🙁": 0.25, "😞": 0},
-    }
-
-    # Get the score mapping based on the selected feedback option
-    scores = score_mappings[feedback_option]
-
-    if feedback:
-        # Get the score from the selected feedback option's score mapping
-        score = scores.get(feedback["score"])
-
-        if score is not None:
-            # Formulate feedback type string incorporating the feedback option
-            # and score value
-            feedback_type_str = f"{feedback_option} {feedback['score']}"
-
-            # Record the feedback with the formulated feedback type string
-            feedback_record = ls_client.create_feedback(
-                run_id,
-                feedback_type_str,
-                score=score,
-                comment=feedback.get("text"),
-            )
-            st.session_state.feedback = {
-                "feedback_id": str(feedback_record.id),
-                "score": score,
-            }
-        else:
-            st.warning("Invalid feedback score.")
+        # Add messages to memory
+        st.session_state["memory"].chat_memory.add_user_message(
+            HumanMessage(content=prompt)
+        )
+        st.session_state["memory"].chat_memory.add_ai_message(
+            AIMessage(content=response)
+        )
+if st.session_state.executor:
+    st.session_state.executor.shutdown(wait=False)