src/gradio_app.py

import os

import pypdfium2 as pdfium
import gradio as gr

from langchain.chat_models import ChatOpenAI
from src.summarization import (
    parallel_summarization,
    PARALLEL_SUMMARIZATION_MAPPING,
)
from src.legal_implications import parallel_legal_implications
from src.mailing_utils import update_email_content, construct_url
import urllib


def _file_render_helper(file):
    pdf = pdfium.PdfDocument(file.name)
    images = []
    for page_index in range(len(pdf)):
        page = pdf.get_page(page_index)
        bitmap = page.render(
            scale=200 / 72,  # 200dpi resolution
            rotation=0,  # no additional rotation
            # ... further rendering options
        )
        images.append(
            (bitmap.to_pil(), f"{file.name.split('/')[-1]} Seite {page_index+1}")
        )
    return images


# Function to render a specific page of a PDF file as an image
def render_file(file):
    images = _file_render_helper(file)
    return gr.update(
        value=images,
    )


def render_files(files):
    all_images = []
    for file in files:
        images = _file_render_helper(file)
        all_images.extend(images)

    return gr.update(
        value=all_images,
    )


def create_email_textbox() -> gr.Textbox:
    return gr.Textbox(
        label="Email Instructions",
        placeholder="Write Email Instructions here.",
        value=(
            "Dear Recipient\n\n"
            "Please find the summarization of the uploaded document below.\n\n"
            "<TEXT_FROM_LLM>\n\n"
            "Kind regards,\n"
            "Your Legal Assistant"
        ),
        lines=9,
    )


def switch_buttons(interactive: bool):
    """This switches the buttons to interactive or not interactive.

    Args:
        interactive (bool): Whether the buttons should be interactive or not.
    """
    return (
        gr.update(
            interactive=interactive,
        ),
        gr.update(
            interactive=interactive,
        ),
        gr.update(
            interactive=interactive,
        ),
    )


def update_email_url(address: str, subject: str, email_instructions_summary: str):
    body = urllib.parse.quote(email_instructions_summary)
    return gr.update(link=construct_url(address=address, subject=subject, body=body))


def update_email_instructions_summary(email_wrapping_content: str, summary: str):
    return gr.update(
        value=update_email_content(
            email_wrapping_content=email_wrapping_content, summary=summary
        )
    )


