Update app.py

app.py

@@ -6,13 +6,13 @@ from datetime import datetime
 from openai import OpenAI
 
 # Streamlit Page Config
-st.set_page_config(page_title="CTrack Document Ai Assistant", layout="wide")
+st.set_page_config(page_title="Ctrack AI Assistant", layout="wide")
 
 # Authentication
 AUTHORIZED_USERS = {
     "[email protected]": "Pass.123",
     "[email protected]": "Pass.123",
-    "test@test,com": "Pass.123"
+    "test@test.com": "Pass.123"
 }
 
 if "authenticated" not in st.session_state:
@@ -33,22 +33,22 @@ if not st.session_state["authenticated"]:
     st.button("Login", on_click=login)
     st.stop()
 
-# Main App
-st.title("CTrack Document Ai Assistant")
-st.caption("Chat with your contract or manage meeting minutes")
+# Page Header
+st.title("🚚 Ctrack Virtual Assistant")
+st.caption("Official AI assistant for Ctrack’s fleet management solutions")
 
-# Load API Key and Assistant IDs from Environment Variables
+# Load OpenAI API Key
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
-ASSISTANT_CONTRACT_ID = os.getenv("ASSISTANT_CONTRACT_ID")
-ASSISTANT_TECHNICAL_ID = os.getenv("ASSISTANT_TECHNICAL_ID")
+CTRACK_ASSISTANT_ID = "asst_ItfaIt2UF2nN1UpOS5xa6l5p"
 
-if not OPENAI_API_KEY or not ASSISTANT_CONTRACT_ID or not ASSISTANT_TECHNICAL_ID:
-    st.error("Missing required environment variables. Please set OPENAI_API_KEY, ASSISTANT_CONTRACT_ID, and ASSISTANT_TECHNICAL_ID.")
+if not OPENAI_API_KEY:
+    st.error("Missing OPENAI_API_KEY environment variable.")
     st.stop()
 
-# Tabs
-tab1, tab2, tab3 = st.tabs(["Contract", "Technical", "Flagged Responses"])
+# Create Tabs
+main_tab, flagged_tab = st.tabs(["Ctrack Assistant", "Flagged Responses"])
 
+# Directory to store flagged responses
 FLAGGED_RESPONSES_DIR = "flagged_responses"
 os.makedirs(FLAGGED_RESPONSES_DIR, exist_ok=True)
 
@@ -60,43 +60,38 @@ def save_flagged_response(user_query, ai_response):
         "response": ai_response,
         "timestamp": datetime.now().isoformat()
     }
-    
     with open(filename, "w") as file:
         json.dump(flagged_data, file, indent=4)
-    
     st.success(f"Response flagged and saved as {filename}.")
 
-# Contract Chat Section
-def contract_chat_section(tab, assistant_id, session_key, input_key):
+def ctrack_chat(tab):
     with tab:
-        st.subheader("Chat")
+        st.subheader("🤖 Chat with Ctrack AI")
         client = OpenAI(api_key=OPENAI_API_KEY)
+        session_key = "ctrack_chat"
 
         if session_key not in st.session_state:
             st.session_state[session_key] = []
 
-        if st.button("Clear Chat", key=f"clear_chat_{session_key}"):
+        if st.button("Clear Chat", key="clear_chat"):
             st.session_state[session_key] = []
             st.rerun()
 
