FIX: integration test fixes (microsoft#1185)

Author: jsong468

MANIFEST.in

@@ -2,5 +2,5 @@ recursive-include pyrit *.json
 recursive-include pyrit *.prompt
 recursive-include pyrit *.yaml
 recursive-include pyrit/datasets/seed_prompts *
-recursive-include pyrit/datasets/score/scorer_evals *
+recursive-include pyrit/datasets/score/config/scorer_evals *
 include pyrit/auxiliary_attacks/gcg/src/Dockerfile

doc/code/executor/promptgen/1_anecdoctor_generator.ipynb

@@ -5,7 +5,11 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.17.3
+#       jupytext_version: 1.17.2
+#   kernelspec:
+#     display_name: pyrit-dev
+#     language: python
+#     name: python3
 # ---
 
 # %% [markdown]
@@ -130,13 +134,13 @@
 from pyrit.executor.promptgen import AnecdoctorContext
 
 
-def visualize_knowledge_graph(kg_result):
+def visualize_knowledge_graph(kg_result: str):
     """
     Parses the knowledge graph result, converts it to a DataFrame, and visualizes it as a graph.
     """
     # 1) Parse as JSON
     clean_output = kg_result.strip("`")
-    clean_output = kg_result.replace("json\n", "")  # Remove "json\n" if present
+    clean_output = clean_output.replace("json\n", "")  # Remove "json\n" if present
     data = json.loads(clean_output)
 
     # 2) Convert to DataFrame

doc/code/executor/promptgen/1_anecdoctor_generator.py

@@ -140,7 +140,6 @@
     "        \"tool_choice\": \"auto\",\n",
     "    },\n",
     "    httpx_client_kwargs={\"timeout\": 60.0},\n",
-    "    # api_version=None,  # this can be uncommented if using api.openai.com\n",
     ")\n",
     "\n",
     "# Build the user prompt\n",
@@ -302,7 +301,6 @@
     "    endpoint=os.getenv(\"PLATFORM_OPENAI_RESPONSES_ENDPOINT\"),\n",
     "    api_key=os.getenv(\"PLATFORM_OPENAI_RESPONSES_KEY\"),\n",
     "    model_name=\"gpt-5\",\n",
-    "    api_version=None,\n",
     "    extra_body_parameters={\"tools\": [grammar_tool], \"tool_choice\": \"required\"},\n",
     "    temperature=1.0,\n",
     ")\n",
@@ -311,7 +309,6 @@
     "    endpoint=os.getenv(\"PLATFORM_OPENAI_RESPONSES_ENDPOINT\"),\n",
     "    api_key=os.getenv(\"PLATFORM_OPENAI_RESPONSES_KEY\"),\n",
     "    model_name=\"gpt-5\",\n",
-    "    api_version=None,\n",
     "    temperature=1.0,\n",
     ")\n",
     "\n",
@@ -331,14 +328,6 @@
     "    if piece.original_value_data_type != \"reasoning\":\n",
     "        print(f\"{idx} | {piece.role}: {piece.original_value}\")"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "8",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {

doc/code/targets/8_openai_responses_target.ipynb

@@ -5,7 +5,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.17.3
+#       jupytext_version: 1.17.2
 # ---
 
 # %% [markdown]
@@ -97,7 +97,6 @@ async def get_current_weather(args):
         "tool_choice": "auto",
     },
     httpx_client_kwargs={"timeout": 60.0},
-    # api_version=None,  # this can be uncommented if using api.openai.com
 )
 
 # Build the user prompt
@@ -203,7 +202,6 @@ async def get_current_weather(args):
     endpoint=os.getenv("PLATFORM_OPENAI_RESPONSES_ENDPOINT"),
     api_key=os.getenv("PLATFORM_OPENAI_RESPONSES_KEY"),
     model_name="gpt-5",
-    api_version=None,
     extra_body_parameters={"tools": [grammar_tool], "tool_choice": "required"},
     temperature=1.0,
 )
@@ -212,7 +210,6 @@ async def get_current_weather(args):
     endpoint=os.getenv("PLATFORM_OPENAI_RESPONSES_ENDPOINT"),
     api_key=os.getenv("PLATFORM_OPENAI_RESPONSES_KEY"),
     model_name="gpt-5",
-    api_version=None,
     temperature=1.0,
 )
 
@@ -231,5 +228,3 @@ async def get_current_weather(args):
 for idx, piece in enumerate(result.message_pieces):
     if piece.original_value_data_type != "reasoning":
         print(f"{idx} | {piece.role}: {piece.original_value}")
-
-# %%

doc/code/targets/8_openai_responses_target.py

@@ -2,6 +2,7 @@
 # Licensed under the MIT license.
 
 import logging
+import os
 from typing import List, Optional
 
 from datasets import load_dataset
