MAINT: Minor pyrit.model module reorganization #1999
Conversation
Move ScenarioResult into pyrit.models.results and ConversationReference / ChatMessage into pyrit.models.messages. These symbols are re-exported from pyrit.models, so the canonical import stays `from pyrit.models import X` and no deprecation shim is needed at the old submodule paths. Switch all external consumers and tests to the short-form package-root import per the style guide, keep same-package imports on the specific module to avoid cycles, and clarify that ChatMessage is the OpenAI Chat Completions wire shape (distinct from the domain Message/MessagePiece). Adds a dedicated Conversation model test. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
rlundeen2
changed the title
| SeedSimulatedConversation, | ||
| SeedType, | ||
| ) | ||
| from pyrit.models.scenario_result import ScenarioRunState |
There was a problem hiding this comment.
why don't we do something like
PyRIT/pyrit/models/__init__.py
Line 215 in c39db20
if not, we should at least label the PR as breaking
There was a problem hiding this comment.
Rich agrees with this comment but it is copilot generated: We went the other direction and removed the deep-path shims entirely instead of adding more. The rule we settled on: if a symbol only moves within pyrit.models and is still re-exported from the package root, there's no shim — from pyrit.models import ScenarioRunState (and every other symbol) keeps working, and that's the only import form we support/document. Warning shims are reserved for symbols that actually leave the package (e.g. the pyrit.memory.storage moves). For consistency I also deleted the older silent shims that were still lingering (attack_result, message, message_piece, strategy_result). The only thing that breaks is the uncanonical deep path from pyrit.models.scenario_result import ..., which was never the documented import.
There was a problem hiding this comment.
(it got messy having so many deprecated shims, and I think when the import mostly stays the same it's likely okay to just move it)
adrian-gavrila
left a comment
adrian-gavrila
left a comment
There was a problem hiding this comment.
Overll looks good, left a couple of small comments to consider
| group_message_pieces_into_conversations, | ||
| sort_message_pieces, | ||
| ) | ||
| from pyrit.models.messages.chat_message import ( |
There was a problem hiding this comment.
Should we add a back-compat shim at every old path? message.py, message_piece.py, attack_result.py, strategy_result.py all still have them.
from pyrit.models import X still resolves for all three so nothing in-repo breaks, but from pyrit.models.scenario_result import ScenarioResult is now a hard ModuleNotFoundError for any downstream code on the old deep path. I think we should either add the 3 shims or add a note about the breaking change?
There was a problem hiding this comment.
Rich agrees with this comment but it is copilot generated: Good catch on the inconsistency. Rather than add three more shims, we removed the deep-path shims altogether — including the four existing silent ones (message, message_piece, attack_result, strategy_result). Principle: internal pyrit.models reshuffles stay shim-free because everything is re-exported from pyrit.models (the canonical import); warning shims are only for symbols that move to a different package. So the only thing that breaks is from pyrit.models.scenario_result import X, which was never a supported path — from pyrit.models import X is unchanged.
Address PR review: the deep-path back-compat shims for purely internal pyrit.models relocations add confusion rather than value, since every symbol stays re-exported from pyrit.models. Delete the lingering silent shims at attack_result.py, message.py, message_piece.py, and strategy_result.py so the single canonical import (from pyrit.models import X) is the only supported form. Warning shims remain only where a symbol genuinely left the package (the pyrit.memory.storage moves). Export ToolCall from pyrit.models so it is reachable via the package root like every other messages symbol, and fold the one-class conversation.py module into conversations.py to remove the near-duplicate module name. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
3 participants
Moving some modules around to where they make sense. Deprecation seemed more confusing than helpful since everything is imported from the base.