app.py

"""Project Task Planner.

Turns a project idea into a structured task plan (tasks, subtasks, owners,
hour estimates) and lets a small instruction-tuned model, running locally on
CPU inside this Space, add project-specific risks and tasks. No external API.
"""

from __future__ import annotations

import csv
import io
import json
import math
import uuid
from copy import deepcopy
from threading import Thread

import streamlit as st

MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
HOURS_PER_PERSON_WEEK = 30  # focused hours, not calendar hours

# ---------------------------------------------------------------------------
# Templates
# ---------------------------------------------------------------------------

BASE_TASKS = [
    {
        "title": "Project setup",
        "priority": "P0",
        "owner": "Tech Lead",
        "description": "Repository, environments and delivery pipeline.",
        "acceptance": ["Repository with branch rules", "CI runs tests on every push", "Dev and staging environments"],
        "subtasks": [("Repository and branch rules", "DevOps", 4), ("CI pipeline with tests", "DevOps", 12), ("Staging environment", "DevOps", 8)],
    },
    {
        "title": "Requirements and scope",
        "priority": "P0",
        "owner": "Product Owner",
        "description": "Agree on what version 1 must do and what it will not do.",
        "acceptance": ["Written scope with non-goals", "Prioritised backlog", "Sign-off from stakeholders"],
        "subtasks": [("Stakeholder interviews", "Product Owner", 8), ("User stories and backlog", "Product Owner", 10)],
    },
]

CLOSING_TASKS = [
    {
        "title": "Testing and quality",
        "priority": "P1",
        "owner": "QA",
        "description": "Automated and manual checks before release.",
        "acceptance": ["Critical paths covered by tests", "No open P0 bugs"],
        "subtasks": [("Automated test suite", "Developer", 16), ("Manual acceptance test", "QA", 8)],
    },
    {
        "title": "Launch and monitoring",
        "priority": "P1",
        "owner": "Tech Lead",
        "description": "Go live safely and know when something breaks.",
        "acceptance": ["Rollback plan written", "Error alerts reach the team", "Backups tested"],
        "subtasks": [("Monitoring and alerts", "DevOps", 8), ("Backup and restore test", "DevOps", 4), ("Release checklist", "Tech Lead", 4)],
    },
]

