Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 71 additions & 1 deletion pyrit/backend/routes/scenarios.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,14 @@
get_scenario_run_service,
)
from pyrit.backend.services.scenario_service import get_scenario_service
from pyrit.models import ScenarioQueueSnapshot, ScenarioResult, ScenarioRunProgress, ScenarioRunState
from pyrit.models import (
ScenarioHierarchyNode,
ScenarioHierarchyProjection,
ScenarioQueueSnapshot,
ScenarioResult,
ScenarioRunProgress,
ScenarioRunState,
)
from pyrit.models.catalog import (
RegisteredScenario,
RunScenarioRequest,
Expand Down Expand Up @@ -335,6 +342,69 @@ async def get_scenario_run_progress( # pyrit-async-suffix-exempt
return progress


@router.get(
"/runs/{scenario_result_id}/hierarchy",
response_model=ScenarioHierarchyProjection,
responses={
404: {"model": ProblemDetail, "description": "Run not found"},
},
)
async def get_scenario_run_hierarchy( # pyrit-async-suffix-exempt
*,
scenario_result_id: IdentifierStr,
max_nodes: int = Query(1000, ge=1, le=5000, description="Maximum nodes to assemble in the hierarchy tree"),
) -> ScenarioHierarchyProjection:
"""
Get the bounded result hierarchy projection for a scenario run.

Returns:
ScenarioHierarchyProjection: Bounded hierarchy projection connecting orchestration parents,
target-facing children, and generic read-only conversations.
"""
service = get_scenario_run_service()
hierarchy = await service.get_run_hierarchy_async(
scenario_result_id=scenario_result_id,
max_nodes=max_nodes,
)
if hierarchy is None:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Scenario run '{scenario_result_id}' not found",
)
return hierarchy


@router.get(
"/runs/{scenario_result_id}/hierarchy/nodes/{attack_result_id}",
response_model=ScenarioHierarchyNode,
responses={
404: {"model": ProblemDetail, "description": "Node or run not found"},
},
)
async def get_scenario_run_hierarchy_node( # pyrit-async-suffix-exempt
*,
scenario_result_id: IdentifierStr,
attack_result_id: IdentifierStr,
) -> ScenarioHierarchyNode:
"""
Get single-node hierarchy details and generic conversation projections without full-run hydration.

Returns:
ScenarioHierarchyNode: The target node with resolved immediate children and conversations.
"""
service = get_scenario_run_service()
node = await service.get_run_hierarchy_node_async(
scenario_result_id=scenario_result_id,
attack_result_id=attack_result_id,
)
if node is None:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Hierarchy node '{attack_result_id}' not found in scenario run '{scenario_result_id}'",
)
return node


@router.post(
"/runs/{scenario_result_id}/cancel",
response_model=ScenarioRunSummary,
Expand Down
Loading