Update agents/software_architect_agent.py

agents/software_architect_agent.py

@@ -10,7 +10,11 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 
-def run(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=512)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+def run(state):
+    prompt = state["input"]
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+    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 {"pm_output": output, "chat_log": new_log}