FEAT: Add --output console | json flag in scenario-results CLI command

pyrit/cli/_cli_args.py

@@ -91,6 +91,44 @@ def parse_scenario_result_view(raw: str) -> ScenarioResultView:
         raise argparse.ArgumentTypeError(f"invalid view '{raw}' (choose from {valid})") from None
 
 
+class OutputFormat(str, Enum):
+    """
+    Serialization format for a ``scenario-results`` render.
+
+    Like ``ScenarioResultView``, this lives in this parse-time-safe module so the
+    argument parsers can reference it without importing ``pydantic``. ``html`` is
+    intentionally omitted until a later phase adds its renderer.
+    """
+
+    #: Human-readable console text (the default).
+    CONSOLE = "console"
+    #: Machine-readable JSON via the payload's ``model_dump_json``.
+    JSON = "json"
+
+
+def parse_output_format(raw: str) -> OutputFormat:
+    """
+    Parse an ``--output`` token into an ``OutputFormat``.
+
+    Used as an argparse ``type=`` so an invalid value produces an error that
+    lists the valid format names (mirroring ``parse_scenario_result_view``).
+
+    Args:
+        raw (str): The raw ``--output`` token.
+
+    Returns:
+        OutputFormat: The matching format.
+
+    Raises:
+        argparse.ArgumentTypeError: If *raw* is not a valid format name.
+    """
+    try:
+        return OutputFormat(raw)
+    except ValueError:
+        valid = ", ".join(fmt.value for fmt in OutputFormat)
+        raise argparse.ArgumentTypeError(f"invalid output '{raw}' (choose from {valid})") from None
+
+
 # ---------------------------------------------------------------------------
 # Pure validators
 # ---------------------------------------------------------------------------
@@ -465,6 +503,14 @@ def add_results_arguments(*, parser: argparse.ArgumentParser) -> None:
         help="Show at most N attacks (ignored for --view overview; defaults to 5 for "
         "--view conversations/full when no --attack-result-ids is given)",
     )
+    group.add_argument(
+        "--output",
+        type=parse_output_format,
+        default=OutputFormat.CONSOLE,
+        metavar="{" + ",".join(fmt.value for fmt in OutputFormat) + "}",
+        help="Output format: 'console' (human-readable, default) or 'json' "
+        "(machine-readable; informational notes go to stderr so stdout stays parseable)",
+    )
 
 
 def build_scenario_results_parser() -> argparse.ArgumentParser:

pyrit/cli/_output.py

