app.py

import streamlit as st
import base64
import os
import json
from dotenv import load_dotenv
from openai import OpenAI

#langchain imports
from langchain.chains import RetrievalQA
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain.chains import create_qa_with_sources_chain
from langchain.chains.combine_documents.stuff import StuffDocumentsChain
from langchain.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from langchain.schema import Document

from langchain.chains.query_constructor.base import AttributeInfo
from langchain.retrievers.self_query.base import SelfQueryRetriever


#from pdf2image import convert_from_path
import io
import json
import re

# Function to encode the image to base64
# def encode_image(image_file):
#     return base64.b64encode(image_file.getvalue()).decode("utf-8")
                
def encode_image(image_files):
    base64_images = []

    for image_file in image_files:
        base64_images.append(base64.b64encode(image_file.getvalue()).decode("utf-8"))

    return base64_images, image_files

st.set_page_config(page_title="Document/Image AI Analyst", layout="centered", initial_sidebar_state="collapsed")
# Streamlit page setup
st.title("Document/Image AI Analyst: `GPT-4 with Vision v2` 👀")
st.write("")
st.write("")
st.write("")

# Guide
st.subheader('What can it do?')
with st.expander('Read more details', expanded=False):
    #st.write("There are various use cases that the AI analyst can do!")
    st.markdown('- It can describe details found on the images. For instance, it can tell the details on an uploaded passport, such as full name, passport number, country, etc.')
    st.markdown('- It can check for validity of images or identification documents. It also checks the legitimacy of documents (if applicable). `Try uploading a suspicious passport picture!`')
    st.markdown("- It can compare multiple documents, such as identifying whether a person's photo is the same as the uploaded personal documents. In a comparison use case, feel free to provide extra info (optional) on what comparison you want to perform.")
    st.markdown('- And anything else! For a simpler, general demo, upload any image and let it describe what it sees!')
    if st.button('Happy prompting and Cheers! 🎈'):
        st.balloons()

# new line space
st.write("")

st.subheader('1. We need an OpenAI API key: ')
api_key = st.text_input('',placeholder='Enter your OpenAI API key', type='password', help="You can find your OpenAI API key here: https://platform.openai.com/api-keys. Or if you are provided with one by your organization.")
st.caption('Never share your OpenAI API key to anyone. Note that usage of your OpenAI API key will be billed to your OpenAI account. Keep in mind that an image analysis costs approximately `$0.04`')


# Initialize messages object
messages = []


# new line space
st.write("")

# File uploader allows user to add their own image
st.subheader('2. Upload Images: ')
uploaded_files = st.file_uploader("", help='Up to five images only.', type=["jpg", "png", "jpeg"], accept_multiple_files=True)


if api_key:
    # Initialize the OpenAI client with the API key
    client = OpenAI(api_key=api_key)    
    # load from disk
    docsearch = Chroma(persist_directory="./knowledge_base", embedding_function=OpenAIEmbeddings())

    metadata_field_info = [
        AttributeInfo(
            name="type",
            description="The type of document",
            type="string",
        ),
        AttributeInfo(
            name="name",
            description="The name of the document",
            type="string",
        ),
        AttributeInfo(
            name="filename",
            description="The file name of the document",
            type="string",
        ),    
    ]

    document_content_description = "Type of a document"
    llm = ChatOpenAI(temperature=0)

    retriever = SelfQueryRetriever.from_llm(
        llm,
        docsearch,
        document_content_description,
        metadata_field_info,
        enable_limit=True,
        search_kwargs={"k": 1}
    )

# This sample for querying
# results = retriever.invoke("requirements in passport")
# st.write("RESULTS: ", results)
#############
if uploaded_files:
    if len(uploaded_files) <= 5:
        st.success("You uploaded " + str(len(uploaded_files)) + " images!", icon="✅")
    elif len(uploaded_files) > 5:
        st.error("More than 5 uploaded images. Please remove.", icon="❌")

    file_names = [file.name for file in uploaded_files]
    print("Filenames: ", file_names)

    for uploaded_file in uploaded_files:
        with st.expander("Uploaded image: `" + uploaded_file.name + "`", expanded = False):
            st.image(uploaded_file, use_column_width=True)
    
