Update coordinator/task_assigner.py

coordinator/task_assigner.py

@@ -2,14 +2,14 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-MODEL_NAME = "facebook/opt-125m"
+MODEL_NAME = "facebook/opt-1.3b"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
 def assign_tasks_with_reasoning(tasks):
     """
-    Use LLM to assign tasks and show reasoning
-    Returns a dict:
+    Use LLM to assign tasks to Frontend/Backend with reasoning
+    Returns dict:
         {
             "reasoning": <llm reasoning>,
             "assignments": {task: agent}
@@ -17,19 +17,21 @@ def assign_tasks_with_reasoning(tasks):
     """
     task_text = "\n".join(f"- {task}" for task in tasks)
     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.
+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
 
-Reasoning and Assignments:
+Tasks:
+{task_text}
 """
     inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=200)
+    outputs = model.generate(**inputs, max_new_tokens=300)
     decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Split reasoning and final assignments
     if "Assignments:" in decoded:
         reasoning, assignments_text = decoded.split("Assignments:", 1)
     else:
@@ -40,4 +42,5 @@ Reasoning and Assignments:
         if "->" in line:
             task, agent = line.split("->")
             assignments[task.strip()] = agent.strip()
+
     return {"reasoning": reasoning.strip(), "assignments": assignments}