Update interim.py

interim.py

@@ -1,8 +1,6 @@
-# ref: https://github.com/kram254/Mixture-of-Agents-running-on-Groq/tree/main
 import streamlit as st
 import json
-import asyncio
-from typing import Union, Iterable, AsyncIterable
+from typing import Iterable
 from moa.agent import MOAgent
 from moa.agent.moa import ResponseChunk
 from streamlit_ace import st_ace
@@ -31,7 +29,7 @@ layer_agent_config_def = {
     },
 }
 
-# Recommended configuration
+# Recommended Configuration
 rec_config = {
     "main_model": "llama3-70b-8192",
     "cycles": 2,
@@ -61,37 +59,41 @@ layer_agent_config_rec = {
     },
 }
 
-# Unified streaming function to handle async and sync responses
-async def stream_or_async_response(messages: Union[Iterable[ResponseChunk], AsyncIterable[ResponseChunk]]):
+def stream_response(messages: Iterable[ResponseChunk]):
     layer_outputs = {}
-
-    async def process_message(message):
+    progress_bar = st.progress(0)
+    total_steps = len(messages)  # Estimate total messages for progress tracking
+    current_step = 0
+
+    for message in messages:
+        current_step += 1
+        progress_bar.progress(current_step / total_steps)
+        
         if message['response_type'] == 'intermediate':
             layer = message['metadata']['layer']
             if layer not in layer_outputs:
                 layer_outputs[layer] = []
             layer_outputs[layer].append(message['delta'])
+            
+            # Real-time rendering for intermediate outputs
+            with st.container():
+                st.markdown(f"**Layer {layer} (In Progress)**")
+                for output in layer_outputs[layer]:
+                    st.markdown(f"- {output}")
+
         else:
+            # Finalize and display accumulated layer outputs
             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()
+                st.markdown(f"### Layer {layer} Final Output")
+                for output in outputs:
+                    st.write(output)
+            layer_outputs = {}  # Reset for next layers
+            
+            # Yield the main agent's output
             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)
+    progress_bar.empty()  # Clear progress bar once done
 
-# Set up the MOAgent
 def set_moa_agent(
     main_model: str = default_config['main_model'],
     cycles: int = default_config['cycles'],
@@ -112,7 +114,6 @@ def set_moa_agent(
         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,
@@ -122,16 +123,13 @@ def set_moa_agent(
         )
 
     del cls_ly_conf
-    del layer_agent_config
 
-# Streamlit app layout
 st.set_page_config(
     page_title="Mixture of Agents",
-    menu_items={
-        'About': "## Groq Mixture-Of-Agents \n Powered by [Groq](https://groq.com)"
-    },
-    layout="wide"
+    layout="wide",
+    menu_items={'About': "## Mixture-of-Agents\nPowered by Groq"}
 )
+
 valid_model_names = [
     'llama3-70b-8192',
     'llama3-8b-8192',
@@ -140,100 +138,79 @@ valid_model_names = [
     'mixtral-8x7b-32768'
 ]
 
-st.markdown("<a href='https://groq.com'><img src='app/static/banner.png' width='500'></a>", unsafe_allow_html=True)
-st.write("---")
-
-# Initialize session state
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
 set_moa_agent()
 
-# Sidebar for configuration
+# Sidebar Configuration
 with st.sidebar:
     st.title("MOA Configuration")
-    with st.form("Agent Configuration", border=False):
+    with st.form("Agent Configuration", clear_on_submit=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
-        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,
-            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',
-            placeholder="Layer Agent Configuration (JSON)",
-            show_gutter=False,
-            wrap=True,
-            auto_update=True
-        )
-
-        if st.form_submit_button("Update Configuration"):
-            try:
-                new_layer_config = json.loads(new_layer_agent_config)
-                set_moa_agent(
-                    main_model=new_main_model,
-                    cycles=new_cycles,
-                    layer_agent_config=new_layer_config,
-                    main_model_temperature=main_temperature,
-                    override=True
-                )
-                st.session_state.messages = []
-                st.success("Configuration updated successfully!")
-            except Exception as e:
-                st.error(f"Error updating configuration: {str(e)}")
+            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!")
+
+        # Config toggling
+        show_advanced = st.checkbox("Show Advanced Configurations")
+        if show_advanced:
+            new_main_model = st.selectbox(
+                "Main Model",
+                valid_model_names,
+                index=valid_model_names.index(st.session_state.main_model)
+            )
+
+            new_cycles = st.number_input(
+                "Number of Layers",
+                min_value=1,
+                max_value=10,
+                value=st.session_state.cycles
+            )
+
+            main_temperature = st.slider(
+                "Main Model Temperature",
+                min_value=0.0,
+                max_value=1.0,
+                value=st.session_state.main_temp,
+                step=0.05
+            )
+
+            new_layer_agent_config = st_ace(
+                value=json.dumps(st.session_state.layer_agent_config, indent=2),
+                language="json",
+                show_gutter=False,
+                wrap=True,
+                auto_update=True
+            )
+            
+            if st.form_submit_button("Update Config"):
+                try:
+                    parsed_config = json.loads(new_layer_agent_config)
+                    set_moa_agent(
+                        main_model=new_main_model,
+                        cycles=new_cycles,
+                        layer_agent_config=parsed_config,
+                        main_model_temperature=main_temperature,
+                        override=True
+                    )
+                    st.session_state.messages = []
+                    st.success("Configuration updated successfully!")
+                except json.JSONDecodeError:
+                    st.error("Invalid JSON in Layer Agent Config.")
+                except Exception as e:
+                    st.error(f"Error updating config: {str(e)}")
 
 # Main app layout
 st.header("Mixture of Agents")
-st.write("This project oversees implementation of Mixture of Agents architecture powered by Groq LLMs.")
-
-# Display current configuration
+st.markdown("Real-time response tracking with intermediate and final results.")
 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',
-        placeholder="Layer Agent Configuration (JSON)",
-        show_gutter=False,
-        wrap=True,
-        readonly=True,
-        auto_update=True
-    )
+    st.json(st.session_state.layer_agent_config)
 
 # Chat interface
 for message in st.session_state.messages:
@@ -241,27 +218,17 @@ for message in st.session_state.messages:
         st.markdown(message["content"])
 
 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})
+    st.session_state.messages.append({"role": "user", "content": query})
+    with st.chat_message("user"):
+        st.markdown(query)
 
-    asyncio.run(handle_query())
+    moa_agent: MOAgent = st.session_state.moa_agent
+    with st.chat_message("assistant"):
+        message_placeholder = st.empty()
+        ast_mess = stream_response(moa_agent.chat(query, output_format="json"))
+        response = st.write_stream(ast_mess)
 
+    st.session_state.messages.append({"role": "assistant", "content": response})
 
-# Add acknowledgment at the bottom
 st.markdown("---")
-st.markdown("""
-###
-This app is based on [Emmanuel M. Ndaliro's work](https://github.com/kram254/Mixture-of-Agents-running-on-Groq/tree/main).
-""")
+st.markdown("Powered by [Groq](https://groq.com).")