def load_summary_section(llm: ChatOpenAI):
    """Load the summary section

    Args:
        llm (ChatOpenAI): Language model.

    Returns:
        gr.Blocks: The summarization section
    """

    with gr.Blocks(
        theme="soft",
    ) as summary_section:
        with gr.Row().style(equal_height=True):
            with gr.Column(scale=1):
                file_upload_summary = gr.File(
                    file_count="single",
                    file_types=[".pdf", ".txt"],
                    label="Upload PDF",
                )
                summary_parallel_button = gr.Button("Summarize", interactive=False)
                clear_button = gr.Button("Clear All")
            with gr.Column(scale=2):
                sections_to_select = [
                    i for i in PARALLEL_SUMMARIZATION_MAPPING.keys() if "II." not in i
                ]
                summary_sections_dropdown = gr.Dropdown(
                    sections_to_select,
                    value=sections_to_select,
                    interactive=True,
                    multiselect=True,
                    label="Sections for Summarization",
                    info="Select the sections you want to include in the summarization.",
                )
        with gr.Row().style(equal_height=True):
            with gr.Column(scale=1):
                summary_show_pdf = gr.Gallery(label="Uploaded PDF").style(
                    object_fit="contain"
                )
            with gr.Column(scale=1):
                summary_output = gr.Textbox(label="Summary", lines=9).style(
                    show_copy_button=True
                )

        with gr.Row().style(equal_height=True):
            with gr.Column(scale=1):
                recipient_email_summary = gr.Textbox(
                    label="Recipient Email", placeholder="Enter Email"
                )
                subject_email_summary = gr.Textbox(
                    label="Subject", placeholder="Enter Subject"
                )
                send_email_button = gr.Button(
                    "TEST email", link="mailto:", interactive=False, variant="primary"
                )
            with gr.Column(scale=3):
                email_instructions_summary = create_email_textbox()

        # Once a file is uploaded, enable the summarization buttons and visualize the uploaded file
        file_upload_summary.upload(
            switch_buttons,
            [gr.State(True)],
            [summary_parallel_button, gr.State(None), gr.State(None)],
            queue=False,
        ).then(fn=render_file, inputs=[file_upload_summary], outputs=[summary_show_pdf])

        summary_parallel_button.click(
            switch_buttons,
            [gr.State(False)],
            [summary_parallel_button, gr.State(None), gr.State(None)],
            queue=False,
        ).then(
            parallel_summarization,
            [file_upload_summary, summary_sections_dropdown, gr.State([llm])],
            [summary_output],
            queue=False,
        ).then(
            fn=update_email_instructions_summary,
            inputs=[
                email_instructions_summary,
                summary_output,
            ],
            outputs=[email_instructions_summary],
            queue=False,
        ).then(
            fn=update_email_url,
            inputs=[
                recipient_email_summary,
                subject_email_summary,
                email_instructions_summary,
            ],
            outputs=[send_email_button],
            queue=False,
        ).then(
            switch_buttons,
            [gr.State(True)],
            [summary_parallel_button, send_email_button, gr.State(None)],
            queue=False,
        ).then

        # The clear button clears the dashboard
        clear_button.click(lambda: None, None, summary_output, queue=False).then(
            lambda: None, None, file_upload_summary, queue=False
        ).then(lambda: None, None, summary_show_pdf, queue=False).then(
            lambda: None, None, send_email_button, queue=False
        ).then(
            lambda: None, None, email_instructions_summary, queue=False
        ).then(
            lambda: None, None, recipient_email_summary, queue=False
        )

        recipient_email_summary.input(
            fn=update_email_url,
            inputs=[
                recipient_email_summary,
                subject_email_summary,
                email_instructions_summary,
            ],
            outputs=[send_email_button],
            queue=False,
        )
        subject_email_summary.input(
            fn=update_email_url,
            inputs=[
                recipient_email_summary,
                subject_email_summary,
                email_instructions_summary,
            ],
            outputs=[send_email_button],
            queue=False,
        )

    return summary_section


def load_legal_implications_section(llm: ChatOpenAI, summarization_kwargs: dict = {}):
    """Load the legal implications section

    Args:
        llm (ChatOpenAI): Language model.
        summarization_kwargs (dict, optional): Keyword arguments for the summarization. Defaults to {}.

    Returns:
        gr.Block: Legal Implications Section
    """
    with gr.Blocks(theme="soft") as legal_implications_section:
        with gr.Row().style(equal_height=True):
            with gr.Column(scale=3):
                file_upload_legal_implications = gr.File(
                    file_count="multiple",
                    file_types=[".pdf", ".txt"],
                    label="Upload PDF",
                )
            with gr.Column(scale=1):
                extract_legal_implications_button = gr.Button(
                    "Extract Legal Implications", interactive=False
                )
                clear_legal_implications_button = gr.Button("Clear All")

        with gr.Row().style(equal_height=True):
            with gr.Column(scale=1):
                legal_implications_show_pdf = gr.Gallery(label="Uploaded PDF").style(
                    object_fit="contain"
                )
            with gr.Column(scale=1):
                legal_implications_output = gr.Textbox(
                    label="Legal Implications", lines=9
                ).style(show_copy_button=True)

        with gr.Row().style(equal_height=True):
            with gr.Column(scale=1):
                recipiant_email_legal_implications = gr.Textbox(
                    label="Recipient Email", placeholder="Enter Email"
                )
                subject_email_legal_implications = gr.Textbox(
                    label="Subject", placeholder="Enter Subject"
                )
                send_email_button_legal_implications = gr.Button(
                    "Open Email", link="mailto:", interactive=False
                )
            with gr.Column(scale=3):
                email_instructions_legal_implications = create_email_textbox()

        # Once a file is uploaded, enable the summarization buttons and visualize the uploaded file
        file_upload_legal_implications.upload(
            switch_buttons,
            [gr.State(True)],
            [extract_legal_implications_button, gr.State(None), gr.State(None)],
            queue=False,
        ).then(
            fn=render_files,
            inputs=[file_upload_legal_implications],
            outputs=[legal_implications_show_pdf],
        )

        extract_legal_implications_button.click(
            switch_buttons,
            [gr.State(False)],
            [extract_legal_implications_button, gr.State(None), gr.State(None)],
            queue=False,
        ).then(
            parallel_legal_implications,
            [
                file_upload_legal_implications,
                gr.State([llm]),
                gr.State(summarization_kwargs),
            ],
            [legal_implications_output],
            queue=False,
        ).then(
            switch_buttons,
            [gr.State(True)],
            [
                extract_legal_implications_button,
                send_email_button_legal_implications,
                gr.State(None),
            ],
            queue=False,
        )
        # The clear button clears the dashboard
        clear_legal_implications_button.click(
            lambda: None, None, legal_implications_output, queue=False
        ).then(lambda: None, None, file_upload_legal_implications, queue=False).then(
            lambda: None, None, legal_implications_show_pdf, queue=False
        ).then(
            lambda: None, None, send_email_button_legal_implications, queue=False
        ).then(
            lambda: None, None, email_instructions_legal_implications, queue=False
        ).then(
            lambda: None, None, recipiant_email_legal_implications, queue=False
        )


