Update app.py

app.py

@@ -5,165 +5,147 @@ from PIL import Image, ImageDraw
 from io import BytesIO
 import fitz  # PyMuPDF
 import time
-import os
 
-# Document categories
+# Configuration - Get API key from Streamlit secrets
+GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
+GEMINI_MODEL = "gemini-2-flash"
 DOCUMENT_TYPES = ["Land Records", "Caste Certificates", "Property Registrations"]
 
-# Set up session state
+# Initialize session state
 def initialize_session_state():
+    if "chat_history" not in st.session_state:
+        st.session_state["chat_history"] = []
     if "processed_doc" not in st.session_state:
-        st.session_state.processed_doc = None
-    if "current_file" not in st.session_state:
-        st.session_state.current_file = None
-    if "gemini_api_key" not in st.session_state:
-        st.session_state.gemini_api_key = None
+        st.session_state["processed_doc"] = None
+    if "doc_preview" not in st.session_state:
+        st.session_state["doc_preview"] = None
+    if "uploaded_file" not in st.session_state:
+        st.session_state["uploaded_file"] = None
 
 # Reset session state
 def reset_session_state():
-    st.session_state.processed_doc = None
-    st.session_state.current_file = None
+    for key in ["chat_history", "processed_doc", "doc_preview", "uploaded_file"]:
+        st.session_state.pop(key, None)
 
-# Encode file to base64
-def encode_file(file):
+# Encode uploaded file to base64
+def encode_file(uploaded_file):
     try:
-        file_content = file.getvalue()
-        return base64.b64encode(file_content).decode('utf-8')
+        file_bytes = uploaded_file.getvalue()
+        
+        if uploaded_file.type == "application/pdf":
+            pdf = fitz.open(stream=BytesIO(file_bytes))
+            page = pdf[0]
+            pix = page.get_pixmap()
+            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+        elif uploaded_file.type.startswith('image/'):
+            img = Image.open(BytesIO(file_bytes))
+        elif uploaded_file.type == "text/plain":
+            text = file_bytes.decode('utf-8')
+            img = Image.new('RGB', (800, 600), color=(73, 109, 137))
+            d = ImageDraw.Draw(img)
+            d.text((10, 10), text, fill=(255, 255, 0))
+        else:
+            st.error("Unsupported file format")
+            return None
+
+        img_byte_arr = BytesIO()
+        img.save(img_byte_arr, format='JPEG')
+        return base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
     except Exception as e:
-        st.error(f"Error encoding file: {str(e)}")
+        st.error(f"File processing error: {str(e)}")
         return None
 
 # Query Gemini API
-def query_gemini(prompt, image_b64):
-    if not st.session_state.gemini_api_key:
-        st.error("API Key is missing. Please enter your Gemini API Key.")
-        return None
+def query_gemini(prompt, image_b64=None):
+    url = f"https://generativelanguage.googleapis.com/v1/models/{GEMINI_MODEL}:generateContent?key={GEMINI_API_KEY}"
+    parts = [{"text": prompt}]
+    if image_b64:
+        parts.append({"inline_data": {"mime_type": "image/jpeg", "data": image_b64}})
     
     try:
-        headers = {
-            "Authorization": f"Bearer {st.session_state.gemini_api_key}",
-            "Content-Type": "application/json"
-        }
-        payload = {
-            "model": "gemini-1.5-pro",  # Change based on the latest API model
-            "prompt": prompt,
-            "image": image_b64
-        }
         response = requests.post(
-            "https://generativelanguage.googleapis.com/v1/models/gemini-1.5-pro:generateContent",  # Corrected URL
-            headers=headers,
-            json=payload
+            url,
+            json={"contents": [{"parts": parts}]},
+            headers={"Content-Type": "application/json"},
+            timeout=30
         )
-        response.raise_for_status()
-        return response.json().get("result", "No response received.")
+        if response.status_code != 200:
+            st.error(f"API Request failed with status code: {response.status_code}")
+            return None
+        
+        data = response.json()
+        if 'error' in data:
+            st.error(f"API Error: {data['error'].get('message', 'Unknown error')}")
+            return None
+        
+        if not data.get('candidates'):
+            st.error("No response candidates found in API response")
+            return None
+            
+        candidate = data['candidates'][0]
+        return candidate.get('content', {}).get('parts', [{}])[0].get('text', 'No response text found')
+    except requests.exceptions.RequestException as e:
+        st.error(f"API Request failed: {str(e)}")
+        return None
     except Exception as e:
-        st.error(f"Gemini API error: {str(e)}")
+        st.error(f"Unexpected error: {str(e)}")
         return None
 
