Create dioscuri.py

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()