Update app.py

app.py

@@ -1,76 +1,67 @@
+
 import os
 import re
-import streamlit as st
 from dotenv import load_dotenv
+load_dotenv()
+
+import gradio as gr
+
 from langchain.agents.openai_assistant import OpenAIAssistantRunnable
+from langchain.schema import HumanMessage, AIMessage
 
-# Load environment variables
-load_dotenv()
-api_key = os.getenv("OPENAI_API_KEY")
-extractor_agent = os.getenv("ASSISTANT_ID_SOLUTION_SPECIFIER_A")
+api_key = os.getenv('OPENAI_API_KEY')
+extractor_agent = os.getenv('ASSISTANT_ID_SOLUTION_SPECIFIER_A')
 
-# Create the assistant
+# Create the assistant. By default, we don't specify a thread_id,
+# so the first call that doesn't pass one will create a new thread.
 extractor_llm = OpenAIAssistantRunnable(
     assistant_id=extractor_agent,
     api_key=api_key,
     as_agent=True
 )
 
-def remove_citation(text: str) -> str:
+# We will store thread_id globally or in a session variable.
+THREAD_ID = None
+
+def remove_citation(text):
     pattern = r"【\d+†\w+】"
     return re.sub(pattern, "📚", text)
 
-# Initialize session state for messages and thread_id
-if "messages" not in st.session_state:
-    st.session_state["messages"] = []
-if "thread_id" not in st.session_state:
-    st.session_state["thread_id"] = None
-
-st.title("Solution Specifier A")
-
-def predict(user_input: str) -> str:
+def predict(message, history):
     """
-    This function calls our OpenAIAssistantRunnable to get a response.
-    If we don't have a thread_id yet, we create a new thread on the first call.
-    Otherwise, we continue the existing thread.
+    Receives the new user message plus the entire conversation history 
+    from Gradio. If no thread_id is set, we create a new thread. 
+    Otherwise we pass the existing thread_id.
     """
-    if st.session_state["thread_id"] is None:
-        response = extractor_llm.invoke({"content": user_input})
-        st.session_state["thread_id"] = response.thread_id
-    else:
-        response = extractor_llm.invoke(
-            {"content": user_input, "thread_id": st.session_state["thread_id"]}
-        )
-    output = response.return_values["output"]
-    return remove_citation(output)
-
-# Display any existing messages (from a previous run or refresh)
-for msg in st.session_state["messages"]:
-    if msg["role"] == "user":
-        with st.chat_message("user"):
-            st.write(msg["content"])
-    else:
-        with st.chat_message("assistant"):
-            st.write(msg["content"])
-
-# Create the chat input widget at the bottom of the page
-user_input = st.chat_input("Type your message here...")
-
-# When the user hits ENTER on st.chat_input
-if user_input:
-    # Add the user message to session state
-    st.session_state["messages"].append({"role": "user", "content": user_input})
+    global THREAD_ID
+    
+    # debug print
+    print("current history:", history)
     
-    # Display the user's message
-    with st.chat_message("user"):
-        st.write(user_input)
-        
-    # Get the assistant's response
-    response_text = predict(user_input)
+    # If history is empty, this means that it is probably a new conversation and therefore the thread shall be reset
+    if not history:
+        THREAD_ID = None
     
-    # Add the assistant response to session state
-    st.session_state["messages"].append({"role": "assistant", "content": response_text})
+    # 1) Decide if we are creating a new thread or continuing the old one
+    if THREAD_ID is None:
+        # No thread_id yet -> this is the first user message
+        response = extractor_llm.invoke({"content": message})
+        THREAD_ID = response.thread_id  # store for subsequent calls
+    else:
+        # We already have a thread_id -> continue that same thread
+        response = extractor_llm.invoke({"content": message, "thread_id": THREAD_ID})
+    
+    # 2) Extract the text output from the response
+    output = response.return_values["output"]
+    non_cited_output = remove_citation(output)
     
-    # Display the assistant's reply
-    with st.chat_message("assistant"):
-        st.write(response_text)
+    # 3) Return the model's text to display in Gradio
+    return non_cited_output
+
+# Create a Gradio ChatInterface using our predict function
+chat = gr.ChatInterface(
+    fn=predict, 
+    title="Solution Specifier A", 
+    #description="Testing threaded conversation"
+)
+chat.launch(share=True)