Update coordinator/task_assigner.py

coordinator/task_assigner.py

@@ -8,28 +8,23 @@ model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
 def assign_tasks_with_reasoning(tasks):
     """
-    Use LLM to assign tasks to Frontend/Backend with reasoning
-    Returns dict:
+    Assigns tasks to Frontend/Backend agents with reasoning.
+    Returns a dict:
         {
-            "reasoning": <llm reasoning>,
+            "reasoning": <str>,
             "assignments": {task: agent}
         }
     """
     task_text = "\n".join(f"- {task}" for task in tasks)
     prompt = f"""
-You are an expert project manager. Assign the following technical tasks to either a Frontend or Backend agent.
-Explain your reasoning step by step.
-Output format:
-Reasoning:
-- Step by step explanation
-Assignments:
-Task -> Agent
-
-Tasks:
+You are an AI assistant. Given the following technical tasks:
 {task_text}
+
+Explain your reasoning step-by-step on whether each task should go to Frontend or Backend agent.
+Then provide the final assignments in the format: Task -> Agent.
 """
     inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=300)
+    outputs = model.generate(**inputs, max_new_tokens=250)
     decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
     if "Assignments:" in decoded:
@@ -37,10 +32,10 @@ Tasks:
     else:
         reasoning, assignments_text = decoded, ""
 
+    # Parse assignments
     assignments = {}
     for line in assignments_text.split("\n"):
         if "->" in line:
             task, agent = line.split("->")
             assignments[task.strip()] = agent.strip()
-
     return {"reasoning": reasoning.strip(), "assignments": assignments}