Update test.py

test.py

@@ -7,6 +7,8 @@ from moa.agent import MOAgent
 from moa.agent.moa import ResponseChunk
 from streamlit_ace import st_ace
 import copy
+import pandas as pd
+import time
 
 # Default configuration
 default_config = {
@@ -24,106 +26,9 @@ layer_agent_config_def = {
         "system_prompt": "Respond with a thought and then your response to the question. {helper_response}",
         "model_name": "gemma-7b-it",
         "temperature": 0.7
-    },
-    "layer_agent_3": {
-        "system_prompt": "You are an expert at logic and reasoning. Always take a logical approach to the answer. {helper_response}",
-        "model_name": "llama3-8b-8192"
-    },
-}
-
-# Recommended configuration
-rec_config = {
-    "main_model": "llama3-70b-8192",
-    "cycles": 2,
-    "layer_agent_config": {}
-}
-
-layer_agent_config_rec = {
-    "layer_agent_1": {
-        "system_prompt": "Think through your response step by step. {helper_response}",
-        "model_name": "llama3-8b-8192",
-        "temperature": 0.1
-    },
-    "layer_agent_2": {
-        "system_prompt": "Respond with a thought and then your response to the question. {helper_response}",
-        "model_name": "llama3-8b-8192",
-        "temperature": 0.2
-    },
-    "layer_agent_3": {
-        "system_prompt": "You are an expert at logic and reasoning. Always take a logical approach to the answer. {helper_response}",
-        "model_name": "llama3-8b-8192",
-        "temperature": 0.4
-    },
-    "layer_agent_4": {
-        "system_prompt": "You are an expert planner agent. Create a plan for how to answer the human's query. {helper_response}",
-        "model_name": "mixtral-8x7b-32768",
-        "temperature": 0.5
-    },
+    }
 }
 
-# Unified streaming function to handle async and sync responses
-async def stream_or_async_response(messages: Union[Iterable[ResponseChunk], AsyncIterable[ResponseChunk]]):
-    layer_outputs = {}
-
-    async def process_message(message):
-        if message['response_type'] == 'intermediate':
-            layer = message['metadata']['layer']
-            if layer not in layer_outputs:
-                layer_outputs[layer] = []
-            layer_outputs[layer].append(message['delta'])
-        else:
-            for layer, outputs in layer_outputs.items():
-                st.write(f"Layer {layer}")
-                cols = st.columns(len(outputs))
-                for i, output in enumerate(outputs):
-                    with cols[i]:
-                        st.expander(label=f"Agent {i+1}", expanded=False).write(output)
-
-            layer_outputs.clear()
-            yield message['delta']
-
-    if isinstance(messages, AsyncIterable):
-        # Process asynchronous messages
-        async for message in messages:
-            await process_message(message)
-    else:
-        # Process synchronous messages
-        for message in messages:
-            await process_message(message)
-
-# Set up the MOAgent
-def set_moa_agent(
-    main_model: str = default_config['main_model'],
-    cycles: int = default_config['cycles'],
-    layer_agent_config: dict[dict[str, any]] = copy.deepcopy(layer_agent_config_def),
-    main_model_temperature: float = 0.1,
-    override: bool = False
-):
-    if override or ("main_model" not in st.session_state):
-        st.session_state.main_model = main_model
-
-    if override or ("cycles" not in st.session_state):
-        st.session_state.cycles = cycles
-
-    if override or ("layer_agent_config" not in st.session_state):
-        st.session_state.layer_agent_config = layer_agent_config
-
-    if override or ("main_temp" not in st.session_state):
-        st.session_state.main_temp = main_model_temperature
-
-    cls_ly_conf = copy.deepcopy(st.session_state.layer_agent_config)
-
-    if override or ("moa_agent" not in st.session_state):
-        st.session_state.moa_agent = MOAgent.from_config(
-            main_model=st.session_state.main_model,
-            cycles=st.session_state.cycles,
-            layer_agent_config=cls_ly_conf,
-            temperature=st.session_state.main_temp
-        )
-
-    del cls_ly_conf
-    del layer_agent_config
-
 # Streamlit app layout
 st.set_page_config(
     page_title="Karios Agents Powered by Groq",
@@ -133,65 +38,114 @@ st.set_page_config(
     },
     layout="wide"
 )
