Upload 11 files

Dockerfile

@@ -0,0 +1,10 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+
+COPY requirements.txt ./
+COPY . /app/
+
+RUN pip install -r requirements.txt
+
+CMD ["streamlit", "run", "app.py", "--server.port", "8080"]

app.py

@@ -0,0 +1,272 @@
+import streamlit as st
+import json
+from typing import Iterable
+from moa.agent import MOAgent
+from moa.agent.moa import ResponseChunk
+from streamlit_ace import st_ace
+import copy
+
+# Default configuration
+default_config = {
+    "main_model": "llama3-70b-8192",
+    "cycles": 3,
+    "layer_agent_config": {}
+}
+
+layer_agent_config_def = {
+    "layer_agent_1": {
+        "system_prompt": "Think through your response step by step. {helper_response}",
+        "model_name": "llama3-8b-8192"
+    },
+    "layer_agent_2": {
+        "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
+    },
+}
+
+
+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['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)
+            
+            # Clear layer outputs for the next iteration
+            layer_outputs = {}
+            
+            # Yield the main agent's output
+            yield message['delta']
+
+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
+    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,
+            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
+
+st.set_page_config(
+    page_title="Mixture-Of-Agents Powered by Groq",
+    page_icon='static/favicon.ico',
+        menu_items={
+        'About': "## Groq Mixture-Of-Agents \n Powered by [Groq](https://groq.com)"
+    },
+    layout="wide"
+)
+valid_model_names = [
+    'llama3-70b-8192',
+    'llama3-8b-8192',
+    'gemma-7b-it',
+    'gemma2-9b-it',
+    '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
+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"):
+            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 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",
+            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
+        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',
+            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 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", anchor=False)
+st.write("A demo of the Mixture of Agents architecture proposed by Together AI, 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(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)",
+        show_gutter=False,
+        wrap=True,
+        readonly=True,
+        auto_update=True
+    )
+
+# Chat interface
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        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})

moa/agent/__init__.py

@@ -0,0 +1 @@
+from .moa import MOAgent

moa/agent/moa.py

