Update app.py

app.py

@@ -1,13 +1,11 @@
-#ref: https://github.com/kram254/Mixture-of-Agents-running-on-Groq/tree/main
+# ref: https://github.com/kram254/Mixture-of-Agents-running-on-Groq/tree/main
 import streamlit as st
 import json
-import threading
-from typing import Iterable
+import asyncio
+from typing import Union, Iterable, AsyncIterable
 from moa.agent import MOAgent
 from moa.agent.moa import ResponseChunk
 from streamlit_ace import st_ace
-import asyncio
-from typing import Union, Iterable, AsyncIterable
 import copy
 
 # Default configuration
@@ -31,11 +29,9 @@ 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,
@@ -65,8 +61,8 @@ layer_agent_config_rec = {
     },
 }
 
-
-async def async_stream_response(messages: Union[Iterable[ResponseChunk], AsyncIterable[ResponseChunk]]):
+# 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):
@@ -82,7 +78,7 @@ async def async_stream_response(messages: Union[Iterable[ResponseChunk], AsyncIt
                 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']
 
@@ -95,6 +91,7 @@ async def async_stream_response(messages: Union[Iterable[ResponseChunk], AsyncIt
         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'],
@@ -104,26 +101,18 @@ 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,
@@ -135,10 +124,11 @@ def set_moa_agent(
     del cls_ly_conf
     del layer_agent_config
 
+# Streamlit app layout
 st.set_page_config(
     page_title="Karios Agents Powered by Groq",
     page_icon='static/favicon.ico',
-        menu_items={
+    menu_items={
         'About': "## Groq Mixture-Of-Agents \n Powered by [Groq](https://groq.com)"
     },
     layout="wide"
@@ -154,8 +144,6 @@ valid_model_names = [
 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 = []
@@ -164,7 +152,6 @@ 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"):
@@ -177,10 +164,9 @@ 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",
@@ -206,8 +192,6 @@ 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',
@@ -229,31 +213,20 @@ 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("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)
+st.header("Karios Agents")
+st.write("This project oversees implementation of Mixture of Agents architecture powered by Groq LLMs.")
 
 # 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(f"**Layer Agents Config**:")
-    new_layer_agent_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("**Layer Agents Config:**")
+    st_ace(
         value=json.dumps(st.session_state.layer_agent_config, indent=2),
         language='json',
         placeholder="Layer Agent Configuration (JSON)",
@@ -269,14 +242,19 @@ for message in st.session_state.messages:
         st.markdown(message["content"])
 
 if query := st.chat_input("Ask a question"):
-    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})
+    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())