Update coordinator/task_assigner.py

coordinator/task_assigner.py

@@ -1,13 +1,27 @@
 # coordinator/task_assigner.py
+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):
     """
-    Minimal assignment logic for demonstration
+    Use LLM to assign tasks to Frontend or Backend agents
     """
-    assigned = {}
-    for task in tasks:
-        if "authentication" in task or "task sharing" in task:
-            assigned[task] = "Backend Agent"
-        else:
-            assigned[task] = "Frontend Agent"
-    return assigned
+    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:"
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(**inputs, max_new_tokens=150)
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Parse into dict
+    assignments = {}
+    for line in decoded.split("\n"):
+        if "->" in line:
+            task, agent = line.split("->")
+            assignments[task.strip()] = agent.strip()
+    return assignments