diff --git a/dioscuri.py b/dioscuri.py new file mode 100644 index 0000000..0ef6671 --- /dev/null +++ b/dioscuri.py @@ -0,0 +1,82 @@ +import streamlit as st +from gemini_chat import GeminiChat +from appendages import get_appended_prompt + +def main(): + st.title("Chat Interface") + + # Initialize GeminiChat + gemini_chat = GeminiChat() + + # Chat history + chat_history_container = st.container() + + # Function to display chat messages + def display_chat_history(messages): + with chat_history_container: + for i, (role, message) in enumerate(messages): + st.markdown(f"**{role}:**") + st.write(message) + st.write("---") + + # Initialize or retrieve chat history from session state + if "chat_history" not in st.session_state: + st.session_state["chat_history"] = [] + + messages = st.session_state["chat_history"] + + # State variable for user input + user_input_state = st.session_state.get("user_input_state", "") + + # User input + user_input = st.text_input("Enter your message:", key="user_input", value=user_input_state) + + # Temperature slider + temperature = st.slider("Temperature", 0.0, 1.0, 0.5, 0.01) + + # Model dropdown + model_options = ["gemini-pro"] + selected_model = st.selectbox("Select Model", model_options, index=3) # Default to "gemini" + + # Toggle for appending encoded string (jailbreak) + append_encoded_string = st.checkbox("Toggle Jailbreak", value=False) + + if user_input: + prompt = user_input + if append_encoded_string: + # Append encoded string to the user input + prompt = user_input + " " + get_appended_prompt() + + max_retries = 5 + retry_count = 0 + response = None + + while retry_count < max_retries: + # Send the prompt to GeminiChat + if selected_model == "gemini-pro": + response = gemini_chat.send_message(prompt, temperature=temperature) + else: + response = gemini_chat.send_message(prompt, temperature=temperature, model=selected_model) + + if response is not None: + break + + retry_count += 1 + print(f"Retrying query ({retry_count}/{max_retries})...") + + if response is None: + st.error("Failed to generate a response after multiple retries.") + + # Update chat history + messages.append(("User", user_input)) + messages.append(("Assistant", response)) + st.session_state["chat_history"] = messages + + # Clear user input state + st.session_state["user_input_state"] = "" + + # Display chat history + display_chat_history(messages) + +if __name__ == "__main__": + main()