app.py

import gradio as gr
import mysql.connector
import os
from langchain_openai import ChatOpenAI
from langchain_core.prompts import (
    ChatPromptTemplate,
    PromptTemplate,
    FewShotPromptTemplate,
)

# Use a pipeline as a high-level helper
from transformers import pipeline

from sentence_transformers import SentenceTransformer, util

# get db info from env vars
db_host = os.environ.get("DB_HOST")
db_user = os.environ.get("DB_USER")
db_pass = os.environ.get("DB_PASS")
db_name = os.environ.get("DB_NAME")

openai_api_key = os.environ.get("OPENAI_API_KEY")

db_connection = mysql.connector.connect(
    host=db_host,
    user=db_user,
    password=db_pass,
    database=db_name,
)

db_cursor = db_connection.cursor()

ORG_ID = 731
AI_PERSON_ID = 11056

potential_labels = []

llm = ChatOpenAI(openai_api_key=openai_api_key, model="gpt-4")

system_prompt = "You are a representative for a local government. A constituent has reached out to you with a question about a local policy. Base your response using the examples below. Be sure to address all points and concerns raised by the constituent. If you do not have enough information to be able to answer the question (you do not see an example that answers the question from the constituent), please make note and another representative will fill in the missing information.\n\n"

examples_prompt = PromptTemplate(
    input_variables=["example"], template="Example:\n\n {example}"
)

ai_response = ""


def get_potential_labels():
    # get potential labels from db
    global potential_labels
    db_connection = mysql.connector.connect(
        host=db_host,
        user=db_user,
        password=db_pass,
        database=db_name,
    )

    db_cursor = db_connection.cursor()
    potential_labels = db_cursor.execute(
        "SELECT message_category_name FROM radmap_frog12.message_categorys"
    )

    potential_labels = db_cursor.fetchall()

    potential_labels = [label[0] for label in potential_labels]

    return potential_labels


potential_labels = get_potential_labels()


# Function to handle the classification
def classify_email_and_generate_response(representative_email, constituent_email):
    potential_labels = get_potential_labels()
    classifier_model = pipeline(
        "zero-shot-classification", model="MoritzLaurer/deberta-v3-large-zeroshot-v1"
    )
    print("classifying email")
    model_out = classifier_model(constituent_email, potential_labels, multi_label=True)
    print("classification complete")
    top_labels = [
        label
        for label, score in zip(model_out["labels"], model_out["scores"])
        if score > 0.95
    ]
    if top_labels == []:
        # Find the index of the highest score
        max_score_index = model_out["scores"].index(max(model_out["scores"]))
        # Return the label with the highest score
        top_labels = [model_out["labels"][max_score_index]]

    labels_with_enough_examples = ["Enforcement", "Financial", "Rules"]
    # see if any of the labels are in labels_with_enough_examples, if so get the messages for that category, else return

    examples = get_similar_messages(constituent_email)

    if representative_email != "":
        current_thread = (
            "Representative message: \n\n"
            + representative_email
            + "\n\nConstituent message: \n\n"
            + constituent_email
        )
    else:
        current_thread = "Constituent message: \n\n" + constituent_email

    prompt = FewShotPromptTemplate(
        examples=examples,
        example_prompt=examples_prompt,
        prefix=system_prompt,
        suffix="Current thread:\n\n {current_thread}\n\nYour response:\n\n",
        input_variables=["current_thread"],
    )

    formatted_prompt = prompt.format(current_thread=current_thread)

    print(formatted_prompt)

    print("Generating GPT4 response")
    import time

    start = time.time()
    ai_out = llm.invoke(formatted_prompt).content

    print("GPT4 response generated in", time.time() - start, "seconds")

    global ai_response
    ai_response = ai_out

    return ", ".join(top_labels), ai_out


def remove_spaces_after_comma(s):
    parts = s.split(",")
    parts = [part.strip() for part in parts]
    return ",".join(parts)