TYPE_TASKS = {
    "E-commerce website": [
        {
            "title": "Product catalog",
            "priority": "P0",
            "owner": "Frontend Team",
            "description": "Browse, search and filter products.",
            "acceptance": ["Product listing and detail pages", "Search", "Filters by category and price"],
            "subtasks": [("Product data model", "Backend Dev", 8), ("Listing and detail pages", "Frontend Dev", 16), ("Search and filters", "Frontend Dev", 20)],
        },
        {
            "title": "Cart and checkout",
            "priority": "P0",
            "owner": "Backend Team",
            "description": "From cart to paid order.",
            "acceptance": ["Cart survives page reload", "Payment provider in test mode", "Order confirmation e-mail"],
            "subtasks": [("Cart logic", "Backend Dev", 12), ("Payment integration", "Backend Dev", 20), ("Order e-mails", "Backend Dev", 6)],
        },
    ],
    "Mobile app": [
        {
            "title": "App structure",
            "priority": "P0",
            "owner": "Mobile Team",
            "description": "Navigation, state and API layer.",
            "acceptance": ["Navigation between core screens", "State management", "API client with error handling"],
            "subtasks": [("Navigation", "Mobile Dev", 12), ("State management", "Mobile Dev", 16), ("API client", "Mobile Dev", 12)],
        },
        {
            "title": "Store release",
            "priority": "P1",
            "owner": "Mobile Team",
            "description": "Get through App Store and Play Store review.",
            "acceptance": ["Store listings and screenshots", "Privacy labels filled in", "Builds approved"],
            "subtasks": [("Store assets", "Designer", 8), ("Privacy and review forms", "Product Owner", 6), ("Release builds", "Mobile Dev", 8)],
        },
    ],
    "SaaS platform": [
        {
            "title": "Accounts and tenants",
            "priority": "P0",
            "owner": "Backend Team",
            "description": "Sign-up, login, organisations and roles.",
            "acceptance": ["Sign-up and login", "Data separated per customer", "Admin and member roles"],
            "subtasks": [("Authentication", "Backend Dev", 16), ("Tenant data separation", "Backend Dev", 16), ("Roles and permissions", "Backend Dev", 12)],
        },
        {
            "title": "Billing",
            "priority": "P1",
            "owner": "Backend Team",
            "description": "Plans, subscriptions and invoices.",
            "acceptance": ["Plans and trial", "Recurring payments", "Invoices downloadable"],
            "subtasks": [("Plans and trial logic", "Backend Dev", 10), ("Subscription payments", "Backend Dev", 16), ("Invoices", "Backend Dev", 8)],
        },
    ],
    "AI/ML solution": [
        {
            "title": "Data pipeline",
            "priority": "P0",
            "owner": "Data Engineer",
            "description": "Reliable, versioned training and inference data.",
            "acceptance": ["Data sources documented", "Data quality checks", "Reproducible train/test split"],
            "subtasks": [("Source inventory and access", "Data Engineer", 8), ("ETL with quality checks", "Data Engineer", 20), ("Train/test split and versioning", "Data Engineer", 6)],
        },
        {
            "title": "Model and evaluation",
            "priority": "P0",
            "owner": "ML Engineer",
            "description": "Baseline first, then a model that beats it on the right metric.",
            "acceptance": ["Simple baseline", "Metric agreed with business", "Model beats baseline on held-out data"],
            "subtasks": [("Baseline", "ML Engineer", 6), ("Model training and tuning", "ML Engineer", 24), ("Evaluation report", "ML Engineer", 8)],
        },
        {
            "title": "Serving and monitoring",
            "priority": "P1",
            "owner": "ML Engineer",
            "description": "Model in production with drift checks.",
            "acceptance": ["Prediction API or batch job", "Prediction logging", "Drift alert"],
            "subtasks": [("Serving", "ML Engineer", 12), ("Prediction logging", "Data Engineer", 6), ("Drift monitoring", "ML Engineer", 8)],
        },
    ],
    "Enterprise system": [
        {
            "title": "Integrations",
            "priority": "P0",
            "owner": "Backend Team",
            "description": "Connect existing systems (ERP, CRM, identity).",
            "acceptance": ["Single sign-on", "Data sync with source systems", "Integration error handling"],
            "subtasks": [("Single sign-on", "Backend Dev", 12), ("ERP/CRM connectors", "Backend Dev", 24), ("Retry and error queue", "Backend Dev", 8)],
        },
        {
            "title": "Compliance and audit",
            "priority": "P1",
            "owner": "Tech Lead",
            "description": "Audit trail and data protection.",
            "acceptance": ["Audit log of changes", "Data retention rules", "Access review"],
            "subtasks": [("Audit log", "Backend Dev", 10), ("Retention and deletion", "Backend Dev", 8), ("Data protection review", "Tech Lead", 6)],
        },
    ],
    "Blockchain platform": [
        {
            "title": "Smart contracts",
            "priority": "P0",
            "owner": "Blockchain Dev",
            "description": "Contracts written, tested and audited.",
            "acceptance": ["Contracts covered by tests", "Deployed to testnet", "External audit done"],
            "subtasks": [("Contract development", "Blockchain Dev", 24), ("Contract tests", "Blockchain Dev", 16), ("External audit", "Tech Lead", 8)],
        },
        {
            "title": "Wallet and front end",
            "priority": "P1",
            "owner": "Frontend Team",
            "description": "Users connect a wallet and use the contracts.",
            "acceptance": ["Wallet connection", "Transaction status shown", "Clear error messages"],
            "subtasks": [("Wallet connection", "Frontend Dev", 10), ("Transaction flows", "Frontend Dev", 16)],
        },
    ],
}

PROJECT_TYPES = list(TYPE_TASKS)


def build_plan(project_type: str) -> list[dict]:
    """Return a fresh plan for a project type: base + type-specific + closing tasks."""
    raw = deepcopy(BASE_TASKS + TYPE_TASKS[project_type] + CLOSING_TASKS)
    plan = []
    for task in raw:
        subtasks = [
            {"id": uuid.uuid4().hex[:8], "title": t, "owner": o, "hours": h, "done": False}
            for t, o, h in task["subtasks"]
        ]
        plan.append(
            {
                "id": uuid.uuid4().hex[:8],
                "title": task["title"],
                "priority": task["priority"],
                "owner": task["owner"],
                "description": task["description"],
                "acceptance": task["acceptance"],
                "hours": sum(s["hours"] for s in subtasks),
                "subtasks": subtasks,
            }
        )
    return plan


def plan_stats(plan: list[dict], team_size: int) -> dict:
    total = sum(t["hours"] for t in plan)
    done = sum(s["hours"] for t in plan for s in t["subtasks"] if s["done"])
    weeks = math.ceil(total / (max(team_size, 1) * HOURS_PER_PERSON_WEEK))
    return {"total": total, "done": done, "weeks": weeks}


# ---------------------------------------------------------------------------
# Exports
# ---------------------------------------------------------------------------

def to_markdown(name: str, project_type: str, plan: list[dict], notes: str = "") -> str:
    lines = [f"# {name}", "", f"Type: {project_type}", ""]
    for t in plan:
        lines.append(f"## {t['title']} ({t['priority']}, {t['hours']} h, {t['owner']})")
        lines.append(t["description"])
        lines.append("")
        lines.append("Acceptance criteria:")
        lines += [f"- {a}" for a in t["acceptance"]]
        lines.append("")
        lines += [f"- [{'x' if s['done'] else ' '}] {s['title']} ({s['owner']}, {s['hours']} h)" for s in t["subtasks"]]
        lines.append("")
    if notes:
        lines += ["## Project-specific notes", "", notes, ""]
    return "\n".join(lines)


