pre-commit

Author: rlundeen2

pyrit/memory/memory_models.py

@@ -1001,6 +1001,9 @@ def __init__(self, *, entry: ScenarioResult) -> None:
 
         Args:
             entry (ScenarioResult): The scenario result object to convert into a database entry.
+
+        Raises:
+            ValueError: If ``entry.objective_target_identifier`` is ``None``.
         """
         self.id = entry.id
         self.scenario_name = entry.scenario_identifier.name

pyrit/models/attack_result.py

@@ -8,19 +8,17 @@
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Optional, TypeVar
-
+from typing import Any, Optional, TypeVar
+
+from pyrit.common.deprecation import print_deprecation_message
+from pyrit.identifiers.atomic_attack_identifier import build_atomic_attack_identifier
+from pyrit.identifiers.component_identifier import ComponentIdentifier
+from pyrit.models.conversation_reference import ConversationReference, ConversationType
+from pyrit.models.message_piece import MessagePiece
+from pyrit.models.retry_event import RetryEvent
+from pyrit.models.score import Score
 from pyrit.models.strategy_result import StrategyResult
 
-if TYPE_CHECKING:
-    from pyrit.identifiers.component_identifier import ComponentIdentifier
-    from pyrit.models.conversation_reference import ConversationReference
-    from pyrit.models.message_piece import MessagePiece
-    from pyrit.models.retry_event import RetryEvent
-    from pyrit.models.score import Score
-
-from pyrit.models.conversation_reference import ConversationType
-
 AttackResultT = TypeVar("AttackResultT", bound="AttackResult")
 
 
@@ -119,8 +117,6 @@ def attack_identifier(self) -> Optional[ComponentIdentifier]:
             Optional[ComponentIdentifier]: The attack strategy identifier, or ``None``.
 
         """
-        from pyrit.common.deprecation import print_deprecation_message
-
         print_deprecation_message(
             old_item="AttackResult.attack_identifier",
             new_item="AttackResult.atomic_attack_identifier or get_attack_strategy_identifier()",
@@ -231,8 +227,6 @@ def to_dict(self) -> dict[str, Any]:
         Returns:
             dict[str, Any]: Serialized payload suitable for REST APIs or persistence.
         """
-        from pyrit.models.conversation_reference import ConversationReference
-
         return {
             "conversation_id": self.conversation_id,
             "objective": self.objective,
@@ -274,12 +268,6 @@ def from_dict(cls, data: dict[str, Any]) -> AttackResult:
         Returns:
             AttackResult: Reconstructed instance.
         """
-        from pyrit.identifiers.component_identifier import ComponentIdentifier
-        from pyrit.models.conversation_reference import ConversationReference
-        from pyrit.models.message_piece import MessagePiece
-        from pyrit.models.retry_event import RetryEvent
-        from pyrit.models.score import Score
-
         return cls(
             conversation_id=data["conversation_id"],
             objective=data["objective"],
@@ -330,16 +318,12 @@ def _add_attack_identifier_compat(cls: type) -> type:
     def wrapped_init(self: Any, *args: Any, **kwargs: Any) -> None:
         attack_identifier = kwargs.pop("attack_identifier", None)
         if attack_identifier is not None:
-            from pyrit.common.deprecation import print_deprecation_message
-
             print_deprecation_message(
                 old_item="AttackResult(attack_identifier=...)",
                 new_item="AttackResult(atomic_attack_identifier=...)",
                 removed_in="0.15.0",
             )
             if kwargs.get("atomic_attack_identifier") is None:
-                from pyrit.identifiers.atomic_attack_identifier import build_atomic_attack_identifier
-
                 kwargs["atomic_attack_identifier"] = build_atomic_attack_identifier(
                     attack_identifier=attack_identifier,
                 )

pyrit/models/message_piece.py

@@ -9,12 +9,13 @@
 from uuid import uuid4
 
 from pyrit.common.deprecation import print_deprecation_message
+from pyrit.identifiers.component_identifier import ComponentIdentifier
+from pyrit.models.data_type_serializer import data_serializer_factory
 from pyrit.models.literals import ChatMessageRole, PromptDataType, PromptResponseError
+from pyrit.models.score import Score
 
 if TYPE_CHECKING:
-    from pyrit.identifiers.component_identifier import ComponentIdentifier
     from pyrit.models.message import Message
-    from pyrit.models.score import Score
 
 Originator = Literal["attack", "converter", "undefined", "scorer"]
 """Deprecated: The Originator type alias will be removed in a future release."""
@@ -218,8 +219,6 @@ async def set_sha256_values_async(self) -> None:
         Note, this method is async due to the blob retrieval. And because of that, we opted
         to take it out of main and setter functions. The disadvantage is that it must be explicitly called.
         """
-        from pyrit.models.data_type_serializer import data_serializer_factory
-
         original_serializer = data_serializer_factory(
             category="prompt-memory-entries",
             data_type=self.original_value_data_type,
@@ -365,9 +364,6 @@ def from_dict(cls, data: dict[str, Any]) -> MessagePiece:
         Returns:
             MessagePiece: Reconstructed instance.
         """
-        from pyrit.identifiers.component_identifier import ComponentIdentifier
-        from pyrit.models.score import Score
-
         return cls(
             id=data.get("id"),
             role=data.get("role", "user"),

tests/unit/models/test_attack_result.py

@@ -8,7 +8,10 @@
 from pyrit.identifiers.atomic_attack_identifier import build_atomic_attack_identifier
 from pyrit.memory.memory_models import AttackResultEntry
 from pyrit.models.attack_result import AttackOutcome, AttackResult
+from pyrit.models.conversation_reference import ConversationReference, ConversationType
+from pyrit.models.message_piece import MessagePiece
 from pyrit.models.retry_event import RetryEvent
+from pyrit.models.score import Score
 
 
 class TestAttackResultDeprecation:
@@ -354,11 +357,6 @@ def test_traceback_truncation(self) -> None:
 
 
 def test_to_dict_from_dict_roundtrip():
-    from pyrit.identifiers.component_identifier import ComponentIdentifier
-    from pyrit.models.conversation_reference import ConversationReference, ConversationType
-    from pyrit.models.message_piece import MessagePiece
-    from pyrit.models.score import Score
-
     scorer_id = ComponentIdentifier(
         class_name="SelfAskTrueFalseScorer",
         class_module="pyrit.score",