# if uploaded_file:
#     # Display the uploaded image
#     with st.expander("Image", expanded = True):
#         st.image(uploaded_file, caption=uploaded_file.name, use_column_width=True)

# new line space
st.write("")

# Toggle for showing additional details input
st.subheader('3. Details about the images:')
show_details = st.toggle("Add details about the images (optional)", value=False)
st.caption('')

if show_details:
    # Text input for additional details about the image, shown only if toggle is True
    additional_details = st.text_area(
        "Add any additional details or context about the image(s) here:",
        placeholder='I am typically able to understand images without context, but feel free to describe what type of analysis you want. For instance, verifying personal documents, checking for falsification or nothing at all (optional)',
        disabled=not show_details
    )


# new line space
st.write("")

# Button to trigger the analysis
st.subheader('4. Analyze! ')
analyze_button = st.button("Analyse the image(s)", type="secondary")
st.caption('')

# Check if an image has been uploaded, if the API key is available, and if the button has been pressed
if uploaded_files is not None and api_key and analyze_button:

    with st.spinner("Analysing the image(s) ..."):
        # Encode the image
        base64_image, filenames = encode_image(uploaded_files)
        #print("base64 images: ", filenames)
    
        #kb_details = ("")

        # Optimized prompt for additional clarity and detail

        prompt_text = (
            "You are a highly knowledgeable student admission document identifier. Evaluate each document and identify what type of admission document it is. "
            "Write your output as a JSON object mapping of the document file name and the type of document you identified, using double-quote, not single-quote. "
            "The format of the output must be enclosed in curly-brackets. "
            f"\n\nYou are provided with these corresponding document file names in order:\n{file_names}"
        )

    
        if show_details and additional_details:
            prompt_text += (
                f"\n\nAdditional Context Provided by the User:\n{additional_details}"
                #f"\n\nGeneral Admissions Guide on how to verify student admission documents:\n{kb_details}"
                
            )
    
        # IF scenarios for images payload for messages var
        if len(uploaded_files) == 1:
            messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt_text},
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image[0]}",
                    },
                ],
            }
            ]
        elif len(uploaded_files) == 2:
            messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt_text},
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image[0]}",
                    },
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image[1]}",
                    },
                ],
            }
            ]
        elif len(uploaded_files) == 3:
            messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt_text},
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image[0]}",
                    },
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image[1]}",
                    },
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image[2]}",
                    },                    
                ],
            }
            ]
        elif len(uploaded_files) == 4:
            messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt_text},
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image[0]}",
                    },
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image[1]}",
                    },
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image[2]}",
                    },                    
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image[3]}",
                    },
                ],
            }
            ]
        elif len(uploaded_files) == 5:
            messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt_text},
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image[0]}",
                    },
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image[1]}",
                    },
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image[2]}",
                    },
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image[3]}",
                    },                                        
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image[4]}",
                    },
                ],
            }
            ]
        elif len(uploaded_files) > 5:
            messages = []
                
        

        # Make the request to the OpenAI API
        try:
            ##################
            # STEP 1: Identify and map filenames to type of document
            full_response = ''
            message_placeholder = st.empty()
            
            for completion in client.chat.completions.create(
                model="gpt-4-vision-preview", messages=messages, 
                max_tokens=1200, stream=True,
            ):
                # Check if there is content to display
                if completion.choices[0].delta.content is not None:
                    full_response += completion.choices[0].delta.content
                    message_placeholder.markdown(full_response + "▌")
            # Final update to placeholder after the stream ends
            #message_placeholder.text(full_response)

            #full_response = '{"filename1": "Give full information about IELTS Test Result without summarizing", "filename2": "Give full information about TOEFL Test results without summarizing"}'
            ##################
            # STEP 2: perform validation via KB retrieval
            # STEP 2.1: Retrieval individual info and concat into an array

            # Clean up format: 
            json_response = re.search(r'\{.*\}', full_response, re.DOTALL).group()
            
            full_response_json = json.loads(json_response)

            # Display identified documents
            # Initialize an empty list to store markdown lines
            markdown_list = []
            
            # Iterate through the dictionary and format each item
            for filename, content in full_response_json.items():
                markdown_list.append(f"- {filename} is a {content}")

            # Join the list into a single string with newline characters
            markdown_string = '\n'.join(markdown_list)
            message_placeholder.markdown("**Identified Documents:**\n\n" + markdown_string)

            ### TO DO : The ABOVE code can be in a function or object.
            # Initialize an empty array to store the full prompt as log 
            retrieved_knowledge = []
            
            # Iterate over each item in the full_response dictionary
            st.subheader("5. Analysis Results")
            for i, (filename, details) in enumerate(full_response_json.items()):

                with st.expander(details):
                    file_content = ""
                    # perform query to KB, details is the query
                    results = retriever.invoke(details)

                    # Step 1: Access the first Document
                    document = results[0]

                    # Step 2: Retrieve the filename from the metadata
                    kb_retrieved_filename = document.metadata['filename']

                    # Step 3: Open the file using the extracted filename
                    with open(kb_retrieved_filename, 'r') as file:
                        file_content = file.read()

                    ### TO DO: NOW perform the actual validation
                    prompt_text_validation = (
                        "You are a highly knowledgeable student admission document image verifier. "
                        "Your task is to examine the image in detail. "
                        #"Provide a comprehensive, factual, and accurate explanation of what each image depict and classify them based on which student admission document it is. "
                        "Verify and compare whether the details in the document image is valid by checking the provided DOCUMENT REQUIREMENTS. "
                        "Highlight key elements on the document image that is valid and invalid, and present your analysis in clear, well-structured markdown format. "
                        "If applicable, identify any falsification, tampering and editing of the image that could potentially mean the document is not valid and was tampered. "
                        #"Assume the reader has a basic understanding of how student admission documents should be."
                        "Lastly, include your final verdict on whether the document is legit or needs further checking. Label as [LEGIT], [NEED FURTHER CHECKING] or [NOT LEGIT]"
                        f"\n\nDOCUMENT REQUIREMENTS that must be met to be considered valid:\n{file_content}"
                    ) 

                    messages = [
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": prompt_text_validation},
                            {
                                "type": "image_url",
                                "image_url": f"data:image/jpeg;base64,{base64_image[i]}",
                            },
                        ],
                    }
                    ]

                    validation_response = ''
                    validation_placeholder = st.empty()
                    
                    prompt_message = ''

                    
                    #prompt_placeholder = st.write("COMPLETE PROMPT", prompt_text_validation)

                    for completion in client.chat.completions.create(
                        model="gpt-4-vision-preview", messages=messages, 
                        max_tokens=1200, stream=True
                    ):
                        # Check if there is content to display
                        if completion.choices[0].delta.content is not None:
                            validation_response += completion.choices[0].delta.content
                            validation_placeholder.markdown(validation_response + "▌")
                    # Final update to validation placeholder after the stream ends
                    
                    #validation_placeholder.text(validation_response)                

                    # Now you can work with file_content
                    #print(file_content)
                    # add data to array
                    retrieved_knowledge.extend([prompt_text_validation, messages]) 

                    validation_placeholder.markdown(validation_response)

            with st.expander("Process Logs: "):
                st.json(retrieved_knowledge) 

            # Display the response in the app
            # st.write(response.choices[0].message.content)
        except Exception as e:
            st.error(f"An error occurred: {e}")
else:
    # Warnings for user action required
    if not uploaded_files and analyze_button:
        st.warning("Please upload at least one image. Up to five.", icon="⚠️")
    if not api_key:
        st.error("Please enter your OpenAI API key.", icon="❌")