Update app.py

app.py

@@ -16,39 +16,38 @@ import base64
 # Set Gemini API Key
 client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
 
-# System prompt
-SYSTEM_PROMPT = (
-    "You are part of an intelligent driver distraction detection system. "
-    "Your role is to analyze documents, images, and videos to help identify and classify driver states, "
-    "including: drowsy, alert, yawning, microsleep, eyes/lips open or closed, and other distractions. "
-    "The user may ask for timestamps of events in video, technical explanations from code/docs, "
-    "or structured summaries. Prioritize your answers for this use case and act as a helpful assistant "
-    "within this domain. "
-    "Additionally, your classification may rely on the State Farm Distracted Driver Detection dataset, "
-    "which defines 10 distraction categories such as texting, phone use, reaching behind, adjusting the radio, "
-    "and safe driving. Use this labeling scheme when interpreting relevant images or video frames."
-)
-
 # --------- File Extractors ---------
 def extract_text_from_docx(docx_file):
-    return "\n".join([para.text for para in Document(docx_file).paragraphs])
+    document = Document(docx_file)
+    return "\n".join([para.text for para in document.paragraphs])
 
 def extract_text_from_csv(csv_file):
-    return pd.read_csv(csv_file).to_string(index=False)
+    df = pd.read_csv(csv_file)
+    return df.to_string(index=False)
 
 def extract_text_from_xlsx(xlsx_file):
-    return pd.read_excel(xlsx_file).to_string(index=False)
+    df = pd.read_excel(xlsx_file)
+    return df.to_string(index=False)
 
 def extract_text_from_pptx(pptx_file):
     prs = Presentation(pptx_file)
-    return "\n".join([shape.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, "text")])
+    text_runs = []
+    for slide in prs.slides:
+        for shape in slide.shapes:
+            if hasattr(shape, "text"):
+                text_runs.append(shape.text)
+    return "\n".join(text_runs)
 
 def extract_text_from_pdf(pdf_file):
     doc = fitz.open(stream=pdf_file.read(), filetype="pdf")
-    return "".join([page.get_text() for page in doc])
+    text = ""
+    for page in doc:
+        text += page.get_text()
+    return text
 
 def extract_text_from_html(html_file):
-    return BeautifulSoup(html_file.read(), "html.parser").get_text()
+    soup = BeautifulSoup(html_file.read(), "html.parser")
+    return soup.get_text()
 
 def extract_text_from_tex(tex_file):
     content = tex_file.read().decode("utf-8")
@@ -57,131 +56,170 @@ def extract_text_from_tex(tex_file):
     return content
 
 def process_image(image_file):
-    return types.Part.from_data(
-        data=image_file.read(),
-        mime_type=image_file.type
-    )
+    image_bytes = image_file.read()
+    encoded_image = base64.b64encode(image_bytes).decode("utf-8")
+    return {
+        "inline_data": {
+            "mime_type": image_file.type,
+            "data": encoded_image
+        }
+    }
 
 def process_video(video_file):
+    video_bytes = video_file.read()
     filetype = video_file.name.split(".")[-1].lower()
     mime_map = {
-        "avi": "video/x-msvideo", "mkv": "video/x-matroska",
-        "flv": "video/x-flv", "wmv": "video/x-ms-wmv",
+        "avi": "video/x-msvideo",
+        "mkv": "video/x-matroska",
+        "flv": "video/x-flv",
+        "wmv": "video/x-ms-wmv",
         "mpeg": "video/mpeg"
     }
     mime_type = mime_map.get(filetype, video_file.type)
-    return types.Part.from_data(data=video_file.read(), mime_type=mime_type)
+    encoded_video = base64.b64encode(video_bytes).decode("utf-8")
+    return {
+        "inline_data": {
+            "mime_type": mime_type,
+            "data": encoded_video
+        }
+    }
 
 def export_conversation_to_pdf(conversation_history):
     pdf_path = "Conversation.pdf"
     doc = SimpleDocTemplate(pdf_path, pagesize=A4)
     styles = getSampleStyleSheet()
     elements = []
+
     for i, (user_q, gemini_a) in enumerate(conversation_history):
         elements.append(Paragraph(f"<b>Q{i+1}:</b> {user_q}", styles["Normal"]))
         elements.append(Spacer(1, 8))
         elements.append(Paragraph(f"<b>A{i+1}:</b> {gemini_a.replace(chr(10), '<br/>')}", styles["Normal"]))
         elements.append(Spacer(1, 16))
+
     doc.build(elements)
     return pdf_path
 
 # --------- Main App ---------
 def main():
