Update moa/agent/moa.py

moa/agent/moa.py

@@ -17,9 +17,10 @@ load_dotenv()
 valid_model_names = Literal[
     'llama3-70b-8192',
     'llama3-8b-8192',
-    'gemma-7b-it',
+    'llama-3.2-3b-preview',
     'gemma2-9b-it',
-    'mixtral-8x7b-32768'
+    'mixtral-8x7b-32768',
+    'llama-3.2-1b-preview'
 ]
 
 class ResponseChunk(TypedDict):
@@ -97,7 +98,7 @@ class MOAgent:
         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_2' : {'system_prompt': SYSTEM_PROMPT, 'model_name': 'llama-3.2-3b-preview'},
                 'layer_agent_3' : {'system_prompt': SYSTEM_PROMPT, 'model_name': 'mixtral-8x7b-32768'}
             }
 
@@ -105,7 +106,7 @@ class MOAgent:
         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'),
+                model_name=value.pop("model_name", 'llama-3.2-3b-preview'),
                 **value
             )
             parallel_chain_map[key] = RunnablePassthrough() | chain