-        for idx, message in enumerate(st.session_state[session_key]):
-            role, content = message["role"], message["content"]
-            if role == "assistant":
-                with st.container():
-                    st.chat_message(role).write(content)
-                    if st.button("🚩 Flag", key=f"flag_{session_key}_{idx}"):
-                        user_query = st.session_state[session_key][idx-1]["content"] if idx > 0 else "Unknown"
-                        save_flagged_response(user_query, content)
-            else:
-                st.chat_message(role).write(content)
-
-        if prompt := st.chat_input("Enter your message:", key=input_key):
+        for idx, msg in enumerate(st.session_state[session_key]):
+            role, content = msg["role"], msg["content"]
+            st.chat_message(role).write(content)
+            if role == "assistant" and st.button("🚩 Flag", key=f"flag_{idx}"):
+                user_query = st.session_state[session_key][idx-1]["content"] if idx > 0 else "Unknown"
+                save_flagged_response(user_query, content)
+
+        if prompt := st.chat_input("Ask about Ctrack's services, fleet solutions, or integrations"):
             st.session_state[session_key].append({"role": "user", "content": prompt})
             st.chat_message("user").write(prompt)
 
             try:
                 thread = client.beta.threads.create()
                 thread_id = thread.id
+
                 client.beta.threads.messages.create(
                     thread_id=thread_id,
                     role="user",
@@ -105,39 +100,37 @@ def contract_chat_section(tab, assistant_id, session_key, input_key):
 
                 run = client.beta.threads.runs.create(
                     thread_id=thread_id,
-                    assistant_id=assistant_id
+                    assistant_id=CTRACK_ASSISTANT_ID
                 )
 
                 while True:
-                    run_status = client.beta.threads.runs.retrieve(thread_id=thread_id, run_id=run.id)
-                    if run_status.status == "completed":
+                    status = client.beta.threads.runs.retrieve(thread_id=thread_id, run_id=run.id)
+                    if status.status == "completed":
                         break
                     time.sleep(1)
 
                 messages = client.beta.threads.messages.list(thread_id=thread_id)
-                assistant_message = messages.data[0].content[0].text.value
-                st.session_state[session_key].append({"role": "assistant", "content": assistant_message})
-                with st.container():
-                    st.chat_message("assistant").write(assistant_message)
-                    if st.button("🚩 Flag", key=f"flag_response_{session_key}_{len(st.session_state[session_key])}"):
-                        save_flagged_response(prompt, assistant_message)
+                response = messages.data[0].content[0].text.value
+                st.session_state[session_key].append({"role": "assistant", "content": response})
+                st.chat_message("assistant").write(response)
+                if st.button("🚩 Flag", key=f"flag_response_{len(st.session_state[session_key])}"):
+                    save_flagged_response(prompt, response)
+
             except Exception as e:
-                st.error(f"Error: {str(e)}")
-
-contract_chat_section(tab1, ASSISTANT_CONTRACT_ID, "contract_messages", "contract_input")
-contract_chat_section(tab2, ASSISTANT_TECHNICAL_ID, "technical_messages", "technical_input")
-
-# Flagged Responses Tab
-with tab3:
-    st.subheader("Flagged Responses")
-    flagged_files = [f for f in os.listdir(FLAGGED_RESPONSES_DIR) if f.endswith(".json")]
-    
-    if flagged_files:
-        selected_file = st.selectbox("Select a flagged response file to download:", flagged_files)
-        
-        if selected_file:
-            with open(os.path.join(FLAGGED_RESPONSES_DIR, selected_file), "r") as file:
-                flagged_responses = file.read()
-            st.download_button("Download Selected Flagged Responses", data=flagged_responses, file_name=selected_file, mime="application/json")
+                st.error(f"OpenAI Error: {str(e)}")
+
+ctrack_chat(main_tab)
+
+# Flagged Response Viewer
+with flagged_tab:
+    st.subheader("🚩 Flagged Responses")
+    files = [f for f in os.listdir(FLAGGED_RESPONSES_DIR) if f.endswith(".json")]
+
+    if files:
+        selected = st.selectbox("Select a flagged response to download:", files)
+        if selected:
+            with open(os.path.join(FLAGGED_RESPONSES_DIR, selected), "r") as file:
+                flagged = file.read()
+            st.download_button("📥 Download", data=flagged, file_name=selected, mime="application/json")
     else:
-        st.info("No flagged responses available.")
+        st.info("No flagged responses yet.")