training/dataset.py

"""
training/dataset.py

TaskCurriculum — generates tasks with progressive difficulty.

Per hackathon guide: "Make success possible early."
Start with short horizons, then gradually remove scaffolding.

Difficulty tiers:
  easy   — single source, 1 year of data, explicit answer in ground truth
  medium — 2 sources, 2 years, partial ground truth
  hard   — 2+ sources, 3 years, schema drift, ambiguous answer
"""

from __future__ import annotations
import random
from typing import Optional


EASY_INSTRUCTIONS = [
    "Find all travel and ride receipts from Gmail in the last 10 days and list them.",
    "Summarize all Uber and Rapido emails from 2023.",
    "What is the total amount spent on shopping invoices in January 2022?",
    "List all tax invoices from Google Drive created in March 2023.",
    "Find emails with the word 'flight' or 'hotel' in 2022.",
    "Read my historical transactions from Google Sheets.",
    "Retrieve the latest 20 financial emails from Gmail.",
    "Sync my recent Gmail transactions to Google Sheets.",
    "Check Google Sheets for the total spend last month.",
    "Find all food delivery receipts in Gmail from last week."
]

MEDIUM_INSTRUCTIONS = [
    "Audit all ride receipts from Gmail between 2022 and 2023, and calculate the total spend by vendor.",
    "Find all travel-related emails and documents from 2022-2023 and summarize total expenses and dates.",
    "Identify any recurring subscriptions and their monthly costs across my email history from 2022-2024.",
    "Summarize my Gmail shopping invoices from 2023 and sum the totals.",
    "Retrieve transactions from Gmail, then check Google Sheets to see if they are already logged.",
    "Sync all new shopping receipts from Gmail to the Google Sheets ledger.",
    "Compare the travel expenses found in Gmail with the records in Google Sheets.",
    "Aggregate my food orders from Gmail and update the summary in Google Sheets."
]

HARD_INSTRUCTIONS = [
    "Perform a full financial audit of my travel and ride footprint from 2022 to 2024, flag any missing receipts, and produce a summary report with exact amounts.",
    "Build a complete breakdown of my financial transactions across categories (rides, travel, shopping) from 2022-2024, identifying the top 5 vendors by spend.",
    "Analyze all receipts, tax invoices, and bookings across Gmail from 2022-2024. Extract all numeric totals and aggregate them by category.",
    "Given my financial history across 2022-2024, estimate total annual ride costs.",
    "Retrieve all financial transactions from Gmail, sync them to Google Sheets, then generate a merged financial dashboard summary.",
    "Audit the Google Sheets ledger against raw Gmail receipts to find discrepancies, and output a final reconciled total spend.",
    "Fetch all unlogged invoices from Gmail, sync them to Sheets, and summarize the top spending categories across both sources."
]


def _make_ground_truth(instruction: str, difficulty: str) -> dict:
    """Generate a plausible ground truth for the task (for reward scoring)."""
    # Since we are training on REAL live data (Gmail/Sheets), we cannot hardcode
    # expected numeric targets. Instead, we return None for answer/numeric_target
    # so the reward function evaluates answer coherence, completeness, and structure.
    if difficulty == "easy":
        return {
            "answer": None,
            "expected_numeric_target": None,
            "expected_sources": ["gmail", "sheets"],
            "expected_steps": 5,
        }
    elif difficulty == "medium":
        return {
            "answer": None,
            "expected_numeric_target": None,
            "expected_sources": ["gmail", "sheets"],
            "expected_steps": 10,
        }
    else:
        return {
            "answer": None,
            "expected_numeric_target": None,
            "expected_sources": ["gmail", "sheets"],
            "expected_steps": 15,
            "schema_drift": True,
        }


class TaskCurriculum:
    """
    Samples tasks from easy → medium → hard based on training progress.
    Implements a simple staged curriculum.
    """

    STAGES = [
        {"difficulty": "easy",   "weight": 1.0, "until_step": 50},
        {"difficulty": "medium", "weight": 1.0, "until_step": 150},
        {"difficulty": "hard",   "weight": 1.0, "until_step": None},
    ]

    def __init__(self, config: dict):
        self.config = config
        self._global_step = 0

    def advance(self, step: int):
        """Update the curriculum based on training step."""
        self._global_step = step

    def sample(self) -> dict:
        """Sample a task at the appropriate difficulty level."""
        difficulty = self._current_difficulty()

        if difficulty == "easy":
            instruction = random.choice(EASY_INSTRUCTIONS)
            sources = ["gmail", "sheets"]
        elif difficulty == "medium":
            instruction = random.choice(MEDIUM_INSTRUCTIONS)
            sources = ["gmail", "sheets"]
        else:
            instruction = random.choice(HARD_INSTRUCTIONS)
            sources = ["gmail", "sheets"]

        return {
            "instruction": instruction,
            "difficulty": difficulty,
            "available_sources": sources,
            "ground_truth": _make_ground_truth(instruction, difficulty),
        }

    def _current_difficulty(self) -> str:
        step = self._global_step
        if step < 50:
            return "easy"
        elif step < 150:
            return "medium"
        else:
            return "hard"