Update app.py

app.py

@@ -11,9 +11,10 @@ from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
 from reportlab.lib.styles import getSampleStyleSheet
 from bs4 import BeautifulSoup
 import re
+import base64
 
 # Set Gemini API Key
-os.environ["GEMINI_API_KEY"] = "AIzaSyAQLUStrjXv2_zrAjssapa_sYXQWOqVeiE"  # Ensure this is set securely in Hugging Face environment
+os.environ["GEMINI_API_KEY"] = "AIzaSyAQLUStrjXv2_zrAjssapa_sYXQWOqVeiE"  # Ensure this is set securely in your environment
 client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
 
 # Function to extract text from different file formats
@@ -56,6 +57,17 @@ def extract_text_from_tex(tex_file):
     content = re.sub(r'\\[a-zA-Z]+', '', content)  # Remove LaTeX commands like \section, \textbf, etc.
     return content
 
+def process_image(image_file):
+    # Convert the image to base64
+    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 export_conversation_to_pdf(conversation_history):
     pdf_path = "conversation_clean.pdf"
     doc = SimpleDocTemplate(pdf_path, pagesize=A4)
@@ -82,8 +94,10 @@ def main():
     if "chat_active" not in st.session_state:
         st.session_state.chat_active = True
 
-    uploaded_file = st.file_uploader("Upload a file (.txt, .docx, .csv, .xlsx, .pptx, .pdf, .html, .htm, .tex):",
-                                     type=["txt", "docx", "csv", "xlsx", "pptx", "pdf", "html", "htm", "tex"])
+    uploaded_file = st.file_uploader(
+        "Upload a file (.txt, .docx, .csv, .xlsx, .pptx, .pdf, .html, .htm, .tex, .jpg, .jpeg, .png):",
+        type=["txt", "docx", "csv", "xlsx", "pptx", "pdf", "html", "htm", "tex", "jpg", "jpeg", "png"]
+    )
 
     if uploaded_file and not st.session_state.document_text:
         filetype = uploaded_file.name.split(".")[-1].lower()
@@ -104,6 +118,11 @@ def main():
                 st.session_state.document_text = extract_text_from_html(uploaded_file)
             elif filetype == "tex":
                 st.session_state.document_text = extract_text_from_tex(uploaded_file)
+            elif filetype in ["jpg", "jpeg", "png"]:
+                image_data = process_image(uploaded_file)
+                st.session_state.document_text = "Image uploaded and processed."
+                st.image(uploaded_file, caption=uploaded_file.name, use_container_width=True)
+                st.session_state.image_data = image_data  # Store the image data for later use
             else:
                 st.error("Unsupported file format.")
         except Exception as e:
@@ -126,30 +145,41 @@ def main():
                         st.success("Chat ended. Reload to start again.")
                     else:
                         with st.spinner("Gemini is thinking..."):
-                            contents = [
-                                types.Content(
-                                    role="user",
-                                    parts=[types.Part.from_text(text=st.session_state.document_text),
-                                           types.Part.from_text(text=user_input),
-                                           types.Part.from_text(text="If needed, search the web for more context.")]
+                            content_blocks = []
+
+                            # Add prior Q&A as part of the context
+                            if st.session_state.conversation:
+                                history_text = "\n\n".join(
+                                    f"Q: {entry[0]}\nA: {entry[1]}"
+                                    for entry in st.session_state.conversation
                                 )
-                            ]
-                            tools = [types.Tool(google_search=types.GoogleSearch())]
-                            generate_config = types.GenerateContentConfig(
-                                tools=tools,
-                                response_mime_type="text/plain",
-                            )
+                                content_blocks.append({"text": f"Previous conversation:\n{history_text}"})
+
+                            # Add file text content
+                            if st.session_state.document_text.strip():
+                                content_blocks.append({"text": f"Context:\n{st.session_state.document_text}"})
 
-                            response_text = ""
-                            for chunk in client.models.generate_content_stream(
+                            # Add image blocks (if image data exists)
+                            if "image_data" in st.session_state:
+                                content_blocks.append(st.session_state.image_data)
+
+                            # Add current question
+                            content_blocks.append({"text": f"Question: {user_input}"})
+
+                            response = client.models.generate_content(
                                 model="gemini-2.0-flash",
-                                contents=contents,
-                                config=generate_config,
-                            ):
-                                response_text += chunk.text
+                                contents=content_blocks
+                            )
+
+                            st.session_state.conversation.append((user_input, response.text))
+                            st.success("💡 Answer:")
+                            st.write(response.text)
 
-                            st.session_state.conversation.append((user_input, response_text))
-                            st.rerun()
+                            # Save to session history
+                            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)