def run_summarization_model_gradio(
    llm: ChatOpenAI,
    share_gradio_via_link: bool = False,
    summarization_kwargs: dict = {},
    run_local: bool = False,
):
    """Run the Summarization assistant with gradio

    Args:
        llm (ChatOpenAI): Language model.
        share_gradio_via_link (bool, optional): Whether to launch the gradio app via a public link. Defaults to False.
        summarization_kwargs (dict, optional): Keyword arguments for the summarization. Defaults to {}.
        run_local (bool, optional): Whether to run the gradio app locally. Defaults to False.

    """
    title = "Legal Document Assistant"
    description = (
        f"This is a proof of concept. Don't upload any confidential documents."
    )

    with gr.Blocks(
        theme="soft",
        title=title,
    ) as webui:
        with gr.Row().style(equal_height=True):
            Header_box = generate_title(title=title, description=description)
        with gr.Tab("Summarize Verdict"):
            load_summary_section(llm=llm)
        with gr.Tab("Legal Implications"):
            load_legal_implications_section(
                llm=llm, summarization_kwargs=summarization_kwargs
            )

    webui.queue()

    if run_local:
        webui.launch(share=share_gradio_via_link)
    else:
        basic_auth = (
            "lenox",
            os.environ["PROTOTYPE_PASSWORD"],
        )
        webui.launch(auth=basic_auth)


def generate_title(title: str, description: str):
    """THIS IS AN EXAMPLE HTML CODE FOR A TITLE"""
    return gr.HTML(
        (
            """
            <html>
            <head>
            <style>
            .awesome-title {
                background: linear-gradient(27deg, rgba(252,92,125,1) 19%, rgba(79,70,229,1) 45%);
                color: #FFFFFF;
                text-align: center;
                padding: 20px;
                border-radius: 10px;
                height: 10%;
                box-shadow: 0px 0px 20px 0px rgba(0,0,0,0.1);
            }
            
            .awesome-title h1 {
                font-family: 'Arial', sans-serif;
                font-size: 32px;
                margin: 0;
            }
            
            .awesome-title p {
                font-family: 'Arial', sans-serif;
                font-size: 16px;
                margin: 10px 0 0;
            }
            </style>
            </head>
            <body>
            """
            + f"""
            <div class="awesome-title">
                <img src="https://aeiljuispo.cloudimg.io/v7/https://cdn-uploads.huggingface.co/production/uploads/64e3560ad5e75575cd641c9c/WpgaNwGDfaTMjsimt8i87.png" alt="Lenox AI" style="height: 80%; position:absolute;">
                <h1>{title}</h1>
                <p>{description}</p>
            </div>
            </body>
            </html>
            """
        ).strip(),
    )