Feat/sentinel v2

services/budapp/budapp/commons/constants.py

@@ -609,6 +609,7 @@ class WorkflowTypeEnum(StrEnum):
     PROMPT_CREATION = auto()
     PROMPT_SCHEMA_CREATION = auto()
     TOOL_CREATION = auto()
+    CUSTOM_PROBE_CREATION = auto()
 
 
 class NotificationType(Enum):

services/budapp/budapp/guardrails/crud.py

@@ -1089,16 +1089,35 @@ async def create_custom_probe_with_rule(
         name: str,
         description: str | None,
         scanner_type: str,
-        model_id: UUID,
+        model_id: UUID | None,
         model_config: dict,
         model_uri: str,
         model_provider_type: str,
         is_gated: bool,
-        project_id: UUID,
         user_id: UUID,
         provider_id: UUID,
+        guard_types: list[str] | None = None,
+        modality_types: list[str] | None = None,
     ) -> GuardrailProbe:
-        """Create a custom probe with a single model-based rule atomically."""
+        """Create a custom probe with a single model-based rule atomically.
+
+        Args:
+            name: Name of the custom probe
+            description: Description of the probe
+            scanner_type: Type of scanner (e.g., "llm")
+            model_id: Optional model ID (can be None if model lookup happens at deployment)
+            model_config: Configuration dictionary for the model
+            model_uri: URI of the model
+            model_provider_type: Provider type for the model
+            is_gated: Whether the model requires gated access
+            user_id: User ID creating the probe
+            provider_id: Provider ID for the probe
+            guard_types: Optional list of guard types (e.g., ["input", "output"])
+            modality_types: Optional list of modality types (e.g., ["text", "image"])
+
+        Returns:
+            The created GuardrailProbe with its rule
+        """
         # Generate URI for uniqueness check
         probe_uri = f"custom.{user_id}.{name.lower().replace(' ', '_')}"
 
@@ -1141,6 +1160,8 @@ async def create_custom_probe_with_rule(
                 is_gated=is_gated,
                 model_config_json=model_config,
                 model_id=model_id,
+                guard_types=guard_types,
+                modality_types=modality_types,
                 created_by=user_id,
                 status=GuardrailStatusEnum.ACTIVE,
             )

services/budapp/budapp/guardrails/guardrail_routes.py

