Update coordinator/task_assigner.py

coordinator/task_assigner.py

@@ -2,26 +2,42 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Reuse the same small OPT model
 MODEL_NAME = "facebook/opt-125m"
-
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
-def assign_tasks(tasks):
+def assign_tasks_with_reasoning(tasks):
     """
-    Use LLM to assign tasks to Frontend or Backend agents
+    Use LLM to assign tasks and show reasoning
+    Returns a dict:
+        {
+            "reasoning": <llm reasoning>,
+            "assignments": {task: agent}
+        }
     """
     task_text = "\n".join(f"- {task}" for task in tasks)
-    prompt = f"Assign the following technical tasks to either Frontend or Backend agent:\n{task_text}\n\nAssignments:"
+    prompt = f"""
+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.
+
+Reasoning and Assignments:
+"""
     inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=150)
+    outputs = model.generate(**inputs, max_new_tokens=200)
     decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Parse into dict
+    # Split reasoning and final assignments
+    if "Assignments:" in decoded:
+        reasoning, assignments_text = decoded.split("Assignments:", 1)
+    else:
+        reasoning, assignments_text = decoded, ""
+
     assignments = {}
-    for line in decoded.split("\n"):
+    for line in assignments_text.split("\n"):
         if "->" in line:
             task, agent = line.split("->")
             assignments[task.strip()] = agent.strip()
-    return assignments
+    return {"reasoning": reasoning.strip(), "assignments": assignments}