Skip to content
Merged
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
35 changes: 18 additions & 17 deletions pyrit/memory/memory_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -401,14 +401,17 @@ def __init__(self, *, entry: Score) -> None:
self.score_rationale = entry.score_rationale
self.score_metadata = entry.score_metadata or {}
normalized_scorer = entry.scorer_class_identifier
# Ensure eval_hash is set before truncation so it survives the DB round-trip
if normalized_scorer.eval_hash is None:
normalized_scorer = normalized_scorer.with_eval_hash(
ScorerEvaluationIdentifier(normalized_scorer).eval_hash
if normalized_scorer is None:
self.scorer_class_identifier = {}
else:
# Ensure eval_hash is set before truncation so it survives the DB round-trip
if normalized_scorer.eval_hash is None:
normalized_scorer = normalized_scorer.with_eval_hash(
ScorerEvaluationIdentifier(normalized_scorer).eval_hash
)
self.scorer_class_identifier = normalized_scorer.model_dump(
context={"max_value_length": MAX_IDENTIFIER_VALUE_LENGTH},
)
self.scorer_class_identifier = normalized_scorer.model_dump(
context={"max_value_length": MAX_IDENTIFIER_VALUE_LENGTH},
)
self.prompt_request_response_id = entry.message_piece_id if entry.message_piece_id else None
self.timestamp = entry.timestamp
# Store in both columns for backward compatibility
Expand Down Expand Up @@ -439,7 +442,7 @@ def get_score(self) -> Score:
score_category=self.score_category,
score_rationale=self.score_rationale,
score_metadata=self.score_metadata,
scorer_class_identifier=scorer_identifier, # type: ignore[ty:invalid-argument-type]
scorer_class_identifier=scorer_identifier,
message_piece_id=self.prompt_request_response_id,
timestamp=_ensure_utc(self.timestamp),
objective=self.objective,
Expand Down Expand Up @@ -975,8 +978,8 @@ class ScenarioResultEntry(Base):
scenario_init_data (dict): Optional initialization parameters used to configure the scenario.
objective_target_identifier (dict): Identifier for the target being evaluated in the scenario.
objective_scorer_identifier (dict): Optional identifier for the scorer used to evaluate results.
scenario_run_state (Literal["CREATED", "IN_PROGRESS", "COMPLETED", "FAILED"]): Current execution state
of the scenario.
scenario_run_state (str): Current execution state of the scenario
(one of CREATED, IN_PROGRESS, COMPLETED, FAILED, CANCELLED).
attack_results_json (str): JSON-serialized dictionary mapping attack names to conversation IDs.
Format: {"attack_name": ["conversation_id1", "conversation_id2", ...]}.
The full AttackResult objects are stored in AttackResultEntries and can be queried by conversation_id.
Expand All @@ -1003,9 +1006,7 @@ class ScenarioResultEntry(Base):
scenario_init_data: Mapped[dict[str, Any] | None] = mapped_column(JSON, nullable=True)
objective_target_identifier: Mapped[dict[str, str]] = mapped_column(JSON, nullable=False)
objective_scorer_identifier: Mapped[dict[str, str] | None] = mapped_column(JSON, nullable=True)
scenario_run_state: Mapped[Literal["CREATED", "IN_PROGRESS", "COMPLETED", "FAILED", "CANCELLED"]] = mapped_column(
String, nullable=False, default="CREATED"
)
scenario_run_state: Mapped[str] = mapped_column(String, nullable=False, default="CREATED")
attack_results_json: Mapped[str] = mapped_column(Unicode, nullable=False)
display_group_map_json: Mapped[str | None] = mapped_column(Unicode, nullable=True)
labels: Mapped[dict[str, str] | None] = mapped_column(JSON, nullable=True)
Expand Down Expand Up @@ -1072,7 +1073,7 @@ def __init__(self, *, entry: ScenarioResult) -> None:
self.attack_results_json = json.dumps(serialized_attack_results)

# Serialize display_group_map if present
self.display_group_map_json = json.dumps(entry._display_group_map) if entry._display_group_map else None
self.display_group_map_json = json.dumps(entry.display_group_map) if entry.display_group_map else None

self.error_message = entry.error_message
self.error_type = entry.error_type
Expand Down Expand Up @@ -1126,14 +1127,14 @@ def get_scenario_result(self) -> ScenarioResult:
attack_results=attack_results,
objective_scorer_identifier=scorer_identifier,
scenario_run_state=self.scenario_run_state,
labels=self.labels,
labels=self.labels or {},
creation_time=self.timestamp,
number_tries=self.number_tries,
completion_time=self.completion_time,
display_group_map=display_group_map,
display_group_map=display_group_map or {},
error_message=self.error_message,
error_type=self.error_type,
metadata=dict(self.scenario_metadata) if self.scenario_metadata else None,
metadata=dict(self.scenario_metadata) if self.scenario_metadata else {},
)

def get_conversation_ids_by_attack_name(self) -> dict[str, list[str]]:
Expand Down
2 changes: 1 addition & 1 deletion pyrit/models/identifiers/component_identifier.py
Original file line number Diff line number Diff line change
Expand Up @@ -168,7 +168,7 @@ class ComponentIdentifier(BaseModel):
#: may have been truncated.
hash: Optional[str] = None
#: Version tag for storage. Not included in the content hash.
pyrit_version: str = Field(default_factory=lambda: pyrit.__version__)
pyrit_version: str = Field(default=pyrit.__version__)
#: Evaluation hash. Computed by EvaluationIdentifier subclasses and attached
#: to the identifier so it survives DB round-trips with truncated params.
eval_hash: Optional[str] = None
Expand Down
Loading
Loading