@@ -84,9 +85,11 @@
 
 
 def fetch_sorry_bench_dataset(
+    *,
     cache_dir: Optional[str] = None,
     categories: Optional[List[str]] = None,
     prompt_style: Optional[str] = None,
+    token: Optional[str] = None,
 ) -> SeedDataset:
     """
     Fetch Sorry-Bench dataset from Hugging Face (updated 2025/03 version).
@@ -97,23 +100,29 @@ def fetch_sorry_bench_dataset(
     Reference: https://arxiv.org/abs/2406.14598
 
     Args:
-        cache_dir: Optional cache directory for Hugging Face datasets
-        categories: Optional list of categories to filter. Full list in:
+        cache_dir (Optional[str]): Optional cache directory for Hugging Face datasets.
+        categories (Optional[List[str]]): Optional list of categories to filter. Full list in:
             https://huggingface.co/datasets/sorry-bench/sorry-bench-202503/blob/main/meta_info.py
-        prompt_style: Optional prompt style to filter. Available styles:
+        prompt_style (Optional[str]): Optional prompt style to filter. Available styles:
             "base", "ascii", "caesar", "slang", "authority_endorsement", etc.
             Default: "base" (only base prompts, no mutations)
             Full list: https://huggingface.co/datasets/sorry-bench/sorry-bench-202503
+        token (Optional[str]): Hugging Face authentication token. If not provided,
+            will attempt to read from HUGGINGFACE_TOKEN environment variable. This is needed for
+            accessing gated datasets on Hugging Face.
 
     Returns:
-        SeedDataset containing Sorry-Bench prompts with harm categories.
+        SeedDataset: SeedDataset containing Sorry-Bench prompts with harm categories.
 
     Raises:
         ValueError: If invalid categories or prompt_style are provided.
     """
     if prompt_style is None:
         prompt_style = "base"
 
+    if token is None:
+        token = os.environ.get("HUGGINGFACE_TOKEN")
+
     if prompt_style not in VALID_PROMPT_STYLES:
         raise ValueError(f"Invalid prompt_style '{prompt_style}'. Must be one of: {', '.join(VALID_PROMPT_STYLES)}")
 
@@ -128,7 +137,7 @@ def fetch_sorry_bench_dataset(
     try:
         source = "sorry-bench/sorry-bench-202503"
         logger.info(f"Loading Sorry-Bench dataset from {source}")
-        data = load_dataset(source, cache_dir=cache_dir)
+        data = load_dataset(source, cache_dir=cache_dir, token=token)
 
         dataset_split = data["train"]
 

pyrit/datasets/sorry_bench_dataset.py

@@ -34,7 +34,6 @@ class OpenAIChatTarget(OpenAIChatTargetBase):
         endpoint (str): The endpoint for the OpenAI API
         model_name (str): The model name for the OpenAI API
         deployment_name (str): For Azure, the deployment name
-        api_version (str): The api version for the OpenAI API
         temperature (float): The temperature for the completion
         max_completion_tokens (int): The maximum number of tokens to be returned by the model.
             The total length of input tokens and generated tokens is limited by

pyrit/prompt_target/openai/openai_chat_target.py

@@ -16,7 +16,7 @@ async def test_openai_responses_gpt5(sqlite_instance):
     args = {
         "endpoint": os.getenv("AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT"),
         "model_name": os.getenv("AZURE_OPENAI_GPT5_MODEL"),
-        "api_key": os.getenv("AZURE_OPENAI_GPT5_RESPONSES_KEY"),
+        "api_key": os.getenv("AZURE_OPENAI_GPT5_KEY"),
     }
 
     target = OpenAIResponseTarget(**args)

tests/integration/targets/test_openai_responses_gpt5.py

@@ -51,7 +51,7 @@ def test_fetch_sorry_bench_dataset_success(self, mock_load_dataset, mock_sorry_b
         mock_load_dataset.return_value = mock_dataset
 
         # Call the function (defaults to base only)
-        result = fetch_sorry_bench_dataset()
+        result = fetch_sorry_bench_dataset(token="mock_token")
 
         # Assertions
         assert isinstance(result, SeedDataset)
@@ -71,7 +71,7 @@ def test_fetch_sorry_bench_dataset_success(self, mock_load_dataset, mock_sorry_b
         assert "base" in prompt1.groups
 
         # Verify load_dataset was called correctly
-        mock_load_dataset.assert_called_once_with("sorry-bench/sorry-bench-202503", cache_dir=None)
+        mock_load_dataset.assert_called_once_with("sorry-bench/sorry-bench-202503", cache_dir=None, token="mock_token")
 
     @patch("pyrit.datasets.sorry_bench_dataset.load_dataset")
     def test_fetch_sorry_bench_dataset_ascii_style(self, mock_load_dataset, mock_sorry_bench_data):
@@ -211,10 +211,12 @@ def test_fetch_sorry_bench_dataset_custom_cache_dir(self, mock_load_dataset, moc
         mock_load_dataset.return_value = mock_dataset
 
         custom_cache = "/custom/cache/path"
-        result = fetch_sorry_bench_dataset(cache_dir=custom_cache)
+        result = fetch_sorry_bench_dataset(cache_dir=custom_cache, token="mock_token")
 
         assert isinstance(result, SeedDataset)
-        mock_load_dataset.assert_called_once_with("sorry-bench/sorry-bench-202503", cache_dir=custom_cache)
+        mock_load_dataset.assert_called_once_with(
+            "sorry-bench/sorry-bench-202503", cache_dir=custom_cache, token="mock_token"
+        )
 
     def test_fetch_sorry_bench_dataset_invalid_prompt_style(self):
         """Test error handling for invalid prompt_style."""