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.")