diff --git a/.github/workflows/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.yml b/.github/workflows/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.yml index 8ad5ebe21..269369344 100644 --- a/.github/workflows/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.yml +++ b/.github/workflows/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.yml @@ -1,28 +1,30 @@ name: 'Test #4 | RAG + CoT | microsoft/Phi-3-mini-4k-instruct' - on: workflow_dispatch: - jobs: test: runs-on: ubuntu-latest + strategy: + # Matrix strategy to run 10 batches of 10 prompts each + matrix: + batch: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + # Prevent all jobs from being cancelled if one fails + fail-fast: false steps: - name: 'checkout' uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 - - name: 'set up Python' uses: actions/setup-python@v3 with: python-version: '3.12' - - name: 'set up Python dependencies' shell: bash run: pip install -r ${{ github.workspace }}/requirements.txt - - - name: 'run text generation tests' + - name: 'run text generation tests - batch ${{ matrix.batch }}' shell: bash + env: + PROMPT_BATCH: ${{ matrix.batch }} run: pytest tests/integration/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.py -s --disable-warnings - - name: Check for changes id: verify-changed-files run: | @@ -31,12 +33,11 @@ jobs: else echo "changed=false" >> $GITHUB_OUTPUT fi - - - name: Commit and push changes + - name: Commit and push changes (batch ${{ matrix.batch }}) if: steps.verify-changed-files.outputs.changed == 'true' run: | git config --local user.email "42450907+lightbroker@users.noreply.github.com" git config --local user.name "Adam Wilson" git add . - git commit -m "Auto-generated files from workflow [skip ci]" + git commit -m "Auto-generated files from workflow batch ${{ matrix.batch }} [skip ci]" git push \ No newline at end of file diff --git a/tests/integration/test_03_malicious_prompts_rag_microsoft_phi_3_mini4k_instruct.py b/tests/integration/test_03_malicious_prompts_rag_microsoft_phi_3_mini4k_instruct.py index 84f064539..64539eb24 100644 --- a/tests/integration/test_03_malicious_prompts_rag_microsoft_phi_3_mini4k_instruct.py +++ b/tests/integration/test_03_malicious_prompts_rag_microsoft_phi_3_mini4k_instruct.py @@ -1,5 +1,7 @@ +import pytest +from src.text_generation.common.model_id import ModelId from src.text_generation.services.nlp.abstract_text_generation_completion_service import AbstractTextGenerationCompletionService -from tests.integration.test_utils import run_prompt_analysis_test +from tests.integration.test_utils import get_prompt_batch, run_prompt_analysis_test @pytest.mark.mitigation @@ -18,7 +20,7 @@ def test_03_malicious_prompts_rag_microsoft_phi_3_mini4k_instruct( model_id=ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT, text_generation_completion_service=text_generation_completion_service, semantic_similarity_service=semantic_similarity_service, - prompts=malicious_prompts, + prompts=get_prompt_batch(malicious_prompts), comparison_texts=prompt_injection_all_example_responses, service_configurator=configure_service ) \ No newline at end of file diff --git a/tests/integration/test_04_malicious_prompts_rag_and_cot_apple_openelm_3b_instruct.py b/tests/integration/test_04_malicious_prompts_rag_and_cot_apple_openelm_3b_instruct.py index a0bfe0af8..99a670efc 100644 --- a/tests/integration/test_04_malicious_prompts_rag_and_cot_apple_openelm_3b_instruct.py +++ b/tests/integration/test_04_malicious_prompts_rag_and_cot_apple_openelm_3b_instruct.py @@ -1,6 +1,6 @@ from src.text_generation.common.model_id import ModelId from src.text_generation.services.nlp.abstract_text_generation_completion_service import AbstractTextGenerationCompletionService -from tests.integration.test_utils import run_prompt_analysis_test +from tests.integration.test_utils import get_prompt_batch, run_prompt_analysis_test import pytest @@ -23,7 +23,7 @@ def test_04_malicious_prompts_rag_and_cot_apple_openelm_3b_instruct( model_id=ModelId.APPLE_OPENELM_3B_INSTRUCT, text_generation_completion_service=text_generation_completion_service, semantic_similarity_service=semantic_similarity_service, - prompts=malicious_prompts[:1], + prompts=get_prompt_batch(malicious_prompts), comparison_texts=prompt_injection_all_example_responses, service_configurator=configure_service ) \ No newline at end of file diff --git a/tests/integration/test_04_malicious_prompts_rag_and_cot_meta_llama_3_2_3b_instruct.py b/tests/integration/test_04_malicious_prompts_rag_and_cot_meta_llama_3_2_3b_instruct.py index 624a87a8b..13e119e0b 100644 --- a/tests/integration/test_04_malicious_prompts_rag_and_cot_meta_llama_3_2_3b_instruct.py +++ b/tests/integration/test_04_malicious_prompts_rag_and_cot_meta_llama_3_2_3b_instruct.py @@ -1,6 +1,6 @@ from src.text_generation.common.model_id import ModelId from src.text_generation.services.nlp.abstract_text_generation_completion_service import AbstractTextGenerationCompletionService -from tests.integration.test_utils import run_prompt_analysis_test +from tests.integration.test_utils import run_prompt_analysis_test, get_prompt_batch import pytest @@ -23,7 +23,7 @@ def test_04_malicious_prompts_rag_and_cot_meta_llama_3_2_3b_instruct( model_id=ModelId.META_LLAMA_3_2_3B_INSTRUCT, text_generation_completion_service=text_generation_completion_service, semantic_similarity_service=semantic_similarity_service, - prompts=malicious_prompts[:1], + prompts=get_prompt_batch(malicious_prompts), comparison_texts=prompt_injection_all_example_responses, service_configurator=configure_service ) \ No newline at end of file diff --git a/tests/integration/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.py b/tests/integration/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.py index 20ad864c3..0a9ede299 100644 --- a/tests/integration/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.py +++ b/tests/integration/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.py @@ -1,9 +1,9 @@ +import os +import pytest + from src.text_generation.common.model_id import ModelId from src.text_generation.services.nlp.abstract_text_generation_completion_service import AbstractTextGenerationCompletionService -from tests.integration.test_utils import run_prompt_analysis_test - - -import pytest +from tests.integration.test_utils import run_prompt_analysis_test, get_prompt_batch @pytest.mark.mitigation @@ -17,13 +17,13 @@ def test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct( return (service .with_chain_of_thought_guidelines() .with_rag_context_guidelines()) - + results = run_prompt_analysis_test( test_id=4, model_id=ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT, text_generation_completion_service=text_generation_completion_service, semantic_similarity_service=semantic_similarity_service, - prompts=malicious_prompts[:1], + prompts=get_prompt_batch(malicious_prompts), comparison_texts=prompt_injection_all_example_responses, service_configurator=configure_service ) \ No newline at end of file diff --git a/tests/integration/test_utils.py b/tests/integration/test_utils.py index 362d1737c..2bc2e38a2 100644 --- a/tests/integration/test_utils.py +++ b/tests/integration/test_utils.py @@ -1,4 +1,5 @@ import inspect +import os from typing import List, Callable from src.text_generation.common.model_id import ModelId from src.text_generation.domain.semantic_similarity_result import SemanticSimilarityResult @@ -9,6 +10,20 @@ from src.text_generation.services.nlp.abstract_text_generation_completion_servic from src.text_generation.services.nlp.text_generation_completion_service import TextGenerationCompletionService + +def get_prompt_batch(prompts, batch_size=10, env_var='PROMPT_BATCH'): + """ + Returns a batch of prompts based on the PROMPT_BATCH environment variable. + Prints batch info for debugging. + """ + batch_num = int(os.getenv(env_var, '1')) + start_idx = (batch_num - 1) * batch_size + end_idx = min(start_idx + batch_size, len(prompts)) + prompt_subset = prompts[start_idx:end_idx] + print(f"Running batch {batch_num}: prompts {start_idx+1}-{end_idx} ({len(prompt_subset)} prompts)") + return prompt_subset + + def run_prompt_analysis_test( test_id: int, model_id: ModelId,