Update app.py

app.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
@@ -29,9 +27,11 @@ layer_agent_config_def = {
         "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
+# Recommended Configuration
+
 rec_config = {
     "main_model": "llama3-70b-8192",
     "cycles": 2,
@@ -61,58 +61,30 @@ 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]]):
-    """
-    Handles both synchronous and asynchronous message streams.
-    Processes each response chunk and manages intermediate outputs.
-    :param messages: Union[Iterable, AsyncIterable]
-    """
-    layer_outputs = {}
-
-    async def process_message(message):
-        """
-        Process an individual message and return the content for streaming.
-        Returns an empty string if the message type is intermediate or invalid.
-        """
-        # Debug: Log the incoming message
-        st.write(f"Processing message: {message}")
 
-        if message.get('response_type') == 'intermediate':
-            layer = message['metadata'].get('layer', "Unknown")
+def stream_response(messages: Iterable[ResponseChunk]):
+    layer_outputs = {}
+    for message in messages:
+        if message['response_type'] == 'intermediate':
+            layer = message['metadata']['layer']
             if layer not in layer_outputs:
                 layer_outputs[layer] = []
-            layer_outputs[layer].append(message.get('delta', ""))
-            return ""  # Intermediate messages don't contribute to the final response
-        elif message.get('response_type') == 'final':
-            # Final message processing
+            layer_outputs[layer].append(message['delta'])
+        else:
+            # 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()
-            return message.get('delta', "")  # Ensure no None value is returned
-        else:
-            # Debug: Log unexpected message types
-            st.write(f"Unexpected message type: {message}")
-            return ""  # Default to empty string for invalid types
-
-    # Check if the input is an async or sync iterable
-    if hasattr(messages, "__aiter__"):  # Asynchronous iterable
-        async for message in messages:
-            resolved_message = await process_message(message)
-            yield resolved_message or ""  # Yield empty string if None
-    elif hasattr(messages, "__iter__"):  # Synchronous iterable
-        for message in messages:
-            resolved_message = process_message(message)  # Do not `await` sync messages
-            yield resolved_message or ""  # Yield empty string if None
-    else:
-        raise TypeError("'messages' must be an Iterable or AsyncIterable.")
-
+                        st.expander(label=f"Agent {i+1}", expanded=False).write(output)
+            
+            # Clear layer outputs for the next iteration
+            layer_outputs = {}
+            
+            # Yield the main agent's output
+            yield message['delta']
 
-# Set up the MOAgent
 def set_moa_agent(
     main_model: str = default_config['main_model'],
     cycles: int = default_config['cycles'],
@@ -122,18 +94,26 @@ def set_moa_agent(
 ):
     if override or ("main_model" not in st.session_state):
         st.session_state.main_model = main_model
+    else:
+        if "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
+    else:
+        if "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
+    else:
+        if "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
+    else:
+        if "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,
@@ -145,10 +125,9 @@ 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={
+        menu_items={
         'About': "## Groq Mixture-Of-Agents \n Powered by [Groq](https://groq.com)"
     },
     layout="wide"
@@ -161,9 +140,6 @@ 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 = []
@@ -172,6 +148,7 @@ set_moa_agent()
 
 # Sidebar for configuration
 with st.sidebar:
+    # config_form = st.form("Agent Configuration", border=False)
     st.title("MOA Configuration")
     with st.form("Agent Configuration", border=False):
         if st.form_submit_button("Use Recommended Config"):
@@ -184,9 +161,10 @@ with st.sidebar:
                 )
                 st.session_state.messages = []
                 st.success("Configuration updated successfully!")
+            except json.JSONDecodeError:
+                st.error("Invalid JSON in Layer Agent Configuration. Please check your input.")
             except Exception as e:
                 st.error(f"Error updating configuration: {str(e)}")
-
         # Main model selection
         new_main_model = st.selectbox(
             "Select Main Model",
@@ -212,6 +190,8 @@ with st.sidebar:
         )
 
         # Layer agent configuration
+        tooltip = "Agents in the layer agent configuration run in parallel _per cycle_. Each layer agent supports all initialization parameters of [Langchain's ChatGroq](https://api.python.langchain.com/en/latest/chat_models/langchain_groq.chat_models.ChatGroq.html) class as valid dictionary fields."
+        st.markdown("Layer Agent Config", help=tooltip)
         new_layer_agent_config = st_ace(
             value=json.dumps(st.session_state.layer_agent_config, indent=2),
             language='json',
@@ -233,20 +213,31 @@ with st.sidebar:
                 )
                 st.session_state.messages = []
                 st.success("Configuration updated successfully!")
+            except json.JSONDecodeError:
+                st.error("Invalid JSON in Layer Agent Configuration. Please check your input.")
             except Exception as e:
                 st.error(f"Error updating configuration: {str(e)}")
 
+    st.markdown("---")
+    st.markdown("""
+    ### Credits
+    - MOA: [Together AI](https://www.together.ai/blog/together-moa)
+    - LLMs: [Groq](https://groq.com/)
+    - Paper: [arXiv:2406.04692](https://arxiv.org/abs/2406.04692)
+    """)
+
 # Main app layout
-st.header("Mixture of Agents")
-st.write("This project oversees implementation of Mixture of Agents architecture powered by Groq LLMs.")
+st.header("Karios Agents", anchor=False)
+st.write("A this project oversees implementation of Mixture of Agents architecture Powered by Groq LLMs.")
+# st.image("./static/moa_groq.svg", caption="Mixture of Agents Workflow", width=1000)
 
 # Display current configuration
 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(
+    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(f"**Layer Agents Config**:")
+    new_layer_agent_config = st_ace(
         value=json.dumps(st.session_state.layer_agent_config, indent=2),
         language='json',
         placeholder="Layer Agent Configuration (JSON)",
@@ -262,53 +253,14 @@ for message in st.session_state.messages:
         st.markdown(message["content"])
 
 if query := st.chat_input("Ask a question"):
-    async def handle_query():
-        # Append the user's message to the session state
-        st.session_state.messages.append({"role": "user", "content": query})
-        with st.chat_message("user"):
-            st.write(query)
-
-        # Access the MOAgent instance
-        moa_agent: MOAgent = st.session_state.moa_agent
-
-        with st.chat_message("assistant"):
-            message_placeholder = st.empty()
-
-            # Get messages from MOAgent (supports async streaming)
-            messages = moa_agent.chat(query, output_format='json')
-
-            try:
-                # Stream and display responses from `stream_or_async_response`
-                final_response = ""
-                async for response in stream_or_async_response(messages):
-                    # Skip None or empty responses
-                    if response is not None and response.strip():
-                        final_response += response
-                        message_placeholder.markdown(final_response)
-                    else:
-                        # Debug: Log skipped empty/None responses
-                        st.write(f"Skipped empty or None response: {response}")
-
-                # Save the final response to session state
-                if final_response.strip():  # Only save non-empty responses
-                    st.session_state.messages.append({"role": "assistant", "content": final_response})
-                else:
-                    st.error("Received no valid response from the assistant.")
-            except Exception as e:
-                st.error(f"Error processing response: {e}")
-                # Debug: Log the exception
-                st.write(f"Exception details: {e}")
-
-    # Run the asynchronous handle_query function
-    asyncio.run(handle_query())
-
-
-
-
-
-# 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.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()
+        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})