-    st.set_page_config(page_title="Driver State Chat with Gemini", layout="wide")
-    st.title("🚗 Chat with Driver Distraction Files using Gemini Flash (Team 18)")
-
-    for key in ["conversation", "documents_text", "chat_active", "chat_history", "image_data", "video_data"]:
-        if key not in st.session_state:
-            st.session_state[key] = [] if key == "conversation" else ""
+    st.set_page_config(page_title="Gemini 2.0 Flash Chat Q&A", layout="wide")
+    st.title("\U0001F4C4 Chat with Your Documents leveraging Internet Using Gemini 2.0 Flash (Team 18)")
+
+    if "conversation" not in st.session_state:
+        st.session_state.conversation = []
+    if "documents_text" not in st.session_state:
+        st.session_state.documents_text = []
+    if "chat_active" not in st.session_state:
+        st.session_state.chat_active = True
+    if "chat_history" not in st.session_state:
+        st.session_state.chat_history = []
 
     uploaded_files = st.file_uploader(
-        "Upload documents, images, or videos:",
-        type=["txt", "docx", "csv", "xlsx", "pptx", "pdf", "html", "htm", "tex",
-              "jpg", "jpeg", "png", "mp4", "webm", "mov", "avi", "mkv", "flv", "wmv", "mpeg"],
+        "Upload files (.txt, .docx, .csv, .xlsx, .pptx, .pdf, .html, .htm, .tex, .jpg, .jpeg, .png, video formats):",
+        type=[
+            "txt", "docx", "csv", "xlsx", "pptx", "pdf", "html", "htm", "tex",
+            "jpg", "jpeg", "png",
+            "mp4", "webm", "mov", "avi", "mkv", "flv", "wmv", "mpeg"
+        ],
         accept_multiple_files=True
     )
 
     if uploaded_files:
-        all_text = ""
-        for f in uploaded_files:
-            ext = f.name.split(".")[-1].lower()
+        all_text_content = ""
+        for uploaded_file in uploaded_files:
+            filetype = uploaded_file.name.split(".")[-1].lower()
             try:
-                if ext == "txt":
-                    all_text += f.read().decode("utf-8") + "\n"
-                elif ext == "docx":
-                    all_text += extract_text_from_docx(f) + "\n"
-                elif ext == "csv":
-                    all_text += extract_text_from_csv(f) + "\n"
-                elif ext == "xlsx":
-                    all_text += extract_text_from_xlsx(f) + "\n"
-                elif ext == "pptx":
-                    all_text += extract_text_from_pptx(f) + "\n"
-                elif ext == "pdf":
-                    all_text += extract_text_from_pdf(f) + "\n"
-                elif ext in ["html", "htm"]:
-                    all_text += extract_text_from_html(f) + "\n"
-                elif ext == "tex":
-                    all_text += extract_text_from_tex(f) + "\n"
-                elif ext in ["jpg", "jpeg", "png"]:
-                    st.session_state.image_data = process_image(f)
-                    st.image(f, caption=f.name, use_container_width=True)
-                    all_text += f"Image {f.name} uploaded.\n"
-                elif ext in ["mp4", "webm", "mov", "avi", "mkv", "flv", "wmv", "mpeg"]:
-                    st.session_state.video_data = process_video(f)
-                    st.video(f)
-                    all_text += f"Video {f.name} uploaded.\n"
+                if filetype == "txt":
+                    all_text_content += uploaded_file.read().decode("utf-8") + "\n"
+                elif filetype == "docx":
+                    all_text_content += extract_text_from_docx(uploaded_file) + "\n"
+                elif filetype == "csv":
+                    all_text_content += extract_text_from_csv(uploaded_file) + "\n"
+                elif filetype == "xlsx":
+                    all_text_content += extract_text_from_xlsx(uploaded_file) + "\n"
+                elif filetype == "pptx":
+                    all_text_content += extract_text_from_pptx(uploaded_file) + "\n"
+                elif filetype == "pdf":
+                    all_text_content += extract_text_from_pdf(uploaded_file) + "\n"
+                elif filetype in ["html", "htm"]:
+                    all_text_content += extract_text_from_html(uploaded_file) + "\n"
+                elif filetype == "tex":
+                    all_text_content += extract_text_from_tex(uploaded_file) + "\n"
+                elif filetype in ["jpg", "jpeg", "png"]:
+                    st.session_state.image_data = process_image(uploaded_file)
+                    all_text_content += "Image uploaded and processed.\n"
+                    st.image(uploaded_file, caption=uploaded_file.name, use_container_width=True)
+                elif filetype in ["mp4", "webm", "mov", "avi", "mkv", "flv", "wmv", "mpeg"]:
+                    st.session_state.video_data = process_video(uploaded_file)
+                    all_text_content += f"Video file '{uploaded_file.name}' uploaded and processed.\n"
+                    st.video(uploaded_file)
+                else:
+                    st.error(f"Unsupported file format: {uploaded_file.name}")
             except Exception as e:
