app.py

app.py

@@ -1,54 +1,62 @@
-from langchain.callbacks.base import BaseCallbackHandler
-from langchain.schema import ChatMessage
-from langchain_openai import ChatOpenAI
-import streamlit as st
-
-st.set_page_config(page_title="Streamlit + Langchain")
-
-st.title("Basic Chatbot with Streamlit and Langchain")
-st.caption("Features text streaming")
-
-
-class StreamHandler(BaseCallbackHandler):
-    def __init__(self, container, text=""):
-        self.container = container
-        self.text = text
-
-    def on_llm_new_token(self, token: str, **kwargs) -> None:
-        self.text += token
-        self.container.markdown(self.text)
-
-
-# Text input to enter OpenAI API key
-with st.sidebar:
-    OPENAI_API_KEY = st.text_input("Enter your OpenAI API Key", type="password")
-
-# Streamlit session state
-if "messages" not in st.session_state:
-    st.session_state["messages"] = [
-        ChatMessage(role="assistant", content="How can I help you?")
-    ]
-# Display all chat messages from session state
-for message in st.session_state.messages:
-    st.chat_message(message.role).write(message.content)
-# If user submits a prompt in the text input, continue
-if prompt := st.chat_input():
-    if not OPENAI_API_KEY:
-        st.error("Please add your OpenAI API key to continue.")
-        st.stop()
-    # Add user's prompt to the chat messages
-    st.session_state.messages.append(ChatMessage(role="user", content=prompt))
-    st.chat_message("user").write(prompt)
-    # Display the assistant's response with langchain query
-    with st.chat_message("assistant"):
-        stream_handler = StreamHandler(st.empty())
-        llm = ChatOpenAI(
-            model="gpt-4o-mini",
-            openai_api_key=OPENAI_API_KEY,
-            streaming=True,
-            callbacks=[stream_handler],
-        )
-        response = llm.invoke(st.session_state.messages)
-        st.session_state.messages.append(
-            ChatMessage(role="assistant", content=response.content)
-        )
+import streamlit as st
+from openai import OpenAI
+
+st.set_page_config(page_title="Streamlit + OpenAI Assistant")
+
+st.title("Basic Chatbot with Streamlit and OpenAI Assistant")
+st.caption("Features text streaming using OpenAI Assistants API")
+
+# Text input to enter OpenAI API key
+with st.sidebar:
+    OPENAI_API_KEY = st.text_input("Enter your OpenAI API Key", type="password")
+
+# Initialize OpenAI client
+if OPENAI_API_KEY:
+    client = OpenAI(api_key=OPENAI_API_KEY)
+else:
+    st.error("Please enter your OpenAI API key to continue.")
+    st.stop()
+
+ASSISTANT_ID = "asst_rd9h8PfYuOmHbkvOF3RTmVfn"
+
+# Streamlit session state
+if "messages" not in st.session_state:
+    st.session_state["messages"] = []
+
+# Display all chat messages from session state
+for message in st.session_state.messages:
+    st.chat_message(message["role"]).write(message["content"])
+
+# User input
+if prompt := st.chat_input():
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    st.chat_message("user").write(prompt)
+
+    # Send message to OpenAI Assistant
+    try:
+        thread = client.beta.threads.create()
+        thread_id = thread.id
+        client.beta.threads.messages.create(
+            thread_id=thread_id,
+            role="user",
+            content=prompt
+        )
+        
+        run = client.beta.threads.runs.create(
+            thread_id=thread_id,
+            assistant_id=ASSISTANT_ID
+        )
+
+        while True:
+            run_status = client.beta.threads.runs.retrieve(thread_id=thread_id, run_id=run.id)
+            if run_status.status == "completed":
+                break
+
+        messages = client.beta.threads.messages.list(thread_id=thread_id)
+        assistant_message = messages.data[0].content[0].text.value
+
+        st.chat_message("assistant").write(assistant_message)
+        st.session_state.messages.append({"role": "assistant", "content": assistant_message})
+
+    except Exception as e:
+        st.error(f"Error: {str(e)}")