def get_similar_messages(constituent_email):
    db_connection = mysql.connector.connect(
        host=db_host,
        user=db_user,
        password=db_pass,
        database=db_name,
    )
    db_cursor = db_connection.cursor()

    messages_for_category = db_cursor.execute(
        "SELECT id, person_id, body FROM radmap_frog12.messages WHERE person_id <> %s AND id IN (SELECT message_id FROM radmap_frog12.message_category_associations)",
        (AI_PERSON_ID,),
    )

    messages_for_category = db_cursor.fetchall()
    embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

    all_message_chains = []

    for message in messages_for_category:
        # TODO: refactor for when integrated with RADMAP
        # if person_id is set
        if message[1] != 0:
            message_chain = "Representative message: \n\n" + message[2] + "\n\n"
            is_representative_turn = False
        else:
            message_chain = "Constituent message: \n\n" + message[2] + "\n\n"
            is_representative_turn = True
            embedding = embedding_model.encode([message[2]])[0]

        next_message_id = message[0]

        while next_message_id:
            next_message = db_cursor.execute(
                "SELECT id, body FROM radmap_frog12.messages WHERE previous_message_id = %s AND person_id <> %s",
                (next_message_id, AI_PERSON_ID),
            )
            next_message = db_cursor.fetchall()
            if not next_message:
                break
            if is_representative_turn:
                message_chain += (
                    "Representative message: \n\n" + next_message[0][1] + "\n\n"
                )
                is_representative_turn = False
            else:
                message_chain += (
                    "Constituent message: \n\n" + next_message[0][1] + "\n\n"
                )
                is_representative_turn = True

                embedding = embedding_model.encode([next_message[0][1]])[0]

            next_message_id = next_message[0][0]

        all_message_chains.append((message_chain, embedding))

        target_embedding = embedding_model.encode([constituent_email])[0]

        # Compute cosine-similarities and keep the top 3 most similar sentences

        top_messages = []

        for message, embedding in all_message_chains:
            cosine_score = util.pytorch_cos_sim(embedding, target_embedding)
            if cosine_score > 0.98:
                continue
            top_messages.append((message, cosine_score))

        top_messages = sorted(top_messages, key=lambda x: x[1], reverse=True)

    return [{"example": message} for message, score in top_messages[0:3]]


# Function to handle saving data
def save_data(orig_user_email, constituent_email, labels, user_response, current_user):
    # save the data to the database
    # orig_user_email should have volley 0
    # constituent_email should have volley 1
    # user_response should have volley 2
    # app_id, org_id, and person_id should be 0
    # subject should be "Email Classification and Response Tracking"
    # body should be the original email

    db_connection = mysql.connector.connect(
        host=db_host,
        user=db_user,
        password=db_pass,
        database=db_name,
    )

    db_cursor = db_connection.cursor()

    if current_user == "Sheryl Springer":
        person_id = 11021
    elif current_user == "Diane Taylor":
        person_id = 11023
    elif current_user == "Ann E. Belyea":
        person_id = 11025
    elif current_user == "Marcelo Mejia":
        person_id = 11027
    elif current_user == "Rishi Vasudeva":
        person_id = 11029
    else:
        return "You need to select a user to save data"

    try:
        first_message_id = 0
        if orig_user_email != "":
            db_cursor.execute(
                "INSERT INTO radmap_frog12.messages (app_id, org_id, person_id, communication_method_id, status_id, subject, body, send_date, message_type, previous_message_id) VALUES (345678, %s, %s, 1, 1, 'Email Classification and Response Tracking', %s, NOW(), 'Email Classification and Response Tracking', %s)",
                (ORG_ID, person_id, orig_user_email, message_id),
            )

            first_message_id = db_cursor.lastrowid

            db_cursor.execute(
                "INSERT INTO radmap_frog12.messages_log (datetime, message_id, app_id, org_id, person_id, communication_method_id, status_id, subject, body, send_date, message_type, previous_message_id) VALUES (NOW(), %s, 345678, %s, %s, 1, 1, 'Email Classification and Response Tracking', %s, NOW(), 'Email Classification and Response Tracking', %s)",
                (first_message_id, ORG_ID, person_id, orig_user_email, 0),
            )

        db_cursor.execute(
            "INSERT INTO radmap_frog12.messages (app_id, org_id, person_id, communication_method_id, status_id, subject, body, send_date, message_type, previous_message_id) VALUES (345678, %s, 0, 1, 1, 'Email Classification and Response Tracking', %s, NOW(), 'Email Classification and Response Tracking', %s)",
            (ORG_ID, constituent_email, first_message_id),
        )

        second_message_id = db_cursor.lastrowid

        db_cursor.execute(
            "INSERT INTO radmap_frog12.messages_log (datetime, message_id, app_id, org_id, person_id, communication_method_id, status_id, subject, body, send_date, message_type, previous_message_id) VALUES (NOW(), %s, 345678, %s, 0, 1, 1, 'Email Classification and Response Tracking', %s, NOW(), 'Email Classification and Response Tracking', %s)",
            (second_message_id, ORG_ID, constituent_email, first_message_id),
        )

        # add the ai response to the database
        db_cursor.execute(
            "INSERT INTO radmap_frog12.messages (app_id, org_id, person_id, communication_method_id, status_id, subject, body, send_date, message_type, previous_message_id) VALUES (345678, %s, %s, 1, 1, 'Email Classification and Response Tracking', %s, NOW(), 'Email Classification and Response Tracking', %s)",
            (ORG_ID, AI_PERSON_ID, ai_response, second_message_id),
        )

        ai_message_id = db_cursor.lastrowid

        db_cursor.execute(
            "INSERT INTO radmap_frog12.messages_log (datetime, message_id, app_id, org_id, person_id, communication_method_id, status_id, subject, body, send_date, message_type, previous_message_id) VALUES (NOW(), %s, 345678, %s, %s, 1, 1, 'Email Classification and Response Tracking', %s, NOW(), 'Email Classification and Response Tracking', %s)",
            (ai_message_id, ORG_ID, AI_PERSON_ID, ai_response, second_message_id),
        )

        db_cursor.execute(
            "INSERT INTO radmap_frog12.messages (app_id, org_id, person_id, communication_method_id, status_id, subject, body, send_date, message_type, previous_message_id) VALUES (345678, %s, %s, 1, 1, 'Email Classification and Response Tracking', %s, NOW(), 'Email Classification and Response Tracking', %s)",
            (ORG_ID, person_id, user_response, second_message_id),
        )

        third_message_id = db_cursor.lastrowid

        db_cursor.execute(
            "INSERT INTO radmap_frog12.messages_log (datetime, message_id, app_id, org_id, person_id, communication_method_id, status_id, subject, body, send_date, message_type, previous_message_id) VALUES (NOW(), %s, 345678, %s, %s, 1, 1, 'Email Classification and Response Tracking', %s, NOW(), 'Email Classification and Response Tracking', %s)",
            (third_message_id, ORG_ID, person_id, user_response, second_message_id),
        )

        # insert a row into the message_categorys_associations table for each valid label in labels with the message_id of the constituent_email
        # if there is a comma, remove all spaces after the comma

        labels = remove_spaces_after_comma(labels)
        labels = labels.split(",")
        for label in labels:
            label_exists = db_cursor.execute(
                "SELECT * FROM radmap_frog12.message_categorys WHERE message_category_name = %s",
                (label,),
            )
            label_exists = db_cursor.fetchall()
            if label_exists:
                message_id = db_cursor.execute(
                    "SELECT id FROM radmap_frog12.messages WHERE body = %s",
                    (constituent_email,),
                )
                message_id = db_cursor.fetchall()

                db_cursor.execute(
                    "INSERT INTO radmap_frog12.message_category_associations (message_id, message_category_id) VALUES (%s, %s)",
                    (message_id[0][0], label_exists[0][0]),
                )

                message_catergory_association_id = db_cursor.lastrowid

                # save to logs
                db_cursor.execute(
                    "INSERT INTO radmap_frog12.message_category_associations_log (datetime, message_category_associations_id, message_id, message_category_id) VALUES (NOW(), %s, %s, %s)",
                    (
                        message_catergory_association_id,
                        message_id[0][0],
                        label_exists[0][0],
                    ),
                )

        db_connection.commit()

        return "Response successfully saved to database"

    except Exception as e:
        print(e)
        db_connection.rollback()
        return "Error saving data to database"


