app/vote.py

from .utils import *
from .config import *
from .models import *
from .db import *
from .init import *

import gradio as gr

# Vote

def upvote_model(model, uname, prompt="", response=""):
    print("Establishing database connection for upvoting.")
    conn = get_db()
    cursor = conn.cursor()
    print(f"Updating upvote count for model: {model}")
    cursor.execute('UPDATE model SET upvote = upvote + 1 WHERE name = ?', (model,))
    if cursor.rowcount == 0:
        print(f"No existing entry found for model '{model}'. Inserting new model with upvote=1 and downvote=0.")
        cursor.execute('INSERT OR REPLACE INTO model (name, upvote, downvote) VALUES (?, 1, 0)', (model,))
    print(f"Inserting vote record: username={uname}, model={model}, vote=1, prompt={prompt}, response={response}")
    cursor.execute('INSERT INTO vote (username, model, vote, prompt, response) VALUES (?, ?, ?, ?, ?)', 
                  (uname, model, 1, prompt, response))
    print("Committing upvote transaction.")
    conn.commit()
    print("Closing cursor after upvoting.")
    cursor.close()

def downvote_model(model, uname, prompt="", response=""):
    print("Establishing database connection for downvoting.")
    conn = get_db()
    cursor = conn.cursor()
    print(f"Updating downvote count for model: {model}")
    cursor.execute('UPDATE model SET downvote = downvote + 1 WHERE name = ?', (model,))
    if cursor.rowcount == 0:
        print(f"No existing entry found for model '{model}'. Inserting new model with upvote=0 and downvote=1.")
        cursor.execute('INSERT OR REPLACE INTO model (name, upvote, downvote) VALUES (?, 0, 1)', (model,))
    print(f"Inserting vote record: username={uname}, model={model}, vote=-1, prompt={prompt}, response={response}")
    cursor.execute('INSERT INTO vote (username, model, vote, prompt, response) VALUES (?, ?, ?, ?, ?)', 
                  (uname, model, -1, prompt, response))
    print("Committing downvote transaction.")
    conn.commit()
    print("Closing cursor after downvoting.")
    cursor.close()

# Battle Mode

def a_is_better(model1, model2, userid, prompt="", response1="", response2=""):
    print("Processing vote: A is better.")
    print(f"Comparing models: {model1} vs {model2}")
    userid = mkuuid(userid)
    print(f"Generated UUID for user: {userid}")
    if model1 and model2:
        print("Establishing database connection for voting.")
        conn = get_db()
        cursor = conn.cursor()
        print(f"Inserting votelog: username={userid}, chosen={model1}, rejected={model2}, is_tie=False")
        cursor.execute('INSERT INTO votelog (username, chosen, rejected, prompt, chosen_response, rejected_response, is_tie) VALUES (?, ?, ?, ?, ?, ?, ?)', 
                      (str(userid), model1, model2, prompt, response1, response2, False))
        if scheduler:
            print("Scheduler detected. Acquiring scheduler lock before committing.")
            with scheduler.lock:
                print("Committing votelog transaction with scheduler lock.")
                conn.commit()
        else:
            print("Committing votelog transaction without scheduler lock.")
            conn.commit()
        print("Closing cursor after logging vote.")
        cursor.close()
        print(f"Upvoting model: {model1}")
        upvote_model(model1, str(userid), prompt, response1)
        print(f"Downvoting model: {model2}")
        downvote_model(model2, str(userid), prompt, response2)
    print("Reloading UI after voting.")
    return reload(model1, model2, userid, chose_a=True)

def b_is_better(model1, model2, userid, prompt="", response1="", response2=""):
    print("Processing vote: B is better.")
    print(f"Comparing models: {model1} vs {model2}")
    userid = mkuuid(userid)
    print(f"Generated UUID for user: {userid}")
    if model1 and model2:
        print("Establishing database connection for voting.")
        conn = get_db()
        cursor = conn.cursor()
        print(f"Inserting votelog: username={userid}, chosen={model2}, rejected={model1}, is_tie=False")
        cursor.execute('INSERT INTO votelog (username, chosen, rejected, prompt, chosen_response, rejected_response, is_tie) VALUES (?, ?, ?, ?, ?, ?, ?)', 
                      (str(userid), model2, model1, prompt, response2, response1, False))
        if scheduler:
            print("Scheduler detected. Acquiring scheduler lock before committing.")
            with scheduler.lock:
                print("Committing votelog transaction with scheduler lock.")
                conn.commit()
        else:
            print("Committing votelog transaction without scheduler lock.")
            conn.commit()
        print("Closing cursor after logging vote.")
        cursor.close()
        print(f"Upvoting model: {model2}")
        upvote_model(model2, str(userid), prompt, response2)
        print(f"Downvoting model: {model1}")
        downvote_model(model1, str(userid), prompt, response1)
    print("Reloading UI after voting.")
    return reload(model1, model2, userid, chose_b=True)

