app.py

import streamlit as st
import base64
import requests
from PIL import Image
from io import BytesIO
import fitz  # PyMuPDF
import os

# Configuration - Replace with your API key
GEMINI_API_KEY = st.secrets["GEMINI_API_KEY"]
GEMINI_MODEL = "gemini-2.0-flash"

DOCUMENT_TYPES = [
    "Insurance Policies", "Explanation of Benefits (EOBs)",
    "Claims (Approved, Denied, or Pending)", "Visit Summaries",
    "Test Results (Lab Reports, Imaging Reports)", 
    "Prescriptions (E-Prescriptions, Handwritten)",
    "Discharge Summaries", "Medical Bills", "Payment Statements",
    "Pharmacy Receipts", "Prior Authorization Requests",
    "Consent Forms", "Referral Letters", "Others"
]

def initialize_session_state():
    """Initialize all session state variables"""
    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 "doc_preview" not in st.session_state:
        st.session_state.doc_preview = None

def encode_file(uploaded_file):
    """Safely encode different file types to base64"""
    try:
        file_bytes = uploaded_file.getvalue()
        
        if uploaded_file.type == "application/pdf":
            # Convert PDF to image using PyMuPDF
            pdf = fitz.open(stream=BytesIO(file_bytes))
            page = pdf[0]  # Get the first page
            pix = page.get_pixmap()
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            img_byte_arr = BytesIO()
            img.save(img_byte_arr, format='JPEG')
            return base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
        
        elif uploaded_file.type in ["image/webp", "image/avif", "image/jpeg", "image/png", "image/gif"]:
            img = Image.open(BytesIO(file_bytes))
            img_byte_arr = BytesIO()
            img.save(img_byte_arr, format='JPEG')
            return base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
        
        elif uploaded_file.type == "text/plain":
            # Convert TXT to image
            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))
            img_byte_arr = BytesIO()
            img.save(img_byte_arr, format='JPEG')
            return base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
        
        return base64.b64encode(file_bytes).decode('utf-8')
    except Exception as e:
        st.error(f"File processing error: {str(e)}")
        return None

def query_gemini(prompt, image_b64=None):
    """Handle Gemini API communication"""
    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:
        response = requests.post(
            url,
            json={"contents": [{"parts": parts}]},
            headers={"Content-Type": "application/json"},
            timeout=30
        )
        response.raise_for_status()
        return response.json()["candidates"][0]["content"]["parts"][0]["text"]
    except Exception as e:
        st.error(f"API Error: {str(e)}")
        return None

def process_document():
    """Handle document processing pipeline"""
    uploaded_file = st.session_state.uploaded_file
    if not uploaded_file:
        return
    
    try:
        with st.spinner("Analyzing document..."):
            # Convert to base64
            image_b64 = encode_file(uploaded_file)
            if not image_b64:
                return

            # Generate preview
            if uploaded_file.type == "application/pdf":
                pdf = fitz.open(stream=BytesIO(uploaded_file.getvalue()))
                page = pdf[0]  # Get the first page
                pix = page.get_pixmap()
                img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                st.session_state.doc_preview = img
            elif uploaded_file.type in ["image/webp", "image/avif", "image/jpeg", "image/png", "image/gif"]:
                st.session_state.doc_preview = Image.open(uploaded_file)
            elif uploaded_file.type == "text/plain":
                text = uploaded_file.getvalue().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))
                st.session_state.doc_preview = img

            # Classify document
            classify_prompt = f"Classify this healthcare document into one of: {DOCUMENT_TYPES}. Respond only with the category name."
            doc_type = query_gemini(classify_prompt, image_b64) or "Others"
            
            # Store results
            st.session_state.processed_doc = {
                "type": doc_type,
                "content": image_b64,
                "summary": query_gemini("Create a detailed structured summary of this healthcare document.", image_b64)
            }
            
    except Exception as e:
        st.error(f"Processing failed: {str(e)}")
        st.session_state.processed_doc = None

def handle_chat_query():
    """Process user chat input"""
    user_input = st.session_state.chat_input
    if not user_input or not st.session_state.processed_doc:
        return
    
    prompt = f"""
    Document Context:
    - Type: {st.session_state.processed_doc['type']}
    - Summary: {st.session_state.processed_doc['summary']}
    
    Question: {user_input}
    
    Answer concisely and factually. If unsure, state "Information not found".
    """
    
    with st.spinner("Generating response..."):
        response = query_gemini(prompt, st.session_state.processed_doc['content']) 
    
    st.session_state.chat_history.append(("user", user_input))
    st.session_state.chat_history.append(("assistant", response or "Could not generate response"))

# UI Layout
def main():
    st.set_page_config(page_title="Healthcare Document Assistant", layout="wide")
    initialize_session_state()

    # Sidebar Section
    with st.sidebar:
        st.header("Document Management")
        
        # Preview above upload button
        if st.session_state.doc_preview:
            st.subheader("Preview")
            st.image(st.session_state.doc_preview, use_container_width=True)
        
        # Upload button
        st.file_uploader(
            "Upload Document",
            type=["pdf", "webp", "avif", "jpg", "jpeg", "png", "gif", "txt"],
            key="uploaded_file",
            on_change=process_document
        )
        
        # Document type below upload button
        if st.session_state.processed_doc:
            st.divider()
            st.subheader("Document Type")
            st.markdown(f"**{st.session_state.processed_doc['type']}**")

    # Main Content
    st.title("Healthcare Document Assistant")
    
    if st.session_state.processed_doc:
        # Document Summary
        st.subheader("Document Summary")
        st.markdown(st.session_state.processed_doc['summary'])
        
        # Chat Interface
        st.divider()
        st.subheader("Document Q&A")
        
        # Chat history
        for role, message in st.session_state.chat_history:
            with st.chat_message(role.capitalize()):
                st.markdown(message)
        
        # Chat input
        st.chat_input(
            "Ask about the document...",
            key="chat_input",
            on_submit=handle_chat_query
        )
    else:
        st.info("Please upload a document to begin analysis")

if __name__ == "__main__":
    main()