def to_csv(plan: list[dict]) -> str:
    buf = io.StringIO()
    writer = csv.writer(buf)
    writer.writerow(["task", "priority", "subtask", "owner", "hours", "done"])
    for t in plan:
        for s in t["subtasks"]:
            writer.writerow([t["title"], t["priority"], s["title"], s["owner"], s["hours"], s["done"]])
    return buf.getvalue()


# ---------------------------------------------------------------------------
# Local model
# ---------------------------------------------------------------------------

@st.cache_resource(show_spinner="Loading the local model (first run only)...")
def load_model():
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    torch.set_num_threads(max(torch.get_num_threads(), 2))
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
    model.eval()
    return tokenizer, model


def build_prompt(name: str, project_type: str, description: str, plan: list[dict]) -> list[dict]:
    task_titles = ", ".join(t["title"] for t in plan)
    user = (
        f"Project: {name}\n"
        f"Type: {project_type}\n"
        f"Description: {description or 'not given'}\n"
        f"Planned tasks: {task_titles}\n\n"
        "Write two short markdown lists.\n"
        "### Risks\n3 risks specific to this project, one line each.\n"
        "### Missing tasks\n3 tasks the plan is missing, one line each with an hour estimate.\n"
        "No introduction, no conclusion."
    )
    return [
        {"role": "system", "content": "You are a pragmatic senior project manager. Be specific and brief."},
        {"role": "user", "content": user},
    ]


def stream_notes(messages: list[dict], max_new_tokens: int = 256):
    from transformers import TextIteratorStreamer

    tokenizer, model = load_model()
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt")
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        temperature=0.6,
        top_p=0.9,
        repetition_penalty=1.1,
    )
    Thread(target=model.generate, kwargs=kwargs, daemon=True).start()
    yield from streamer


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

def render_plan(plan: list[dict]) -> None:
    for t in plan:
        with st.expander(f"{t['priority']} · {t['title']} · {t['hours']} h · {t['owner']}"):
            left, right = st.columns([3, 2])
            with left:
                st.write(t["description"])
                st.markdown("**Acceptance criteria**")
                st.markdown("\n".join(f"- {a}" for a in t["acceptance"]))
            with right:
                st.markdown("**Subtasks**")
                for s in t["subtasks"]:
                    s["done"] = st.checkbox(f"{s['title']} ({s['owner']}, {s['hours']} h)", value=s["done"], key=s["id"])


def main() -> None:
    st.set_page_config(page_title="Project Task Planner", page_icon="📋", layout="wide")
    st.title("📋 Project Task Planner")
    st.caption(
        f"Turns a project idea into tasks, owners and hour estimates. "
        f"Optional: a small model ({MODEL_ID}) running on this Space's CPU adds project-specific risks. No external API."
    )

    with st.form("project"):
        c1, c2, c3 = st.columns([3, 2, 1])
        name = c1.text_input("Project name", placeholder="e.g. Online shop for a Zurich bakery")
        project_type = c2.selectbox("Project type", PROJECT_TYPES)
        team_size = c3.number_input("Team size", min_value=1, max_value=50, value=3)
        description = st.text_area("Short description (optional, used by the model)", height=80)
        submitted = st.form_submit_button("Generate plan", type="primary")

    if submitted:
        if not name.strip():
            st.warning("Please enter a project name.")
        else:
            st.session_state.update(
                plan=build_plan(project_type),
                name=name.strip(),
                project_type=project_type,
                description=description.strip(),
                notes="",
            )

    plan = st.session_state.get("plan")
    if not plan:
        st.info("Enter a project name and click **Generate plan**.")
        return

    stats = plan_stats(plan, team_size)
    m1, m2, m3 = st.columns(3)
    m1.metric("Estimated effort", f"{stats['total']} h")
    m2.metric(f"Duration with {team_size} people", f"~{stats['weeks']} weeks")
    m3.metric("Done", f"{stats['done']} h")
    st.progress(stats["done"] / stats["total"] if stats["total"] else 0.0)

    st.subheader(f"Plan: {st.session_state['name']}")
    render_plan(plan)

    st.subheader("Project-specific risks and missing tasks")
    if st.button("Ask the local model", help="Runs on CPU; takes about 20-60 seconds."):
        messages = build_prompt(st.session_state["name"], st.session_state["project_type"], st.session_state["description"], plan)
        try:
            st.session_state["notes"] = st.write_stream(stream_notes(messages))
        except Exception as exc:  # model download or generation failed
            st.error(f"The model could not run: {exc}")
    elif st.session_state.get("notes"):
        st.markdown(st.session_state["notes"])

    st.subheader("Export")
    d1, d2, d3 = st.columns(3)
    md = to_markdown(st.session_state["name"], st.session_state["project_type"], plan, st.session_state.get("notes", ""))
    d1.download_button("Markdown", md, file_name="plan.md", mime="text/markdown")
    d2.download_button("CSV", to_csv(plan), file_name="plan.csv", mime="text/csv")
    d3.download_button("JSON", json.dumps(plan, indent=2), file_name="plan.json", mime="application/json")


if __name__ == "__main__":
    main()