Update agents/software_engineer_agent.py

agents/software_engineer_agent.py

@@ -19,10 +19,8 @@ def run(state: dict) -> dict:
     output_ids = model.generate(input_ids, max_new_tokens=3000)
     output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
-    ai_message = {"role": "assistant", "content": output}
-
     return {
-        "messages": [ai_message],
+        "messages": [AIMessage(content=output)],
         "chat_log": state["chat_log"] + [{"role": "Software Engineer", "content": output}],
         "dev_output": output,
     }