MAINT: Fixed doscstrings for pyrit/analytics+auth (microsoft#1192)

Co-authored-by: Roman Lutz <romanlutz13@gmail.com>

Author: riyosha

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/{analytics,auth,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,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/__init__.py" = ["D104"]
 
 [tool.ruff.lint.pydocstyle]

pyrit/analytics/__init__.py

@@ -1,6 +1,9 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+"""Analytics module for PyRIT conversation and result analysis."""
+
+
 from pyrit.analytics.conversation_analytics import ConversationAnalytics
 from pyrit.analytics.result_analysis import analyze_results, AttackStats
 from pyrit.analytics.text_matching import (

pyrit/analytics/conversation_analytics.py

@@ -18,7 +18,7 @@ class ConversationAnalytics:
 
     def __init__(self, *, memory_interface: MemoryInterface):
         """
-        Initializes the ConversationAnalytics with a memory interface for data access.
+        Initialize the ConversationAnalytics with a memory interface for data access.
 
         Args:
             memory_interface (MemoryInterface): An instance of MemoryInterface for accessing conversation data.
@@ -29,7 +29,7 @@ def get_prompt_entries_with_same_converted_content(
         self, *, chat_message_content: str
     ) -> list[ConversationMessageWithSimilarity]:
         """
-        Retrieves chat messages that have the same converted content.
+        Retrieve chat messages that have the same converted content.
 
         Args:
             chat_message_content (str): The content of the chat message to find similar messages for.
@@ -58,7 +58,7 @@ def get_similar_chat_messages_by_embedding(
         self, *, chat_message_embedding: list[float], threshold: float = 0.8
     ) -> list[EmbeddingMessageWithSimilarity]:
         """
-        Retrieves chat messages that are similar to the given embedding based on cosine similarity.
+        Retrieve chat messages that are similar to the given embedding based on cosine similarity.
 
         Args:
             chat_message_embedding (List[float]): The embedding of the chat message to find similar messages for.

pyrit/analytics/result_analysis.py

@@ -10,6 +10,8 @@
 
 @dataclass
 class AttackStats:
+    """Statistics for attack analysis results."""
+
     success_rate: Optional[float]
     total_decided: int
     successes: int

pyrit/auth/__init__.py

@@ -2,7 +2,7 @@
 # Licensed under the MIT license.
 
 """
-This module contains authentication functionality for a variety of services.
+Authentication functionality for a variety of services.
 """
 
 from pyrit.auth.authenticator import Authenticator

pyrit/auth/authenticator.py

@@ -14,8 +14,20 @@ class Authenticator(abc.ABC):
 
     @abstractmethod
     def refresh_token(self) -> str:
+        """
+        Refresh the authentication token.
+
+        Returns:
+            str: The refreshed authentication token.
+        """
         raise NotImplementedError("refresh_token method not implemented")
 
     @abstractmethod
     def get_token(self) -> str:
+        """
+        Get the current authentication token.
+
+        Returns:
+            str: The current authentication token.
+        """
         raise NotImplementedError("get_token method not implemented")

pyrit/auth/azure_auth.py

@@ -31,11 +31,21 @@ class AzureAuth(Authenticator):
     _token_scope: str
 
     def __init__(self, token_scope: str, tenant_id: str = ""):
+        """
+        Initialize Azure authentication.
+
+        Args:
+            token_scope (str): The token scope for authentication.
+            tenant_id (str, optional): The tenant ID. Defaults to "".
+        """
         self._tenant_id = tenant_id
         self._token_scope = token_scope
         self._set_default_token()
 
     def _set_default_token(self) -> None:
+        """
+        Set up default Azure credentials and retrieve access token.
+        """
         self.azure_creds = DefaultAzureCredential()
         self.access_token = self.azure_creds.get_token(self._token_scope)
         self.token = self.access_token.token
@@ -45,8 +55,7 @@ def refresh_token(self) -> str:
         Refresh the access token if it is expired.
 
         Returns:
-            A token
-
+            str: A token
         """
         curr_epoch_time_in_ms = int(time.time()) * 1_000
         access_token_epoch_expiration_time_in_ms = int(self.access_token.expires_on) * 1_000
@@ -63,13 +72,23 @@ def get_token(self) -> str:
         """
         Get the current token.
 
-        Returns: The current token
-
+        Returns:
+            str: current token
         """
         return self.token
 
 
 def get_access_token_from_azure_cli(*, scope: str, tenant_id: str = ""):
+    """
+    Get access token from Azure CLI.
+
+    Args:
+        scope (str): The scope to request.
+        tenant_id (str, optional): The tenant ID. Defaults to "".
+
+    Returns:
+        str: The access token.
+    """
     try:
         credential = AzureCliCredential(tenant_id=tenant_id)
         token = credential.get_token(scope)
@@ -90,7 +109,7 @@ def get_access_token_from_azure_msi(*, client_id: str, scope: str):
         scope (str): The scope to request
 
     Returns:
-        Authentication token
+        str: Authentication token
     """
     try:
         credential = ManagedIdentityCredential(client_id=client_id)
@@ -103,15 +122,15 @@ def get_access_token_from_azure_msi(*, client_id: str, scope: str):
 
 def get_access_token_from_msa_public_client(*, client_id: str, scope: str):
     """
-    Uses MSA account to connect to an AOAI endpoint via interactive login. A browser window
+    Use MSA account to connect to an AOAI endpoint via interactive login. A browser window
     will open and ask for login credentials.
 
     Args:
         client_id (str): The client ID of the service
         scope (str): The scope to request
 
     Returns:
-        Authentication token
+        str: Authentication token
     """
     try:
         app = msal.PublicClientApplication(client_id)
@@ -124,11 +143,14 @@ def get_access_token_from_msa_public_client(*, client_id: str, scope: str):
 
 def get_access_token_from_interactive_login(scope: str) -> str:
     """
-    Connects to an OpenAI endpoint with an interactive login from Azure. A browser window will
+    Connect to an OpenAI endpoint with an interactive login from Azure. A browser window will
     open and ask for login credentials.  The token will be scoped for Azure Cognitive services.
 
+    Args:
+        scope (str): The scope to request
+
     Returns:
-        Authentication token
+        str: Authentication token
     """
     try:
         token_provider = get_bearer_token_provider(InteractiveBrowserCredential(), scope)
@@ -143,7 +165,7 @@ def get_token_provider_from_default_azure_credential(scope: str) -> Callable[[],
     Connect to an AOAI endpoint via default Azure credential.
 
     Returns:
-        Authentication token provider
+        Callable[[], str]: Authentication token provider
     """
     try:
         token_provider = get_bearer_token_provider(DefaultAzureCredential(), scope)
@@ -161,7 +183,7 @@ def get_default_scope(endpoint: str) -> str:
         endpoint (str): The endpoint to get the scope for.
 
     Returns:
-        The default scope for the given endpoint.
+        str: The default scope for the given endpoint.
     """
     try:
         parsed_uri = urlparse(endpoint)
@@ -184,7 +206,7 @@ def get_speech_config(resource_id: Union[str, None], key: Union[str, None], regi
         region (str): The region to get the token for.
 
     Returns:
-        The speech config based on passed in args
+        speechsdk.SpeechConfig: The speech config based on passed in args
 
     Raises:
         ModuleNotFoundError: If azure.cognitiveservices.speech is not installed.
@@ -223,6 +245,9 @@ def get_speech_config_from_default_azure_credential(resource_id: str, region: st
 
     Returns:
         The speech config for the given resource ID and region.
+
+    Raises:
+        ModuleNotFoundError: If azure.cognitiveservices.speech is not installed.
     """
     try:
         import azure.cognitiveservices.speech as speechsdk  # noqa: F811

pyrit/auth/azure_storage_auth.py

@@ -22,7 +22,7 @@ class AzureStorageAuth:
     @staticmethod
     async def get_user_delegation_key(blob_service_client: BlobServiceClient) -> UserDelegationKey:
         """
-        Retrieves a user delegation key valid for one day.
+        Retrieve a user delegation key valid for one day.
 
         Args:
             blob_service_client (BlobServiceClient): An instance of BlobServiceClient to interact
@@ -43,13 +43,16 @@ async def get_user_delegation_key(blob_service_client: BlobServiceClient) -> Use
     @staticmethod
     async def get_sas_token(container_url: str) -> str:
         """
-        Generates a SAS token for the specified blob using a user delegation key.
+        Generate a SAS token for the specified blob using a user delegation key.
 
         Args:
             container_url (str): The URL of the Azure Blob Storage container.
 
         Returns:
             str: The generated SAS token.
+
+        Raises:
+            ValueError: If container_url is empty or invalid.
         """
         if not container_url:
             raise ValueError(