app.py

import streamlit as st
from main import benchmark_model_multithreaded, benchmark_model_sequential
from prompts import questions as predefined_questions
import requests
import pandas as pd

# Set the title in the browser tab
st.set_page_config(page_title="Aidan Bench - Generator")

st.title("Aidan Bench - Generator")

# API Key Inputs with Security and User Experience Enhancements
st.warning("Please keep your API keys secure and confidential. This app does not store or log your API keys.")

if "open_router_key" not in st.session_state:
    st.session_state.open_router_key = ""
if "openai_api_key" not in st.session_state:
    st.session_state.openai_api_key = ""

open_router_key = st.text_input("Enter your Open Router API Key:", type="password", value=st.session_state.open_router_key)
openai_api_key = st.text_input("Enter your OpenAI API Key:", type="password", value=st.session_state.openai_api_key)

if st.button("Confirm API Keys"):
    if open_router_key and openai_api_key:
        st.session_state.open_router_key = open_router_key
        st.session_state.openai_api_key = openai_api_key
        st.success("API keys confirmed!")
    else:
        st.warning("Please enter both API keys.")

# Access API keys from session state
if st.session_state.open_router_key and st.session_state.openai_api_key:
    # Fetch models from OpenRouter API
    try:
        response = requests.get("https://openrouter.ai/api/v1/models")
        response.raise_for_status()  # Raise an exception for bad status codes
        all_models = response.json()["data"]
        # Sort models alphabetically by their ID
        all_models.sort(key=lambda model: model["id"])

        # --- Create dictionaries for easy model lookup ---
        models_by_id = {model["id"]: model for model in all_models}
        judge_models = [model["id"] for model in all_models if "gpt" in model["id"]]
        judge_models.sort()

        model_names = list(models_by_id.keys())
    except requests.exceptions.RequestException as e:
        st.error(f"Error fetching models from OpenRouter API: {e}")
        model_names = []  # Provide an empty list if API call fails
        judge_models = []

    # Model Selection
    if model_names:
        model_name = st.selectbox("Select a Contestant Model", model_names)
        # --- Display pricing for the selected model ---
        selected_model = models_by_id.get(model_name)
        if selected_model:
            pricing_info = selected_model.get('pricing', {})
            prompt_price = float(pricing_info.get("prompt", 0)) * 1000000
            completion_price = float(pricing_info.get("completion", 0)) * 1000000

            # Display pricing information with increased precision
            st.write(f"**Prompt Pricing:** ${prompt_price:.2f}/Million tokens (if applicable)")
            st.write(f"**Completion Pricing:** ${completion_price:.2f}/Million tokens")
        else:
            st.write("**Pricing:** N/A")
    else:
        st.error("No models available. Please check your API connection.")
        st.stop()

    # Judge Model Selection
    if judge_models:
        judge_model_name = st.selectbox("Select a Judge Model", judge_models)
        # --- Display pricing for the selected judge model ---
        selected_judge_model = models_by_id.get(judge_model_name)
        if selected_judge_model:
            pricing_info = selected_judge_model.get('pricing', {})
            prompt_price = float(pricing_info.get("prompt", 0)) * 1000000
            completion_price = float(pricing_info.get("completion", 0)) * 1000000

            # Display pricing information with increased precision
            st.write(f"**Prompt Pricing:** ${prompt_price:.2f}/Million tokens (if applicable)")
            st.write(f"**Completion Pricing:** ${completion_price:.2f}/Million tokens")
        else:
            st.write("**Pricing:** N/A")
    else:
        st.error("No judge models available. Please check your API connection.")
        st.stop()


    # Initialize session state for user_questions and predefined_questions
    if "user_questions" not in st.session_state:
        st.session_state.user_questions = []

    # Threshold Sliders
    st.sidebar.subheader("Threshold Sliders")
    coherence_threshold = st.sidebar.slider("Coherence Threshold (0-5):", 0, 5, 3)
    novelty_threshold = st.sidebar.slider("Novelty Threshold (0-1):", 0.0, 1.0, 0.1)

    st.sidebar.subheader("Temp Sliders")
    temp_threshold = st.sidebar.slider("Temperature (0-2):", 0.0, 2.0, 1.0)
    top_p = st.sidebar.slider("Top P (0-1):", 0.0, 1.0, 1.0)

    # Workflow Selection
    workflow = st.radio("Select Workflow:", ["Use Predefined Questions", "Use User-Defined Questions"])

    # Handle Predefined Questions
    if workflow == "Use Predefined Questions":
        st.header("Question Selection")
        # Multiselect for predefined questions
        selected_questions = st.multiselect(
            "Select questions to benchmark:",
            predefined_questions,
            predefined_questions  # Select all by default
        )

    # Handle User-Defined Questions
    elif workflow == "Use User-Defined Questions":
        st.header("Question Input")

        # Input for adding a new question
        new_question = st.text_input("Enter a new question:")
        if st.button("Add Question") and new_question:
            new_question = new_question.strip()  # Remove leading/trailing whitespace
            if new_question and new_question not in st.session_state.user_questions:
                st.session_state.user_questions.append(new_question)  # Append to session state
                st.success(f"Question '{new_question}' added successfully.")
            else:
                st.warning("Question already exists or is empty!")

        # Display multiselect with updated user questions
        selected_questions = st.multiselect(
            "Select your custom questions:",
            options=st.session_state.user_questions,
            default=st.session_state.user_questions
        )

    # Display selected questions
    st.write("Selected Questions:", selected_questions)

    # Choose execution mode
    execution_mode = st.radio("Execution Mode:", ["Sequential", "Multithreaded"])

    # If multithreaded, allow user to configure thread pool size
    if execution_mode == "Multithreaded":
        max_threads = st.slider("Maximum Number of Threads:", 1, 10, 4)  # Default to 4 threads
    else:
        max_threads = None  # For sequential mode



    # Benchmark Execution
    if st.button("Start Benchmark"):
        if not selected_questions:
            st.warning("Please select at least one question.")
        else:
            num_questions = len(selected_questions)
            results = []

            # Stop button (not implemented yet)
            stop_button = st.button("Stop Benchmark")

            # Benchmarking logic using the chosen execution mode
            if execution_mode == "Sequential":
                question_results = benchmark_model_sequential(model_name, selected_questions, st.session_state.open_router_key, st.session_state.openai_api_key,judge_model_name,coherence_threshold,novelty_threshold,temp_threshold,top_p)
            else:  # Multithreaded
                question_results = benchmark_model_multithreaded(model_name, selected_questions, st.session_state.open_router_key, st.session_state.openai_api_key, max_threads, judge_model_name, coherence_threshold,novelty_threshold,temp_threshold,top_p)

            results.extend(question_results)

            # Display results in a table
            st.write("Results:")
            results_table = []
            for result in results:
                for answer in result["answers"]:
                    results_table.append({
                        "Question": result["question"],
                        "Answer": answer,
                        "Contestant Model": model_name,
                        "Judge Model": judge_model_name,
                        "Coherence Score": result["coherence_score"],
                        "Novelty Score": result["novelty_score"]
                    })
            st.table(results_table)

            df = pd.DataFrame(results_table)  # Create a Pandas DataFrame from the results
            csv = df.to_csv(index=False).encode('utf-8')  # Convert DataFrame to CSV
            st.download_button(
                label="Export Results as CSV",
                data=csv,
                file_name="benchmark_results.csv",
                mime='text/csv'
                )

            if stop_button:
                st.warning("Partial results displayed due to interruption.")
            else:
                st.success("Benchmark completed!")

else:
    st.warning("Please confirm your API keys first.")