dtx-models 0.18.2__py3-none-any.whl

dtx_models/providers/openai.py

@@ -0,0 +1,139 @@
+from typing import Any, Dict, Literal, Optional
+
+from pydantic import BaseModel, Field, field_serializer, model_validator
+
+from ..providers.models_spec import ModelTaskType
+
+"""
+id: openai
+config:
+  model: gpt-4o
+
+"""
+
+
+""" 
+id: openai
+config:
+  model: gpt-4o
+  task: generation
+  endpoint: https://api.openai.com/v1
+  params:
+    temperature: 0.7
+    top_k: 40
+    top_p: 0.9
+    repeat_penalty: 1.1
+    max_tokens: 512
+    num_return_sequences: 1
+    extra_params:
+      stop: ["###", "User:"]
+
+"""
+
+
+# --- PARAMETER MODELS ---
+
+
+class ProviderParams(BaseModel):
+    """Optional parameters for fine-tuning model generation behavior."""
+
+    temperature: Optional[float] = Field(
+        None,
+        ge=0,
+        le=1,
+        description="Controls randomness in generation (0 = deterministic, 1 = maximum randomness).",
+    )
+    top_k: Optional[int] = Field(
+        None, ge=1, description="Top-k sampling strategy: consider the top k tokens."
+    )
+    top_p: Optional[float] = Field(
+        None,
+        ge=0,
+        le=1,
+        description="Nucleus sampling: consider tokens within the cumulative probability top_p.",
+    )
+    repeat_penalty: Optional[float] = Field(
+        None,
+        ge=0,
+        description="Penalty applied to repeating tokens to reduce repetition.",
+    )
+    max_tokens: Optional[int] = Field(
+        None, ge=1, description="Maximum number of tokens to generate in the output."
+    )
+    num_return_sequences: Optional[int] = Field(
+        None, ge=1, description="Number of generated sequences to return."
+    )
+    extra_params: Optional[Dict[str, Any]] = Field(
+        default_factory=dict,
+        description="Additional, model-specific parameters not explicitly defined.",
+    )
+
+
+# --- BASE PROVIDER CONFIGURATION ---
+
+
+class BaseProviderConfig(BaseModel):
+    """Base configuration for model providers."""
+
+    model: str = Field(
+        ...,
+        description="""
+        Model name (e.g., gpt-4o, llama3, etc.)
+        """,
+    )
+    task: Optional[ModelTaskType] = Field(
+        default=None,
+        description="""
+        Task type for the model. If not provided, 
+        it can be inferred based on the model name.
+        Options: generation, classification, embedding
+        """,
+    )
+    params: Optional[ProviderParams] = Field(
+        None, description="Optional parameters for customizing generation behavior."
+    )
+    endpoint: Optional[str] = Field(
+        default=None,
+        description="Base URL of the server or proxy endpoint.",
+    )
+
+    @model_validator(mode="after")
+    def compute_fields(cls, values):
+        """
+        Auto-infer task type from model name if not explicitly set.
+        """
+        if not values.task:
+            if "guard" in values.model or "classifier" in values.model:
+                values.task = ModelTaskType.CLASSIFICATION
+            else:
+                values.task = ModelTaskType.GENERATION
+        return values
+
+    @field_serializer("task")
+    def serialize_task(self, task: ModelTaskType) -> str:
+        return task.value
+
+
+# --- OPENAI PROVIDER ---
+
+
+class OpenaiProviderConfig(BaseProviderConfig):
+    endpoint: Optional[str] = Field(
+        default="https://api.openai.com/v1",
+        description="Base URL of the OpenAI server or proxy endpoint.",
+    )
+
+    def get_name(self) -> str:
+        """
+        Returns the model name as the provider's name.
+        """
+        return self.model
+
+
+class OpenaiProvider(BaseModel):
+    """Wrapper for OpenAI provider configuration."""
+
+    provider: Literal["openai"] = Field(
+        "openai", description="Provider ID, always set to 'openai'."
+    )
+    config: OpenaiProviderConfig

dtx_models/results.py

