Fixed aggregator prompt.

.gitignore

@@ -1 +1,2 @@
-env/
+env/
+client_secret.json

app.py

@@ -54,7 +54,7 @@ PROVIDER_TO_AVATAR_MAP = {
     "anthropic://claude-3-haiku-20240307": "data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHdpZHRoPSIxZW0iIGhlaWdodD0iMWVtIiB2aWV3Qm94PSIwIDAgMjQgMjQiPjxwYXRoIGZpbGw9ImN1cnJlbnRDb2xvciIgZD0iTTE3LjMwNCAzLjU0MWgtMy42NzJsNi42OTYgMTYuOTE4SDI0Wm0tMTAuNjA4IDBMMCAyMC40NTloMy43NDRsMS4zNy0zLjU1M2g3LjAwNWwxLjM2OSAzLjU1M2gzLjc0NEwxMC41MzYgMy41NDFabS0uMzcxIDEwLjIyM0w4LjYxNiA3LjgybDIuMjkxIDUuOTQ1WiIvPjwvc3ZnPg==",
 }
 
-AGGREGATORS = ["openai://gpt-4", "openai://gpt-3.5-turbo"]
+AGGREGATORS = ["openai://gpt-4o-mini", "openai://gpt-4o"]
 
 
 def anthropic_streamlit_streamer(stream):
@@ -147,6 +147,10 @@ def get_llm_response(model_identifier, prompt):
         return None
 
 
+def get_response_key(model):
+    return model + ".response"
+
+
 # Main Streamlit App
 def main():
     st.set_page_config(
@@ -226,22 +230,30 @@ def main():
                             stream = google_streamlit_streamer(stream)
                         elif model.startswith("together"):
                             stream = together_streamlit_streamer(stream)
-                        message_placeholder.write_stream(stream)
+                        st.session_state[get_response_key(model)] = (
+                            message_placeholder.write_stream(stream)
+                        )
 
             # Constructing the aggregator prompt
             aggregator_prompt = f"User prompt: {prompt}\n\n"
-            aggregator_prompt += "Responses from other LLMs:\n"
+            aggregator_prompt += "Responses from other LLMs:\n\n"
             aggregator_prompt += "\n".join(
                 [
-                    f"{model}: {st.session_state.get(model, '')}"
+                    f"{model}: {st.session_state.get(get_response_key(model))} \n\n"
                     for model in selected_models
                 ]
             )
             aggregator_prompt += "\n\nPlease provide an aggregated response."
 
+            with st.expander("Aggregator Prompt"):
+                st.write(aggregator_prompt)
+
             # Fetching and streaming response from the aggregator
             st.write(f"Aggregated response from {selected_aggregator}:")
-            with st.chat_message(selected_aggregator):
+            with st.chat_message(
+                selected_aggregator,
+                avatar=PROVIDER_TO_AVATAR_MAP[selected_aggregator],
+            ):
                 message_placeholder = st.empty()
                 aggregator_stream = get_llm_response(
                     selected_aggregator, aggregator_prompt
@@ -257,6 +269,11 @@ if __name__ == "__main__":
     main()
 
 
+# Fix the aggregator step.
+# Add a judging step.
+# Add visualizations.
+
+
 # import streamlit as st
 # from components import llm_council_selector