Spaces:

crodri
/

vfrplan-agent

Paused

Retry a 502/504 from the model gateway instead of failing the turn

aaaf5af 3 months ago

14.1 kB

"""The agent loop: model → tool calls → model → ... → answer.
Written out by hand rather than taken from a framework, because the whole
value of this screen is watching the loop happen: which tool, with which
arguments, what came back, what the model made of it. A framework that owns
the loop also owns what you are allowed to see of it.
Every stage yields an event; main.py forwards them to the browser as SSE.
Nothing here writes to the transport, so the loop stays testable on its own.
"""

from __future__ import annotations

import json
import re
from collections.abc import AsyncIterator
from pathlib import Path
from typing import Any

from .config import AppConfig
from .llm import LLMError, chat_completion, extract_message, extract_reasoning
from .mcp_tools import mcp_session, result_to_text, to_openai_tools
from .wikipedia_tool import WIKIPEDIA_TOOL_NAME, WIKIPEDIA_TOOL_SCHEMA, wikipedia_lookup

PROMPT_PATH = Path(__file__).resolve().parent.parent / "PROMPT.md"

# Fallback only. The real rules live in PROMPT.md, shipped alongside this app
# and copied unmodified from the vfrplan project — it is what stops the model
# inventing figures it does not have. If the file goes missing the agent still
# runs, but the safety-critical parts of its behaviour are gone, so the
# fallback says so out loud rather than pretending nothing happened.
FALLBACK_PROMPT = """You are a VFR flight planning assistant for Europe.
Use the provided tools for every aeronautical fact. Never invent aerodrome
data, weather, NOTAMs or performance figures. Report every assumption and
warning the tools return. Ask the pilot for missing inputs instead of
guessing them. Fuel, weather minima and go/no-go remain the pilot's decisions.
(Note: PROMPT.md was not found, so these are reduced fallback instructions.)"""


def system_prompt() -> str:
    try:
        text = PROMPT_PATH.read_text(encoding="utf-8").strip()
        return text or FALLBACK_PROMPT
    except OSError:
        return FALLBACK_PROMPT


def _event(kind: str, **data: Any) -> dict[str, Any]:
    return {"type": kind, **data}


# Matches an ISO date, optionally followed by a time and/or a trailing 'Z' /
# numeric UTC offset, e.g. "2026-08-18", "2026-08-18T11:00:00" or
# "2026-08-14T09:00:00+00:00".
_DATETIME_RE = re.compile(
    r"^(?P<date>\d{4}-\d{2}-\d{2})"
    r"(?:[T ](?P<time>\d{2}:\d{2}(?::\d{2})?)(?:\.\d+)?)?"
    r"(?:Z|[+-]\d{2}:?\d{2})?$"
)


def _sanitize_arguments(
    arguments: dict[str, Any], schema: dict[str, Any] | None = None
) -> dict[str, Any]:
    """Reshape ISO-datetime-looking string arguments to match a tool's own schema.
    Two plan_route failures have been seen, both the same root cause under a
    different face. First "Invalid isoformat string:
    '2026-08-14T09:00:00+00:00'" — a trailing offset its parser does not
    accept. Then, after that fix, "Invalid isoformat string:
    '2026-08-18T11:00:00'" on a string with no offset at all, which is
    exactly the error Python's own date.fromisoformat() raises when handed a
    full datetime: the tool has a *date-only* parameter, and the model was
    sending it a combined date+time. A per-tool string field can't be
    hardcoded for this — plan_route's real parameter names are behind the
    private vfrplan-mcp Space and were not confirmed live — so this reads
    each tool's own declared JSON Schema `format` ("date"/"time"/anything
    else) for the argument key in question and reshapes the value to match:
    date-only down to just the date, time-only down to just the time, and
    everything else keeps date+time together with any tz suffix dropped
    (PROMPT.md §2 has every time treated as local by design, so an offset is
    never meaningful here regardless). With no schema, or no format on that
    property, the old date+time-joined-minus-tz behaviour applies.
    """
    properties = (schema or {}).get("properties") or {} if isinstance(schema, dict) else {}
    sanitized = dict(arguments)
    for key, value in arguments.items():
        if not isinstance(value, str):
            continue
        match = _DATETIME_RE.match(value.strip())
        if not match:
            continue
        date_part = match.group("date")
        time_part = match.group("time")
        fmt = (properties.get(key) or {}).get("format")
        if fmt == "date":
            sanitized[key] = date_part
        elif fmt == "time":
            if time_part:
                sanitized[key] = time_part
        else:
            sanitized[key] = f"{date_part}T{time_part}" if time_part else date_part
    return sanitized