@@ -16,7 +16,12 @@
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
-    from pyrit.cli._results import AttacksTablePayload, ConversationsPayload, TranscriptMessage
+    from pyrit.cli._results import (
+        AttacksTablePayload,
+        ConversationsPayload,
+        ResultsPayload,
+        TranscriptMessage,
+    )
     from pyrit.models import ScenarioResult
     from pyrit.models.catalog import (
         RegisteredInitializer,
@@ -457,7 +462,7 @@ def print_attacks_table(*, payload: AttacksTablePayload) -> None:
         print(f"       technique: {row.atomic_attack_name}")
         print(f"       objective: {row.objective}")
 
-    shown = len(payload.rows)
+    shown = payload.shown
     if shown < payload.total:
         print(f"\nShowing {shown} of {payload.total} attacks (use --limit to change).")
     else:
@@ -487,7 +492,7 @@ def print_conversations(*, payload: ConversationsPayload) -> None:
         print(f"       objective: {convo.objective}")
         _print_transcript(messages=convo.messages)
 
-    shown = len(payload.conversations)
+    shown = payload.shown
     if shown < payload.total:
         print(f"\nShowing {shown} of {payload.total} attacks (use --limit or --attack-result-ids to change).")
     else:
@@ -514,6 +519,50 @@ def _print_transcript(*, messages: list[TranscriptMessage]) -> None:
                 print(_wrap(text=f"rationale: {message.score.rationale}", indent="           "))
 
 
+def print_results_json(*, payload: ResultsPayload) -> None:
+    """
+    Serialize a results payload as JSON to stdout.
+
+    Emits the payload's own ``model_dump_json`` so the JSON always matches its
+    fields; informational notes are routed to stderr elsewhere so this stdout
+    document stays parseable.
+
+    Args:
+        payload (ResultsPayload): The view payload to serialize.
+    """
+    print(payload.model_dump_json(indent=2))
+
+
+async def print_results_console_async(*, result: ScenarioResult, payload: ResultsPayload) -> None:
+    """
+    Render a results payload to the console, dispatching on its concrete type.
+
+    The ``overview`` payload keeps the rich framework pretty printer (which reads
+    *result* directly), so its typed payload is used only by the JSON format; the
+    other views render from *payload*.
+
+    Args:
+        result (ScenarioResult): The scenario result, used by the overview printer.
+        payload (ResultsPayload): The payload built for the requested view.
+    """
+    from pyrit.cli._results import (
+        AttacksTablePayload,
+        ConversationsPayload,
+        FullPayload,
+        ScenarioOverviewPayload,
+    )
+
+    if isinstance(payload, ScenarioOverviewPayload):
+        await print_scenario_result_async(result=result)
+    elif isinstance(payload, AttacksTablePayload):
+        print_attacks_table(payload=payload)
+    elif isinstance(payload, ConversationsPayload):
+        print_conversations(payload=payload)
+    elif isinstance(payload, FullPayload):
+        print_attacks_table(payload=payload.attacks)
+        print_conversations(payload=payload.conversations)
+
+
 # ---------------------------------------------------------------------------
 # Scenario run history
 # ---------------------------------------------------------------------------

pyrit/cli/_results.py

@@ -16,9 +16,10 @@
 
 from __future__ import annotations
 
+import sys
 from typing import TYPE_CHECKING, Any
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, computed_field
 
 from pyrit.cli._cli_args import ScenarioResultView
 
@@ -33,6 +34,36 @@
 _DEFAULT_HEAVY_VIEW_LIMIT = 5
 
 
+class GroupSummary(BaseModel):
+    """One display group's totals and success rate in the ``overview`` view."""
+
+    name: str
+    total_results: int
+    success_rate: int
+
+
+class ScenarioOverviewPayload(BaseModel):
+    """
+    The ``overview`` view: scenario-level aggregates and per-group success rates.
+
+    Captures the numbers the console pretty printer computes inline, but as a
+    typed object so ``--output json`` can emit the same aggregates the console
+    shows. The console path still renders via the framework pretty printer while
+    JSON serializes this payload, so the two compute independently; the aggregates
+    are kept in lockstep by ``test_overview_payload_matches_framework_printer``.
+    """
+
+    scenario_result_id: str
+    scenario_name: str
+    scenario_version: int
+    pyrit_version: str
+    total_techniques: int
+    total_attack_results: int
+    overall_success_rate: int
+    unique_objectives: int
+    groups: list[GroupSummary] = Field(default_factory=list)
+
+
 class AttackRow(BaseModel):
     """A single attack result rendered as one row of the attacks table."""
 
@@ -49,14 +80,21 @@ class AttacksTablePayload(BaseModel):
     The ``attacks`` view: one row per attack result in a scenario run.
 
     ``total`` is the number of attacks that matched the selection before
-    ``--limit`` was applied; ``len(rows)`` is how many are actually included.
-    Exposing both lets any renderer show a "showing N of M" note.
+    ``--limit`` was applied; ``shown`` (== ``len(rows)``) is how many are actually
+    included. Exposing both lets any renderer — or a JSON consumer — detect when
+    the list was truncated.
     """
 
     scenario_result_id: str
     rows: list[AttackRow] = Field(default_factory=list)
     total: int = 0
 
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def shown(self) -> int:
+        """Rows actually included, i.e. ``total`` minus any ``--limit`` truncation."""
+        return len(self.rows)
+
 
 class TranscriptScore(BaseModel):
     """
@@ -98,14 +136,40 @@ class ConversationsPayload(BaseModel):
     The ``conversations`` view: the main-conversation transcript per attack.
 
     ``total`` is the number of attacks that matched the selection before
-    ``--limit`` was applied; ``len(conversations)`` is how many are actually
-    included. Exposing both lets any renderer show a "showing N of M" note.
+    ``--limit`` was applied; ``shown`` (== ``len(conversations)``) is how many are
+    actually included. Exposing both lets any renderer — or a JSON consumer —
+    detect when the list was truncated.
     """
 
     scenario_result_id: str
     conversations: list[AttackConversation] = Field(default_factory=list)
     total: int = 0
 
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def shown(self) -> int:
+        """Transcripts actually included, i.e. ``total`` minus any ``--limit`` truncation."""
+        return len(self.conversations)
+
+
+class FullPayload(BaseModel):
+    """
+    The ``full`` view: the ``attacks`` table together with the ``conversations``.
+
+    Bundling both sub-payloads into one model lets ``--output json`` emit a
+    single well-formed document (rather than two concatenated ones) while the
+    console renderer prints the table then the transcripts.
+    """
+
+    scenario_result_id: str
+    attacks: AttacksTablePayload
+    conversations: ConversationsPayload
+
+
+#: Any of the per-view payloads a ``scenario-results`` render can produce. Named
+#: so the builder, the JSON renderer, and both front-ends share one type.
+ResultsPayload = ScenarioOverviewPayload | AttacksTablePayload | ConversationsPayload | FullPayload
+
 
 def resolve_view(*, view: ScenarioResultView | None) -> ScenarioResultView:
     """
@@ -155,20 +219,61 @@ def apply_view_limit_policy(
     """
     if view is ScenarioResultView.OVERVIEW:
         if limit is not None:
-            print("Note: --limit has no effect with --view overview; ignoring it.")
+            print("Note: --limit has no effect with --view overview; ignoring it.", file=sys.stderr)
         return None
     if view in (ScenarioResultView.CONVERSATIONS, ScenarioResultView.FULL):
         if limit is None and not attack_result_ids:
             print(
                 f"Note: no --attack-result-ids or --limit given; showing at most "
                 f"{_DEFAULT_HEAVY_VIEW_LIMIT} conversations. Pass --limit or "
-                "--attack-result-ids to see more."
+                "--attack-result-ids to see more.",
+                file=sys.stderr,
             )
             return _DEFAULT_HEAVY_VIEW_LIMIT
         return limit
     return limit
 
 
+def build_overview_payload(*, result: ScenarioResult, scenario_result_id: str) -> ScenarioOverviewPayload:
+    """
+    Build the ``overview`` payload from an already-fetched scenario result.
+
+    Reproduces the console pretty printer's "Overall Statistics" and "Per-Group
+    Breakdown" by calling the same ``ScenarioResult`` aggregate methods. The
+    console printer computes the same numbers independently, so
+    ``test_overview_payload_matches_framework_printer`` locks the two in sync and
+    fails if the printer's math changes.
+
+    Args:
+        result (ScenarioResult): The scenario result to aggregate.
+        scenario_result_id (str): The run id, echoed back on the payload.
+
+    Returns:
+        ScenarioOverviewPayload: The scenario-level aggregates and per-group rates.
+    """
+    from pyrit.models import AttackOutcome
+
+    groups: list[GroupSummary] = []
+    for group_name, group_results in result.get_display_groups().items():
+        total_group = len(group_results)
+        successful = sum(1 for attack_result in group_results if attack_result.outcome == AttackOutcome.SUCCESS)
+        success_rate = int((successful / total_group) * 100) if total_group else 0
+        groups.append(GroupSummary(name=group_name, total_results=total_group, success_rate=success_rate))
+
+    total_attack_results = sum(len(results) for results in result.attack_results.values())
+    return ScenarioOverviewPayload(
+        scenario_result_id=scenario_result_id,
+        scenario_name=result.scenario_name,
+        scenario_version=result.scenario_version,
+        pyrit_version=result.pyrit_version,
+        total_techniques=len(result.get_techniques_used()),
+        total_attack_results=total_attack_results,
+        overall_success_rate=result.objective_achieved_rate(),
+        unique_objectives=len(result.get_objectives()),
+        groups=groups,
+    )
+
+
 def build_attacks_table_payload(
     *,
     result: ScenarioResult,
@@ -279,6 +384,76 @@ async def build_conversations_payload_async(
     )
 
 
+async def build_results_payload_async(
+    *,
+    view: ScenarioResultView,
+    result: ScenarioResult,
+    client: PyRITApiClient,
+    scenario_result_id: str,
+    attack_result_ids: list[str] | None = None,
+    limit: int | None = None,
+) -> ResultsPayload:
+    """
+    Build the payload for *view*, dispatching to the per-view builder.
+
+    A single entry point so both front-ends ask for "the payload for this view"
+    without special-casing which one: the ``overview`` / ``attacks`` builders are
+    pure and local, while ``conversations`` / ``full`` fetch each attack's
+    transcript, and this hides that async/sync split behind one ``await``.
+
+    Args:
+        view (ScenarioResultView): The resolved view to build.
+        result (ScenarioResult): The already-fetched scenario result.
+        client (PyRITApiClient): Client used by the transcript-fetching views.
+        scenario_result_id (str): The run id, echoed back on the payload.
+        attack_result_ids (list[str] | None): Restrict to these attack ids.
+            Defaults to None (all attacks).
+        limit (int | None): Effective attack cap from ``apply_view_limit_policy``.
+            Defaults to None.
+
+    Returns:
+        ResultsPayload: The payload matching *view*.
+    """
+    if view is ScenarioResultView.OVERVIEW:
+        return build_overview_payload(result=result, scenario_result_id=scenario_result_id)
+
+    if view is ScenarioResultView.ATTACKS:
+        return build_attacks_table_payload(
+            result=result,
+            scenario_result_id=scenario_result_id,
+            attack_result_ids=attack_result_ids,
+            limit=limit,
+        )
+
+    if view is ScenarioResultView.CONVERSATIONS:
+        return await build_conversations_payload_async(
+            result=result,
+            client=client,
+            scenario_result_id=scenario_result_id,
+            attack_result_ids=attack_result_ids,
+            limit=limit,
+        )
+
+    attacks = build_attacks_table_payload(
+        result=result,
+        scenario_result_id=scenario_result_id,
+        attack_result_ids=attack_result_ids,
+        limit=limit,
+    )
+    conversations = await build_conversations_payload_async(
+        result=result,
+        client=client,
+        scenario_result_id=scenario_result_id,
+        attack_result_ids=attack_result_ids,
+        limit=limit,
+    )
+    return FullPayload(
+        scenario_result_id=scenario_result_id,
+        attacks=attacks,
+        conversations=conversations,
+    )
+
+
 def _select_attacks(*, result: ScenarioResult, attack_result_ids: list[str] | None) -> list[tuple[str, AttackResult]]:
     """
     Return ``(atomic_attack_name, attack_result)`` pairs, optionally id-filtered.