app.py

app.py

@@ -1,15 +1,9 @@
 import os
-import gradio as gr
-import subprocess
 import streamlit as st
-from openai import OpenAI
 import time
-import PyPDF2
-from docx import Document
-from datetime import datetime
-from pydub import AudioSegment
-from io import BytesIO
 import json
+from datetime import datetime
+from openai import OpenAI
 
 # Streamlit Page Config
 st.set_page_config(page_title="Schlager ContractAi", layout="wide")
@@ -53,9 +47,7 @@ if not OPENAI_API_KEY:
 # Tabs
 tab1, tab2 = st.tabs(["Contract", "Technical"])
 
-SUPPORTED_AUDIO_FORMATS = (".mp3", ".wav", ".m4a")
-SUPPORTED_TEXT_FORMATS = (".txt", ".docx", ".csv", ".xlsx", ".pdf")
-
+# Function to save flagged responses
 def save_flagged_response(user_query, ai_response):
     flagged_data = {
         "query": user_query,
@@ -70,7 +62,6 @@ def save_flagged_response(user_query, ai_response):
 def contract_chat_section(tab, assistant_id, session_key, input_key):
     with tab:
         st.subheader("Chat")
-
         client = OpenAI(api_key=OPENAI_API_KEY)
 
         if session_key not in st.session_state:
@@ -80,15 +71,15 @@ def contract_chat_section(tab, assistant_id, session_key, input_key):
             st.session_state[session_key] = []
             st.rerun()
 
-        for message in st.session_state[session_key]:
+        for idx, message in enumerate(st.session_state[session_key]):
             role, content = message["role"], message["content"]
             if role == "assistant":
                 col1, col2 = st.columns([5, 1])
                 with col1:
                     st.chat_message(role).write(content)
                 with col2:
-                    if st.button("🚩 Flag", key=f"flag_{session_key}_{len(st.session_state[session_key])}"):
-                        save_flagged_response(st.session_state[session_key][-2]["content"], content)
+                    if st.button("🚩 Flag", key=f"flag_{session_key}_{idx}"):
+                        save_flagged_response(st.session_state[session_key][idx-1]["content"], content)
             else:
                 st.chat_message(role).write(content)
 
@@ -123,8 +114,9 @@ def contract_chat_section(tab, assistant_id, session_key, input_key):
             except Exception as e:
                 st.error(f"Error: {str(e)}")
 
+# Assign Assistant IDs
 ASSISTANT_CONTRACT_ID = "asst_rd9h8PfYuOmHbkvOF3RTmVfn"
 ASSISTANT_TECHNICAL_ID = "asst_xizNZBCJuy4TqdjqjwkxbAki"
 
 contract_chat_section(tab1, ASSISTANT_CONTRACT_ID, "contract_messages", "contract_input")
-contract_chat_section(tab2, ASSISTANT_TECHNICAL_ID, "technical_messages", "technical_input")
+contract_chat_section(tab2, ASSISTANT_TECHNICAL_ID, "technical_messages", "technical_input")