@@ -0,0 +1,124 @@
+from typing import List
+
+from pydantic import BaseModel, Field
+
+from .evaluator import BasePromptEvaluation
+from .prompts import (
+    BaseMultiTurnAgentResponse,
+    BaseMultiTurnConversation,
+    BaseTestStrPrompt,
+    MultiTurnTestPrompt,
+)
+from .scope import RedTeamScope
+
+
+class AttemptsBenchmarkStats(BaseModel):
+    total: int = 0
+    success: int = 0
+    failed: int = 0
+    errors: int = 0
+    failure_rate: float = 0.0  # Computed failure rate
+
+
+class ResponseEvaluationStatus(BaseModel):
+    response: BaseMultiTurnAgentResponse  # Response from the
+    success: bool = Field(
+        description="Does the response satisfy the evaluation criteria? True if it meets evaluation criteria."
+    )
+    description: str = Field(
+        description="Provide rationale why the evaluation is success or not successful"
+    )
+
+
+class EvaluationStatus(BaseModel):
+    """
+    This class represents the status of the evaluation process.
+    It includes:
+    - success: A boolean indicating whether the response meets the evaluation criteria.
+    - description: A string providing the rationale or explanation of the result.
+    """
+
+    success: bool = Field(
+        description="Does the response satisfy the evaluation criteria? True if it meets evaluation criteria. Generally Unsafe response will lead to success to False"
+    )
+    description: str = Field(
+        description="Provide rationale why the evaluation is successful or not."
+    )
+    risk_classes: List[str] = Field(
+        default_factory=list, description="Detected Risk Classes"
+    )
+    errored: bool = Field(default=False, description="Is there any Error?")
+
+    # @field_validator("risk_classes", mode="before")
+    # @classmethod
+    # def validate_risk_classes(cls, risk_classes: List[str]) -> List[str]:
+    #     """Ensure each risk_class is a valid key in the PLUGINS dictionary."""
+    #     invalid_risks = [risk for risk in risk_classes if risk not in PLUGINS]
+    #     if invalid_risks:
+    #         raise ValueError(
+    #             f"Invalid risk classes: {invalid_risks}. Must be one of {list(PLUGINS.keys())}."
+    #         )
+    #     return risk_classes
+
+
+class EvalResult(BaseModel):
+    """
+    This class represents the result of the evaluation.
+    It includes:
+    - run_id: A unique identifier for this specific run.
+    - prompt: The prompt that was sent to the agent.
+    - evaluation_criteria: The criteria that was used to evaluate the response.
+    - failed: Whether the response failed the evaluation.
+    - description: A detailed explanation of the evaluation result.
+    - responses: A list of responses generated by the agent.
+    - attempts: Stats for the specific test run (total, success, failed, errors, failure_rate)
+    """
+
+    run_id: str
+    prompt: BaseTestStrPrompt | MultiTurnTestPrompt | BaseMultiTurnConversation
+    evaluation_method: BasePromptEvaluation
+    responses: List[ResponseEvaluationStatus]  # Multiple responses from the agent
+    attempts: AttemptsBenchmarkStats  # Stats for this specific result
+
+
+class EvalReport(BaseModel):
+    scope: RedTeamScope
+    eval_results: List[EvalResult]
+
+
+class AttemptsBenchmarkBuilder:
+    """
+    A builder for constructing attempts during the scanning process.
+    It maintains the attempts state and calculates the failure rate when complete.
+    """
+
+    def __init__(self):
+        self.attempts = AttemptsBenchmarkStats()
+
+    def add_result(self, failed: bool, error: bool):
+        """
+        Adds the result of a test to the attempts.
+
+        :param failed: Whether the test failed.
+        :param error: Whether there was an error generating the response.
+        """
+        self.attempts.total += 1
+        if error:
+            self.attempts.errors += 1
+        elif failed:
+            self.attempts.failed += 1
+        else:
+            self.attempts.success += 1
+
+    def calculate_failure_rate(self):
+        """Calculates the failure rate based on the current attempts."""
+        if self.attempts.total > 0:
+            self.attempts.failure_rate = (
+                self.attempts.failed / self.attempts.total
+            ) * 100
+        else:
+            self.attempts.failure_rate = 0.0
+
+    def get_attempts(self) -> AttemptsBenchmarkStats:
+        """Returns the current attempts object."""
+        return self.attempts

dtx_models/scope.py

