Planner-Executor Agents with Self-Critique: Plans, Validators, Replanning and Stopping Rules

"""A trip planner for Ferrow Travel: a planner writes the searches, code runs them, a solver picks the trip, and a
critic checks it. You write the functions marked TODO, one step at a time."""
import datetime as dt
import json
import re

import harness as H
import travel as T

MAX_STEPS = 12


class PlanError(ValueError):
    pass


def task_text(t):
    """The request as the travel desk receives it."""
    return (f"{t['traveller']} lives near {t['home']} airport and must attend an event in {t['city']} from "
            f"{t['event_start'].replace('T', ' ')} to {t['event_end'].replace('T', ' ')} (local time). Travel policy: "
            f"arrive at least 2 hours before the event starts; fly home on a flight leaving at least 90 minutes after "
            f"it ends (the same day or the next); fly direct, or change once at FRA or AMS with at least 60 minutes "
            f"between flights on the same day; stay every night between arriving and flying home in one hotel at most "
            f"{t['max_km']} km from the venue; keep flights plus hotel within {t['budget']} EUR in total.")


def first_json(reply, opener):
    """The first JSON value in reply that starts with opener ("[" or "{"), or None."""
    closer = "]" if opener == "[" else "}"
    m = re.search(re.escape(opener) + r".*" + re.escape(closer), reply or "", re.S)
    if not m:
        return None
    try:
        return json.loads(m.group(0))
    except ValueError:
        return None


# ---------- Step 1: plan the searches ----------

def plan_prompt(task, feedback=None):
    """Ask for every search needed, as a JSON list of {"id", "tool", "args"}. With feedback (what is missing), ask
    only for the extra searches."""
    # TODO (Step 1): the request (task_text), the tools with their parameters, what the policy allows, and the
    # JSON format; with feedback, ask only for the extra searches.
    raise NotImplementedError("Step 1: write plan_prompt()")


def parse_plan(reply):
    """The plan as a list of steps. Raise PlanError unless it is a non-empty JSON list of at most MAX_STEPS dicts, each
    with a unique string "id", a "tool" in T.TOOLS and "args" with exactly that tool's T.PARAMS, all strings."""
    # TODO (Step 1): first_json(reply, "["), then check each step; raise PlanError with the reason.
    raise NotImplementedError("Step 1: write parse_plan()")


# ---------- Step 2: run the plan, then solve ----------

def execute(plan):
    """{step id: the tool's result}, or {"error": message} for a step whose tool raised."""
    raise NotImplementedError("execute() arrives in Step 2")


def evidence(plan, results):
    """Compact text for the solver, where every option has an id "<step id>.<n>" (n from 1):
    "s1 search_journeys LIS->OSL 2026-10-12: 2 journeys" then "  s1.1 TP857 2026-10-12 + LH748 2026-10-12 | 08:55-22:04
    EUR 299"; "s5 search_hotels OSL: 5 hotels" then "  s5.1 Grand OSL EUR 104/night 1.2 km"; a failed step as
    "s2 ...: error <message>"."""
    raise NotImplementedError("evidence() arrives in Step 2")


def option(ref, plan, results):
    """The journey or hotel dict an id like "s1.3" points to, or None."""
    sid, _, n = str(ref).rpartition(".")
    r = results.get(sid)
    if not isinstance(r, list) or not n.isdigit() or not 1 <= int(n) <= len(r):
        return None
    return r[int(n) - 1]


def build(choice, plan, results):
    """(itinerary, problems) from the solver's choice {"outbound": id, "return": id, "hotel": id or null}.
    The itinerary: {"outbound": legs, "return": legs, "hotel": name or None, "checkin": arrival date or None,
    "checkout": return date or None, "total": both journeys' prices + nightly x nights}, where nights is the days
    from the outbound date to the return date. Problems name each id that does not point to the right kind of
    option ("outbound: s9.1 is not a journey in the search results"); with any problem the itinerary is None."""
    raise NotImplementedError("build() arrives in Step 2")