@@ -31,6 +31,7 @@
 from budapp.commons.schemas import ErrorResponse, PaginatedSuccessResponse, SuccessResponse
 from budapp.guardrails.crud import GuardrailsDeploymentDataManager
 from budapp.guardrails.schemas import (
+    CustomProbeWorkflowRequest,
     GuardrailCustomProbeCreate,
     GuardrailCustomProbeDetailResponse,
     GuardrailCustomProbeResponse,
@@ -57,6 +58,7 @@
     TagsListResponse,
 )
 from budapp.guardrails.services import (
+    GuardrailCustomProbeService,
     GuardrailDeploymentWorkflowService,
     GuardrailProbeRuleService,
     GuardrailProfileDeploymentService,
@@ -1013,6 +1015,54 @@ async def add_guardrail_deployment_workflow(
         ).to_http_response()
 
 
+@router.post(
+    "/custom-probe-workflow",
+    responses={
+        status.HTTP_500_INTERNAL_SERVER_ERROR: {
+            "model": ErrorResponse,
+            "description": "Service is unavailable due to server error",
+        },
+        status.HTTP_400_BAD_REQUEST: {
+            "model": ErrorResponse,
+            "description": "Service is unavailable due to client error",
+        },
+        status.HTTP_200_OK: {
+            "model": RetrieveWorkflowDataResponse,
+            "description": "Successfully add custom probe workflow",
+        },
+    },
+    description="Add custom probe workflow",
+)
+@require_permissions(permissions=[PermissionEnum.MODEL_MANAGE])
+async def add_custom_probe_workflow(
+    current_user: Annotated[User, Depends(get_current_active_user)],
+    session: Annotated[Session, Depends(get_session)],
+    request: CustomProbeWorkflowRequest,
+) -> Union[RetrieveWorkflowDataResponse, ErrorResponse]:
+    """Add custom probe workflow.
+
+    Multi-step workflow for creating custom probes:
+    - Step 1: Select probe type (llm_policy, etc.) - auto-derives model_uri, scanner_type
+    - Step 2: Configure policy (PolicyConfig)
+    - Step 3: Probe metadata + trigger_workflow=true creates probe
+    """
+    try:
+        db_workflow = await GuardrailCustomProbeService(session).add_custom_probe_workflow(
+            current_user_id=current_user.id,
+            request=request,
+        )
+
+        return await WorkflowService(session).retrieve_workflow_data(db_workflow.id)
+    except ClientException as e:
+        logger.exception(f"Failed to add custom probe workflow: {e}")
+        return ErrorResponse(code=e.status_code, message=e.message).to_http_response()
+    except Exception as e:
+        logger.exception(f"Failed to add custom probe workflow: {e}")
+        return ErrorResponse(
+            code=status.HTTP_500_INTERNAL_SERVER_ERROR, message="Failed to add custom probe workflow"
+        ).to_http_response()
+
+
 # Deployment endpoints
 
 

services/budapp/budapp/guardrails/schemas.py

@@ -60,6 +60,18 @@ class ModelDeploymentStatus(str, Enum):
     DELETING = "deleting"
 
 
+class CustomProbeTypeEnum(str, Enum):
+    """Available custom probe type options.
+
+    Each option maps to a specific model_uri, scanner_type, handler, and provider.
+    """
+
+    LLM_POLICY = "llm_policy"
+    # Future extensions:
+    # CLASSIFIER = "classifier"
+    # REGEX = "regex"
+
+
 class GuardrailModelStatus(BaseModel):
     """Status of a model required by guardrail rules."""
 
@@ -373,6 +385,8 @@ class GuardrailCustomProbeResponse(BaseModel):
     model_id: UUID4 | None = None
     model_uri: str | None = None
     model_config_json: dict | None = None
+    guard_types: list[str] | None = None
+    modality_types: list[str] | None = None
     status: str
     created_at: datetime
     modified_at: datetime
@@ -397,6 +411,8 @@ def extract_rule_data(cls, data: Any) -> Any:
                 "model_id": getattr(rule, "model_id", None),
                 "model_uri": getattr(rule, "model_uri", None),
                 "model_config_json": getattr(rule, "model_config_json", None),
+                "guard_types": getattr(rule, "guard_types", None),
+                "modality_types": getattr(rule, "modality_types", None),
                 "status": data.status,
                 "created_at": data.created_at,
                 "modified_at": data.modified_at,
@@ -745,6 +761,81 @@ class GuardrailDeploymentWorkflowSteps(BaseModel):
     pending_profile_data: dict | None = None
 
 
+class CustomProbeWorkflowRequest(BaseModel):
+    """Custom probe workflow request schema (multi-step).
+
+    Similar to GuardrailDeploymentWorkflowRequest but for creating custom probes.
+    Follows the probe-first pattern where the probe is created with model_uri only,
+    and model_id gets assigned later during deployment (or immediately if model is already onboarded).
+
+    Workflow Steps:
+    - Step 1: Select probe type (llm_policy, etc.) - system auto-derives model_uri, scanner_type, etc.
+    - Step 2: Configure policy (PolicyConfig)
+    - Step 3: Probe metadata + trigger_workflow=true creates probe
+    """
+
+    # Workflow management
+    workflow_id: UUID4 | None = None
+    workflow_total_steps: int | None = None  # Should be 3 for new workflows
+    step_number: int = Field(..., gt=0)
+    trigger_workflow: bool = False
+
+    # Step 1: Probe type selection
+    probe_type_option: CustomProbeTypeEnum | None = None
+
+    # Step 2: Policy configuration
+    policy: PolicyConfig | None = None
+
+    # Step 3: Probe metadata
+    name: str | None = None
+    description: str | None = None
+    guard_types: list[str] | None = None
+    modality_types: list[str] | None = None
+
+    @model_validator(mode="after")
+    def validate_fields(self) -> "CustomProbeWorkflowRequest":
+        """Validate workflow request fields.
+
+        Either workflow_id OR workflow_total_steps must be provided, but not both.
+        """
+        if self.workflow_id is None and self.workflow_total_steps is None:
+            raise ValueError("workflow_total_steps is required when workflow_id is not provided")
+
+        if self.workflow_id is not None and self.workflow_total_steps is not None:
+            raise ValueError("workflow_total_steps and workflow_id cannot be provided together")
+
+        return self
+
+
+class CustomProbeWorkflowSteps(BaseModel):
+    """Custom probe workflow step data schema.
+
+    Tracks accumulated data across workflow steps for custom probe creation.
+    """
+
+    # Step 1 data
+    probe_type_option: CustomProbeTypeEnum | None = None
+    # Auto-derived from probe_type_option
+    model_uri: str | None = None
+    scanner_type: str | None = None
+    handler: str | None = None
+    model_provider_type: str | None = None
+
+    # Step 2 data
+    policy: dict | None = None  # PolicyConfig as dict
+
+    # Step 3 data
+    name: str | None = None
+    description: str | None = None
+    guard_types: list[str] | None = None
+    modality_types: list[str] | None = None
+
+    # Result data (after trigger_workflow)
+    probe_id: UUID4 | None = None
+    model_id: UUID4 | None = None  # Assigned if model exists
+    workflow_execution_status: dict | None = None
+
+
 class BudSentinelConfig(BaseModel):
     """BudSentinel config."""