diff --git a/build_scripts/export_adversarial_benchmark_result.py b/build_scripts/export_adversarial_benchmark_result.py index cf7cd13c6d..6941258997 100644 --- a/build_scripts/export_adversarial_benchmark_result.py +++ b/build_scripts/export_adversarial_benchmark_result.py @@ -5,15 +5,12 @@ import argparse import asyncio -import contextlib import csv import json from collections import Counter, defaultdict from pathlib import Path from typing import Any -from pyrit.cli._output import print_attacks_table -from pyrit.cli._results import build_attacks_table_payload from pyrit.memory import CentralMemory from pyrit.models import ScenarioResult from pyrit.output.scenario_result.pretty import PrettyScenarioResultMemoryPrinter @@ -46,16 +43,37 @@ async def _write_overview_async(*, result: ScenarioResult, output_dir: Path) -> await printer.write_async(result) -def _write_attacks(*, result: ScenarioResult, output_dir: Path) -> None: +def _attack_rows(*, result: ScenarioResult) -> list[dict[str, Any]]: + """Build machine-readable per-attack rows from the embedded attack results.""" + rows: list[dict[str, Any]] = [] + for atomic_attack_name, attacks in result.attack_results.items(): + for attack in attacks: + score = attack.last_score + score_value = None + if score is not None: + score_value = score.score_value if score.score_value is not None else score.status.value + rows.append( + { + "attack_result_id": attack.attack_result_id, + "atomic_attack_name": atomic_attack_name, + "objective": attack.objective, + "outcome": attack.outcome.value, + "executed_turns": attack.executed_turns, + "score_value": score_value, + } + ) + return rows + + +async def _write_attacks_async(*, result: ScenarioResult, output_dir: Path) -> None: """Write machine-readable and console-style partial attack tables.""" - payload = build_attacks_table_payload( - result=result, - scenario_result_id=str(result.id), - ) - (output_dir / "attacks.json").write_text(payload.model_dump_json(indent=2), encoding="utf-8") - with open(output_dir / "attacks.txt", "w", encoding="utf-8") as output: - with contextlib.redirect_stdout(output): - print_attacks_table(payload=payload) + rows = _attack_rows(result=result) + document = {"scenario_result_id": str(result.id), "rows": rows, "total": len(rows)} + (output_dir / "attacks.json").write_text(json.dumps(document, indent=2), encoding="utf-8") + + sink = FileSink(path=output_dir / "attacks.txt") + printer = PrettyScenarioResultMemoryPrinter(sink=sink, enable_colors=False) + await printer.write_async(result, view="attacks") def _build_technique_metrics(*, result: ScenarioResult) -> list[dict[str, Any]]: @@ -143,7 +161,7 @@ async def _export_async(*, scenario_result_id: str, output_dir: Path) -> None: result = await _load_result_async(scenario_result_id=scenario_result_id) output_dir.mkdir(parents=True, exist_ok=True) await _write_overview_async(result=result, output_dir=output_dir) - await asyncio.to_thread(_write_attacks, result=result, output_dir=output_dir) + await _write_attacks_async(result=result, output_dir=output_dir) await asyncio.to_thread(_write_technique_metrics, result=result, output_dir=output_dir) diff --git a/pyrit/cli/_output.py b/pyrit/cli/_output.py index 0883ee5d46..3c3fa36030 100644 --- a/pyrit/cli/_output.py +++ b/pyrit/cli/_output.py @@ -16,7 +16,7 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: - from pyrit.cli._results import AttacksTablePayload, ConversationsPayload, TranscriptMessage + from pyrit.cli.api_client import PyRITApiClient from pyrit.models import ScenarioResult from pyrit.models.catalog import ( RegisteredInitializer, @@ -425,94 +425,67 @@ async def print_scenario_result_async(*, result: ScenarioResult) -> None: "undetermined": None, } -# Per-role transcript colors, mirroring PrettyConversationPrinter's palette so the -# thin-client transcript reads like the framework's own conversation output. -_ROLE_COLORS = { - "user": "blue", - "assistant": "yellow", - "system": "magenta", -} - - -def print_attacks_table(*, payload: AttacksTablePayload) -> None: - """ - Print the per-attack table for a scenario run. - Args: - payload (AttacksTablePayload): The rows to render plus the pre-limit total. +async def print_conversations_async( + *, + result: ScenarioResult, + client: PyRITApiClient, + scenario_result_id: str, + attack_result_ids: list[str] | None = None, + limit: int | None = None, +) -> None: """ - if not payload.rows: - print(f"\nNo attack results found for scenario {payload.scenario_result_id}.") - return - - _header(f"Attack Results — scenario {payload.scenario_result_id}") - for index, row in enumerate(payload.rows, start=1): - outcome = row.outcome.upper() - score = row.score_value if row.score_value is not None else "—" - _cprint( - f" {index}. [{outcome}] turns={row.executed_turns} score={score}", - color=_OUTCOME_COLORS.get(row.outcome), - bold=True, - ) - print(f" id: {row.attack_result_id}") - print(f" technique: {row.atomic_attack_name}") - print(f" objective: {row.objective}") - - shown = len(payload.rows) - if shown < payload.total: - print(f"\nShowing {shown} of {payload.total} attacks (use --limit to change).") - else: - print(f"\nTotal attacks: {payload.total}") - + Print each attack's main-conversation transcript, rendered by the framework. -def print_conversations(*, payload: ConversationsPayload) -> None: - """ - Print the per-attack main-conversation transcripts for a scenario run. + Reuses ``pyrit.output``'s conversation printer via a REST-backed source, so the + CLI transcript matches the framework's own conversation output. The per-attack + fetch loop is gated by *limit* (network calls, not just rendered rows). Args: - payload (ConversationsPayload): The transcripts to render plus the - pre-limit total. - """ - if not payload.conversations: - print(f"\nNo conversations found for scenario {payload.scenario_result_id}.") + result (ScenarioResult): The scenario result whose attacks to inspect. + client (PyRITApiClient): Client used to fetch each conversation's messages. + scenario_result_id (str): The run id, echoed in the header. + attack_result_ids (list[str] | None): Restrict to these attack ids. Defaults to None. + limit (int | None): Maximum number of attacks to fetch and render. Defaults to None. + """ + from pyrit.cli._results import _objective_scorer_key, _select_attacks + from pyrit.cli._sources import RestApiConversationSource + from pyrit.output.conversation.pretty import PrettyConversationPrinter + + selected = _select_attacks(result=result, attack_result_ids=attack_result_ids) + total = len(selected) + if limit is not None: + selected = selected[:limit] + + if not selected: + print(f"\nNo conversations found for scenario {scenario_result_id}.") return - _header(f"Conversations — scenario {payload.scenario_result_id}") - for index, convo in enumerate(payload.conversations, start=1): + objective_hash, objective_class = _objective_scorer_key(result=result) + _header(f"Conversations — scenario {scenario_result_id}") + for index, (atomic_attack_name, attack_result) in enumerate(selected, start=1): _cprint( - f" {index}. [{convo.outcome.upper()}] {convo.atomic_attack_name}", - color=_OUTCOME_COLORS.get(convo.outcome), + f" {index}. [{attack_result.outcome.value.upper()}] {atomic_attack_name}", + color=_OUTCOME_COLORS.get(attack_result.outcome.value), bold=True, ) - print(f" id: {convo.attack_result_id}") - print(f" objective: {convo.objective}") - _print_transcript(messages=convo.messages) + print(f" id: {attack_result.attack_result_id}") + print(f" objective: {attack_result.objective}") + source = RestApiConversationSource( + client=client, + attack_result_id=attack_result.attack_result_id, + objective_hash=objective_hash, + objective_class=objective_class, + ) + messages = await source.get_messages_async(conversation_id=attack_result.conversation_id) + printer = PrettyConversationPrinter(source=source) + print(await printer.render_async(messages, include_scores=True)) - shown = len(payload.conversations) - if shown < payload.total: - print(f"\nShowing {shown} of {payload.total} attacks (use --limit or --attack-result-ids to change).") + shown = len(selected) + if shown < total: + print(f"\nShowing {shown} of {total} attacks (use --limit or --attack-result-ids to change).") else: - print(f"\nTotal attacks: {payload.total}") - - -def _print_transcript(*, messages: list[TranscriptMessage]) -> None: - """Print one attack's ordered messages with their optional scores.""" - if not messages: - print(" (no messages)") - return - for message in messages: - _cprint( - f" [{message.role.upper()}] (turn {message.turn})", - color=_ROLE_COLORS.get(message.role.lower()), - bold=True, - ) - print(_wrap(text=message.text, indent=" ")) - if message.score is not None: - value = message.score.value if message.score.value is not None else "—" - label = f"SCORE [{message.score.scorer}]" if message.score.scorer else "SCORE" - _cprint(f" {label}: {value}", color="magenta", bold=True) - if message.score.rationale: - print(_wrap(text=f"rationale: {message.score.rationale}", indent=" ")) + print(f"\nTotal attacks: {total}") # --------------------------------------------------------------------------- diff --git a/pyrit/cli/_results.py b/pyrit/cli/_results.py index 571fc885a3..9e1bb54c99 100644 --- a/pyrit/cli/_results.py +++ b/pyrit/cli/_results.py @@ -2,29 +2,23 @@ # Licensed under the MIT license. """ -Typed payloads and builders for the ``scenario-results`` command. +View resolution, ``--limit`` policy, and attack selection for the +``scenario-results`` command. -A *view* selects the data (one of these payloads); a *format* serializes it. -Keeping the payload a Pydantic model makes it the single source of truth: the -console renderer reads it today, and ``--output json`` will serialize the same -object in a later phase, so every format stays consistent. - -This module imports ``pydantic`` and is therefore loaded only from deferred -(post-parse) call sites, never on the CLI ``--help`` path. The lightweight -``ScenarioResultView`` enum lives in ``pyrit.cli._cli_args`` for that reason. +Rendering is delegated to ``pyrit.output`` (the scenario, attacks, and conversation +printers); this module holds only the CLI-side flag policy and the shared +attack-selection helpers. ``ScenarioResultView`` lives in ``pyrit.cli._cli_args`` +so the argument parsers can reference it cheaply. """ from __future__ import annotations -from typing import TYPE_CHECKING, Any - -from pydantic import BaseModel, Field +from typing import TYPE_CHECKING from pyrit.cli._cli_args import ScenarioResultView if TYPE_CHECKING: - from pyrit.cli.api_client import PyRITApiClient - from pyrit.models import AttackResult, ScenarioResult, Score + from pyrit.models import AttackResult, ScenarioResult #: Default cap on how many attacks the expensive views (``conversations`` / #: ``full``) render when the user gives neither ``--attack-result-ids`` nor @@ -33,80 +27,6 @@ _DEFAULT_HEAVY_VIEW_LIMIT = 5 -class AttackRow(BaseModel): - """A single attack result rendered as one row of the attacks table.""" - - attack_result_id: str - atomic_attack_name: str - objective: str - outcome: str - executed_turns: int - score_value: str | None = None - - -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. - """ - - scenario_result_id: str - rows: list[AttackRow] = Field(default_factory=list) - total: int = 0 - - -class TranscriptScore(BaseModel): - """ - The objective (top-level) score attached to a transcript message. - - A response is often scored by several scorers (refusal, objective, - composite, ...), but only the scenario's objective scorer determines - success, so the transcript surfaces just that one. ``scorer`` names it so a - bare ``true``/``false`` isn't ambiguous about which dimension it measures. - """ - - scorer: str | None = None - value: str | None = None - rationale: str | None = None - - -class TranscriptMessage(BaseModel): - """One message in a conversation transcript (role, turn, text, objective score).""" - - role: str - turn: int - text: str - score: TranscriptScore | None = None - - -class AttackConversation(BaseModel): - """The main-conversation transcript for a single attack result.""" - - attack_result_id: str - atomic_attack_name: str - objective: str - outcome: str - conversation_id: str - messages: list[TranscriptMessage] = Field(default_factory=list) - - -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. - """ - - scenario_result_id: str - conversations: list[AttackConversation] = Field(default_factory=list) - total: int = 0 - - def resolve_view(*, view: ScenarioResultView | None) -> ScenarioResultView: """ Resolve an optional ``--view`` value to a concrete view. @@ -169,132 +89,6 @@ def apply_view_limit_policy( return limit -def _row_score_value(*, score: Score | None) -> str | None: - """ - Render a score for the attacks table. - - Args: - score (Score | None): The attack's last score, if it had one. - - Returns: - str | None: The score value, its status when undetermined, or None. - """ - if score is None: - return None - # An undetermined score has no value, so name its status instead. - return score.score_value if score.score_value is not None else score.status.value - - -def build_attacks_table_payload( - *, - result: ScenarioResult, - scenario_result_id: str, - attack_result_ids: list[str] | None = None, - limit: int | None = None, -) -> AttacksTablePayload: - """ - Build the ``attacks`` payload from an already-fetched scenario result. - - Every ``AttackResult`` is already embedded in *result* (grouped by atomic - attack name), so no extra server calls are needed. ``--limit`` is applied - here, on the payload, rather than in a renderer, so that all output formats - honor it identically. - - Args: - result (ScenarioResult): The full scenario result to read attacks from. - scenario_result_id (str): The run id, echoed back on the payload. - attack_result_ids (list[str] | None): When provided, keep only attacks - whose id is in this set. Defaults to None (all attacks). - limit (int | None): Maximum number of rows to include. Defaults to None. - - Returns: - AttacksTablePayload: The rows plus the pre-limit total. - """ - selected = _select_attacks(result=result, attack_result_ids=attack_result_ids) - total = len(selected) - if limit is not None: - selected = selected[:limit] - - rows = [ - AttackRow( - attack_result_id=attack_result.attack_result_id, - atomic_attack_name=atomic_attack_name, - objective=attack_result.objective, - outcome=attack_result.outcome.value, - executed_turns=attack_result.executed_turns, - score_value=_row_score_value(score=attack_result.last_score), - ) - for atomic_attack_name, attack_result in selected - ] - return AttacksTablePayload(scenario_result_id=scenario_result_id, rows=rows, total=total) - - -async def build_conversations_payload_async( - *, - result: ScenarioResult, - client: PyRITApiClient, - scenario_result_id: str, - attack_result_ids: list[str] | None = None, - limit: int | None = None, -) -> ConversationsPayload: - """ - Build the ``conversations`` payload, fetching each attack's main transcript. - - Unlike the ``attacks`` builder (all data embedded, pure and sync), this view - needs one message fetch per attack, so the builder owns the I/O loop. Attack - selection and ``--limit`` are applied *before* fetching, so the effective - limit gates the number of network calls — not just the rendered rows — and - both front-ends share that guard. - - Args: - result (ScenarioResult): The full scenario result to read attacks from. - client (PyRITApiClient): Client used to fetch each conversation's messages. - scenario_result_id (str): The run id, echoed back on the payload. - attack_result_ids (list[str] | None): When provided, keep only attacks - whose id is in this set. Defaults to None (all attacks). - limit (int | None): Maximum number of attacks to fetch. Defaults to None. - - Returns: - ConversationsPayload: The per-attack transcripts plus the pre-limit total. - """ - selected = _select_attacks(result=result, attack_result_ids=attack_result_ids) - total = len(selected) - if limit is not None: - selected = selected[:limit] - - objective_hash, objective_class = _objective_scorer_key(result=result) - conversations: list[AttackConversation] = [] - for atomic_attack_name, attack_result in selected: - response = await client.get_conversation_messages_async( - attack_result_id=attack_result.attack_result_id, - conversation_id=attack_result.conversation_id, - ) - messages = [ - _message_to_transcript( - message=message, - objective_hash=objective_hash, - objective_class=objective_class, - ) - for message in response.get("messages", []) - ] - conversations.append( - AttackConversation( - attack_result_id=attack_result.attack_result_id, - atomic_attack_name=atomic_attack_name, - objective=attack_result.objective, - outcome=attack_result.outcome.value, - conversation_id=attack_result.conversation_id, - messages=messages, - ) - ) - - return ConversationsPayload( - scenario_result_id=scenario_result_id, - conversations=conversations, - total=total, - ) - - 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. @@ -320,53 +114,6 @@ def _select_attacks(*, result: ScenarioResult, attack_result_ids: list[str] | No return selected -def _message_to_transcript( - *, - message: dict[str, Any], - objective_hash: str | None, - objective_class: str | None, -) -> TranscriptMessage: - """ - Map one raw ``MessageView`` payload into a ``TranscriptMessage``. - - Args: - message (dict[str, Any]): A single message from the - ``ConversationMessagesResponse`` payload. - objective_hash (str | None): The objective scorer's identity hash, used - to select which of the message's scores to surface. - objective_class (str | None): The objective scorer's class name, used as - a fallback match when the hash is unavailable. - - Returns: - TranscriptMessage: The role, turn, joined text, and objective score. - """ - pieces: list[dict[str, Any]] = message.get("message_pieces") or [] - return TranscriptMessage( - role=str(message.get("role", "")), - turn=int(message.get("turn_number", 0) or 0), - text=_join_piece_text(pieces=pieces), - score=_select_objective_score( - pieces=pieces, - objective_hash=objective_hash, - objective_class=objective_class, - ), - ) - - -def _join_piece_text(*, pieces: list[dict[str, Any]]) -> str: - """ - Join the text of a message's pieces, preferring the converted value. - - Args: - pieces (list[dict[str, Any]]): The message's piece payloads. - - Returns: - str: The non-empty piece values joined by newlines. - """ - parts = [str(value) for piece in pieces if (value := piece.get("converted_value") or piece.get("original_value"))] - return "\n".join(parts) - - def _objective_scorer_key(*, result: ScenarioResult) -> tuple[str | None, str | None]: """ Extract the scenario objective scorer's ``(hash, class_name)`` match key. @@ -386,56 +133,3 @@ def _objective_scorer_key(*, result: ScenarioResult) -> tuple[str | None, str | if identifier is None: return None, None return identifier.hash, identifier.class_name - - -def _select_objective_score( - *, - pieces: list[dict[str, Any]], - objective_hash: str | None, - objective_class: str | None, -) -> TranscriptScore | None: - """ - Pick the objective (top-level) score from a message's pieces. - - A response carries several scores (refusal, objective, composite, ...); only - the objective scorer's verdict reflects attack success. Match it by identity - hash (exact), falling back to class name, rather than by list order — which - is undocumented and would arbitrarily surface an auxiliary sub-score. - - Args: - pieces (list[dict[str, Any]]): The message's piece payloads. - objective_hash (str | None): The objective scorer's identity hash. - objective_class (str | None): The objective scorer's class name (fallback). - - Returns: - TranscriptScore | None: The objective score, or ``None`` if not found. - """ - fallback: TranscriptScore | None = None - for piece in pieces: - for score in piece.get("scores") or []: - identifier = score.get("scorer_class_identifier") or {} - if objective_hash and identifier.get("hash") == objective_hash: - return _to_transcript_score(score=score) - if objective_class and fallback is None and score.get("scorer_type") == objective_class: - fallback = _to_transcript_score(score=score) - return fallback - - -def _to_transcript_score(*, score: dict[str, Any]) -> TranscriptScore: - """ - Map one raw ``ScoreView`` payload into a ``TranscriptScore``. - - Args: - score (dict[str, Any]): A single score payload from a message piece. - - Returns: - TranscriptScore: The scorer name, value, and rationale. - """ - value = score.get("score_value") - if value is None: - value = score.get("status") - return TranscriptScore( - scorer=score.get("scorer_type") or None, - value=str(value) if value is not None else None, - rationale=score.get("score_rationale") or None, - ) diff --git a/pyrit/cli/_sources.py b/pyrit/cli/_sources.py new file mode 100644 index 0000000000..20bcc7cc3b --- /dev/null +++ b/pyrit/cli/_sources.py @@ -0,0 +1,167 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +CLI-side data sources for the ``pyrit.output`` printers. + +The thin REST client can't use the framework's ``CentralMemory``-backed sources, +so this module supplies a ``ConversationSource`` that hydrates ``pyrit.models`` +objects from the ``/messages`` view JSON and serves the objective score inline +(no extra endpoint). It lives here — not in ``pyrit.output`` — so the framework +output layer never imports ``pyrit.cli``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from pyrit.models import Message, MessagePiece, Score + +if TYPE_CHECKING: + from pyrit.cli.api_client import PyRITApiClient + + +def _only_known(model_cls: type, data: dict[str, Any]) -> dict[str, Any]: + """ + Keep only the fields the domain model declares (drops view-only extras). + + Args: + model_cls (type): The pydantic model whose fields to keep. + data (dict[str, Any]): The raw view JSON. + + Returns: + dict[str, Any]: The subset of ``data`` whose keys are model fields. + """ + return {key: value for key, value in data.items() if key in model_cls.model_fields} + + +class RestApiConversationSource: + """ + ``ConversationSource`` backed by the REST api client (thin CLI path). + + ``get_messages_async`` hydrates ``pyrit.models`` messages from the + ``/messages`` view JSON and, in the same pass, captures the **objective** + score per piece (matched by the scenario's objective scorer identity) so + ``get_scores_async`` needs no additional endpoint. Only the objective score is + surfaced — the response also carries auxiliary sub-scores (refusal, etc.) + whose verdicts read backwards relative to attack success. + """ + + def __init__( + self, + *, + client: PyRITApiClient, + attack_result_id: str, + objective_hash: str | None = None, + objective_class: str | None = None, + ) -> None: + """ + Args: + client (PyRITApiClient): Transport used to fetch the conversation messages. + attack_result_id (str): The attack whose conversation is fetched. + objective_hash (str | None): The objective scorer's identity hash, used + to select which score to surface. Defaults to None. + objective_class (str | None): The objective scorer's class name, used as + a fallback match when the hash is unavailable. Defaults to None. + """ + self._client = client + self._attack_result_id = attack_result_id + self._objective_hash = objective_hash + self._objective_class = objective_class + self._scores_by_piece: dict[str, list[Score]] = {} + + async def get_messages_async(self, *, conversation_id: str) -> list[Message]: + """ + Fetch and hydrate the conversation's messages, capturing objective scores. + + Args: + conversation_id (str): The conversation whose messages to fetch. + + Returns: + list[Message]: The hydrated messages in order. + """ + response = await self._client.get_conversation_messages_async( + attack_result_id=self._attack_result_id, + conversation_id=conversation_id, + ) + self._scores_by_piece = {} + messages: list[Message] = [] + for message_json in response.get("messages", []): + role = message_json.get("role") + pieces: list[MessagePiece] = [] + for piece_json in message_json.get("message_pieces") or []: + piece = self._hydrate_piece(piece_json=piece_json, message_role=role) + pieces.append(piece) + objective = self._select_objective_score(piece_json=piece_json) + if objective is not None: + self._scores_by_piece[str(piece.id)] = [objective] + if pieces: + messages.append(Message(message_pieces=pieces)) + return messages + + async def get_scores_async(self, *, prompt_ids: list[str]) -> list[Score]: + """ + Return the objective scores captured for the given piece ids. + + Args: + prompt_ids (list[str]): The message-piece ids to fetch scores for. + + Returns: + list[Score]: The objective scores for those pieces (empty if none). + """ + return [score for prompt_id in prompt_ids for score in self._scores_by_piece.get(prompt_id, [])] + + def _hydrate_piece(self, *, piece_json: dict[str, Any], message_role: str | None) -> MessagePiece: + """ + Hydrate one ``MessagePieceView`` payload into a domain ``MessagePiece``. + + Args: + piece_json (dict[str, Any]): A single piece from the view JSON. + message_role (str | None): The enclosing message's role, used when the + piece omits its own (keeps ``Message``'s role invariant satisfied). + + Returns: + MessagePiece: The hydrated piece. + """ + data = _only_known(MessagePiece, piece_json) + if not data.get("role"): + data["role"] = message_role or "user" + if not data.get("original_value"): + data["original_value"] = piece_json.get("converted_value") or piece_json.get("original_value") or "" + return MessagePiece.model_validate(data) + + def _select_objective_score(self, *, piece_json: dict[str, Any]) -> Score | None: + """ + Pick the objective score from a piece's scores, matched by scorer identity. + + Matches by identity hash first, then class name, rather than list order. + + Args: + piece_json (dict[str, Any]): A single piece from the view JSON. + + Returns: + Score | None: The hydrated objective score, or None when absent. + """ + fallback: Score | None = None + for score_json in piece_json.get("scores") or []: + identifier = score_json.get("scorer_class_identifier") or {} + if self._objective_hash and identifier.get("hash") == self._objective_hash: + return self._hydrate_score(score_json=score_json) + if self._objective_class and fallback is None and score_json.get("scorer_type") == self._objective_class: + fallback = self._hydrate_score(score_json=score_json) + return fallback + + def _hydrate_score(self, *, score_json: dict[str, Any]) -> Score | None: + """ + Hydrate one ``ScoreView`` payload into a domain ``Score``, best-effort. + + Args: + score_json (dict[str, Any]): A single score from the view JSON. + + Returns: + Score | None: The hydrated score, or None if it can't be validated. + """ + try: + return Score.model_validate(_only_known(Score, score_json)) + except Exception: + return None diff --git a/pyrit/cli/pyrit_scan.py b/pyrit/cli/pyrit_scan.py index bfcf6bb296..23481149ea 100644 --- a/pyrit/cli/pyrit_scan.py +++ b/pyrit/cli/pyrit_scan.py @@ -821,10 +821,9 @@ async def _handle_results_async(*, client: Any, parsed_args: Namespace) -> int: from pyrit.cli._cli_args import ScenarioResultView from pyrit.cli._results import ( apply_view_limit_policy, - build_attacks_table_payload, - build_conversations_payload_async, resolve_view, ) + from pyrit.output import output_scenario_attacks_async scenario_result_id = parsed_args.scenario_result_id view = resolve_view(view=parsed_args.view) @@ -841,18 +840,16 @@ async def _handle_results_async(*, client: Any, parsed_args: Namespace) -> int: return 0 if view in (ScenarioResultView.ATTACKS, ScenarioResultView.FULL): - attacks_payload = build_attacks_table_payload( - result=result, - scenario_result_id=scenario_result_id, + await output_scenario_attacks_async( + result, attack_result_ids=parsed_args.attack_result_ids, limit=limit, ) - _output.print_attacks_table(payload=attacks_payload) if view is ScenarioResultView.ATTACKS: return 0 try: - conversations_payload = await build_conversations_payload_async( + await _output.print_conversations_async( result=result, client=client, scenario_result_id=scenario_result_id, @@ -862,7 +859,6 @@ async def _handle_results_async(*, client: Any, parsed_args: Namespace) -> int: except Exception as exc: _print_cli_exception(exc=exc) return 1 - _output.print_conversations(payload=conversations_payload) return 0 diff --git a/pyrit/cli/pyrit_shell.py b/pyrit/cli/pyrit_shell.py index 80ea4419b8..8b7e9045ab 100644 --- a/pyrit/cli/pyrit_shell.py +++ b/pyrit/cli/pyrit_shell.py @@ -620,13 +620,12 @@ def do_scenario_results(self, arg: str) -> None: import shlex from pyrit.cli._cli_args import ScenarioResultView, build_scenario_results_parser - from pyrit.cli._output import print_attacks_table, print_conversations, print_scenario_result_async + from pyrit.cli._output import print_conversations_async, print_scenario_result_async from pyrit.cli._results import ( apply_view_limit_policy, - build_attacks_table_payload, - build_conversations_payload_async, resolve_view, ) + from pyrit.output import output_scenario_attacks_async try: tokens = shlex.split(arg) @@ -663,19 +662,19 @@ def do_scenario_results(self, arg: str) -> None: return if view in (ScenarioResultView.ATTACKS, ScenarioResultView.FULL): - attacks_payload = build_attacks_table_payload( - result=result, - scenario_result_id=parsed.scenario_result_id, - attack_result_ids=parsed.attack_result_ids, - limit=limit, + self._run_async( + output_scenario_attacks_async( + result, + attack_result_ids=parsed.attack_result_ids, + limit=limit, + ) ) - print_attacks_table(payload=attacks_payload) if view is ScenarioResultView.ATTACKS: return try: - conversations_payload = self._run_async( - build_conversations_payload_async( + self._run_async( + print_conversations_async( result=result, client=self._api_client, scenario_result_id=parsed.scenario_result_id, @@ -686,7 +685,6 @@ def do_scenario_results(self, arg: str) -> None: except Exception as exc: print(f"Error: {exc}") return - print_conversations(payload=conversations_payload) def do_print_scenario(self, arg: str) -> None: """ diff --git a/pyrit/output/__init__.py b/pyrit/output/__init__.py index 1d72304af8..fb223e16ef 100644 --- a/pyrit/output/__init__.py +++ b/pyrit/output/__init__.py @@ -25,6 +25,7 @@ output_attack_async, output_conversation_async, output_scenario_async, + output_scenario_attacks_async, output_score_async, output_scorer_async, ) @@ -38,6 +39,7 @@ "output_attack_async": "pyrit.output.helpers", "output_conversation_async": "pyrit.output.helpers", "output_scenario_async": "pyrit.output.helpers", + "output_scenario_attacks_async": "pyrit.output.helpers", "output_score_async": "pyrit.output.helpers", "output_scorer_async": "pyrit.output.helpers", "PrinterBase": "pyrit.output.base", diff --git a/pyrit/output/attack_result/markdown.py b/pyrit/output/attack_result/markdown.py index ce49b70db3..76f1b81d6c 100644 --- a/pyrit/output/attack_result/markdown.py +++ b/pyrit/output/attack_result/markdown.py @@ -6,7 +6,7 @@ from pyrit.models import AttackResult, ConversationType, Message, Score from pyrit.output.attack_result.base import AttackResultPrinterBase -from pyrit.output.conversation.markdown import MarkdownConversationPrinter +from pyrit.output.conversation.markdown import MarkdownConversationMemoryPrinter, MarkdownConversationPrinter from pyrit.output.score.markdown import MarkdownScorePrinter from pyrit.output.sink import Sink @@ -52,7 +52,7 @@ def __init__( super().__init__(sink=sink) self._display_inline = display_inline self._score_printer = score_printer or MarkdownScorePrinter(sink=sink) - self._conversation_printer = conversation_printer or MarkdownConversationPrinter( + self._conversation_printer = conversation_printer or MarkdownConversationMemoryPrinter( sink=sink, score_printer=self._score_printer, blur_images=blur_images, diff --git a/pyrit/output/attack_result/pretty.py b/pyrit/output/attack_result/pretty.py index c31381fd88..9e92859a60 100644 --- a/pyrit/output/attack_result/pretty.py +++ b/pyrit/output/attack_result/pretty.py @@ -9,7 +9,7 @@ from pyrit.models import AttackOutcome, AttackResult, ConversationType, Message, Score from pyrit.output._formatting import _PrettyPrinterMixin from pyrit.output.attack_result.base import AttackResultPrinterBase -from pyrit.output.conversation.pretty import PrettyConversationPrinter +from pyrit.output.conversation.pretty import PrettyConversationMemoryPrinter, PrettyConversationPrinter from pyrit.output.score.pretty import PrettyScorePrinter from pyrit.output.sink import Sink @@ -59,7 +59,7 @@ def __init__( self._score_printer = score_printer or PrettyScorePrinter( sink=sink, width=width, indent_size=indent_size, enable_colors=enable_colors ) - self._conversation_printer = conversation_printer or PrettyConversationPrinter( + self._conversation_printer = conversation_printer or PrettyConversationMemoryPrinter( sink=sink, width=width, indent_size=indent_size, diff --git a/pyrit/output/conversation/base.py b/pyrit/output/conversation/base.py index 9eaf74df92..082b0409b8 100644 --- a/pyrit/output/conversation/base.py +++ b/pyrit/output/conversation/base.py @@ -4,7 +4,7 @@ import json from abc import abstractmethod -from pyrit.models import Message, MessagePiece, Score +from pyrit.models import Message, MessagePiece from pyrit.output.base import PrinterBase @@ -12,8 +12,8 @@ class ConversationPrinterBase(PrinterBase): """ Abstract base class for printing conversation message histories. - Subclasses implement data-fetching methods (``_get_scores_async``, - ``_display_image_async``) and rendering via ``render_async``. + Data access goes through an injected ``ConversationSource``; subclasses + provide only rendering via ``render_async`` (and image display, if any). """ _REASONING_RENDER_WARNING = "⚠ WARNING: Reasoning summary failed to render; conversation is intact." @@ -106,18 +106,6 @@ def _extract_reasoning_summary(reasoning_value: str) -> str: return "\n".join(parts) - @abstractmethod - async def _get_scores_async(self, *, prompt_ids: list[str]) -> list[Score]: - """ - Fetch scores for given prompt piece IDs. - - Args: - prompt_ids (list[str]): The message piece IDs to fetch scores for. - - Returns: - list[Score]: The scores associated with the given piece IDs. - """ - async def _display_image_async(self, piece: MessagePiece) -> None: """ Display an image from a message piece. No-op by default. diff --git a/pyrit/output/conversation/markdown.py b/pyrit/output/conversation/markdown.py index 2b1ab74b8f..b0aa812575 100644 --- a/pyrit/output/conversation/markdown.py +++ b/pyrit/output/conversation/markdown.py @@ -6,8 +6,9 @@ import os from pathlib import Path -from pyrit.models import Message, MessagePiece, Score +from pyrit.models import Message, MessagePiece from pyrit.output.conversation.base import ConversationPrinterBase +from pyrit.output.conversation.source import ConversationSource, MemoryConversationSource from pyrit.output.score.markdown import MarkdownScorePrinter from pyrit.output.sink import Sink @@ -25,6 +26,7 @@ class MarkdownConversationPrinter(ConversationPrinterBase): def __init__( self, *, + source: ConversationSource, sink: Sink | None = None, score_printer: MarkdownScorePrinter | None = None, blur_images: bool = False, @@ -35,6 +37,7 @@ def __init__( Initialize the markdown conversation printer. Args: + source (ConversationSource): Data source used to fetch inline scores. sink (Sink | None): Output sink. Defaults to StdoutSink(). score_printer (MarkdownScorePrinter | None): Score printer for inline score rendering. Defaults to a new MarkdownScorePrinter with matching sink. @@ -52,6 +55,7 @@ def __init__( directory using the original basename plus ``_blurred.png``. """ super().__init__(sink=sink) + self._source = source self._score_printer = score_printer or MarkdownScorePrinter(sink=sink) self._blur_images = blur_images self._blur_radius = blur_radius @@ -430,7 +434,7 @@ async def _format_message_scores_async(self, *, pieces: list[MessagePiece]) -> l """ lines: list[str] = [] for piece in pieces: - scores = await self._get_scores_async(prompt_ids=[str(piece.id)]) + scores = await self._source.get_scores_async(prompt_ids=[str(piece.id)]) if scores: lines.append("\n##### Scores\n") lines.extend(self._score_printer._format_score(score, indent="") for score in scores) @@ -469,15 +473,13 @@ def __init__( Defaults to None (sibling of the original). """ super().__init__( + source=MemoryConversationSource(), sink=sink, score_printer=score_printer, blur_images=blur_images, blur_radius=blur_radius, blurred_dir=blurred_dir, ) - from pyrit.memory import CentralMemory - - self._memory = CentralMemory.get_memory_instance() async def render_async( self, @@ -500,12 +502,3 @@ async def render_async( return await super().render_async( messages, include_scores=include_scores, include_reasoning_summaries=include_reasoning_summaries ) - - async def _get_scores_async(self, *, prompt_ids: list[str]) -> list[Score]: - """ - Fetch scores from CentralMemory. - - Returns: - list[Score]: The scores. - """ - return list(self._memory.get_prompt_scores(prompt_ids=prompt_ids)) diff --git a/pyrit/output/conversation/pretty.py b/pyrit/output/conversation/pretty.py index 4cad444872..2580ba2912 100644 --- a/pyrit/output/conversation/pretty.py +++ b/pyrit/output/conversation/pretty.py @@ -6,9 +6,10 @@ from colorama import Fore, Style -from pyrit.models import Message, MessagePiece, Score +from pyrit.models import Message, MessagePiece from pyrit.output._formatting import _PrettyPrinterMixin from pyrit.output.conversation.base import ConversationPrinterBase +from pyrit.output.conversation.source import ConversationSource, MemoryConversationSource from pyrit.output.score.pretty import PrettyScorePrinter from pyrit.output.sink import Sink @@ -19,13 +20,14 @@ class PrettyConversationPrinter(_PrettyPrinterMixin, ConversationPrinterBase): """ Pretty printer for conversation message histories with ANSI-colored formatting. - Contains all formatting logic. Subclasses implement ``_get_scores_async`` - and ``_display_image_async`` for data fetching. + Contains all formatting logic; scores are fetched through the injected + ``ConversationSource``. """ def __init__( self, *, + source: ConversationSource, sink: Sink | None = None, width: int = 100, indent_size: int = 2, @@ -38,6 +40,7 @@ def __init__( Initialize the pretty conversation printer. Args: + source (ConversationSource): Data source used to fetch inline scores. sink (Sink | None): Output sink. Defaults to StdoutSink(). width (int): Maximum width for text wrapping. Defaults to 100. indent_size (int): Number of spaces for indentation. Defaults to 2. @@ -51,6 +54,7 @@ def __init__( Defaults to 20. """ super().__init__(sink=sink) + self._source = source self._width = width self._indent = " " * indent_size self._enable_colors = enable_colors @@ -161,7 +165,7 @@ async def render_async( image_pieces.append(piece) if include_scores: - scores = await self._get_scores_async(prompt_ids=[str(piece.id)]) + scores = await self._source.get_scores_async(prompt_ids=[str(piece.id)]) if scores: lines.append("\n") lines.append(self._format_colored(f"{self._indent}📊 Scores:", Style.DIM, Fore.MAGENTA)) @@ -293,6 +297,7 @@ def __init__( Defaults to 20. """ super().__init__( + source=MemoryConversationSource(), sink=sink, width=width, indent_size=indent_size, @@ -301,9 +306,6 @@ def __init__( blur_images=blur_images, blur_radius=blur_radius, ) - from pyrit.memory import CentralMemory - - self._memory = CentralMemory.get_memory_instance() async def render_async( self, @@ -327,15 +329,6 @@ async def render_async( messages, include_scores=include_scores, include_reasoning_summaries=include_reasoning_summaries ) - async def _get_scores_async(self, *, prompt_ids: list[str]) -> list[Score]: - """ - Fetch scores from CentralMemory. - - Returns: - list[Score]: The scores. - """ - return list(self._memory.get_prompt_scores(prompt_ids=prompt_ids)) - async def _display_image_async(self, piece: MessagePiece) -> None: """ Display an image from a message piece in notebook environments. diff --git a/pyrit/output/conversation/source.py b/pyrit/output/conversation/source.py new file mode 100644 index 0000000000..4f16e003ba --- /dev/null +++ b/pyrit/output/conversation/source.py @@ -0,0 +1,63 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +if TYPE_CHECKING: + from pyrit.models import Message, Score + + +@runtime_checkable +class ConversationSource(Protocol): + """ + The data a conversation printer needs, decoupled from where it comes from. + + ``MemoryConversationSource`` (CentralMemory) backs the framework/notebook path; + a REST-backed source backs the thin CLI client. Printers depend only on this + Protocol, so ``pyrit.output`` needs no knowledge of either backend and never + imports ``pyrit.cli``. + """ + + async def get_messages_async(self, *, conversation_id: str) -> list[Message]: + """Return the ordered messages for a conversation.""" + ... + + async def get_scores_async(self, *, prompt_ids: list[str]) -> list[Score]: + """Return the scores attached to the given message-piece ids (empty if none).""" + ... + + +class MemoryConversationSource: + """``ConversationSource`` backed by ``CentralMemory`` (framework / notebook path).""" + + def __init__(self) -> None: + """Resolve the process-wide memory instance (deferred import).""" + from pyrit.memory import CentralMemory + + self._memory = CentralMemory.get_memory_instance() + + async def get_messages_async(self, *, conversation_id: str) -> list[Message]: + """ + Return the ordered messages for a conversation from memory. + + Args: + conversation_id (str): The conversation to read. + + Returns: + list[Message]: The conversation's messages in order. + """ + return list(self._memory.get_conversation_messages(conversation_id=conversation_id)) + + async def get_scores_async(self, *, prompt_ids: list[str]) -> list[Score]: + """ + Return the scores attached to the given message-piece ids from memory. + + Args: + prompt_ids (list[str]): The message-piece ids to fetch scores for. + + Returns: + list[Score]: The scores for those pieces (empty if none). + """ + return list(self._memory.get_prompt_scores(prompt_ids=prompt_ids)) diff --git a/pyrit/output/helpers.py b/pyrit/output/helpers.py index f95e1ea8d2..aa75b3f0bc 100644 --- a/pyrit/output/helpers.py +++ b/pyrit/output/helpers.py @@ -117,6 +117,30 @@ async def output_scenario_async( await printer.write_async(result) +async def output_scenario_attacks_async( + result: ScenarioResult, + *, + attack_result_ids: list[str] | None = None, + limit: int | None = None, + sink: Sink | None = None, +) -> None: + """ + Print a compact per-attack table for a scenario result. + + Complements ``output_scenario_async`` (which prints the aggregate overview) by + listing individual attack results: id, technique, objective, outcome, and score. + + Args: + result (ScenarioResult): The scenario result whose attacks to list. + attack_result_ids (list[str] | None): Restrict to these attack ids. Defaults to None. + limit (int | None): Maximum number of attacks to show. Defaults to None. + sink (Sink | None): Output sink. Defaults to StdoutSink. + """ + resolved_sink = sink or get_default_sink(StdoutSink) + printer = PrettyScenarioResultMemoryPrinter(sink=resolved_sink) + await printer.write_async(result, view="attacks", attack_result_ids=attack_result_ids, limit=limit) + + async def output_scorer_async( *, scorer_identifier: ComponentIdentifier, diff --git a/pyrit/output/scenario_result/base.py b/pyrit/output/scenario_result/base.py index 1ad2422ef7..8500049ab7 100644 --- a/pyrit/output/scenario_result/base.py +++ b/pyrit/output/scenario_result/base.py @@ -2,10 +2,15 @@ # Licensed under the MIT license. from abc import abstractmethod +from typing import Literal from pyrit.models import ScenarioResult from pyrit.output.base import PrinterBase +# Which projection of a scenario result to render: the aggregate overview or the +# compact per-attack table. +ScenarioView = Literal["overview", "attacks"] + class ScenarioResultPrinterBase(PrinterBase): """ @@ -16,12 +21,25 @@ class ScenarioResultPrinterBase(PrinterBase): """ @abstractmethod - async def render_async(self, result: ScenarioResult) -> str: + async def render_async( + self, + result: ScenarioResult, + *, + view: ScenarioView = "overview", + attack_result_ids: list[str] | None = None, + limit: int | None = None, + ) -> str: """ - Render a scenario result summary and return it as a string. + Render a scenario result and return it as a string. Args: - result (ScenarioResult): The scenario result to summarize. + result (ScenarioResult): The scenario result to render. + view (ScenarioView): Which projection to render — the aggregate ``"overview"`` + or the per-attack ``"attacks"`` table. Defaults to ``"overview"``. + attack_result_ids (list[str] | None): For the ``"attacks"`` view, restrict to + these attack ids. Ignored by other views. Defaults to None. + limit (int | None): For the ``"attacks"`` view, the maximum number of attacks + to show. Ignored by other views. Defaults to None. Returns: str: The rendered scenario result text. diff --git a/pyrit/output/scenario_result/pretty.py b/pyrit/output/scenario_result/pretty.py index 1bdbf3a289..b00b6dd13c 100644 --- a/pyrit/output/scenario_result/pretty.py +++ b/pyrit/output/scenario_result/pretty.py @@ -5,12 +5,19 @@ from colorama import Fore, Style -from pyrit.models import AttackOutcome, ScenarioResult +from pyrit.models import AttackOutcome, AttackResult, ScenarioResult from pyrit.output._formatting import _PrettyPrinterMixin -from pyrit.output.scenario_result.base import ScenarioResultPrinterBase +from pyrit.output.scenario_result.base import ScenarioResultPrinterBase, ScenarioView from pyrit.output.scorer.base import ScorerPrinterBase from pyrit.output.sink import Sink +# A successful attack is a failure for the defender, so success is shown in red. +_ATTACK_OUTCOME_COLORS = { + AttackOutcome.SUCCESS: Fore.RED, + AttackOutcome.FAILURE: Fore.GREEN, + AttackOutcome.UNDETERMINED: Fore.YELLOW, +} + class PrettyScenarioResultPrinter(_PrettyPrinterMixin, ScenarioResultPrinterBase): """ @@ -117,12 +124,25 @@ def _get_rate_color(self, rate: int) -> str: return str(Fore.CYAN) return str(Fore.GREEN) - async def render_async(self, result: ScenarioResult) -> str: + async def render_async( + self, + result: ScenarioResult, + *, + view: ScenarioView = "overview", + attack_result_ids: list[str] | None = None, + limit: int | None = None, + ) -> str: """ - Render the scenario result summary and return it as a string. + Render a scenario result and return it as a string. Args: - result (ScenarioResult): The scenario result to summarize. + result (ScenarioResult): The scenario result to render. + view (ScenarioView): Which projection to render — the aggregate ``"overview"`` + or the per-attack ``"attacks"`` table. Defaults to ``"overview"``. + attack_result_ids (list[str] | None): For the ``"attacks"`` view, restrict to + these attack ids. Ignored by the overview. Defaults to None. + limit (int | None): For the ``"attacks"`` view, the maximum number of attacks + to show. Ignored by the overview. Defaults to None. Returns: str: The rendered scenario result text. @@ -131,6 +151,9 @@ async def render_async(self, result: ScenarioResult) -> str: ValueError: If the result has an ``objective_scorer_identifier`` but no scorer printer is configured. """ + if view == "attacks": + return self._render_attacks(result, attack_result_ids=attack_result_ids, limit=limit) + parts: list[str] = [] lines: list[str] = [] @@ -224,6 +247,81 @@ async def render_async(self, result: ScenarioResult) -> str: return "".join(parts) + def _render_attacks( + self, + result: ScenarioResult, + *, + attack_result_ids: list[str] | None = None, + limit: int | None = None, + ) -> str: + """ + Render a compact per-attack table for the scenario's results. + + Reads the ``AttackResult`` objects embedded in *result* (no fetching), so + the framework and the thin CLI client render attacks identically. + + Args: + result (ScenarioResult): The scenario result whose attacks to list. + attack_result_ids (list[str] | None): When provided, keep only attacks + whose id is in this set. Defaults to None (all attacks). + limit (int | None): Maximum number of attacks to show. Defaults to None. + + Returns: + str: The rendered attacks table. + """ + id_filter = set(attack_result_ids) if attack_result_ids else None + selected = [ + (atomic_attack_name, attack) + for atomic_attack_name, attacks in result.attack_results.items() + for attack in attacks + if id_filter is None or attack.attack_result_id in id_filter + ] + total = len(selected) + if limit is not None: + selected = selected[:limit] + + lines: list[str] = [self._render_section_header("Attack Results")] + if not selected: + lines.append(self._format_colored(f"{self._indent}No attack results.", Fore.YELLOW)) + return "".join(lines) + + for index, (name, attack) in enumerate(selected, start=1): + color = _ATTACK_OUTCOME_COLORS.get(attack.outcome, Fore.CYAN) + lines.append("\n") + lines.append( + self._format_colored( + f"{self._indent}{index}. [{attack.outcome.value.upper()}] " + f"turns={attack.executed_turns} score={self._attack_score(attack)}", + Style.BRIGHT, + color, + ) + ) + lines.append(self._format_colored(f"{self._indent * 2}id: {attack.attack_result_id}", Fore.CYAN)) + lines.append(self._format_colored(f"{self._indent * 2}technique: {name}", Fore.CYAN)) + lines.append(self._format_colored(f"{self._indent * 2}objective: {attack.objective}", Fore.CYAN)) + + shown = len(selected) + footer = f"Showing {shown} of {total} attacks." if shown < total else f"Total attacks: {total}" + lines.append("\n") + lines.append(self._format_colored(f"{self._indent}{footer}", Fore.GREEN)) + return "".join(lines) + + @staticmethod + def _attack_score(attack: AttackResult) -> str: + """ + Return the attack's last score value (or status when undetermined). + + Args: + attack: The attack result to read the score from. + + Returns: + str: The score value, its status, or a dash when there is no score. + """ + score = attack.last_score + if score is None: + return "—" + return score.score_value if score.score_value is not None else score.status.value + class PrettyScenarioResultMemoryPrinter(PrettyScenarioResultPrinter): """ @@ -267,14 +365,27 @@ def __init__( ) self._scorer_printer = scorer_printer - async def render_async(self, result: ScenarioResult) -> str: + async def render_async( + self, + result: ScenarioResult, + *, + view: ScenarioView = "overview", + attack_result_ids: list[str] | None = None, + limit: int | None = None, + ) -> str: """ - Render the scenario result summary and return it as a string. + Render the scenario result and return it as a string. Args: - result (ScenarioResult): The scenario result to summarize. + result (ScenarioResult): The scenario result to render. + view (ScenarioView): Which projection to render — the aggregate ``"overview"`` + or the per-attack ``"attacks"`` table. Defaults to ``"overview"``. + attack_result_ids (list[str] | None): For the ``"attacks"`` view, restrict to + these attack ids. Ignored by the overview. Defaults to None. + limit (int | None): For the ``"attacks"`` view, the maximum number of attacks + to show. Ignored by the overview. Defaults to None. Returns: str: The rendered scenario result text. """ - return await super().render_async(result) + return await super().render_async(result, view=view, attack_result_ids=attack_result_ids, limit=limit) diff --git a/tests/unit/cli/test_output.py b/tests/unit/cli/test_output.py index c361d213dd..0c1b5de765 100644 --- a/tests/unit/cli/test_output.py +++ b/tests/unit/cli/test_output.py @@ -668,160 +668,98 @@ async def test_print_scenario_result_async_accepts_real_scenario_result(): # --------------------------------------------------------------------------- -# print_attacks_table +# print_conversations_async (reuses the framework conversation printer) # --------------------------------------------------------------------------- -def _attacks_payload(*, rows, total): - from pyrit.cli._results import AttackRow, AttacksTablePayload +class _FakeMessagesClient: + def __init__(self, by_conversation=None): + self._by_conversation = by_conversation or {} - return AttacksTablePayload( - scenario_result_id="SID", - rows=[AttackRow(**row) for row in rows], - total=total, - ) + async def get_conversation_messages_async(self, *, attack_result_id, conversation_id): + return self._by_conversation.get(conversation_id, {"messages": []}) -def test_print_attacks_table_empty(capsys): - _output.print_attacks_table(payload=_attacks_payload(rows=[], total=0)) - out = capsys.readouterr().out - assert "No attack results found" in out - assert "SID" in out +def _piece(*, role, text, scores=None): + piece = {"role": role, "sequence": 0, "conversation_id": "conv-1", "original_value": text, "converted_value": text} + if scores is not None: + piece["scores"] = scores + return piece -def test_print_attacks_table_renders_rows(capsys): - rows = [ - { - "attack_result_id": "aid-1", - "atomic_attack_name": "tech_a", - "objective": "extract secrets", - "outcome": "success", - "executed_turns": 3, - "score_value": "0.9", - } - ] - _output.print_attacks_table(payload=_attacks_payload(rows=rows, total=1)) - out = capsys.readouterr().out - assert "aid-1" in out - assert "tech_a" in out - assert "extract secrets" in out - assert "SUCCESS" in out - assert "0.9" in out - assert "Total attacks: 1" in out - +def _result_with_attacks(attacks, *, objective_scorer=None): + from pyrit.models import AttackOutcome, AttackResult -def test_print_attacks_table_shows_truncation_note(capsys): - rows = [ - { - "attack_result_id": "aid-1", - "atomic_attack_name": "tech_a", - "objective": "obj", - "outcome": "failure", - "executed_turns": 1, - "score_value": None, - } - ] - # total (5) exceeds shown rows (1) -> "showing N of M" note. - _output.print_attacks_table(payload=_attacks_payload(rows=rows, total=5)) - out = capsys.readouterr().out - assert "Showing 1 of 5" in out - assert "—" in out # missing score placeholder + attack_results = { + name: [AttackResult(conversation_id=cid, objective=obj, outcome=AttackOutcome.SUCCESS) for cid, obj in items] + for name, items in attacks.items() + } + return make_scenario_result(attack_results=attack_results, objective_scorer_identifier=objective_scorer) -# --------------------------------------------------------------------------- -# print_conversations -# --------------------------------------------------------------------------- +async def test_print_conversations_async_empty(capsys): + result = _result_with_attacks({}) + await _output.print_conversations_async(result=result, client=_FakeMessagesClient(), scenario_result_id="SID") + out = capsys.readouterr().out + assert "No conversations found" in out + assert "SID" in out -def _conversations_payload(*, conversations, total): - from pyrit.cli._results import AttackConversation, ConversationsPayload, TranscriptMessage, TranscriptScore - - built = [] - for convo in conversations: - messages = [ - TranscriptMessage( - role=message["role"], - turn=message["turn"], - text=message["text"], - score=( - TranscriptScore( - scorer=message["score"][0], - value=message["score"][1], - rationale=message["score"][2], +async def test_print_conversations_async_renders_messages_and_objective_score(capsys): + from pyrit.models import ComponentIdentifier + + objective = ComponentIdentifier(class_name="ObjScorer", class_module="tests.unit.mocks") + result = _result_with_attacks({"tech_a": [("conv-1", "extract secrets")]}, objective_scorer=objective) + attack = next(iter(result.attack_results["tech_a"])) + response = { + "messages": [ + {"role": "user", "turn_number": 0, "message_pieces": [_piece(role="user", text="please comply")]}, + { + "role": "assistant", + "turn_number": 1, + "message_pieces": [ + _piece( + role="assistant", + text="sure thing", + scores=[ + { + "score_value": "true", + "score_type": "true_false", + "score_rationale": "clearly harmful", + "scorer_type": "ObjScorer", + "scorer_class_identifier": {"hash": objective.hash}, + } + ], ) - if message.get("score") - else None - ), - ) - for message in convo["messages"] + ], + }, ] - built.append( - AttackConversation( - attack_result_id=convo["attack_result_id"], - atomic_attack_name=convo["atomic_attack_name"], - objective=convo["objective"], - outcome=convo["outcome"], - conversation_id=convo["conversation_id"], - messages=messages, - ) - ) - return ConversationsPayload(scenario_result_id="SID", conversations=built, total=total) - - -def test_print_conversations_empty(capsys): - _output.print_conversations(payload=_conversations_payload(conversations=[], total=0)) - out = capsys.readouterr().out - assert "No conversations found" in out - assert "SID" in out + } + client = _FakeMessagesClient({"conv-1": response}) + await _output.print_conversations_async(result=result, client=client, scenario_result_id="SID") -def test_print_conversations_renders_messages_and_score(capsys): - conversations = [ - { - "attack_result_id": "aid-1", - "atomic_attack_name": "tech_a", - "objective": "extract secrets", - "outcome": "success", - "conversation_id": "conv-1", - "messages": [ - {"role": "user", "turn": 0, "text": "please comply", "score": None}, - { - "role": "assistant", - "turn": 1, - "text": "sure thing", - "score": ("TrueFalseCompositeScorer", "0.9", "clearly harmful"), - }, - ], - } - ] - _output.print_conversations(payload=_conversations_payload(conversations=conversations, total=1)) out = capsys.readouterr().out - assert "aid-1" in out + assert attack.attack_result_id in out assert "extract secrets" in out assert "USER" in out - assert "ASSISTANT" in out assert "please comply" in out - assert "0.9" in out + assert "sure thing" in out + # The objective score renders in the framework's own style. + assert "Scores" in out assert "clearly harmful" in out - assert "TrueFalseCompositeScorer" in out assert "Total attacks: 1" in out -def test_print_conversations_shows_truncation_note(capsys): - conversations = [ - { - "attack_result_id": "aid-1", - "atomic_attack_name": "tech_a", - "objective": "obj", - "outcome": "failure", - "conversation_id": "conv-1", - "messages": [], - } - ] - _output.print_conversations(payload=_conversations_payload(conversations=conversations, total=5)) +async def test_print_conversations_async_truncation_note(capsys): + result = _result_with_attacks({"tech_a": [(f"conv-{i}", f"obj-{i}") for i in range(5)]}) + + await _output.print_conversations_async( + result=result, client=_FakeMessagesClient(), scenario_result_id="SID", limit=1 + ) + out = capsys.readouterr().out assert "Showing 1 of 5" in out - assert "(no messages)" in out # --------------------------------------------------------------------------- diff --git a/tests/unit/cli/test_results.py b/tests/unit/cli/test_results.py index 8420d98d7a..68f9e6d812 100644 --- a/tests/unit/cli/test_results.py +++ b/tests/unit/cli/test_results.py @@ -6,8 +6,6 @@ shared argument parser (``pyrit.cli._results`` and ``pyrit.cli._cli_args``). """ -import uuid - import pytest from pyrit.cli._cli_args import ( @@ -16,63 +14,9 @@ build_scenario_results_parser, ) from pyrit.cli._results import ( - _to_transcript_score, apply_view_limit_policy, - build_attacks_table_payload, - build_conversations_payload_async, resolve_view, ) -from pyrit.models import AttackOutcome, AttackResult, ComponentIdentifier, Score -from unit.mocks import make_scenario_result - - -def _scorer_id(class_name): - return ComponentIdentifier(class_name=class_name, class_module="tests.unit.mocks") - - -def _attack(*, outcome=AttackOutcome.SUCCESS, objective="obj", turns=1, with_score=False): - attack = AttackResult( - conversation_id=str(uuid.uuid4()), - objective=objective, - outcome=outcome, - executed_turns=turns, - ) - if with_score: - attack.last_score = Score( - score_value="0.9", - score_type="float_scale", - message_piece_id=str(uuid.uuid4()), - ) - return attack - - -def _result(attack_results, *, objective_scorer=None): - return make_scenario_result(attack_results=attack_results, objective_scorer_identifier=objective_scorer) - - -def test_to_transcript_score_uses_status_for_null_value(): - score = _to_transcript_score( - score={ - "score_value": None, - "status": "undetermined", - "score_type": "TrueFalseCompositeScorer", - } - ) - - assert score.value == "undetermined" - - -class _FakeMessagesClient: - """A minimal stand-in for ``PyRITApiClient.get_conversation_messages_async``.""" - - def __init__(self, responses=None): - self._responses = responses or {} - self.calls: list[tuple[str, str]] = [] - - async def get_conversation_messages_async(self, *, attack_result_id, conversation_id): - self.calls.append((attack_result_id, conversation_id)) - return self._responses.get(conversation_id, {"messages": []}) - # --------------------------------------------------------------------------- # ScenarioResultView @@ -119,69 +63,6 @@ def test_limit_policy_noop_when_no_limit(capsys): assert capsys.readouterr().out == "" -# --------------------------------------------------------------------------- -# build_attacks_table_payload -# --------------------------------------------------------------------------- - - -def test_builder_includes_all_attacks_grouped_by_atomic_name(): - result = _result( - { - "tech_a": [_attack(objective="a1"), _attack(objective="a2")], - "tech_b": [_attack(objective="b1")], - } - ) - payload = build_attacks_table_payload(result=result, scenario_result_id="SID") - assert payload.scenario_result_id == "SID" - assert payload.total == 3 - assert len(payload.rows) == 3 - assert {row.atomic_attack_name for row in payload.rows} == {"tech_a", "tech_b"} - - -def test_builder_maps_outcome_and_score(): - result = _result( - { - "tech_a": [ - _attack(outcome=AttackOutcome.SUCCESS, turns=4, with_score=True), - _attack(outcome=AttackOutcome.FAILURE, with_score=False), - ] - } - ) - payload = build_attacks_table_payload(result=result, scenario_result_id="SID") - scored, unscored = payload.rows[0], payload.rows[1] - assert scored.outcome == "success" - assert scored.executed_turns == 4 - assert scored.score_value == "0.9" - assert unscored.outcome == "failure" - assert unscored.score_value is None - - -def test_builder_filters_by_attack_result_ids(): - keep = _attack(objective="keep") - drop = _attack(objective="drop") - result = _result({"tech_a": [keep, drop]}) - payload = build_attacks_table_payload( - result=result, - scenario_result_id="SID", - attack_result_ids=[keep.attack_result_id], - ) - assert payload.total == 1 - assert payload.rows[0].attack_result_id == keep.attack_result_id - - -def test_builder_limit_caps_rows_but_total_is_pre_limit(): - result = _result({"tech_a": [_attack() for _ in range(5)]}) - payload = build_attacks_table_payload(result=result, scenario_result_id="SID", limit=2) - assert payload.total == 5 - assert len(payload.rows) == 2 - - -def test_builder_handles_no_attacks(): - payload = build_attacks_table_payload(result=_result({}), scenario_result_id="SID") - assert payload.total == 0 - assert payload.rows == [] - - # --------------------------------------------------------------------------- # Shared argument parser # --------------------------------------------------------------------------- @@ -269,141 +150,3 @@ def test_limit_policy_heavy_view_respects_attack_ids(capsys): effective = apply_view_limit_policy(view=ScenarioResultView.CONVERSATIONS, limit=None, attack_result_ids=["a"]) assert effective is None assert capsys.readouterr().out == "" - - -async def test_build_conversations_payload_selects_objective_score_by_hash(): - # A response carries several scores; only the objective scorer's verdict should - # surface — matched by identity hash, not by position in the list. - objective = _scorer_id("TrueFalseCompositeScorer") - auxiliary = _scorer_id("SelfAskRefusalScorer") - attack = _attack(objective="a1") - result = _result({"tech_a": [attack]}, objective_scorer=objective) - messages_payload = { - "conversation_id": attack.conversation_id, - "messages": [ - {"role": "user", "turn_number": 0, "message_pieces": [{"converted_value": "hi"}]}, - { - "role": "assistant", - "turn_number": 1, - "message_pieces": [ - { - "converted_value": "there", - "scores": [ - { - "score_value": "true", - "score_rationale": "refused", - "scorer_type": "SelfAskRefusalScorer", - "scorer_class_identifier": {"hash": auxiliary.hash}, - }, - { - "score_value": "false", - "score_rationale": "complied", - "scorer_type": "TrueFalseCompositeScorer", - "scorer_class_identifier": {"hash": objective.hash}, - }, - ], - } - ], - }, - ], - } - client = _FakeMessagesClient({attack.conversation_id: messages_payload}) - - payload = await build_conversations_payload_async(result=result, client=client, scenario_result_id="SID") - - assert payload.scenario_result_id == "SID" - assert payload.total == 1 - convo = payload.conversations[0] - assert convo.attack_result_id == attack.attack_result_id - assert convo.atomic_attack_name == "tech_a" - assert [message.text for message in convo.messages] == ["hi", "there"] - assert convo.messages[0].score is None - # The objective (composite) score is surfaced, not the first (refusal) score. - assert convo.messages[1].score.value == "false" - assert convo.messages[1].score.rationale == "complied" - assert convo.messages[1].score.scorer == "TrueFalseCompositeScorer" - - -async def test_build_conversations_payload_falls_back_to_class_name(): - # No identity hash on the stored score — fall back to matching the class name. - objective = _scorer_id("TrueFalseCompositeScorer") - attack = _attack(objective="a1") - result = _result({"tech_a": [attack]}, objective_scorer=objective) - messages_payload = { - "conversation_id": attack.conversation_id, - "messages": [ - { - "role": "assistant", - "turn_number": 1, - "message_pieces": [ - { - "converted_value": "there", - "scores": [ - {"score_value": "true", "scorer_type": "SelfAskRefusalScorer"}, - {"score_value": "false", "scorer_type": "TrueFalseCompositeScorer"}, - ], - } - ], - }, - ], - } - client = _FakeMessagesClient({attack.conversation_id: messages_payload}) - - payload = await build_conversations_payload_async(result=result, client=client, scenario_result_id="SID") - - score = payload.conversations[0].messages[0].score - assert score.value == "false" - assert score.scorer == "TrueFalseCompositeScorer" - - -async def test_build_conversations_payload_no_objective_scorer_yields_no_score(): - # Without a scenario objective scorer there's no canonical score to surface. - attack = _attack(objective="a1") - result = _result({"tech_a": [attack]}) - messages_payload = { - "conversation_id": attack.conversation_id, - "messages": [ - { - "role": "assistant", - "turn_number": 1, - "message_pieces": [ - {"converted_value": "there", "scores": [{"score_value": "true", "scorer_type": "X"}]} - ], - }, - ], - } - client = _FakeMessagesClient({attack.conversation_id: messages_payload}) - - payload = await build_conversations_payload_async(result=result, client=client, scenario_result_id="SID") - - assert payload.conversations[0].messages[0].score is None - - -async def test_build_conversations_payload_filters_by_ids(): - keep = _attack(objective="keep") - drop = _attack(objective="drop") - result = _result({"tech_a": [keep, drop]}) - client = _FakeMessagesClient() - - payload = await build_conversations_payload_async( - result=result, - client=client, - scenario_result_id="SID", - attack_result_ids=[keep.attack_result_id], - ) - - assert payload.total == 1 - assert client.calls == [(keep.attack_result_id, keep.conversation_id)] - - -async def test_build_conversations_payload_limit_gates_fetch(): - attacks = [_attack(objective=f"o{i}") for i in range(4)] - result = _result({"tech_a": attacks}) - client = _FakeMessagesClient() - - payload = await build_conversations_payload_async(result=result, client=client, scenario_result_id="SID", limit=2) - - assert payload.total == 4 - assert len(payload.conversations) == 2 - # --limit caps the number of message fetches, not just the rendered rows. - assert len(client.calls) == 2 diff --git a/tests/unit/cli/test_sources.py b/tests/unit/cli/test_sources.py new file mode 100644 index 0000000000..99b246bef0 --- /dev/null +++ b/tests/unit/cli/test_sources.py @@ -0,0 +1,142 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Unit tests for the CLI-side ``RestApiConversationSource`` (pyrit.cli._sources).""" + +import uuid + +from pyrit.cli._sources import RestApiConversationSource +from pyrit.models import Message, Score + + +class _FakeClient: + def __init__(self, response): + self._response = response + self.calls: list[tuple[str, str]] = [] + + async def get_conversation_messages_async(self, *, attack_result_id, conversation_id): + self.calls.append((attack_result_id, conversation_id)) + return self._response + + +def _piece(*, role, text, scores=None): + piece = { + "id": str(uuid.uuid4()), + "role": role, + "sequence": 0, + "conversation_id": "conv-1", + "original_value": text, + "converted_value": text, + # view-only extras that must be dropped on hydration: + "original_value_url": None, + "converted_value_mime_type": "text/plain", + "converted_filename": None, + } + if scores is not None: + piece["scores"] = scores + return piece + + +def _response(messages): + return {"conversation_id": "conv-1", "messages": messages} + + +async def test_get_messages_hydrates_domain_messages(): + response = _response( + [ + {"role": "user", "turn_number": 0, "message_pieces": [_piece(role="user", text="hello")]}, + {"role": "assistant", "turn_number": 1, "message_pieces": [_piece(role="assistant", text="there")]}, + ] + ) + source = RestApiConversationSource(client=_FakeClient(response), attack_result_id="aid-1") + + messages = await source.get_messages_async(conversation_id="conv-1") + + assert all(isinstance(message, Message) for message in messages) + assert [message.get_piece().converted_value for message in messages] == ["hello", "there"] + assert [message.api_role for message in messages] == ["user", "assistant"] + + +async def test_objective_score_selected_by_hash_and_served_by_piece_id(): + objective = [ + { + "score_value": "false", + "score_type": "true_false", + "score_rationale": "complied", + "scorer_type": "SelfAskRefusalScorer", + "scorer_class_identifier": {"hash": "AUX"}, + }, + { + "score_value": "true", + "score_type": "true_false", + "score_rationale": "achieved", + "scorer_type": "TrueFalseCompositeScorer", + "scorer_class_identifier": {"hash": "OBJ"}, + }, + ] + assistant_piece = _piece(role="assistant", text="there", scores=objective) + response = _response([{"role": "assistant", "turn_number": 1, "message_pieces": [assistant_piece]}]) + source = RestApiConversationSource(client=_FakeClient(response), attack_result_id="aid-1", objective_hash="OBJ") + + messages = await source.get_messages_async(conversation_id="conv-1") + piece_id = str(messages[0].get_piece().id) + scores = await source.get_scores_async(prompt_ids=[piece_id]) + + assert len(scores) == 1 + assert isinstance(scores[0], Score) + # The objective (composite) score is surfaced, not the first (refusal) one. + assert scores[0].score_value == "true" + assert scores[0].score_rationale == "achieved" + + +async def test_objective_score_falls_back_to_class_name(): + scores_json = [ + {"score_value": "true", "score_type": "true_false", "scorer_type": "MyObjective"}, + ] + response = _response( + [ + { + "role": "assistant", + "turn_number": 1, + "message_pieces": [_piece(role="assistant", text="x", scores=scores_json)], + } + ] + ) + source = RestApiConversationSource( + client=_FakeClient(response), attack_result_id="aid-1", objective_class="MyObjective" + ) + + messages = await source.get_messages_async(conversation_id="conv-1") + scores = await source.get_scores_async(prompt_ids=[str(messages[0].get_piece().id)]) + + assert [score.score_value for score in scores] == ["true"] + + +async def test_no_objective_scorer_yields_no_scores(): + scores_json = [{"score_value": "true", "score_type": "true_false", "scorer_type": "X"}] + response = _response( + [ + { + "role": "assistant", + "turn_number": 1, + "message_pieces": [_piece(role="assistant", text="x", scores=scores_json)], + } + ] + ) + source = RestApiConversationSource(client=_FakeClient(response), attack_result_id="aid-1") + + messages = await source.get_messages_async(conversation_id="conv-1") + scores = await source.get_scores_async(prompt_ids=[str(messages[0].get_piece().id)]) + + assert scores == [] + + +async def test_view_only_fields_are_dropped_on_hydration(): + # Pieces carry view-only keys (urls, mime, filenames) that MessagePiece forbids; + # hydration must strip them rather than raise. + response = _response([{"role": "user", "turn_number": 0, "message_pieces": [_piece(role="user", text="hi")]}]) + source = RestApiConversationSource(client=_FakeClient(response), attack_result_id="aid-1") + + messages = await source.get_messages_async(conversation_id="conv-1") + + assert messages[0].get_piece().converted_value == "hi" diff --git a/tests/unit/output/scenario_result/test_pretty.py b/tests/unit/output/scenario_result/test_pretty.py index 6169d0df06..1ff5bbb023 100644 --- a/tests/unit/output/scenario_result/test_pretty.py +++ b/tests/unit/output/scenario_result/test_pretty.py @@ -240,3 +240,50 @@ async def test_write_async_sort_is_stable_for_ties(patch_central_database, capsy await sorting_printer.write_async(result) # Tied 100% groups retain their original relative order; 0% group goes last. assert _group_order(capsys.readouterr().out) == ["first_success", "second_success", "fail"] + + +# --- attacks view --- + + +async def test_render_attacks_lists_each_attack(printer): + a1 = _attack_result(objective="obj-1") + a2 = _attack_result(outcome=AttackOutcome.FAILURE, objective="obj-2") + result = _scenario_result(attack_results={"tech_a": [a1], "tech_b": [a2]}) + + text = await printer.render_async(result, view="attacks") + + assert "Attack Results" in text + assert "obj-1" in text + assert "obj-2" in text + assert a1.attack_result_id in text + assert a2.attack_result_id in text + assert "tech_a" in text + assert "tech_b" in text + assert "Total attacks: 2" in text + + +async def test_render_attacks_limit_truncates(printer): + attacks = [_attack_result(objective=f"o{i}") for i in range(3)] + result = _scenario_result(attack_results={"tech_a": attacks}) + + text = await printer.render_async(result, view="attacks", limit=1) + + assert "Showing 1 of 3 attacks" in text + + +async def test_render_attacks_filters_by_ids(printer): + keep = _attack_result(objective="keep") + drop = _attack_result(objective="drop") + result = _scenario_result(attack_results={"tech_a": [keep, drop]}) + + text = await printer.render_async(result, view="attacks", attack_result_ids=[keep.attack_result_id]) + + assert "keep" in text + assert "drop" not in text + assert "Total attacks: 1" in text + + +async def test_render_attacks_empty(printer): + # Build directly (the _scenario_result helper substitutes a default for {}). + text = await printer.render_async(make_scenario_result(attack_results={}), view="attacks") + assert "No attack results" in text diff --git a/tests/unit/output/test_blur_images.py b/tests/unit/output/test_blur_images.py index e5c24b85ed..a62d337cd0 100644 --- a/tests/unit/output/test_blur_images.py +++ b/tests/unit/output/test_blur_images.py @@ -10,16 +10,24 @@ from PIL import Image -from pyrit.models import MessagePiece, Score +from pyrit.models import MessagePiece from pyrit.output.conversation.markdown import MarkdownConversationPrinter from pyrit.output.conversation.pretty import PrettyConversationMemoryPrinter -class _ConcreteMarkdown(MarkdownConversationPrinter): - async def _get_scores_async(self, *, prompt_ids: list[str]) -> list[Score]: +class _NullConversationSource: + async def get_messages_async(self, *, conversation_id: str) -> list: + return [] + + async def get_scores_async(self, *, prompt_ids: list[str]) -> list: return [] +class _ConcreteMarkdown(MarkdownConversationPrinter): + def __init__(self, **kwargs) -> None: + super().__init__(source=_NullConversationSource(), **kwargs) + + def _make_image_bytes(*, multicolor: bool = True) -> bytes: image = Image.new("RGB", (32, 32), color=(0, 200, 0)) if multicolor: