Update agents/software_architect_agent.py

agents/software_architect_agent.py

@@ -10,12 +10,17 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 
-def run(state):
-    """Designs system architecture from project plan."""
-    prompt = state["input"]
+def run(state: dict) -> dict:
+    """Software Architect designs overall system architecture"""
+    messages = state["messages"]
+    prompt = messages[-1]["content"]
+
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
     output_ids = model.generate(input_ids, max_new_tokens=512)
     output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
-    new_log = state["chat_log"] + [{"role": "Software Architect", "content": output}]
-    return {"arch_output": output, "chat_log": new_log}
+    return {
+        "messages": [{"role": "Software Architect", "content": output}],
+        "chat_log": state["chat_log"] + [{"role": "Software Architect", "content": output}],
+        "arch_output": output,
+    }