freight/scripts/generate_freight_dataset.py

#!/usr/bin/env python3
"""
freight/scripts/generate_freight_dataset.py — Full-Spectrum Multi-Intent US Freight SFT Dataset Generator.

Trains LoadETA across all 5 operational pillars of commercial trucking:
1. Spot Rate Negotiation & Booking (FMCSA verification, Rate Floor, Equipment, Driver info, Rate Con)
2. Check Calls & In-Transit Tracking (Miles out, MacroPoint ELD link, ETA updates, weather/traffic delays)
3. Accessorials, Detention & Lumpers (Detention >2hrs, BOL in/out timestamps, EFS/Comchek codes, TONU)
4. Billing, POD & Factoring Inquiries (Signed BOL/POD submission, Notice of Assignment, QuickPay status)
5. Equipment, Safety, HOS & Regulatory Q&A (IFTA, 11h/14h clock, 80k GVWR axle limits, Kingpin laws, MC vs DOT)

Strict Speech & TTS Hygiene:
- Digit-by-digit pronunciation for IDs and MC numbers
- Natural spoken phrasing for dollar amounts
- Full state name pronunciation
- No company suffixes (LLC, INC)
- No robotic markdown, bullets, or emojis

Usage:
    python3 freight/scripts/generate_freight_dataset.py --count 500 --output freight/data/freight_negotiation_sample.jsonl
"""

import os
import sys
import json
import random
import argparse
from typing import List, Dict, Any

# ─── Seed Data Matrix ───────────────────────────────────────────────────────────

LANES = [
    {"origin": "Chicago, Illinois", "dest": "Atlanta, Georgia", "miles": 715, "min_rpm": 2.65, "tolls": "Moderate on Interstate 65", "highway": "Interstate 65 South"},
    {"origin": "Los Angeles, California", "dest": "Dallas, Texas", "miles": 1435, "min_rpm": 2.20, "tolls": "Low on Interstate 10", "highway": "Interstate 10 East"},
    {"origin": "Laredo, Texas", "dest": "Chicago, Illinois", "miles": 1380, "min_rpm": 2.80, "tolls": "Moderate on Interstate 35", "highway": "Interstate 35 North"},
    {"origin": "Allentown, Pennsylvania", "dest": "Richmond, Virginia", "miles": 260, "min_rpm": 3.40, "tolls": "High on Pennsylvania Turnpike", "highway": "Interstate 95 South"},
    {"origin": "Savannah, Georgia", "dest": "Memphis, Tennessee", "miles": 585, "min_rpm": 2.50, "tolls": "Low on Interstate 16", "highway": "Interstate 16 West"},
    {"origin": "Seattle, Washington", "dest": "Denver, Colorado", "miles": 1305, "min_rpm": 2.45, "tolls": "Low on Interstate 90", "highway": "Interstate 90 East"},
    {"origin": "Houston, Texas", "dest": "Atlanta, Georgia", "miles": 790, "min_rpm": 2.55, "tolls": "Low on Interstate 10", "highway": "Interstate 10 East"},
    {"origin": "Indianapolis, Indiana", "dest": "Orlando, Florida", "miles": 965, "min_rpm": 2.70, "tolls": "Moderate on Interstate 75", "highway": "Interstate 75 South"},
    {"origin": "Kansas City, Missouri", "dest": "Columbus, Ohio", "miles": 620, "min_rpm": 2.75, "tolls": "Low on Interstate 70", "highway": "Interstate 70 East"},
    {"origin": "Elizabeth, New Jersey", "dest": "Greensboro, North Carolina", "miles": 510, "min_rpm": 3.10, "tolls": "High on New Jersey Turnpike", "highway": "Interstate 95 South"},
    {"origin": "Phoenix, Arizona", "dest": "Ontario, California", "miles": 340, "min_rpm": 2.90, "tolls": "Low on Interstate 10", "highway": "Interstate 10 West"},
    {"origin": "Detroit, Michigan", "dest": "Nashville, Tennessee", "miles": 530, "min_rpm": 2.85, "tolls": "Moderate on Interstate 75", "highway": "Interstate 75 South"},
]