-# Process the document
-def process_document(file):
+# Process the uploaded document
+def process_document():
+    if not st.session_state.uploaded_file:
+        st.error("Please upload a document first.")
+        return
     try:
         with st.spinner("Analyzing document..."):
-            image_b64 = encode_file(file)
+            image_b64 = encode_file(st.session_state.uploaded_file)
             if not image_b64:
                 return
 
-            # Store preview image
-            if file.type == "application/pdf":
-                pdf = fitz.open(stream=BytesIO(file.getvalue()))
-                page = pdf[0]
-                pix = page.get_pixmap()
-                st.session_state.doc_preview = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-            elif file.type.startswith('image/'):
-                st.session_state.doc_preview = Image.open(BytesIO(file.getvalue()))
-            else:
-                st.error("Unsupported file format")
-                return
-
-            # Classify document
             classify_prompt = f"Classify this document into one of these categories: {', '.join(DOCUMENT_TYPES)}. Respond only with the category name."
             doc_type = query_gemini(classify_prompt, image_b64)
             
-            # Extract details
-            extract_prompt = """Extract and organize all important details from this document including:
+            extract_prompt = """Extract key details including:
             - Names
             - Dates
             - Identification numbers
             - Locations
-            - Key terms
-            Format as a bullet-point list with clear headings."""
+            Format as a bullet-point list."""
             details = query_gemini(extract_prompt, image_b64)
-
-            # Verify authenticity
-            verify_prompt = "Analyze this document for signs of tampering or forgery. Provide verification results."
+            
+            verify_prompt = "Analyze this document for signs of tampering. Provide verification status."
             verification = query_gemini(verify_prompt, image_b64)
 
             st.session_state.processed_doc = {
                 "type": doc_type or "Unclassified",
                 "details": details or "No details extracted",
                 "verification": verification or "Verification failed",
-                "preview": st.session_state.doc_preview
             }
-
             st.success("Document processing complete!")
             time.sleep(1)
-
     except Exception as e:
         st.error(f"Document processing failed: {str(e)}")
         st.session_state.processed_doc = None
 
-# Main application
+# Main app function
 def main():
     st.set_page_config(page_title="DocVerify AI", layout="wide")
     initialize_session_state()
     
-    # Sidebar Controls
-    with st.sidebar:
-        st.header("Document Controls")
-
-        # API Key input
-        st.session_state.gemini_api_key = st.text_input("Enter Gemini API Key", type="password")
-
-        uploaded_file = st.file_uploader(
-            "Upload Document",
-            type=["pdf", "jpg", "jpeg", "png"],
-            key="uploaded_file"
-        )
-        
-        if st.button("Process Document"):
-            if uploaded_file:
-                st.session_state.current_file = uploaded_file
-                process_document(uploaded_file)
-            else:
-                st.error("Please select a document to process.")
-        
-        if st.button("New Document"):
-            reset_session_state()
-            st.rerun()
-        
-        if st.session_state.processed_doc:
-            st.divider()
-            st.subheader("Document Summary")
-            st.markdown(f"**Type:** {st.session_state.processed_doc['type']}")
-            st.markdown(f"**Verification Status:**\n{st.session_state.processed_doc['verification']}")
-
-    # Main Interface
-    st.title("📄 Automated Document Verifier")
+    st.sidebar.header("Document Controls")
+    st.sidebar.file_uploader("Upload Document", type=["pdf", "jpg", "jpeg", "png", "txt"], key="uploaded_file", on_change=process_document)
     
-    if st.session_state.processed_doc and 'preview' in st.session_state.processed_doc:
-        col1, col2 = st.columns([1, 2])
-        with col1:
-            st.subheader("Document Preview")
-            st.image(st.session_state.processed_doc['preview'], use_column_width=True)
-        
-        with col2:
-            st.subheader("Extracted Details")
-            st.markdown(st.session_state.processed_doc['details'])
-            
-            st.subheader("Verification Analysis")
-            st.markdown(st.session_state.processed_doc['verification'])
-    else:
-        st.info("Please upload a document and click 'Process Document' to start verification.")
+    if st.sidebar.button("New Document"):
+        reset_session_state()
+        st.rerun()
+    
+    st.title("DocVerify AI - Document Analysis")
+    
+    if st.session_state.processed_doc:
+        st.subheader("Document Summary")
+        st.markdown(f"**Type:** {st.session_state.processed_doc['type']}")
+        st.markdown(f"**Verification Status:** {st.session_state.processed_doc['verification']}")
+        st.text_area("Extracted Details", st.session_state.processed_doc['details'], height=200)
 
 if __name__ == "__main__":
     main()