+
 valid_model_names = [
     'llama3-70b-8192',
     'llama3-8b-8192',
-    'gemma-7b-it',
-    'gemma2-9b-it',
-    'mixtral-8x7b-32768'
+    'gemma-7b-it'
 ]
 
-st.markdown("<a href='https://groq.com'><img src='app/static/banner.png' width='500'></a>", unsafe_allow_html=True)
-st.write("---")
+# Caching function
+@st.cache_data
+def cached_chat_response(query, model, config):
+    """Cache responses to minimize redundant processing."""
+    moa_agent = MOAgent.from_config(**config)
+    return moa_agent.chat(query, output_format='json')
+
+# Functions to reset, export, and import configurations
+def reset_session():
+    """Reset session state to clear messages and restart."""
+    st.session_state.messages = []
+    st.experimental_rerun()
+
+def export_config():
+    """Allow the user to download the current configuration as a JSON file."""
+    config_data = {
+        "main_model": st.session_state.main_model,
+        "cycles": st.session_state.cycles,
+        "layer_agent_config": st.session_state.layer_agent_config,
+        "main_temp": st.session_state.main_temp,
+    }
+    st.download_button(
+        label="Download Config",
+        data=json.dumps(config_data, indent=4),
+        file_name="moa_config.json",
+        mime="application/json",
+    )
+
+def import_config(uploaded_file):
+    """Upload and apply a configuration from a JSON file."""
+    try:
+        config_data = json.load(uploaded_file)
+        set_moa_agent(
+            main_model=config_data['main_model'],
+            cycles=config_data['cycles'],
+            layer_agent_config=config_data['layer_agent_config'],
+            main_model_temperature=config_data['main_temp'],
+            override=True
+        )
+        st.success("Configuration imported successfully!")
+        st.experimental_rerun()
+    except Exception as e:
+        st.error(f"Error importing configuration: {str(e)}")
 
 # Initialize session state
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-set_moa_agent()
+if "main_model" not in st.session_state:
+    st.session_state.main_model = default_config['main_model']
 
-# Sidebar for configuration
+if "cycles" not in st.session_state:
+    st.session_state.cycles = default_config['cycles']
+
+if "layer_agent_config" not in st.session_state:
+    st.session_state.layer_agent_config = copy.deepcopy(layer_agent_config_def)
+
+if "main_temp" not in st.session_state:
+    st.session_state.main_temp = 0.1
+
+if "moa_agent" not in st.session_state:
+    st.session_state.moa_agent = MOAgent.from_config(
+        main_model=st.session_state.main_model,
+        cycles=st.session_state.cycles,
+        layer_agent_config=st.session_state.layer_agent_config,
+        temperature=st.session_state.main_temp
+    )
+
+# Sidebar
 with st.sidebar:
     st.title("MOA Configuration")
-    with st.form("Agent Configuration", border=False):
-        if st.form_submit_button("Use Recommended Config"):
-            try:
-                set_moa_agent(
-                    main_model=rec_config['main_model'],
-                    cycles=rec_config['cycles'],
-                    layer_agent_config=layer_agent_config_rec,
-                    override=True
-                )
-                st.session_state.messages = []
-                st.success("Configuration updated successfully!")
-            except Exception as e:
-                st.error(f"Error updating configuration: {str(e)}")
-
-        # Main model selection
+    if st.button("Reset Session"):
+        reset_session()
+    
+    if st.button("Export Configuration"):
+        export_config()
+
+    uploaded_file = st.file_uploader("Upload Config", type=["json"])
+    if uploaded_file:
+        import_config(uploaded_file)
+
+    # Configuration forms
+    with st.form("Agent Configuration"):
         new_main_model = st.selectbox(
             "Select Main Model",
             options=valid_model_names,
             index=valid_model_names.index(st.session_state.main_model)
         )
