import streamlit as st
from PIL import Image
import google.generativeai as genai
import os

gemini_api_key = genai.configure(api_key = os.environ.get("Google_API_KEY"))
model = genai.GenerativeModel('gemini-1.5-flash')

input_prompts = """
You are an expert Invoice Entity Extractor and you are expert in understanding Invoices. 
We will upload an image as Invoice and you will have to answer any question based on the uploaded invoice.
"""

def response_gemini(input, image, prompt):
  response = model.generate_content([input, image[0], prompt])
  return response.text

st.title("Invoice Entity Extractor")

uploaded_file = st.file_uploader("Choose an invoice image", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
  image = Image.open(uploaded_file)
  st.image(image, caption='Uploaded Invoice', use_column_width=True)

  question = st.text_input("Ask a question about the invoice:")
  if st.button("Extract"):
    if question:
      answer = response_gemini(input_prompts, [image], question)
      #st.write("**Answer:**", answer)
      st.markdown(f"**Answer:** {answer}")
    else:
      st.warning("Please enter a question.")