def tie_vote(model1, model2, userid, prompt="", response1="", response2=""):
    print("Processing vote: Tie.")
    print(f"Comparing models: {model1} vs {model2}")
    userid = mkuuid(userid)
    print(f"Generated UUID for user: {userid}")
    
    # Log the tie vote
    print("Establishing database connection for tie vote.")
    conn = get_db()
    cursor = conn.cursor()
    print(f"Inserting votelog: username={userid}, chosen={model1}, rejected={model2}, is_tie=True")
    cursor.execute('INSERT INTO votelog (username, chosen, rejected, prompt, chosen_response, rejected_response, is_tie) VALUES (?, ?, ?, ?, ?, ?, ?)', 
                  (str(userid), model1, model2, prompt, response1, response2, True))
    if scheduler:
        with scheduler.lock:
            conn.commit()
    else:
        conn.commit()
    cursor.close()
    
    # Upvote both models
    print(f"Upvoting both models in tie: {model1} and {model2}")
    upvote_model(model1, str(userid), prompt, response1)
    upvote_model(model2, str(userid), prompt, response2)
    
    print("Reloading UI after voting.")
    return reload(model1, model2, userid, is_tie=True)

def prefer_both_vote(model1, model2, userid, prompt="", response1="", response2=""):
    print("Processing vote: Prefer Both.")
    print(f"Comparing models: {model1} vs {model2}")
    userid = mkuuid(userid)
    print(f"Generated UUID for user: {userid}")
    
    # Log the prefer both vote
    print("Establishing database connection for prefer both vote.")
    conn = get_db()
    cursor = conn.cursor()
    print(f"Inserting votelog: username={userid}, chosen={model1}, rejected={model2}, is_tie=True")
    cursor.execute('INSERT INTO votelog (username, chosen, rejected, prompt, chosen_response, rejected_response, is_tie) VALUES (?, ?, ?, ?, ?, ?, ?)', 
                  (str(userid), model1, model2, prompt, response1, response2, True))
    if scheduler:
        with scheduler.lock:
            conn.commit()
    else:
        conn.commit()
    cursor.close()
    
    # Upvote both models
    print(f"Upvoting both models: {model1} and {model2}")
    upvote_model(model1, str(userid), prompt, response1)
    upvote_model(model2, str(userid), prompt, response2)
    
    print("Reloading UI after voting.")
    return reload(model1, model2, userid, is_prefer_both=True)

def prefer_none_vote(model1, model2, userid, prompt="", response1="", response2=""):
    print("Processing vote: Prefer None.")
    print(f"Comparing models: {model1} vs {model2}")
    userid = mkuuid(userid)
    print(f"Generated UUID for user: {userid}")
    
    # Log the prefer none vote
    print("Establishing database connection for prefer none vote.")
    conn = get_db()
    cursor = conn.cursor()
    print(f"Inserting votelog: username={userid}, chosen={model1}, rejected={model2}, is_tie=True")
    cursor.execute('INSERT INTO votelog (username, chosen, rejected, prompt, chosen_response, rejected_response, is_tie) VALUES (?, ?, ?, ?, ?, ?, ?)', 
                  (str(userid), model1, model2, prompt, response1, response2, True))
    if scheduler:
        with scheduler.lock:
            conn.commit()
    else:
        conn.commit()
    cursor.close()
    
    # Downvote both models
    print(f"Downvoting both models: {model1} and {model2}")
    downvote_model(model1, str(userid), prompt, response1)
    downvote_model(model2, str(userid), prompt, response2)
    
    print("Reloading UI after voting.")
    return reload(model1, model2, userid, is_prefer_none=True)

# Reload

def reload(chosenmodel1=None, chosenmodel2=None, userid=None, chose_a=False, chose_b=False, is_prefer_both=False, is_prefer_none=False):
    out = [
        gr.update(interactive=False),  # a_better
        gr.update(interactive=False),  # b_better
        gr.update(interactive=False),  # prefer_both
        gr.update(interactive=False),  # prefer_none
        gr.update(value=f"Selected: {chosenmodel1}" if (chose_a or is_prefer_both) else 
                       f"Rejected: {chosenmodel1}" if is_prefer_none else chosenmodel1, 
                 interactive=False, 
                 visible=True),  # model1_name
        gr.update(value=f"Selected: {chosenmodel2}" if (chose_b or is_prefer_both) else 
                       f"Rejected: {chosenmodel2}" if is_prefer_none else chosenmodel2, 
                 interactive=False, 
                 visible=True)   # model2_name
    ]
    return out