-
-        # Cycles input
         new_cycles = st.number_input(
             "Number of Layers",
             min_value=1,
             max_value=10,
             value=st.session_state.cycles
         )
-
-        # Main Model Temperature
         main_temperature = st.number_input(
             label="Main Model Temperature",
-            value=0.1,
+            value=st.session_state.main_temp,
             min_value=0.0,
             max_value=1.0,
             step=0.1
         )
-
-        # Layer agent configuration
         new_layer_agent_config = st_ace(
             value=json.dumps(st.session_state.layer_agent_config, indent=2),
             language='json',
@@ -204,19 +158,22 @@ with st.sidebar:
         if st.form_submit_button("Update Configuration"):
             try:
                 new_layer_config = json.loads(new_layer_agent_config)
-                set_moa_agent(
+                st.session_state.main_model = new_main_model
+                st.session_state.cycles = new_cycles
+                st.session_state.main_temp = main_temperature
+                st.session_state.layer_agent_config = new_layer_config
+                st.session_state.moa_agent = MOAgent.from_config(
                     main_model=new_main_model,
                     cycles=new_cycles,
                     layer_agent_config=new_layer_config,
-                    main_model_temperature=main_temperature,
-                    override=True
+                    temperature=main_temperature
                 )
-                st.session_state.messages = []
                 st.success("Configuration updated successfully!")
+                st.experimental_rerun()
             except Exception as e:
                 st.error(f"Error updating configuration: {str(e)}")
 
-# Main app layout
+# Main layout
 st.header("Mixture of Agents")
 st.write("This project oversees implementation of Mixture of Agents architecture powered by Groq LLMs.")
 
@@ -225,7 +182,6 @@ with st.expander("Current MOA Configuration", expanded=False):
     st.markdown(f"**Main Model**: `{st.session_state.main_model}`")
     st.markdown(f"**Main Model Temperature**: `{st.session_state.main_temp:.1f}`")
     st.markdown(f"**Layers**: `{st.session_state.cycles}`")
-    st.markdown("**Layer Agents Config:**")
     st_ace(
         value=json.dumps(st.session_state.layer_agent_config, indent=2),
         language='json',
@@ -236,29 +192,38 @@ with st.expander("Current MOA Configuration", expanded=False):
         auto_update=True
     )
 
-# Chat interface
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
+# Model comparison
+selected_models = st.multiselect(
+    "Select Models for Comparison",
+    valid_model_names,
+    default=[st.session_state.main_model]
+)
 
 if query := st.chat_input("Ask a question"):
-    async def handle_query():
-        st.session_state.messages.append({"role": "user", "content": query})
-        with st.chat_message("user"):
-            st.write(query)
-
-        moa_agent: MOAgent = st.session_state.moa_agent
-
-        with st.chat_message("assistant"):
-            message_placeholder = st.empty()
-            messages = moa_agent.chat(query, output_format='json')
-            async for response in stream_or_async_response(messages):
-                message_placeholder.markdown(response)
-
-        st.session_state.messages.append({"role": "assistant", "content": response})
-
-    asyncio.run(handle_query())
-
+    st.write(f"Query: {query}")
+    results = {}
+    for model in selected_models:
+        try:
+            st.session_state.moa_agent.set_model(model)
+            results[model] = cached_chat_response(query, model, st.session_state.layer_agent_config)
+        except Exception as e:
+            st.error(f"Agent {model} failed: {e}")
+
+    for model, response in results.items():
+        st.subheader(f"Response from {model}")
+        st.write(response)
+
+# Progress bar for layers
+progress_bar = st.progress(0)
+total_layers = st.session_state.cycles
+
+async def process_message(message):
+    try:
+        # Simulate layer processing
+        current_layer = int(message['metadata']['layer'])
+        progress_bar.progress(int((current_layer / total_layers) * 100))
+    except Exception as e:
+        st.error(f"Agent failed: {e}")
 
 # Add acknowledgment at the bottom
 st.markdown("---")