app.py

import os
import base64
import json

import gradio as gr

from typing_extensions import TypedDict

from openai import OpenAI
from langchain_openai import ChatOpenAI

from langgraph.graph import StateGraph, START, END
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_community.document_loaders import PyPDFLoader

from dotenv import load_dotenv

load_dotenv()

client = OpenAI(
    api_key=os.environ["OPENAI_API_KEY"],
    base_url=os.environ["OPENAI_BASE_URL"]
)

llm = ChatOpenAI(
    api_key=os.environ["OPENAI_API_KEY"],
    base_url=os.environ["OPENAI_BASE_URL"],
    model="gpt-4o-mini",
    temperature=0
)

employee_name = 'John Doe'
type_of_expense = 'Restaurant'

company_policy_file_path = 'Company Policy on Expense Claims.pdf'
loader = PyPDFLoader(company_policy_file_path)
company_policy_document = loader.load()


class State(TypedDict):
    image_path: str
    extracted_text: str
    categorized_text: str
    relevant_company_policy: str
    verified_text: str
    revised_calculation: str
    final_output: str

def generate_data_uri(jpg_file_path):

    with open(jpg_file_path, 'rb') as image_file:
        image_data = image_file.read()
    
    # Encode the binary image data to base64
    base64_encoded_data = base64.b64encode(image_data).decode('utf-8')
    
    # Construct the data URI
    data_uri = f"data:image/png;base64,{base64_encoded_data}"
    
    return data_uri

def text_extractor(state: State):
    """
    This function extracts text from an image using OpenAI's GPT-4o mini model.
    """

    text_extraction_system_message = """
        You are an expert in extracting the text in images.
        Extract the following details from the bill presented in the input.
        - Date of bill
        - Bill No
        - Restaurant Name and Address
        - Items ordered quantity and price
        - Tax and Charges
        - Total amount

        Do not output anything except the above details in your output.
        """

    text_extraction_prompt = [
        {
            'role': 'system',
            'content': text_extraction_system_message
        },
        {
            'role': 'user',
            'content': [
                {'type': "image_url", "image_url": {'url': generate_data_uri(state['image_path'])}}
            ]
        }
    ]


    response = client.chat.completions.create(
        model='gpt-4o-mini',
        messages=text_extraction_prompt,
        temperature=0
    )

    extracted_text = response.choices[0].message.content

    return {'extracted_text': extracted_text}

def categorizer(state: State):

    categorization_system_message = """
        You are an expert accountant tasked to categorize the items ordered in the bill.
        Categorize the items STRICTLY into the following categories: Alcoholic Drinks, Non-Alcoholic Drinks and Food.
        Remember to categorize the items into one of the three categories only. Do not use new categories. 
        Present your output as a JSON with the following fields: 
        [{'item': '<name of the item>', 'category': '<category assigned>', 'quantity': '<quantity>', 'price': '<price>'}, ... and so on]        
        
        Do not output anything except the above fields in your JSON output.
        Do not delimit the JSON with any extra tags (e.g., ``` or ```JSON).
        """
    
    categorization_prompt = [
        SystemMessage(content=categorization_system_message),
        HumanMessage(content=state['extracted_text'])
    ]

    categorized_text = llm.invoke(categorization_prompt)

    return {'categorized_text': categorized_text.content}


def verifier(state: State):

    for document in company_policy_document:
        if document.page_content.find(f'{type_of_expense}') != -1: 
            relevant_company_policy = document.page_content
            
    verification_system_message = """
        You are an expert accountant tasked to verify the bill details against the provided company policy.
        Verify the items in the submitted bill against the company policy presented below.

        Present your output in the following JSON format after removing the items inthat are not aligned with the company policy.
        [{'item': '<name of the item>', 'category': '<category assigned>', 'quantity': '<quantity>', 'price': '<price>'}, ... and so on]        
        
        Do not output anything except the above details in your JSON output.
        Do not delimit the JSON with any extra tags (e.g., ``` or ```JSON).
        """ 
        
    verification_prompt = [
        SystemMessage(content=verification_system_message + f"\n Company Policy: \n{relevant_company_policy}"),
        HumanMessage(content=state['categorized_text'])
    ]

    verified_text = llm.invoke(verification_prompt)

    return {'verified_text': verified_text.content, 'relevant_company_policy': relevant_company_policy}