@@ -0,0 +1,181 @@
+"""
+Langchain agent
+"""
+from typing import Generator, Dict, Optional, Literal, TypedDict, List
+from dotenv import load_dotenv
+
+from langchain_groq import ChatGroq
+from langchain.memory import ConversationBufferMemory
+from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.messages import BaseMessage
+from langchain_core.runnables import RunnablePassthrough, RunnableLambda, RunnableSerializable
+from langchain_core.output_parsers import StrOutputParser
+
+from .prompts import SYSTEM_PROMPT, REFERENCE_SYSTEM_PROMPT
+
+load_dotenv()
+valid_model_names = Literal[
+    'llama3-70b-8192',
+    'llama3-8b-8192',
+    'gemma-7b-it',
+    'gemma2-9b-it',
+    'mixtral-8x7b-32768'
+]
+
+class ResponseChunk(TypedDict):
+    delta: str
+    response_type: Literal['intermediate', 'output']
+    metadata: Dict = {}
+
+
+class MOAgent:
+    def __init__(
+        self,
+        main_agent: RunnableSerializable[Dict, str],
+        layer_agent: RunnableSerializable[Dict, Dict],
+        reference_system_prompt: Optional[str] = None,
+        cycles: Optional[int] = None,
+        chat_memory: Optional[ConversationBufferMemory] = None
+    ) -> None:
+        self.reference_system_prompt = reference_system_prompt or REFERENCE_SYSTEM_PROMPT
+        self.main_agent = main_agent
+        self.layer_agent = layer_agent
+        self.cycles = cycles or 1
+        self.chat_memory = chat_memory or ConversationBufferMemory(
+            memory_key="messages",
+            return_messages=True
+        )
+
+    @staticmethod
+    def concat_response(
+        inputs: Dict[str, str],
+        reference_system_prompt: Optional[str] = None
+    ):
+        reference_system_prompt = reference_system_prompt or REFERENCE_SYSTEM_PROMPT
+
+        responses = ""
+        res_list = []
+        for i, out in enumerate(inputs.values()):
+            responses += f"{i}. {out}\n"
+            res_list.append(out)
+
+        formatted_prompt = reference_system_prompt.format(responses=responses)
+        return {
+            'formatted_response': formatted_prompt,
+            'responses': res_list
+        }
+
+    @classmethod
+    def from_config(
+        cls,
+        main_model: Optional[valid_model_names] = 'llama3-70b-8192',
+        system_prompt: Optional[str] = None,
+        cycles: int = 1,
+        layer_agent_config: Optional[Dict] = None,
+        reference_system_prompt: Optional[str] = None,
+        **main_model_kwargs
+    ):
+        reference_system_prompt = reference_system_prompt or REFERENCE_SYSTEM_PROMPT
+        system_prompt = system_prompt or SYSTEM_PROMPT
+        layer_agent = MOAgent._configure_layer_agent(layer_agent_config)
+        main_agent = MOAgent._create_agent_from_system_prompt(
+            system_prompt=system_prompt,
+            model_name=main_model,
+            **main_model_kwargs
+        )
+        return cls(
+            main_agent=main_agent,
+            layer_agent=layer_agent,
+            reference_system_prompt=reference_system_prompt,
+            cycles=cycles
+        )
+
+    @staticmethod
+    def _configure_layer_agent(
+        layer_agent_config: Optional[Dict] = None
+    ) -> RunnableSerializable[Dict, Dict]:
+        if not layer_agent_config:
+            layer_agent_config = {
+                'layer_agent_1' : {'system_prompt': SYSTEM_PROMPT, 'model_name': 'llama3-8b-8192'},
+                'layer_agent_2' : {'system_prompt': SYSTEM_PROMPT, 'model_name': 'gemma-7b-it'},
+                'layer_agent_3' : {'system_prompt': SYSTEM_PROMPT, 'model_name': 'mixtral-8x7b-32768'}
+            }
+
+        parallel_chain_map = dict()
+        for key, value in layer_agent_config.items():
+            chain = MOAgent._create_agent_from_system_prompt(
+                system_prompt=value.pop("system_prompt", SYSTEM_PROMPT), 
+                model_name=value.pop("model_name", 'llama3-8b-8192'),
+                **value
+            )
+            parallel_chain_map[key] = RunnablePassthrough() | chain
+        
+        chain = parallel_chain_map | RunnableLambda(MOAgent.concat_response)
+        return chain
+
+    @staticmethod
+    def _create_agent_from_system_prompt(
+        system_prompt: str = SYSTEM_PROMPT,
+        model_name: str = "llama3-8b-8192",
+        **llm_kwargs
+    ) -> RunnableSerializable[Dict, str]:
+        prompt = ChatPromptTemplate.from_messages([
+            ("system", system_prompt),
+            MessagesPlaceholder(variable_name="messages", optional=True),
+            ("human", "{input}")
+        ])
+
+        assert 'helper_response' in prompt.input_variables
+        llm = ChatGroq(model=model_name, **llm_kwargs)
+        
+        chain = prompt | llm | StrOutputParser()
+        return chain
+
+    def chat(
+        self, 
+        input: str,
+        messages: Optional[List[BaseMessage]] = None,
+        cycles: Optional[int] = None,
+        save: bool = True,
+        output_format: Literal['string', 'json'] = 'string'
+    ) -> Generator[str | ResponseChunk, None, None]:
+        cycles = cycles or self.cycles
+        llm_inp = {
+            'input': input,
+            'messages': messages or self.chat_memory.load_memory_variables({})['messages'],
+            'helper_response': ""
+        }
+        for cyc in range(cycles):
+            layer_output = self.layer_agent.invoke(llm_inp)
+            l_frm_resp = layer_output['formatted_response']
+            l_resps = layer_output['responses']
+            
+            llm_inp = {
+                'input': input,
+                'messages': self.chat_memory.load_memory_variables({})['messages'],
+                'helper_response': l_frm_resp
+            }
+
+            if output_format == 'json':
+                for l_out in l_resps:
+                    yield ResponseChunk(
+                        delta=l_out,
+                        response_type='intermediate',
+                        metadata={'layer': cyc + 1}
+                    )
+
+        stream = self.main_agent.stream(llm_inp)
+        response = ""
+        for chunk in stream:
+            if output_format == 'json':
+                    yield ResponseChunk(
+                        delta=chunk,
+                        response_type='output',
+                        metadata={}
+                    )
+            else:
+                yield chunk
+            response += chunk
+
+        if save:
+            self.chat_memory.save_context({'input': input}, {'output': response})

moa/agent/prompts.py

@@ -0,0 +1,15 @@
+SYSTEM_PROMPT = """\
+You are a personal assistant that is helpful.
+
+{helper_response}\
+"""
+
+REFERENCE_SYSTEM_PROMPT = """\
+You have been provided with a set of responses from various open-source models to the latest user query. 
+Your task is to synthesize these responses into a single, high-quality response. 
+It is crucial to critically evaluate the information provided in these responses, recognizing that some of it may be biased or incorrect. 
+Your response should not simply replicate the given answers but should offer a refined, accurate, and comprehensive reply to the instruction. 
+Ensure your response is well-structured, coherent, and adheres to the highest standards of accuracy and reliability.
+Responses from models:
+{responses}
+"""

moa/main.py

@@ -0,0 +1,20 @@
+from agent import MOAgent
+
+# Configure agent
+layer_agent_config = {
+    'layer_agent_1' : {'system_prompt': "Think through your response with step by step {helper_response}", 'model_name': 'llama3-8b-8192'},
+    'layer_agent_2' : {'system_prompt': "Respond with a thought and then your response to the question {helper_response}", 'model_name': 'gemma-7b-it'},
+    'layer_agent_3' : {'model_name': 'llama3-8b-8192'},
+    'layer_agent_4' : {'model_name': 'gemma-7b-it'},
+    'layer_agent_5' : {'model_name': 'llama3-8b-8192'},
+}
+agent = MOAgent.from_config(
+    main_model='mixtral-8x7b-32768',
+    layer_agent_config=layer_agent_config
+)
+
+while True:
+    inp = input("\nAsk a question: ")
+    stream = agent.chat(inp, output_format='json')
+    for chunk in stream:
+        print(chunk)

requirements.txt

@@ -0,0 +1,6 @@
+langchain>=0.2.5
+langchain_groq>=0.1.5
+streamlit>=1.36.0
+watchdog>=4.0.1
+python-dotenv>=1.0.1
+streamlit-ace>=0.1.1