def _parse_arguments(raw: Any) -> tuple[dict[str, Any], str | None]:
    """Normalise tool-call arguments into a dict.
    The API says this is a JSON string, but providers vary: some send a dict
    already, some send "" for a no-argument call, and some send malformed
    JSON. A parse failure is returned rather than raised so the loop can hand
    it back to the model as a tool error and let it retry — which works far
    more often than aborting the turn.
    """
    if isinstance(raw, dict):
        return raw, None
    if raw in (None, "", "null"):
        return {}, None
    try:
        parsed = json.loads(raw)
    except (TypeError, ValueError) as exc:
        return {}, f"arguments were not valid JSON: {exc}"
    if not isinstance(parsed, dict):
        return {}, f"arguments must be a JSON object, got {type(parsed).__name__}"
    return parsed, None


async def run_agent(
    config: AppConfig,
    user_message: str,
    history: list[dict[str, str]] | None = None,
) -> AsyncIterator[dict[str, Any]]:
    """Run one turn, yielding events as they happen.
    Event kinds consumed by the frontend:
      status       — a line for the activity strip
      tools_loaded — the tool inventory, once, after connecting
      reasoning    — the model's private thinking for this step
      thinking     — assistant prose emitted alongside tool calls
      tool_call    — a call about to be made, with its arguments
      tool_result  — what came back
      final        — the answer
      error        — the turn failed; carries a human-readable explanation
    """
    yield _event("status", message="Connecting to the vfrplan MCP server…")

    try:
        async with mcp_session(config.mcp) as session:
            tool_list = (await session.list_tools()).tools
            # wikipedia_lookup is not an MCP tool: it is the one local, non-AIP
            # exception PROMPT.md §1 carves out, used only to find an ICAO code
            # find_aerodromes could not (see §3).
            tools = to_openai_tools(tool_list) + [WIKIPEDIA_TOOL_SCHEMA]
            # Keyed by tool name so _sanitize_arguments can look up, per call,
            # what that specific tool's own JSON Schema says a date/time
            # argument should look like (see _sanitize_arguments docstring).
            schema_by_tool = {t["function"]["name"]: t["function"].get("parameters") for t in tools}
            yield _event(
                "tools_loaded",
                names=sorted(tool.name for tool in tool_list) + [WIKIPEDIA_TOOL_NAME],
                count=len(tool_list) + 1,
            )

            messages: list[dict[str, Any]] = [{"role": "system", "content": system_prompt()}]
            for turn in history or []:
                role = turn.get("role")
                content = turn.get("content")
                if role in ("user", "assistant") and content:
                    messages.append({"role": role, "content": content})
            messages.append({"role": "user", "content": user_message})

            for step in range(1, config.max_steps + 1):
                yield _event("status", message=f"Step {step} — asking the model…", step=step)

                # chat_completion() retries a 502/504 internally (see llm.py) —
                # a big briefing's later steps can take the model over a
                # minute, right where an intermediary gateway is prone to
                # give up before this app's own, much longer, timeout does.
                # It has no way to yield mid-call, so it reports each retry
                # through this callback instead; queue them and flush after,
                # whether the call goes on to succeed or ultimately fail, so
                # the pilot sees why a step took a while rather than a silent
                # stall followed by either an answer or an error.
                retry_events: list[dict[str, Any]] = []

                def _note_retry(attempt: int, status: int, _step: int = step) -> None:
                    retry_events.append(
                        _event(
                            "status",
                            step=_step,
                            message=(
                                f"Model gateway returned {status} — retrying "
                                f"(attempt {attempt + 1})…"
                            ),
                        )
                    )

                try:
                    response = await chat_completion(
                        config.llm, messages, tools, on_retry=_note_retry
                    )
                except LLMError as exc:
                    for ev in retry_events:
                        yield ev
                    yield _event("error", message=str(exc))
                    return

                for ev in retry_events:
                    yield ev

                message = extract_message(response)

                reasoning = extract_reasoning(message)
                if reasoning:
                    yield _event("reasoning", step=step, content=reasoning)

                tool_calls = message.get("tool_calls") or []
                content = (message.get("content") or "").strip()

                if not tool_calls:
                    # No tool call means the model considers itself done. An
                    # empty answer at that point is a failure, not a result:
                    # silently showing a blank bubble would read as "nothing
                    # to report", which for a flight briefing is the worst
                    # possible way to be wrong.
                    if content:
                        yield _event("final", content=content)
                    else:
                        yield _event(
                            "error",
                            message=(
                                "The model finished without producing an answer. "
                                "This usually means it is not handling tool results "
                                "well — try a different model, or rephrase the request."
                            ),
                        )
                    return

                if content:
                    yield _event("thinking", step=step, content=content)

                # Echo the assistant turn back verbatim: the tool results that
                # follow are only valid as replies to these exact call ids.
                messages.append(
                    {
                        "role": "assistant",
                        "content": message.get("content") or "",
                        "tool_calls": tool_calls,
                    }
                )

                for call in tool_calls:
                    function = call.get("function") or {}
                    name = function.get("name") or "unknown_tool"
                    call_id = call.get("id") or name
                    arguments, parse_error = _parse_arguments(function.get("arguments"))
                    if not parse_error:
                        arguments = _sanitize_arguments(arguments, schema_by_tool.get(name))

                    yield _event("tool_call", step=step, name=name, arguments=arguments)

                    if parse_error:
                        failure = f"Tool call rejected: {parse_error}"
                        yield _event("tool_result", step=step, name=name, content=failure, ok=False)
                        messages.append(
                            {"role": "tool", "tool_call_id": call_id, "content": failure}
                        )
                        continue

                    try:
                        if name == WIKIPEDIA_TOOL_NAME:
                            # Local tool, not routed through the MCP session —
                            # see the comment where `tools` is built above.
                            text = await wikipedia_lookup(arguments.get("query", ""))
                            ok = True
                        else:
                            result = await session.call_tool(name, arguments)
                            text = result_to_text(result)
                            ok = not getattr(result, "isError", False)
                    except Exception as exc:  # noqa: BLE001 — surfaced to the model, not swallowed
                        text = f"The tool call failed: {type(exc).__name__}: {exc}"
                        ok = False

                    yield _event("tool_result", step=step, name=name, content=text, ok=ok)
                    messages.append({"role": "tool", "tool_call_id": call_id, "content": text})

            yield _event(
                "error",
                message=(
                    f"Stopped after {config.max_steps} steps without a final answer. "
                    "Raise AGENT_MAX_STEPS if this request genuinely needs more, or "
                    "check the tool results above for a call that keeps failing."
                ),
            )

    except Exception as exc:  # noqa: BLE001 — the turn is over either way; say why
        yield _event(
            "error",
            message=(
                f"Could not reach the vfrplan MCP server at {config.mcp.url}: "
                f"{type(exc).__name__}: {exc}\n\n"
                "Check that the vfrplan-mcp Space is running and that the token "
                "in use has read access to it."
            ),
        )

	"""The agent loop: model → tool calls → model → ... → answer.

	Written out by hand rather than taken from a framework, because the whole
	value of this screen is watching the loop happen: which tool, with which
	arguments, what came back, what the model made of it. A framework that owns
	the loop also owns what you are allowed to see of it.

	Every stage yields an event; main.py forwards them to the browser as SSE.
	Nothing here writes to the transport, so the loop stays testable on its own.
	"""

	from __future__ import annotations

	import json
	import re
	from collections.abc import AsyncIterator
	from pathlib import Path
	from typing import Any

	from .config import AppConfig
	from .llm import LLMError, chat_completion, extract_message, extract_reasoning
	from .mcp_tools import mcp_session, result_to_text, to_openai_tools
	from .wikipedia_tool import WIKIPEDIA_TOOL_NAME, WIKIPEDIA_TOOL_SCHEMA, wikipedia_lookup

	PROMPT_PATH = Path(__file__).resolve().parent.parent / "PROMPT.md"

	# Fallback only. The real rules live in PROMPT.md, shipped alongside this app
	# and copied unmodified from the vfrplan project — it is what stops the model
	# inventing figures it does not have. If the file goes missing the agent still
	# runs, but the safety-critical parts of its behaviour are gone, so the
	# fallback says so out loud rather than pretending nothing happened.
	FALLBACK_PROMPT = """You are a VFR flight planning assistant for Europe.
	Use the provided tools for every aeronautical fact. Never invent aerodrome
	data, weather, NOTAMs or performance figures. Report every assumption and
	warning the tools return. Ask the pilot for missing inputs instead of
	guessing them. Fuel, weather minima and go/no-go remain the pilot's decisions.
	(Note: PROMPT.md was not found, so these are reduced fallback instructions.)"""


	def system_prompt() -> str:
	try:
	text = PROMPT_PATH.read_text(encoding="utf-8").strip()
	return text or FALLBACK_PROMPT
	except OSError:
	return FALLBACK_PROMPT


	def _event(kind: str, **data: Any) -> dict[str, Any]:
	return {"type": kind, **data}


	# Matches an ISO date, optionally followed by a time and/or a trailing 'Z' /
	# numeric UTC offset, e.g. "2026-08-18", "2026-08-18T11:00:00" or
	# "2026-08-14T09:00:00+00:00".
	_DATETIME_RE = re.compile(
	r"^(?P<date>\d{4}-\d{2}-\d{2})"
	r"(?:[T ](?P<time>\d{2}:\d{2}(?::\d{2})?)(?:\.\d+)?)?"
	r"(?:Z\|[+-]\d{2}:?\d{2})?$"
	)


	def _sanitize_arguments(
	arguments: dict[str, Any], schema: dict[str, Any] \| None = None
	) -> dict[str, Any]:
	"""Reshape ISO-datetime-looking string arguments to match a tool's own schema.

	Two plan_route failures have been seen, both the same root cause under a
	different face. First "Invalid isoformat string:
	'2026-08-14T09:00:00+00:00'" — a trailing offset its parser does not
	accept. Then, after that fix, "Invalid isoformat string:
	'2026-08-18T11:00:00'" on a string with no offset at all, which is
	exactly the error Python's own date.fromisoformat() raises when handed a
	full datetime: the tool has a date-only parameter, and the model was
	sending it a combined date+time. A per-tool string field can't be
	hardcoded for this — plan_route's real parameter names are behind the
	private vfrplan-mcp Space and were not confirmed live — so this reads
	each tool's own declared JSON Schema `format` ("date"/"time"/anything
	else) for the argument key in question and reshapes the value to match:
	date-only down to just the date, time-only down to just the time, and
	everything else keeps date+time together with any tz suffix dropped
	(PROMPT.md §2 has every time treated as local by design, so an offset is
	never meaningful here regardless). With no schema, or no format on that
	property, the old date+time-joined-minus-tz behaviour applies.
	"""
	properties = (schema or {}).get("properties") or {} if isinstance(schema, dict) else {}
	sanitized = dict(arguments)
	for key, value in arguments.items():
	if not isinstance(value, str):
	continue
	match = _DATETIME_RE.match(value.strip())
	if not match:
	continue
	date_part = match.group("date")
	time_part = match.group("time")
	fmt = (properties.get(key) or {}).get("format")
	if fmt == "date":
	sanitized[key] = date_part
	elif fmt == "time":
	if time_part:
	sanitized[key] = time_part
	else:
	sanitized[key] = f"{date_part}T{time_part}" if time_part else date_part
	return sanitized


	def _parse_arguments(raw: Any) -> tuple[dict[str, Any], str \| None]:
	"""Normalise tool-call arguments into a dict.

	The API says this is a JSON string, but providers vary: some send a dict
	already, some send "" for a no-argument call, and some send malformed
	JSON. A parse failure is returned rather than raised so the loop can hand
	it back to the model as a tool error and let it retry — which works far
	more often than aborting the turn.
	"""
	if isinstance(raw, dict):
	return raw, None
	if raw in (None, "", "null"):
	return {}, None
	try:
	parsed = json.loads(raw)
	except (TypeError, ValueError) as exc:
	return {}, f"arguments were not valid JSON: {exc}"
	if not isinstance(parsed, dict):
	return {}, f"arguments must be a JSON object, got {type(parsed).__name__}"
	return parsed, None


	async def run_agent(
	config: AppConfig,
	user_message: str,
	history: list[dict[str, str]] \| None = None,
	) -> AsyncIterator[dict[str, Any]]:
	"""Run one turn, yielding events as they happen.

	Event kinds consumed by the frontend:
	status — a line for the activity strip
	tools_loaded — the tool inventory, once, after connecting
	reasoning — the model's private thinking for this step
	thinking — assistant prose emitted alongside tool calls
	tool_call — a call about to be made, with its arguments
	tool_result — what came back
	final — the answer
	error — the turn failed; carries a human-readable explanation
	"""
	yield _event("status", message="Connecting to the vfrplan MCP server…")

	try:
	async with mcp_session(config.mcp) as session:
	tool_list = (await session.list_tools()).tools
	# wikipedia_lookup is not an MCP tool: it is the one local, non-AIP
	# exception PROMPT.md §1 carves out, used only to find an ICAO code
	# find_aerodromes could not (see §3).
	tools = to_openai_tools(tool_list) + [WIKIPEDIA_TOOL_SCHEMA]
	# Keyed by tool name so _sanitize_arguments can look up, per call,
	# what that specific tool's own JSON Schema says a date/time
	# argument should look like (see _sanitize_arguments docstring).
	schema_by_tool = {t["function"]["name"]: t["function"].get("parameters") for t in tools}
	yield _event(
	"tools_loaded",
	names=sorted(tool.name for tool in tool_list) + [WIKIPEDIA_TOOL_NAME],
	count=len(tool_list) + 1,
	)

	messages: list[dict[str, Any]] = [{"role": "system", "content": system_prompt()}]
	for turn in history or []:
	role = turn.get("role")
	content = turn.get("content")
	if role in ("user", "assistant") and content:
	messages.append({"role": role, "content": content})
	messages.append({"role": "user", "content": user_message})

	for step in range(1, config.max_steps + 1):
	yield _event("status", message=f"Step {step} — asking the model…", step=step)

	# chat_completion() retries a 502/504 internally (see llm.py) —
	# a big briefing's later steps can take the model over a
	# minute, right where an intermediary gateway is prone to
	# give up before this app's own, much longer, timeout does.
	# It has no way to yield mid-call, so it reports each retry
	# through this callback instead; queue them and flush after,
	# whether the call goes on to succeed or ultimately fail, so
	# the pilot sees why a step took a while rather than a silent
	# stall followed by either an answer or an error.
	retry_events: list[dict[str, Any]] = []

	def _note_retry(attempt: int, status: int, _step: int = step) -> None:
	retry_events.append(
	_event(
	"status",
	step=_step,
	message=(
	f"Model gateway returned {status} — retrying "
	f"(attempt {attempt + 1})…"
	),
	)
	)

	try:
	response = await chat_completion(
	config.llm, messages, tools, on_retry=_note_retry
	)
	except LLMError as exc:
	for ev in retry_events:
	yield ev
	yield _event("error", message=str(exc))
	return

	for ev in retry_events:
	yield ev

	message = extract_message(response)

	reasoning = extract_reasoning(message)
	if reasoning:
	yield _event("reasoning", step=step, content=reasoning)

	tool_calls = message.get("tool_calls") or []
	content = (message.get("content") or "").strip()

	if not tool_calls:
	# No tool call means the model considers itself done. An
	# empty answer at that point is a failure, not a result:
	# silently showing a blank bubble would read as "nothing
	# to report", which for a flight briefing is the worst
	# possible way to be wrong.
	if content:
	yield _event("final", content=content)
	else:
	yield _event(
	"error",
	message=(
	"The model finished without producing an answer. "
	"This usually means it is not handling tool results "
	"well — try a different model, or rephrase the request."
	),
	)
	return

	if content:
	yield _event("thinking", step=step, content=content)

	# Echo the assistant turn back verbatim: the tool results that
	# follow are only valid as replies to these exact call ids.
	messages.append(
	{
	"role": "assistant",
	"content": message.get("content") or "",
	"tool_calls": tool_calls,
	}
	)

	for call in tool_calls:
	function = call.get("function") or {}
	name = function.get("name") or "unknown_tool"
	call_id = call.get("id") or name
	arguments, parse_error = _parse_arguments(function.get("arguments"))
	if not parse_error:
	arguments = _sanitize_arguments(arguments, schema_by_tool.get(name))

	yield _event("tool_call", step=step, name=name, arguments=arguments)

	if parse_error:
	failure = f"Tool call rejected: {parse_error}"
	yield _event("tool_result", step=step, name=name, content=failure, ok=False)
	messages.append(
	{"role": "tool", "tool_call_id": call_id, "content": failure}
	)
	continue

	try:
	if name == WIKIPEDIA_TOOL_NAME:
	# Local tool, not routed through the MCP session —
	# see the comment where `tools` is built above.
	text = await wikipedia_lookup(arguments.get("query", ""))
	ok = True
	else:
	result = await session.call_tool(name, arguments)
	text = result_to_text(result)
	ok = not getattr(result, "isError", False)
	except Exception as exc: # noqa: BLE001 — surfaced to the model, not swallowed
	text = f"The tool call failed: {type(exc).__name__}: {exc}"
	ok = False

	yield _event("tool_result", step=step, name=name, content=text, ok=ok)
	messages.append({"role": "tool", "tool_call_id": call_id, "content": text})

	yield _event(
	"error",
	message=(
	f"Stopped after {config.max_steps} steps without a final answer. "
	"Raise AGENT_MAX_STEPS if this request genuinely needs more, or "
	"check the tool results above for a call that keeps failing."
	),
	)

	except Exception as exc: # noqa: BLE001 — the turn is over either way; say why
	yield _event(
	"error",
	message=(
	f"Could not reach the vfrplan MCP server at {config.mcp.url}: "
	f"{type(exc).__name__}: {exc}\n\n"
	"Check that the vfrplan-mcp Space is running and that the token "
	"in use has read access to it."
	),
	)