invoice_generator.py

import base64
import json
import os
from PIL import Image
from tqdm import tqdm
import subprocess
import textwrap
import random
import fitz
import io
import uuid
from openai import AzureOpenAI, Client


def encode_image(image_path: str):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")


def generate_accident_description(client: Client, image_path: str):
    base64_image = encode_image(image_path)

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": """
                        You will be provided with an image of a car accident.
                        If the provided photo is not a photo of a car accident or a damaged car (the picture should look like "First Notice of Loss" photo), you must write "Irrelevant." and nothing else.
                        If a car is completely destroyed and it is not possible to repair it, you must write "Irrelevant." and nothing else.
                        Otherwise, you must return two things:
                        1) Describe in detail all the damages that you see on the car (if there are more than one, choose the most significant ones). Try to add as many details about each damage as possible.
                        2) Come up with a story about how the accident happened. This story is written by a person who contacts the insurance service regarding the accident. The story should be plausible and consistent with the damages you see on the car. Write 2-3 sentences. Write it in simple words, like it was written py a person who is not a professional in car accidents.
                        As a result, you should return a json dictionary with the following keys: "damage_description" and "accident_story".
                        Remember, that you must return "Irrelevant." if the photo is not relevant.
                        DO NOT WRITE ANY ADDITIONAL COMMENTS.
                        """,
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}",
                        },
                    },
                ],
            }
        ],
        max_tokens=1024,
        temperature=0.6,
    )

    message = response.choices[0].message.content
    if "Irrelevant" in message or len(message.split()) < 12:
        return None
    message = message.replace("```json", "").replace("```", "")
    return json.loads(message)


def generate_invoice_file(
    client: Client, image_path: str, meta_info: dict, damage_description: str
):
    base64_image = encode_image(image_path)

    prompt = """
    Given an image of car accident and details of damage, generate repair invoice, which includes the cost of repair, the parts that need to be replaced and the labor cost/hours.
    The invoice should be a standard form typical for Bavaria Direct (write in German).
    Do no write contact details or any additional information, write ONLY repair and cost information.
    The invoice should include a list of items (not tables).
    Each item MUST include fields: Beschädigtes Teil, "Teilkosten" (in EUR), "Arbeitsstunden" (number - hours), "Arbeitskosten" (in EUR/Stunde), "Gesamtkosten" (in EUR).
    Check each calculation and make sure it is correct.
    
    Accident details:
    The accident happened in {0}, year: {1}.
    Damage description: 
    {2}
    """.format(
        meta_info["location"],
        meta_info["year"],
        damage_description,
    )

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}",
                        },
                    },
                ],
            }
        ],
        max_tokens=1024,
        temperature=0.3,
    )

    return response.choices[0].message.content


def compile_latex(invoice_table, template, output_folder, output_pdf_name):
    tex_file_path = os.path.join(output_folder, output_pdf_name.replace(".pdf", ".tex"))
    output_pdf_path = os.path.join(output_folder, output_pdf_name)

    with open(template, "r") as f:
        latex_template = f.read()

    latex_content = latex_template.replace("=========", invoice_table)

    with open(tex_file_path, "w") as f:
        f.write(latex_content)

    try:
        subprocess.run(
            [
                # "xelatex",
                # "-interaction=batchmode",
                "pdflatex",
                "-interaction=nonstopmode",
                "-output-directory",
                output_folder,
                tex_file_path,
            ],
            stdout=subprocess.DEVNULL,
            check=True,
        )
        if not os.path.exists(output_pdf_path):
            print("PDF generation failed.")
    except subprocess.CalledProcessError as e:
        print(f"Error in LaTeX compilation: {e}")

    extensions_to_remove = [".aux", ".out", ".log", ".tex"]
    for ext in extensions_to_remove:
        file_to_remove = os.path.join(
            output_folder, f'{output_pdf_name.replace(".pdf", "")}{ext}'
        )
        if os.path.exists(file_to_remove):
            os.remove(file_to_remove)


def embed_invoice_in_template(client: Client, invoice: str):
    prompt = (
        r"""
    Given an invoice, create a latex table and write the information about Gesamtsumme below the table.
    Do not add new sections or change the structure of the template (e.g. add new rows or columns).
    All you can change is the content of existing table's cells and the text below the table.
    In the column "Beschädigtes Teil" you should insert the corresponding text from the invoice.
    In the columns "Teilkosten", "Arbeitsstunden", "Arbeitskosten", "Gesamtkosten" you should insert the corresponding INTEGER numbers from the invoice.
    The invoice may contain additional information, but you should only use the information that is relevant to the table.
    The invoice's field may look like "80 EUR/Stunde", but you must insert only the number "80" into the table.
    YOU CAN NOT CHANGE THE COLUMN NAMES OR ADD NEW COLUMNS.
    YOU OUTPUT SHOULD BE A LATEX TABLE WITHOUT ANY ADDITIONAL COMMENTS.
    
    Table template:
    ```latex
    \textbf{Beschädigtes Teil} & \textbf{Teilkosten} & \textbf{Arbeitsstunden} & \textbf{Arbeitskosten} & \textbf{Gesamtkosten} \\
    .... & .... & .... & .... & .... \\
    ...
    \hline
    \end{tabular}
    \vspace{2cm}
    \newline
    \textbf{Gesamtsumme:} .... EUR
    ```
    
    If there are more than 3 words in the "Beschädigtes Teil" field, you should add a new line between third and fourth word (new line is \\).
    
    Invoice:
    
    """
        + invoice
    )

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.1,
    )

    return (
        response.choices[0].message.content.replace("```latex", "").replace("```", "")
    )


def complete_pdf(pdf_name: str, image_name: str, accident_story: str):
    doc = fitz.open(pdf_name)

    with Image.open(image_name) as img:
        img_byte_arr = io.BytesIO()
        img.save(img_byte_arr, format='PNG')
        img_byte_arr.seek(0)

    doc.insert_page(1)
    image_rect = fitz.Rect(50, 50, 400, 400)
    second_page = doc[1]
    second_page.insert_image(image_rect, stream=img_byte_arr.read())
    wrapped_text = textwrap.fill(accident_story, width=80)
    text_position = fitz.Point(50, 420)
    second_page.insert_text(text_position, wrapped_text, fontsize=12)

    temp_output_name = "temp.pdf"
    doc.save("temp.pdf")
    doc.close()

    os.replace(temp_output_name, pdf_name)


def generate_invoice(image_path, output_file, template, output_folder):
    os.makedirs(output_folder, exist_ok=True)
    client = AzureOpenAI(
        api_key=os.environ["AZURE_API_KEY"],
        api_version=os.environ["AZURE_API_VERSION"],
        azure_endpoint=os.environ["AZURE_ENDPOINT"],
    )
    meta_info = {
        "location": "Munich",
        "year": 2022,
    }
    print("Generating description...")
    description = generate_accident_description(client, image_path)
    if not description:
        print(f"Image {image_path} is irrelevant.")
        return 0

    assert "damage_description" in description, "damage_description not found"
    assert "accident_story" in description, "accident_story not found"

    print("Generating invoice...")
    invoice = generate_invoice_file(client, image_path, meta_info, description["damage_description"])
    invoice_table = embed_invoice_in_template(client, invoice)

    print("Compiling LaTeX...")
    compile_latex(invoice_table, template, output_folder, output_file)
    complete_pdf(
        f"{output_folder}/{output_file}", image_path, description["accident_story"]
    )
    return 1