# read auth from env vars
auth_username = os.environ.get("AUTH_USERNAME")
auth_password = os.environ.get("AUTH_PASSWORD")

# Define your username and password pairs
auth = (auth_username, auth_password)

# Start building the Gradio interface
# Start building the Gradio interface with two columns
with gr.Blocks(theme=gr.themes.Soft()) as app:
    with gr.Row():
        gr.Markdown("## Campaign Messaging Assistant")

    with gr.Row():
        with gr.Column():
            current_user = gr.Dropdown(
                label="Current User",
                choices=[
                    "Sheryl Springer",
                    "Ann E. Belyea",
                    "Marcelo Mejia",
                    "Rishi Vasudeva",
                    "Diane Taylor",
                ],
            )

            email_labels_input = gr.Markdown(
                "## Message Category Library\n ### " + ", ".join(potential_labels),
            )

            original_email_input = gr.TextArea(
                placeholder="Enter the original email sent by you",
                label="Your Original Email (if any)",
            )

            spacer1 = gr.Label(visible=False)

            constituent_response_input = gr.TextArea(
                placeholder="Enter the incoming message",
                label="Incoming Message (may be a response to original email)",
                lines=15,
            )

            classify_button = gr.Button("Process Message", variant="primary")

        with gr.Column():
            classification_output = gr.TextArea(
                label="Message Categories (modify as needed, but only use categories from Library on left). Separate categories with commas",
                lines=1,
                interactive=True,
            )

            spacer2 = gr.Label(visible=False)

            user_response_input = gr.TextArea(
                placeholder="Enter your response to the constituent",
                label="Suggested Response (modify as needed)",
                lines=25,
            )

            save_button = gr.Button("Save Response", variant="primary")
            save_output = gr.Label(label="Backend Response")

    # Define button actions
    classify_button.click(
        fn=classify_email_and_generate_response,
        inputs=[original_email_input, constituent_response_input],
        outputs=[classification_output, user_response_input],
    )

    save_button.click(
        fn=save_data,
        inputs=[
            original_email_input,
            constituent_response_input,
            classification_output,
            user_response_input,
            current_user,
        ],
        outputs=save_output,
    )

# Launch the app
app.launch(auth=auth, share=True)