-                st.error(f"Error processing {f.name}: {e}")
-        st.session_state.documents_text = all_text
+                st.error(f"Failed to extract/process {uploaded_file.name}: {e}")
+
+        st.session_state.documents_text = all_text_content
 
     if st.session_state.documents_text:
-        st.markdown("### 💬 Chat")
-        for q, a in st.session_state.conversation:
-            st.markdown(f"**You:** {q}")
-            st.markdown(f"**Gemini:** {a}")
+        st.markdown("### \U0001F4AC Conversation")
+        for user_q, gemini_a in st.session_state.conversation:
+            st.markdown(f"**You:** {user_q}")
+            st.markdown(f"**Gemini:**\n\n{gemini_a}")
 
         if st.session_state.chat_active:
-            with st.form("chat_form", clear_on_submit=True):
-                user_input = st.text_input("Ask a question about the uploaded files:")
-                submitted = st.form_submit_button("Send")
-                if submitted and user_input:
+            with st.form(key="chat_form", clear_on_submit=True):
+                user_input = st.text_input("Ask a question about the documents (type 'exit' to stop):")
+                submit = st.form_submit_button("Send")
+
+                if submit and user_input:
                     if user_input.strip().lower() == "exit":
                         st.session_state.chat_active = False
-                        st.success("Chat ended.")
+                        st.success("Chat ended. Reload to start again.")
                     else:
                         with st.spinner("Gemini is thinking..."):
-                            messages = [
-                                types.Content(role="user", parts=[types.Part(text=SYSTEM_PROMPT)])
-                            ]
+                            content_blocks = []
 
                             if st.session_state.conversation:
-                                conv = "\n".join([f"Q: {q}\nA: {a}" for q, a in st.session_state.conversation])
-                                messages.append(types.Content(role="user", parts=[types.Part(text=conv)]))
+                                history_text = "\n\n".join(
+                                    f"Q: {entry[0]}\nA: {entry[1]}"
+                                    for entry in st.session_state.conversation
+                                )
+                                content_blocks.append({"text": f"Previous conversation:\n{history_text}"})
 
-                            if st.session_state.documents_text:
-                                messages.append(types.Content(role="user", parts=[types.Part(text=st.session_state.documents_text)]))
+                            if st.session_state.documents_text.strip():
+                                content_blocks.append({"text": f"Context:\n{st.session_state.documents_text}"})
 
-                            if st.session_state.image_data:
-                                messages.append(types.Content(role="user", parts=[st.session_state.image_data]))
+                            if "image_data" in st.session_state:
+                                content_blocks.append(st.session_state.image_data)
 
-                            if st.session_state.video_data:
-                                messages.append(types.Content(role="user", parts=[st.session_state.video_data]))
+                            if "video_data" in st.session_state:
+                                content_blocks.append(st.session_state.video_data)
 
-                            messages.append(types.Content(role="user", parts=[types.Part(text=user_input)]))
+                            content_blocks.append({"text": f"Question: {user_input}"})
 
-                            response = client.generate_content(
+                            response = client.models.generate_content(
                                 model="gemini-2.0-flash",
-                                contents=messages
+                                contents=content_blocks,
+                                config={"tools": [{"google_search": {}}]}
                             )
 
                             st.session_state.conversation.append((user_input, response.text))
-                            st.success("💡 Answer:")
+                            st.success("\U0001F4A1 Answer:")
                             st.write(response.text)
 
-    if st.session_state.conversation:
-        pdf_path = export_conversation_to_pdf(st.session_state.conversation)
-        with open(pdf_path, "rb") as f:
-            st.download_button("⬇️ Export Conversation as PDF", f, "Conversation.pdf", "application/pdf")
+                            st.session_state.chat_history.append({
+                                "question": user_input,
+                                "answer": response.text
+                            })
+
+        if st.session_state.conversation:
+            pdf_path = export_conversation_to_pdf(st.session_state.conversation)
+            with open(pdf_path, "rb") as f:
+                st.download_button(
+                    label="\U0001F4E5 Export Conversation as PDF",
+                    data=f,
+                    file_name="Conversation.pdf",
+                    mime="application/pdf"
+                )
 
 if __name__ == "__main__":
-    main()
+    main()