def estimator(state: State):

    total_bill = 0
    total_taxes_and_charges = 0

    for item in json.loads(state['verified_text']):
        total_bill += float(item['quantity']) * float(item['price'])

    total_taxes_and_charges = total_bill * 0.10 + total_bill * 0.025 + total_bill * 0.025 + total_bill * 0.20        
    
    revised_calculation = {
        'taxes_and_charges': total_taxes_and_charges,
        'total_amount': total_bill + total_taxes_and_charges
    }

    return {'revised_calculation': revised_calculation}


def formatter(state: State):
    final_output_system_message = """ 
        You are an expert accountant tasked to generate the expense claim report.
        Generate the expense claim report based on the calculated total amount to be reimbursed and other details available to you.
        
        The details of the fields needed for the report are present in the input.
        
        These are:
        - Employee Name:
        - Original Bill:
        - Verified items ordered quantity and price:
        - Total amount to be reimbursed:
        - Tax and Charges: 
        
        Use only the details from the input to generate the report.

        Present your output in the following markdown format.
        # Expense Claim Report
        
        ## Employee Name: <Insert Employee Name>
        ## Date: <Insert Date from original bill>
        ## Bill No: <Insert Bill No from original bill>
        ## Restaurant Name and Address: <Insert Restaurant Name and Address fromm original bill>
        ## Items ordered quantity and price (<arrange in a table format from verified list of items>):
            |Item|Quantity|Price|
                ...
            ... 
        ### Tax and Charges: <enter the tax amount from calculated amounts>
        ### Total amount to be reimbursed: <enter the total from calculated amounts>

        Do not output anything except the above details in your output.
        Do not delimit the output with any extra tags (e.g., ```).
        """
    
    input = f"""
    Employee Name: {employee_name}

    ---

    Original Bill:

    {state['extracted_text']}

    ---

    Verified items ordered quantity and price:
    {state['verified_text']}

    ---
    Calculated amounts:
    Taxes and Charges: {state['revised_calculation']['taxes_and_charges']}
    Total amount to be reimbursed: {state['revised_calculation']['total_amount']}

    """
    
    final_output_prompt = [
        SystemMessage(content=final_output_system_message),
        HumanMessage(content=input)
    ]

    final_output = llm.invoke(final_output_prompt)

    return {'final_output': final_output.content}

    
def claim_generator(input_bill_path):

    workflow = StateGraph(State)
    
    workflow.add_node("text_extractor", text_extractor)
    workflow.add_node("categorizer", categorizer)
    workflow.add_node("verifier", verifier)
    workflow.add_node("estimator", estimator)
    workflow.add_node("formatter", formatter)
    
    workflow.add_edge(START, "text_extractor")
    workflow.add_edge("text_extractor", "categorizer")
    workflow.add_edge("categorizer", "verifier")
    workflow.add_edge("verifier", "estimator")
    workflow.add_edge("estimator", "formatter")
    workflow.add_edge("formatter", END)
    
    chain = workflow.compile()
    
    output = chain.invoke({'image_path': input_bill_path})
    
    return output['final_output']


demo = gr.Interface(
    fn=claim_generator,
    inputs=gr.Image(type="filepath", label="Upload your image"),
    outputs=gr.Textbox(label="Expense Report", show_copy_button=True),
    title="Expense Report Extractor",
    description="This web API presents an interface to extract an expense claim report based on a submitted bill.",
    examples='images',
    cache_examples=False,
    theme=gr.themes.Base(),
    concurrency_limit=16
)

demo.queue()
demo.launch(auth=("demouser", os.getenv('PASSWD')), ssr_mode=False)