MAINT: Fix docstrings for pyrit/chat_message_normalizer and pyrit/cli (microsoft#1196)

riyosha · web-flow · commit 48ebd5c1ac23 · 2025-11-20T14:22:14.000-08:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -251,7 +251,7 @@ extend-select = [
 # Temporary ignores for pyrit/ subdirectories until issue #1176
 # https://github.com/Azure/PyRIT/issues/1176 is fully resolved
 # TODO: Remove these ignores once the issues are fixed
-"pyrit/{auxiliary_attacks,chat_message_normalizer,cli,embedding,exceptions,executor,memory,models,prompt_converter,prompt_normalizer,prompt_target,scenarios,score,setup,ui}/**/*.py" = ["D101", "D102", "D103", "D104", "D105", "D106", "D107", "D401", "D404", "D417", "D418", "DOC102", "DOC201", "DOC202", "DOC402", "DOC501"]
+"pyrit/{auxiliary_attacks,exceptions,executor,memory,models,prompt_converter,prompt_normalizer,prompt_target,scenarios,score,setup,ui}/**/*.py" = ["D101", "D102", "D103", "D104", "D105", "D106", "D107", "D401", "D404", "D417", "D418", "DOC102", "DOC201", "DOC202", "DOC402", "DOC501"]
 "pyrit/__init__.py" = ["D104"]
 
 [tool.ruff.lint.pydocstyle]
diff --git a/pyrit/chat_message_normalizer/__init__.py b/pyrit/chat_message_normalizer/__init__.py
@@ -2,7 +2,7 @@
 # Licensed under the MIT license.
 
 """
-This module contains the functionality to normalize chat messages into compatible formats for targets.
+Functionality to normalize chat messages into compatible formats for targets.
 """
 
 from pyrit.chat_message_normalizer.chat_message_normalizer import ChatMessageNormalizer
diff --git a/pyrit/chat_message_normalizer/chat_message_nop.py b/pyrit/chat_message_normalizer/chat_message_nop.py
@@ -12,7 +12,7 @@ class ChatMessageNop(ChatMessageNormalizer[list[ChatMessage]]):
 
     def normalize(self, messages: list[ChatMessage]) -> list[ChatMessage]:
         """
-        Returns the same list as was passed in.
+        Return the same list as was passed in.
 
         Args:
             messages (list[ChatMessage]): The list of messages to normalize.
diff --git a/pyrit/chat_message_normalizer/chat_message_normalizer.py b/pyrit/chat_message_normalizer/chat_message_normalizer.py
@@ -10,23 +10,30 @@
 
 
 class ChatMessageNormalizer(abc.ABC, Generic[T]):
+    """
+    Abstract base class for normalizing chat messages for model or target compatibility.
+    """
+
     @abc.abstractmethod
     def normalize(self, messages: list[ChatMessage]) -> T:
         """
-        Normalizes the list of chat messages into a compatible format for the model or target.
+        Normalize the list of chat messages into a compatible format for the model or target.
         """
 
     @staticmethod
     def squash_system_message(messages: list[ChatMessage], squash_function) -> list[ChatMessage]:
         """
-        Combines the system message into the first user request.
+        Combine the system message into the first user request.
 
         Args:
             messages: The list of chat messages.
             squash_function: The function to combine the system message with the user message.
 
         Returns:
             The list of chat messages with squashed system messages.
+
+        Raises:
+            ValueError: If the chat message list is empty.
         """
         if not messages:
             raise ValueError("ChatMessage list cannot be empty")
diff --git a/pyrit/chat_message_normalizer/chat_message_normalizer_chatml.py b/pyrit/chat_message_normalizer/chat_message_normalizer_chatml.py
@@ -40,6 +40,9 @@ def from_chatml(content: str) -> list[ChatMessage]:
 
         Returns:
             list[ChatMessage]: The list of chat messages.
+
+        Raises:
+            ValueError: If the input content is invalid.
         """
         messages: list[ChatMessage] = []
         matches = list(re.finditer(r"<\|im_start\|>(.*?)<\|im_end\|>", content, re.DOTALL | re.MULTILINE))
diff --git a/pyrit/chat_message_normalizer/chat_message_normalizer_tokenizer.py b/pyrit/chat_message_normalizer/chat_message_normalizer_tokenizer.py
@@ -9,14 +9,14 @@
 
 class ChatMessageNormalizerTokenizerTemplate(ChatMessageNormalizer[str]):
     """
-    This class enables you to apply the chat template stored in a Hugging Face tokenizer
+    Enable application of the chat template stored in a Hugging Face tokenizer
     to a list of chat messages. For more details, see
     https://huggingface.co/docs/transformers/main/en/chat_templating.
     """
 
     def __init__(self, tokenizer: PreTrainedTokenizerBase):
         """
-        Initializes an instance of the ChatMessageNormalizerTokenizerTemplate class.
+        Initialize an instance of the ChatMessageNormalizerTokenizerTemplate class.
 
         Args:
             tokenizer (PreTrainedTokenizerBase): A Hugging Face tokenizer.
@@ -25,7 +25,7 @@ def __init__(self, tokenizer: PreTrainedTokenizerBase):
 
     def normalize(self, messages: list[ChatMessage]) -> str:
         """
-        Applies the chat template stored in the tokenizer to a list of chat messages.
+        Apply the chat template stored in the tokenizer to a list of chat messages.
 
         Args:
             messages (list[ChatMessage]): A list of ChatMessage objects.
diff --git a/pyrit/chat_message_normalizer/generic_system_squash.py b/pyrit/chat_message_normalizer/generic_system_squash.py
@@ -6,9 +6,13 @@
 
 
 class GenericSystemSquash(ChatMessageNormalizer[list[ChatMessage]]):
+    """
+    Normalizer that combines the first system message with the first user message using generic instruction tags.
+    """
+
     def normalize(self, messages: list[ChatMessage]) -> list[ChatMessage]:
         """
-        Returns the first system message combined with the first user message.
+        Return the first system message combined with the first user message.
 
         The format of the result uses generic instruction tags.
 
@@ -28,7 +32,7 @@ def combine_system_user_message(
         system_message: ChatMessage, user_message: ChatMessage, msg_type: ChatMessageRole = "user"
     ) -> ChatMessage:
         """
-        Combines the system message with the user message.
+        Combine the system message with the user message.
 
         Args:
             system_message (str): The system message.
diff --git a/pyrit/cli/__init__.py b/pyrit/cli/__init__.py
@@ -2,6 +2,6 @@
 # Licensed under the MIT License.
 
 """
-This module provides command-line interface functionalities for PyRIT.
+Module to provide command-line interface functionalities for PyRIT.
 The CLI module is currently experimental.
 """
diff --git a/pyrit/cli/frontend_core.py b/pyrit/cli/frontend_core.py
@@ -30,8 +30,11 @@
 
     # Create a dummy termcolor module for fallback
     class termcolor:  # type: ignore
+        """Dummy termcolor fallback for colored printing if termcolor is not installed."""
+
         @staticmethod
         def cprint(text: str, color: str = None, attrs: list = None) -> None:  # type: ignore
+            """Print text without color."""
             print(text)
 
 
@@ -471,7 +474,7 @@ def validate_integer(value: str, *, name: str = "value", min_value: Optional[int
 
 def _argparse_validator(validator_func: Callable[..., Any]) -> Callable[[Any], Any]:
     """
-    Decorator to convert ValueError to argparse.ArgumentTypeError.
+    Adapt a validator to argparse by converting ValueError to ArgumentTypeError.
 
     This decorator adapts our keyword-only validators for use with argparse's type= parameter.
     It handles two challenges:
@@ -498,6 +501,9 @@ def _argparse_validator(validator_func: Callable[..., Any]) -> Callable[[Any], A
         - Accepts a single positional argument (for argparse compatibility)
         - Calls validator_func with the correct keyword argument
         - Raises ArgumentTypeError instead of ValueError
+
+    Raises:
+        ValueError: If validator_func has no parameters.
     """
     import inspect
 
diff --git a/pyrit/cli/pyrit_scan.py b/pyrit/cli/pyrit_scan.py
@@ -15,7 +15,12 @@
 
 
 def parse_args(args=None) -> Namespace:
-    """Parse command-line arguments for the PyRIT scanner."""
+    """
+    Parse command-line arguments for the PyRIT scanner.
+
+    Returns:
+        Namespace: Parsed command-line arguments.
+    """
     parser = ArgumentParser(
         prog="pyrit_scan",
         description="""PyRIT Scanner - Run security scenarios against AI systems
@@ -120,7 +125,12 @@ def parse_args(args=None) -> Namespace:
 
 
 def main(args=None) -> int:
-    """Main entry point for the PyRIT scanner CLI."""
+    """
+    Start the PyRIT scanner CLI.
+
+    Returns:
+        int: Exit code (0 for success, 1 for error).
+    """
     print("Starting PyRIT...")
     sys.stdout.flush()
 
diff --git a/pyrit/cli/pyrit_shell.py b/pyrit/cli/pyrit_shell.py
@@ -374,7 +374,12 @@ def do_help(self, arg):
             super().do_help(arg)
 
     def do_exit(self, arg):
-        """Exit the shell. Aliases: quit, q."""
+        """
+        Exit the shell. Aliases: quit, q.
+
+        Returns:
+            bool: True to exit the shell.
+        """
         print("\nGoodbye!")
         return True
 
@@ -390,11 +395,21 @@ def do_clear(self, arg):
     do_EOF = do_exit  # Ctrl+D on Unix, Ctrl+Z on Windows
 
     def emptyline(self) -> bool:
-        """Don't repeat last command on empty line."""
+        """
+        Don't repeat last command on empty line.
+
+        Returns:
+            bool: False to prevent repeating the last command.
+        """
         return False
 
     def default(self, line):
-        """Handle unknown commands and convert hyphens to underscores."""
+        """
+        Handle unknown commands and convert hyphens to underscores.
+
+        Returns:
+            None
+        """
         # Try converting hyphens to underscores for command lookup
         parts = line.split(None, 1)
         if parts:
@@ -411,7 +426,12 @@ def default(self, line):
 
 
 def main():
-    """Entry point for pyrit_shell."""
+    """
+    Entry point for pyrit_shell.
+
+    Returns:
+        int: Exit code.
+    """
     import argparse
 
     parser = argparse.ArgumentParser(
diff --git a/pyrit/embedding/__init__.py b/pyrit/embedding/__init__.py
@@ -1,6 +1,9 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+"""Embedding module for PyRI to provide Azure and OpenAI text embedding classes."""
+
+
 from pyrit.embedding.azure_text_embedding import AzureTextEmbedding
 from pyrit.embedding.openai_text_embedding import OpenAiTextEmbedding
 
diff --git a/pyrit/embedding/azure_text_embedding.py b/pyrit/embedding/azure_text_embedding.py
@@ -11,6 +11,10 @@
 
 
 class AzureTextEmbedding(_TextEmbedding):
+    """
+    Provide text embedding functionality using Azure OpenAI services.
+    """
+
     API_KEY_ENVIRONMENT_VARIABLE: str = "AZURE_OPENAI_EMBEDDING_KEY"
     ENDPOINT_URI_ENVIRONMENT_VARIABLE: str = "AZURE_OPENAI_EMBEDDING_ENDPOINT"
     DEPLOYMENT_ENVIRONMENT_VARIABLE: str = "AZURE_OPENAI_EMBEDDING_DEPLOYMENT"
@@ -35,6 +39,9 @@ def __init__(
                         Defaults to environment variable.
             api_version: The API version to use. Defaults to "2024-02-01".
             use_entra_auth: Whether to use Entra authentication. Defaults to False.
+
+        Raises:
+            ValueError: If using Entra ID auth and an api_key is also provided.
         """
         endpoint = default_values.get_required_value(
             env_var_name=self.ENDPOINT_URI_ENVIRONMENT_VARIABLE, passed_value=endpoint
diff --git a/pyrit/embedding/openai_text_embedding.py b/pyrit/embedding/openai_text_embedding.py
@@ -7,12 +7,15 @@
 
 
 class OpenAiTextEmbedding(_TextEmbedding):
+    """
+    Provides text embedding functionality using OpenAI services.
+    """
+
     def __init__(self, *, model: str, api_key: str) -> None:
         """
         Generate embedding using OpenAI API.
 
         Args:
-            api_version: The API version to use
             model: The model to use
             api_key: The API key to use
         """