def solve_prompt(task, evidence_text, feedback=None):
    """Ask the solver to choose one outbound journey, one return journey and a hotel by their ids; with feedback (a
    list of problems with the last choice), ask it to fix them."""
    fix = ""
    if feedback:
        fix = "\nYour previous choice had these problems. Fix all of them:\n" + "\n".join(f"- {p}" for p in feedback) + "\n"
    return f"""Choose one trip that follows every rule of the travel policy from the search results below.

Request: {task_text(task)}

Search results (each option has an id such as s1.2):
{evidence_text}
{fix}
Check each rule before you answer: the outbound journey lands at least 2 hours before the start, the return journey
leaves at least 90 minutes after the end, the hotel is close enough (only needed if the return is on a later date
than the arrival), and journeys plus hotel nights fit the budget.
Reply with only JSON: {{"outbound": "<journey id>", "return": "<journey id>", "hotel": "<hotel id, or null if no night is needed>"}}
If no combination in these results can follow the policy, reply {{"need": "<which searches are missing>"}}."""


def parse_choice(reply):
    """The choice dict from the solver's reply, or None."""
    it = first_json(reply, "{")
    return it if isinstance(it, dict) else None


# ---------- Step 3: check it ----------

def at(date, hhmm):
    return dt.datetime.fromisoformat(f"{date}T{hhmm}")


def check_legs(legs, origin, destination, label):
    """(problems, flight records) for a list of "<flight> <date>" legs that should go from origin to destination."""
    problems, recs = [], []
    for leg in legs if isinstance(legs, list) else []:
        parts = leg.split() if isinstance(leg, str) else []
        f = T.flight(parts[0], parts[1]) if len(parts) == 2 else None
        if f is None:
            problems.append(f"{label}: unknown flight {leg!r}")
        recs.append(f)
    if problems:
        return problems, []
    if not recs:
        return [f"{label}: no flights"], []
    if len(recs) > 2:
        problems.append(f"{label}: {len(recs)} flights; at most one change is allowed")
    if recs[0]["origin"] != origin or recs[-1]["destination"] != destination:
        problems.append(f"{label}: goes from {recs[0]['origin']} to {recs[-1]['destination']}, not {origin} to {destination}")
    for a, b in zip(recs, recs[1:]):
        if b["origin"] != a["destination"]:
            problems.append(f"{label}: {b['flight']} leaves {b['origin']}, but {a['flight']} lands at {a['destination']}")
        elif a["destination"] not in T.HUBS:
            problems.append(f"{label}: changes at {a['destination']}; changes are only allowed at FRA or AMS")
        elif b["date"] != a["date"] or at(b["date"], b["dep"]) - at(a["date"], a["arr"]) < dt.timedelta(minutes=60):
            problems.append(f"{label}: less than 60 minutes to change from {a['flight']} to {b['flight']} at {a['destination']}")
    return problems, recs


def validate(task, it):
    """Every way the itinerary breaks the policy, as a list of short strings ([] if it follows every rule):
    the legs (check_legs), arrival at least 2 h before the start, the first return flight leaving at least 90 min
    after the end, the hotel (exists in the city, within max_km, checkin on the arrival date and checkout on the
    return date, or no hotel when no night is needed), the stated total against the real cost (flights plus
    nightly x nights, 0.5 EUR tolerance), and the real cost against the budget."""
    raise NotImplementedError("validate() arrives in Step 3")


# ---------- Step 4: critique, revise, and know when to stop ----------

def solve(task, critic="validator", max_rounds=4):
    """Plan, run the searches, then up to max_rounds solver answers. After each answer: {"need": ...} means replan
    with that as feedback and add the new searches (counts as a round); otherwise build() the itinerary. build()'s
    problems (ids that point nowhere) are always fed back; with no critic they stop the run as "unreadable".
    Otherwise get problems from the critic ("validator": validate(); "self": self_critique(); None: none) and stop at
    "no_problems", "no_progress" (the same problems as the round before), or "max_rounds".
    Returns {"itinerary", "valid" (validate() of the final itinerary is []), "rounds", "stop", "llm_calls"}."""
    raise NotImplementedError("solve() arrives in Step 4")


# ---------- Step 5: can the model check itself? ----------

def self_critique(task, evidence_text, it):
    """Ask the model to check the itinerary against the request and the search results, like a colleague would:
    a JSON list of problems, [] if there are none. Anything unreadable counts as no problems."""
    raise NotImplementedError("self_critique() arrives in Step 5")


def compare(results):
    """{critic: {"valid" (share of tasks), "llm_calls" (mean), "rounds" (mean)}} from {critic: [solve() results]}."""
    raise NotImplementedError("compare() arrives in Step 5")