EQUIPMENT_TYPES = [
    {"type": "53ft Reefer", "commodity": "Frozen Poultry", "temp": "minus ten degrees continuous", "extra_cost": 300},
    {"type": "53ft Reefer", "commodity": "Fresh Organic Produce", "temp": "thirty-four degrees pre-cooled", "extra_cost": 250},
    {"type": "53ft Dry Van", "commodity": "Palletized Consumer Goods", "temp": "ambient", "extra_cost": 0},
    {"type": "53ft Dry Van", "commodity": "Retail Automotive Parts", "temp": "ambient", "extra_cost": 0},
    {"type": "48ft Flatbed", "commodity": "Structural Steel", "temp": "ambient", "extra_cost": 200},
    {"type": "Stepdeck", "commodity": "Construction Machinery", "temp": "ambient", "extra_cost": 400},
]

BROKER_COMPANIES = [
    {"name": "Apex Logistics", "mc": "782914", "fraud": False},
    {"name": "Summit Freight", "mc": "419820", "fraud": False},
    {"name": "FastTrack Logistics", "mc": "993812", "fraud": False},
    {"name": "Redwood Transportation", "mc": "624108", "fraud": False},
    {"name": "Echo Global Direct", "mc": "512940", "fraud": False},
    {"name": "Ghost Lane", "mc": "1588231", "fraud": True},
]

DRIVERS = [
    {"name": "John Kovacs", "phone": "three one two, five five five, zero one nine eight", "truck": "Truck seven zero four", "trailer": "Trailer five five two"},
    {"name": "Marcus Vance", "phone": "four zero four, five five five, zero two eight one", "truck": "Truck eight one nine", "trailer": "Trailer nine zero one"},
    {"name": "David Miller", "phone": "two one four, five five five, zero three seven two", "truck": "Truck five one two", "trailer": "Trailer three eight zero"},
    {"name": "Sergey Petrov", "phone": "seven one three, five five five, zero four six three", "truck": "Truck six three zero", "trailer": "Trailer four one nine"},
]

SYSTEM_PROMPT = (
    "You are LoadETA, an autonomous commercial freight dispatch and operations voice agent on a live phone call. "
    "Time is money: keep responses short, professional, and natural for voice synthesis (one to two sentences per turn). "
    "Do not use markdown bolding, bullet points, or robotic filler. "
    "Call tools for live lookups, updates, and bookings. Follow strict speech hygiene: pronounce numbers digit-by-digit for IDs "
    "and MCs, naturally for dollar amounts, full state names, and strip legal suffixes like LLC and INC."
)


def number_to_words_usd(amount: int) -> str:
    if amount >= 1000 and amount % 100 == 0:
        hundreds = amount // 100
        return f"{hundreds} hundred dollars"
    elif amount >= 1000 and amount % 50 == 0:
        hundreds = amount // 100
        rem = amount % 100
        return f"{hundreds} {rem} dollars"
    return f"{amount} dollars"


def format_mc_digits(mc_str: str) -> str:
    digit_words = {
        "0": "zero", "1": "one", "2": "two", "3": "three", "4": "four",
        "5": "five", "6": "six", "7": "seven", "8": "eight", "9": "nine"
    }
    return " ".join(digit_words.get(c, c) for c in mc_str)


# ─── Intent 1: Spot Rate Negotiation & Booking ─────────────────────────────────

