Upload app.py

app.py

@@ -1,15 +1,13 @@
-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",
@@ -24,33 +22,24 @@ with open("styles/styles.css") as css:
 
 
 # Get runnable and memory
-def initialize_runnable_and_memory():
+@st.cache_resource(show_spinner=False)
+def get_runnable_and_memory():
     try:
-        return get_runnable(model="gpt-4o", temperature=0)
+        return get_runnable(model="gpt-4-turbo", temperature=0)
     except Exception:
         st.warning(ERROR_MESSAGE, icon="🙁")
         st.stop()
 
 
-# Get the ThreadPoolExecutor
-if "executor" not in st.session_state:
-    st.session_state.executor = ThreadPoolExecutor(max_workers=4)
+chain, memory = get_runnable_and_memory()
 
-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
-    )
 
-# 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()
+# 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()
 
-# Other session state variables
 if "messages" not in st.session_state:
     st.session_state["messages"] = []
 
@@ -96,14 +85,14 @@ user_input = st.chat_input(
     "Ask ELLA...", disabled=st.session_state["disable_chat_input"]
 )
 
-if user_input and user_input.strip():
+if user_input:
     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
 
-# Render chat messages
+# Display previous messages
 user_avatar = "images/user.png"
 ai_avatar = "images/tall-tree-logo.png"
 for msg in st.session_state["messages"]:
@@ -111,19 +100,26 @@ for msg in st.session_state["messages"]:
     with st.chat_message(msg.role, avatar=avatar):
         st.markdown(msg.content)
 
-# 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
+# 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
     with st.chat_message("assistant", avatar=ai_avatar):
         message_placeholder = st.empty()
-        try:
-            response = ""
-            with st.spinner(" "):
-                for chunk in st.session_state["runnable"].stream({"message": prompt}):
-                    response += chunk
-                    message_placeholder.markdown(response + "|")
 
-            message_placeholder.markdown(response)
+        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)
         except openai.BadRequestError:
             st.warning(ERROR_MESSAGE, icon="🙁")
             st.stop()
@@ -131,17 +127,60 @@ if "runnable" in st.session_state and prompt:
             st.warning(ERROR_MESSAGE, icon="🙁")
             st.stop()
 
-        # Add response to the message history
+        # 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
         st.session_state["messages"].append(
-            ChatMessage(role="assistant", content=response)
+            ChatMessage(role="assistant", content=partial_message)
         )
 
-        # 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)
+
+# 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.")