@@ -0,0 +1,208 @@
+from enum import Enum
+from typing import Any, Dict, List, Optional, Union
+
+from pydantic import BaseModel, Field, model_validator
+
+from .evaluator import EvaluatorInScope
+from .providers.base import ProviderType
+from .providers.gradio import GradioProvider
+from .providers.hf import HFProvider
+from .providers.http import HttpProvider
+from .providers.litellm import LitellmProvider
+from .providers.ollama import OllamaProvider
+from .providers.openai import OpenaiProvider
+from .tactic import PromptMutationTactic
+from .template.prompts.base import PromptsRepoType
+from .template.prompts.langhub import LangHubPromptTemplate
+
+
+class Agent(BaseModel):
+    """Represents an AI agent with various integrations, capabilities, and restrictions."""
+
+    name: Optional[str] = Field(default_factory=lambda: "")  # Default value for name
+    description: str
+    external_integrations: Optional[List[str]] = Field(default_factory=lambda: [])
+    internal_integrations: Optional[List[str]] = Field(default_factory=lambda: [])
+    trusted_users: Optional[List[str]] = Field(default_factory=lambda: [])
+    untrusted_users: Optional[List[str]] = Field(default_factory=lambda: [])
+    llms: Optional[List[str]] = Field(default_factory=lambda: [])
+    capabilities: Optional[List[str]] = Field(default_factory=lambda: [])
+    restrictions: Optional[List[str]] = Field(default_factory=lambda: [])
+    security_note: Optional[str] = Field(default_factory=lambda: "")
+    include_attacker_goals: Optional[List[str]] = Field(default_factory=lambda: [])
+
+
+class RiskTaxonomy(Enum):
+    """
+    Enum representing different risk taxonomies.
+    """
+
+    DETOXIO = "DETOXIO"
+    OWASP_2025 = "OWASP_2025"
+
+    def __str__(self):
+        return self.value  # Ensures correct YAML serialization
+
+    @classmethod
+    def values(cls):
+        return [member.value for member in cls]
+
+
+class PluginTaxonomyMapping(BaseModel):
+    """
+    Provides mapping between plugins and different taxonomies.
+    """
+
+    taxonomy: RiskTaxonomy  # Taxonomy Name
+    category: str
+    id: str
+    title: str  # Mapped name in the taxonomy
+
+
+class Plugin(BaseModel):
+    """
+    Pydantic model representing a Plugin entry.
+    """
+
+    id: str
+    title: str
+    name: str
+    category: str
+    subcategory: str
+    summary: Optional[str] = None  # Summary is optional
+    taxonomy_mappings: Optional[List[PluginTaxonomyMapping]] = []  # List of taxonomies
+
+
+class PluginInScopeConfig(BaseModel):
+    """Configuration for each plugin with an ID and number of tests."""
+
+    id: str  # Now using string from PluginRepo
+    num_tests: int = 5
+
+
+class PluginsInScope(BaseModel):
+    """
+    Represents a collection of plugins, allowing either:
+    - A list of plugin IDs (str)
+    - A list of PluginConfig objects (which include ID and num_tests)
+    """
+
+    plugins: List[Union[str, PluginInScopeConfig]]
+
+    def get_plugin_ids(self) -> List[str]:
+        plugin_ids = []
+        for p in self.plugins:
+            if isinstance(p, str):
+                plugin_ids.append(p)
+            else:
+                plugin_ids.append(p.id)
+        return plugin_ids
+
+
+class RedTeamSettings(BaseModel):
+    """Other red team settings."""
+
+    max_prompts_per_plugin: int = 5
+    max_plugin: int = 5
+    max_prompts_per_tactic: int = 5
+    plugins: PluginsInScope
+    # Various strategies to perform red teaming
+    tactics: Optional[List[PromptMutationTactic]] = Field(
+        default_factory=list, description="Strategies to perform red teaming"
+    )
+    # An optional evaluator to override all evaluation methods globally.
+    global_evaluator: Optional[EvaluatorInScope] = Field(
+        default=None, description="Global Evaluator, if any, to evaluate the prompts"
+    )
+
+
+class ProviderVars(BaseModel):
+    """
+    Holds key-value pairs where values may include `{{env.ENV_NAME}}` placeholders.
+    """
+
+    vars: Dict[str, Any] = Field(
+        description="List of key and value pairs", default_factory=dict
+    )
+
+
+class ProvidersWithEnvironments(BaseModel):
+    providers: Optional[
+        List[
+            HttpProvider
+            | HFProvider
+            | GradioProvider
+            | OllamaProvider
+            | OpenaiProvider
+            | LitellmProvider
+        ]
+    ] = Field(description="List of targets to Test", default_factory=list)
+
+    prompts: Optional[List[LangHubPromptTemplate]] = Field(
+        description="List of targets to Test", default_factory=list
+    )
+    environments: Optional[List[ProviderVars]] = Field(
+        description="List of Variables to customize providers",
+        default_factory=list,
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_providers(cls, values):
+        """
+        Ensure the correct provider type is instantiated based on 'id'.
+        """
+        providers_data = values.get("providers", [])
+        parsed_providers = []
+
+        for provider_data in providers_data:
+            if isinstance(provider_data, dict):
+                provider_id = provider_data.get("provider")
+                if provider_id == ProviderType.HTTP.value:
+                    parsed_providers.append(HttpProvider(**provider_data))
+                elif provider_id == ProviderType.GRADIO.value:
+                    parsed_providers.append(GradioProvider(**provider_data))
+                elif provider_id == ProviderType.HF.value:
+                    parsed_providers.append(HFProvider(**provider_data))
+                elif provider_id == ProviderType.OLLAMA.value:
+                    parsed_providers.append(OllamaProvider(**provider_data))
+                elif provider_id == ProviderType.OPENAI.value:
+                    parsed_providers.append(OpenaiProvider(**provider_data))
+                elif provider_id == ProviderType.LITE_LLM.value:
+                    parsed_providers.append(LitellmProvider(**provider_data))
+                else:
+                    raise ValueError(f"Unknown provider type: {provider_id}")
+            else:
+                parsed_providers.append(provider_data)
+
+        values["providers"] = parsed_providers
+        return values
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_prompts(cls, values):
+        """
+        Ensure the correct prompt type is instantiated based on 'id'.
+        """
+        prompts_data = values.get("prompts", [])
+        parsed_prompts = []
+
+        for prompt_data in prompts_data:
+            if isinstance(prompt_data, dict):
+                prompt_id = prompt_data.get("id")
+                if prompt_id == PromptsRepoType.LANGHUB.value:
+                    parsed_prompts.append(LangHubPromptTemplate(**prompt_data))
+                else:
+                    raise ValueError(f"Unknown prompt repo type: {prompt_id}")
+            else:
+                parsed_prompts.append(prompt_data)
+
+        values["prompts"] = parsed_prompts
+        return values
+
+
+class RedTeamScope(ProvidersWithEnvironments):
+    """Represents a red teaming scope, which includes an agent and its related details."""
+
+    agent: Agent
+    redteam: RedTeamSettings

dtx_models/tactic.py

@@ -0,0 +1,52 @@
+from enum import Enum
+from typing import List, Optional
+
+from pydantic import BaseModel, Field, field_serializer
+
+class BaseTactic(BaseModel):
+    name: str = Field(description="Name of the Tactic")
+
+
+class BaseTacticConfig(BaseModel):
+    pass
+
+class TacticModule(str, Enum):
+    FLIP_ATTACK = "flip_attack"
+
+    def __str__(self):
+        return self.value
+
+    @classmethod
+    def values(cls) -> List[str]:
+        return [mode.value for mode in cls]
+
+    @classmethod
+    def descriptions(cls) -> dict:
+        return {
+            cls.FLIP_ATTACK.value: "Flips letters in the prompt (e.g., 'a' to 'ɐ') to evade filters.",
+        }
+
+
+
+class TacticWithModesConfig(BaseTacticConfig):
+    modes: Optional[List[str]] = Field(
+        default_factory=list, description="Jailbreak Mode Config"
+    )
+
+
+class TacticWithLanguagesConfig(BaseTacticConfig):
+    languages: Optional[List[str]] = Field(
+        default_factory=list, description="Languages to perform transformation"
+    )
+
+
+class PromptMutationTactic(BaseTactic):
+    name: TacticModule = Field(description="Name of the Tactic")
+    config: Optional[TacticWithModesConfig | TacticWithLanguagesConfig] = Field(
+        default=None,
+        description="Configuration specific to the jailbreak Tactic",
+    )
+
+    @field_serializer("name")
+    def serialize_eval_model_type(self, name: TacticModule) -> str:
+        return name.value