wehub-resource-sync
253 lines
7.8 KiB
Python
253 lines
7.8 KiB
Python
import os
|
|
import requests
|
|
import streamlit as st
|
|
import sys
|
|
import yaml
|
|
from dotenv import load_dotenv, find_dotenv
|
|
|
|
sys.path.append("../../../aisuite")
|
|
from aisuite.client import Client
|
|
|
|
# Configure Streamlit to use wide mode and hide the top streamlit menu
|
|
st.set_page_config(layout="wide", menu_items={})
|
|
# Add heading with padding
|
|
st.markdown(
|
|
"<div style='padding-top: 1rem;'><h2 style='text-align: center; color: #ffffff;'>Chat & Compare LLM responses</h2></div>",
|
|
unsafe_allow_html=True,
|
|
)
|
|
st.markdown(
|
|
"""
|
|
<style>
|
|
/* Apply default font size globally */
|
|
html, body, [class*="css"] {
|
|
font-size: 14px !important;
|
|
}
|
|
|
|
/* Style for Reset button focus */
|
|
button[data-testid="stButton"][aria-label="Reset Chat"]:focus {
|
|
border-color: red !important;
|
|
box-shadow: 0 0 0 2px red !important;
|
|
}
|
|
</style>
|
|
""",
|
|
unsafe_allow_html=True,
|
|
)
|
|
st.markdown(
|
|
"""
|
|
<style>
|
|
/* Hide Streamlit's default top bar */
|
|
#MainMenu {visibility: hidden;}
|
|
header {visibility: hidden;}
|
|
footer {visibility: hidden;}
|
|
|
|
/* Remove top padding/margin */
|
|
.block-container {
|
|
padding-top: 0rem;
|
|
padding-bottom: 0rem;
|
|
margin-top: 0rem;
|
|
}
|
|
|
|
/* Remove padding from the app container */
|
|
.appview-container {
|
|
padding-top: 0rem;
|
|
}
|
|
|
|
/* Custom CSS for scrollable chat container */
|
|
.chat-container {
|
|
height: 650px;
|
|
overflow-y: auto !important;
|
|
background-color: #1E1E1E;
|
|
border: 1px solid #333;
|
|
border-radius: 10px;
|
|
padding: 20px;
|
|
margin: 10px 0;
|
|
}
|
|
|
|
/* Ensure the container takes full width */
|
|
.stMarkdown {
|
|
width: 100%;
|
|
}
|
|
|
|
/* Style for chat messages to ensure they're visible */
|
|
.chat-message {
|
|
margin: 10px 0;
|
|
padding: 10px;
|
|
}
|
|
|
|
#text_area_1 {
|
|
min-height: 20px !important;
|
|
}
|
|
</style>
|
|
""",
|
|
unsafe_allow_html=True,
|
|
)
|
|
|
|
# Load configuration and initialize aisuite client
|
|
with open("config.yaml", "r") as file:
|
|
config = yaml.safe_load(file)
|
|
configured_llms = config["llms"]
|
|
load_dotenv(find_dotenv())
|
|
client = Client()
|
|
|
|
|
|
# Function to display chat history
|
|
def display_chat_history(chat_history, model_name):
|
|
for message in chat_history:
|
|
role_display = "User" if message["role"] == "user" else model_name
|
|
role = "user" if message["role"] == "user" else "assistant"
|
|
if role == "user":
|
|
with st.chat_message(role, avatar="👤"):
|
|
st.write(message["content"])
|
|
else:
|
|
with st.chat_message(role, avatar="🤖"):
|
|
st.write(message["content"])
|
|
|
|
|
|
# Helper function to query each LLM
|
|
def query_llm(model_config, chat_history):
|
|
print(f"Querying {model_config['name']} with {chat_history}")
|
|
try:
|
|
model = model_config["provider"] + ":" + model_config["model"]
|
|
response = client.chat.completions.create(model=model, messages=chat_history)
|
|
print(
|
|
f"Response from {model_config['name']}: {response.choices[0].message.content}"
|
|
)
|
|
return response.choices[0].message.content
|
|
except Exception as e:
|
|
st.error(f"Error querying {model_config['name']}: {e}")
|
|
return "Error with LLM response."
|
|
|
|
|
|
# Initialize session states
|
|
if "chat_history_1" not in st.session_state:
|
|
st.session_state.chat_history_1 = []
|
|
if "chat_history_2" not in st.session_state:
|
|
st.session_state.chat_history_2 = []
|
|
if "is_processing" not in st.session_state:
|
|
st.session_state.is_processing = False
|
|
if "use_comparison_mode" not in st.session_state:
|
|
st.session_state.use_comparison_mode = False
|
|
|
|
# Top Section - Controls
|
|
col1, col2 = st.columns([1, 2])
|
|
with col1:
|
|
st.session_state.use_comparison_mode = st.checkbox("Comparison Mode", value=True)
|
|
|
|
# Move LLM selection below comparison mode checkbox - now in columns
|
|
llm_col1, llm_col2 = st.columns(2)
|
|
with llm_col1:
|
|
selected_model_1 = st.selectbox(
|
|
"Choose LLM Model 1",
|
|
[llm["name"] for llm in configured_llms],
|
|
key="model_1",
|
|
index=0 if configured_llms else 0,
|
|
)
|
|
with llm_col2:
|
|
if st.session_state.use_comparison_mode:
|
|
selected_model_2 = st.selectbox(
|
|
"Choose LLM Model 2",
|
|
[llm["name"] for llm in configured_llms],
|
|
key="model_2",
|
|
index=1 if len(configured_llms) > 1 else 0,
|
|
)
|
|
|
|
# Display Chat Histories first, always
|
|
# Middle Section - Display Chat Histories
|
|
if st.session_state.use_comparison_mode:
|
|
col1, col2 = st.columns(2)
|
|
with col1:
|
|
chat_container = st.container(height=500)
|
|
with chat_container:
|
|
display_chat_history(st.session_state.chat_history_1, selected_model_1)
|
|
with col2:
|
|
chat_container = st.container(height=500)
|
|
with chat_container:
|
|
display_chat_history(st.session_state.chat_history_2, selected_model_2)
|
|
else:
|
|
chat_container = st.container(height=500)
|
|
with chat_container:
|
|
display_chat_history(st.session_state.chat_history_1, selected_model_1)
|
|
|
|
# Bottom Section - User Input
|
|
st.markdown("<div style='height: 20px;'></div>", unsafe_allow_html=True)
|
|
|
|
col1, col2, col3 = st.columns([6, 1, 1])
|
|
with col1:
|
|
user_query = st.text_area(
|
|
label="Enter your query",
|
|
label_visibility="collapsed",
|
|
placeholder="Enter your query...",
|
|
key="query_input",
|
|
height=70,
|
|
)
|
|
|
|
|
|
# CSS for aligning buttons with the bottom of the text area
|
|
st.markdown(
|
|
"""
|
|
<style>
|
|
/* Adjust the container of the buttons to align at the bottom */
|
|
.stButton > button {
|
|
margin-top: 35px !important; /* Adjust the margin to align */
|
|
}
|
|
|
|
/* Align buttons and "Processing..." text to the bottom of the text area */
|
|
.button-container {
|
|
margin-top: 42px !important;
|
|
text-align: center; /* Center-aligns "Processing..." */
|
|
}
|
|
</style>
|
|
""",
|
|
unsafe_allow_html=True,
|
|
)
|
|
|
|
with col2:
|
|
send_button = False # Initialize send_button
|
|
if st.session_state.is_processing:
|
|
st.markdown(
|
|
"<div class='button-container'>Processing... ⏳</div>",
|
|
unsafe_allow_html=True,
|
|
)
|
|
else:
|
|
send_button = st.button("Send Query", use_container_width=True)
|
|
|
|
with col3:
|
|
if st.button("Reset Chat", use_container_width=True):
|
|
st.session_state.chat_history_1 = []
|
|
st.session_state.chat_history_2 = []
|
|
st.rerun()
|
|
|
|
# Handle send button click and processing
|
|
if send_button and user_query and not st.session_state.is_processing:
|
|
# Set processing state
|
|
st.session_state.is_processing = True
|
|
|
|
# Append user's message to chat histories first
|
|
st.session_state.chat_history_1.append({"role": "user", "content": user_query})
|
|
if st.session_state.use_comparison_mode:
|
|
st.session_state.chat_history_2.append({"role": "user", "content": user_query})
|
|
|
|
st.rerun()
|
|
|
|
# Handle the actual processing
|
|
if st.session_state.is_processing and user_query:
|
|
# Query the selected LLM(s)
|
|
model_config_1 = next(
|
|
llm for llm in configured_llms if llm["name"] == selected_model_1
|
|
)
|
|
response_1 = query_llm(model_config_1, st.session_state.chat_history_1)
|
|
st.session_state.chat_history_1.append({"role": "assistant", "content": response_1})
|
|
|
|
if st.session_state.use_comparison_mode:
|
|
model_config_2 = next(
|
|
llm for llm in configured_llms if llm["name"] == selected_model_2
|
|
)
|
|
response_2 = query_llm(model_config_2, st.session_state.chat_history_2)
|
|
st.session_state.chat_history_2.append(
|
|
{"role": "assistant", "content": response_2}
|
|
)
|
|
|
|
# Reset processing state
|
|
st.session_state.is_processing = False
|
|
st.rerun()
|