def generate_negotiation_scenario(idx: int) -> Dict[str, Any]:
    lane = random.choice(LANES)
    equip = random.choice(EQUIPMENT_TYPES)
    broker = random.choice(BROKER_COMPANIES)
    driver = random.choice(DRIVERS)
    weight = random.randint(34, 44) * 1000

    base_rate = int(lane["miles"] * lane["min_rpm"] + equip["extra_cost"])
    base_rate = (base_rate // 50) * 50
    broker_offer = int(base_rate * 0.85)
    broker_offer = (broker_offer // 50) * 50
    target_counter = int(base_rate * 1.06)
    target_counter = (target_counter // 50) * 50

    broker_offer_spoken = number_to_words_usd(broker_offer)
    target_counter_spoken = number_to_words_usd(target_counter)

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    turn1 = (
        f"Hey, this is Mike from {broker['name']}, Motor Carrier Number {format_mc_digits(broker['mc'])}. "
        f"I have a {equip['type']} load picking up in {lane['origin']} going to {lane['dest']}, "
        f"{weight:,} pounds of {equip['commodity']}. Can you cover this today for {broker_offer_spoken}?"
    )
    messages.append({"role": "user", "content": turn1})

    tool_call_1 = f"call_fmcsa_{idx:04d}"
    tool_call_2 = f"call_rate_{idx:04d}"
    messages.append({
        "role": "assistant",
        "content": None,
        "tool_calls": [
            {
                "id": tool_call_1,
                "type": "function",
                "function": {
                    "name": "fmcsa_verify_mc",
                    "arguments": json.dumps({"mc_number": broker["mc"], "broker_name": broker["name"]}),
                },
            },
            {
                "id": tool_call_2,
                "type": "function",
                "function": {
                    "name": "calculate_rate_floor",
                    "arguments": json.dumps({
                        "origin": lane["origin"],
                        "destination": lane["dest"],
                        "equipment_type": equip["type"],
                        "mileage": lane["miles"],
                        "weight": weight,
                    }),
                },
            },
        ],
    })

    if broker["fraud"]:
        messages.append({"role": "tool", "tool_call_id": tool_call_1, "content": json.dumps({"status": "FLAGGED", "active": False, "reason": "Unverified Authority"})})
        messages.append({"role": "tool", "tool_call_id": tool_call_2, "content": json.dumps({"rate_floor": base_rate})})
        messages.append({"role": "assistant", "content": f"Mike, our compliance system flagged Motor Carrier Number {format_mc_digits(broker['mc'])} for inactive authority. We cannot accept loads with your company. Have a safe day."})
        return {"id": f"freight_intent1_{idx:05d}", "intent": "rate_negotiation", "messages": messages}

    messages.append({"role": "tool", "tool_call_id": tool_call_1, "content": json.dumps({"status": "CLEAN", "active": True, "credit_score": 95})})
    messages.append({"role": "tool", "tool_call_id": tool_call_2, "content": json.dumps({"rate_floor": base_rate, "recommended_target": target_counter})})

    messages.append({
        "role": "assistant",
        "content": (
            f"Motor Carrier Number {format_mc_digits(broker['mc'])} is verified. "
            f"For {lane['miles']} miles on {lane['origin']} to {lane['dest']}, {broker_offer_spoken} is below our operating floor. "
            f"We can roll on this for {target_counter_spoken} all in, with two hours free detention and seventy-five dollars an hour after."
        ),
    })

    messages.append({"role": "user", "content": f"I can do {target_counter_spoken} if you can guarantee {driver['name']} picks up on time. What is the driver and truck info?"})

    tool_call_3 = f"call_book_{idx:04d}"
    messages.append({
        "role": "assistant",
        "content": (
            f"Confirmed. We assign {driver['truck']} with driver {driver['name']}, cell phone {driver['phone']}. "
            f"Send the rate confirmation to dispatch at load eta dot com."
        ),
        "tool_calls": [
            {
                "id": tool_call_3,
                "type": "function",
                "function": {
                    "name": "book_load_offer",
                    "arguments": json.dumps({
                        "broker_name": broker["name"],
                        "mc_number": broker["mc"],
                        "rate_agreed": target_counter,
                        "origin": lane["origin"],
                        "destination": lane["dest"],
                        "equipment": equip["type"],
                        "driver_name": driver["name"],
                        "driver_phone": driver["phone"],
                        "truck_number": driver["truck"],
                        "detention_rate_per_hr": 75,
                        "dispatch_email": "dispatch@loadeta.com",
                    }),
                },
            }
        ],
    })

    messages.append({"role": "tool", "tool_call_id": tool_call_3, "content": json.dumps({"status": "BOOKED", "confirmation_code": f"LE-{random.randint(100000, 999999)}"})})
    messages.append({"role": "user", "content": "Rate confirmation sent to your email. Thank you."})
    messages.append({"role": "assistant", "content": "Load is confirmed on our end. We will sign and return the rate confirmation shortly. Have a great day."})

    return {"id": f"freight_intent1_{idx:05d}", "intent": "rate_negotiation", "messages": messages}


# ─── Intent 2: Check Calls & In-Transit Tracking ───────────────────────────────

def generate_check_call_scenario(idx: int) -> Dict[str, Any]:
    lane = random.choice(LANES)
    broker = random.choice(BROKER_COMPANIES)
    driver = random.choice(DRIVERS)
    miles_out = random.randint(45, 180)
    hours_out = round(miles_out / 50.0, 1)

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    turn1 = (
        f"Hi, this is Sarah from {broker['name']} tracking department calling on load {random.randint(80000, 89999)}. "
        f"Can I get a location update on {driver['name']} heading to {lane['dest']}?"
    )
    messages.append({"role": "user", "content": turn1})

    tool_call_1 = f"call_track_{idx:04d}"
    messages.append({
        "role": "assistant",
        "content": None,
        "tool_calls": [
            {
                "id": tool_call_1,
                "type": "function",
                "function": {
                    "name": "get_truck_route_distance",
                    "arguments": json.dumps({"driver_name": driver["name"], "truck_number": driver["truck"], "destination": lane["dest"]}),
                },
            }
        ],
    })

    messages.append({
        "role": "tool",
        "tool_call_id": tool_call_1,
        "content": json.dumps({
            "current_location": f"Mile marker {random.randint(120, 240)} on {lane['highway']}",
            "miles_to_destination": miles_out,
            "estimated_eta": f"{hours_out} hours",
            "eld_status": "Driving, 4 hours available on 11 hour clock",
            "temperature_status": "Set at minus ten degrees, actual minus ten degrees",
        }),
    })

    messages.append({
        "role": "assistant",
        "content": (
            f"Hello Sarah. {driver['name']} is currently rolling on {lane['highway']}, approximately {miles_out} miles out from {lane['dest']}. "
            f"Estimated arrival time is in two and a half hours, well ahead of the delivery window."
        ),
    })

    messages.append({"role": "user", "content": "Awesome. Is the MacroPoint tracking active on your driver cell phone?"})
    messages.append({
        "role": "assistant",
        "content": f"Yes, MacroPoint is active on cell phone {driver['phone']} with location sharing enabled. Everything is smooth on transit.",
    })

    return {"id": f"freight_intent2_{idx:05d}", "intent": "check_call_tracking", "messages": messages}


# ─── Intent 3: Accessorials, Detention & Lumpers ────────────────────────────────

def generate_detention_lumper_scenario(idx: int) -> Dict[str, Any]:
    lane = random.choice(LANES)
    broker = random.choice(BROKER_COMPANIES)
    driver = random.choice(DRIVERS)
    lumper_fee = random.choice([175, 240, 285, 320])
    lumper_fee_spoken = number_to_words_usd(lumper_fee)

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    turn1 = (
        f"Hi, this is LoadETA dispatch calling {broker['name']} regarding {driver['truck']} on load {random.randint(91000, 99000)}. "
        f"{driver['name']} has been at the receiver in {lane['dest']} for two hours and forty-five minutes and the lumper service requires {lumper_fee_spoken}."
    )
    messages.append({"role": "assistant", "content": turn1})

    messages.append({
        "role": "user",
        "content": f"Okay, does your driver have a signed lumper receipt or breakdown slip? And when did your driver arrive on site?",
    })

    tool_call_1 = f"call_detention_{idx:04d}"
    messages.append({
        "role": "assistant",
        "content": None,
        "tool_calls": [
            {
                "id": tool_call_1,
                "type": "function",
                "function": {
                    "name": "submit_accessorial_claim",
                    "arguments": json.dumps({
                        "truck_number": driver["truck"],
                        "claim_type": "detention_and_lumper",
                        "lumper_amount": lumper_fee,
                        "hours_waiting": 2.75,
                        "in_time": "08:00 AM",
                        "appointment_time": "08:00 AM",
                    }),
                },
            }
        ],
    })

    messages.append({
        "role": "tool",
        "tool_call_id": tool_call_1,
        "content": json.dumps({"status": "SUBMITTED", "qualifies_detention": True, "billable_detention_hours": 1, "detention_usd": 75}),
    })

    messages.append({
        "role": "assistant",
        "content": (
            f"Yes, {driver['name']} checked in at zero eight hundred sharp for a scheduled eight hundred appointment with in-times marked on the Bill of Lading. "
            f"Please issue an EFS code or Comchek for {lumper_fee_spoken} so the unloader can release the truck, and update the rate con with one hour detention."
        ),
    })

    messages.append({"role": "user", "content": f"Understood. The EFS express code for {lumper_fee_spoken} is four eight one zero nine nine. Please have driver get the stamped receipt."})
    messages.append({"role": "assistant", "content": f"Received code four eight one zero nine nine. The stamped receipt and signed Bill of Lading will be uploaded immediately after unload. Thank you."})

    return {"id": f"freight_intent3_{idx:05d}", "intent": "detention_lumper", "messages": messages}


# ─── Intent 4: Billing, POD & Factoring ────────────────────────────────────────

def generate_billing_factoring_scenario(idx: int) -> Dict[str, Any]:
    broker = random.choice(BROKER_COMPANIES)
    load_num = random.randint(70000, 79999)

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    turn1 = f"Hi, this is accounting from {broker['name']}. We received your paperwork for load {load_num}, but we need the factoring Notice of Assignment before processing payment."
    messages.append({"role": "user", "content": turn1})

    tool_call_1 = f"call_factoring_{idx:04d}"
    messages.append({
        "role": "assistant",
        "content": None,
        "tool_calls": [
            {
                "id": tool_call_1,
                "type": "function",
                "function": {
                    "name": "verify_factoring_status",
                    "arguments": json.dumps({"load_number": load_num, "broker_name": broker["name"]}),
                },
            }
        ],
    })

    messages.append({
        "role": "tool",
        "tool_call_id": tool_call_1,
        "content": json.dumps({
            "factoring_company": "Triumph Financial",
            "noa_on_file": True,
            "remit_email": "remit@triumphbancorp.com",
            "signed_pod_uploaded": True,
        }),
    })

    messages.append({
        "role": "assistant",
        "content": (
            f"Hello. Our account is factored through Triumph Financial. "
            f"The Notice of Assignment and clean signed Proof of Delivery were transmitted to your accounts payable department. "
            f"Please verify remittance to Triumph Financial at remit at triumph bancorp dot com."
        ),
    })

    messages.append({"role": "user", "content": "Got it, I see the Triumph Notice of Assignment in the document portal now. Payment will be released on two day QuickPay."})
    messages.append({"role": "assistant", "content": "Confirmed. Thank you for resolving this quickly. Have a great day."})

    return {"id": f"freight_intent4_{idx:05d}", "intent": "billing_factoring", "messages": messages}


# ─── Intent 5: Equipment, Regulations, HOS & General Knowledge ────────────────

def generate_general_knowledge_scenario(idx: int) -> Dict[str, Any]:
    topics = [
        {
            "user": "Hey, can you explain the Hours of Service driving clock limits for our driver heading to Texas?",
            "reply": "Under Federal Motor Carrier Safety Administration rules, property-carrying drivers have an eleven hour maximum driving window after ten consecutive hours off duty. They must also take a thirty minute rest break before reaching eight hours of driving time.",
        },
        {
            "user": "What is the maximum legal gross vehicle weight without over-dimensional permits on interstate highways?",
            "reply": "The federal maximum gross vehicle weight is eighty thousand pounds without special permits, with axle weight limits of twelve thousand pounds on steer axles, thirty-four thousand pounds on drive tandems, and thirty-four thousand pounds on trailer tandems.",
        },
        {
            "user": "What does Motor Carrier Number mean versus Department of Transportation number when booking loads?",
            "reply": "A Department of Transportation number tracks safety and vehicle inspection compliance, while a Motor Carrier Number grants operating authority for interstate commerce and hauling regulated freight for hire.",
        },
        {
            "user": "How do Kingpin to rear axle bridge laws affect fifty-three foot trailers entering California?",
            "reply": "In California, the distance from the kingpin to the center of the rear trailer axle cannot exceed forty feet. Drivers must slide their trailer tandems forward to remain legal before crossing the state border.",
        },
        {
            "user": "What is the difference between continuous run and start-stop mode on refrigerated trailers?",
            "reply": "Continuous mode maintains a constant airflow and steady temperature required for frozen loads like meat and ice cream, while start-stop mode cycles the refrigeration unit on and off to save fuel on resilient freight.",
        },
    ]

    t = random.choice(topics)
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": t["user"]},
        {"role": "assistant", "content": t["reply"]},
    ]
    return {"id": f"freight_intent5_{idx:05d}", "intent": "freight_general_knowledge", "messages": messages}


# ─── Main Generator ─────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Generate Full-Spectrum Freight SFT Dataset")
    parser.add_argument("--count", type=int, default=300, help="Total number of dialogues to generate")
    parser.add_argument("--output", default="freight/data/freight_negotiation_sample.jsonl", help="Output JSONL path")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    args = parser.parse_args()

    random.seed(args.seed)
    os.makedirs(os.path.dirname(args.output), exist_ok=True)

    # Balanced Intent Distribution:
    # 35% Negotiations & Booking
    # 25% Check Calls & In-Transit Tracking
    # 15% Accessorials & Detention/Lumper
    # 15% Billing & Factoring
    # 10% General Knowledge & Regulations
    intent_weights = [
        ("negotiation", generate_negotiation_scenario, 0.35),
        ("check_call", generate_check_call_scenario, 0.25),
        ("detention_lumper", generate_detention_lumper_scenario, 0.15),
        ("billing_factoring", generate_billing_factoring_scenario, 0.15),
        ("general_knowledge", generate_general_knowledge_scenario, 0.10),
    ]

    print(f"Generating {args.count} multi-intent commercial freight conversations...")
    counts_by_intent = {}

    with open(args.output, "w", encoding="utf-8") as f:
        for i in range(1, args.count + 1):
            r = random.random()
            cum = 0.0
            chosen_fn = generate_negotiation_scenario
            intent_name = "negotiation"
            for name, fn, weight in intent_weights:
                cum += weight
                if r <= cum:
                    chosen_fn = fn
                    intent_name = name
                    break

            scenario = chosen_fn(i)
            counts_by_intent[intent_name] = counts_by_intent.get(intent_name, 0) + 1
            f.write(json.dumps(scenario) + "\n")

    print(f"\nGenerated dataset saved to {args.output}")
    print("Intent Distribution Breakdown:")
    for name, cnt in counts_by_intent.items():
        print(f"  - {name}: {cnt} dialogues ({round(cnt/args.count*100, 1)}%)")


if __name__ == "__main__":
    main()