coordinator/task_parser.py

# coordinator/task_parser.py
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

MODEL_NAME = "facebook/opt-125m"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

def parse_brief_with_reasoning(brief: str):
    """
    Use LLM to extract tasks and show reasoning step-by-step
    Returns a dict:
        {
            "reasoning": <llm reasoning>,
            "tasks": [list of tasks]
        }
    """
    prompt = f"""
You are an AI assistant. A user gives you a project brief. 
Explain your reasoning step-by-step about what technical tasks you would extract from this brief.
Then provide the final list of tasks.

Project Brief:
{brief}

Reasoning and Tasks:
"""
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=200)
    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Split reasoning and final tasks
    if "Tasks:" in decoded:
        reasoning, tasks_text = decoded.split("Tasks:", 1)
    else:
        reasoning, tasks_text = decoded, ""

    tasks = [line.strip("- ").strip() for line in tasks_text.split("\n") if line.strip()]
    return {"reasoning": reasoning.strip(), "tasks": tasks}