diff --git a/.github/workflows/test_04_malicious_prompts_rag_and_cot_apple_openelm_3b_instruct.yml b/.github/workflows/test_04_malicious_prompts_rag_and_cot_apple_openelm_3b_instruct.yml index bb9ab7f66..24c9e5bb0 100644 --- a/.github/workflows/test_04_malicious_prompts_rag_and_cot_apple_openelm_3b_instruct.yml +++ b/.github/workflows/test_04_malicious_prompts_rag_and_cot_apple_openelm_3b_instruct.yml @@ -21,8 +21,7 @@ jobs: - name: 'run text generation tests' shell: bash - working-directory: tests/integration - run: pytest test_04_malicious_prompts_rag_and_cot_apple_openelm_3b_instruct.py -s --disable-warnings + run: pytest tests/integration/test_04_malicious_prompts_rag_and_cot_apple_openelm_3b_instruct.py -s --disable-warnings - name: Check for changes id: verify-changed-files diff --git a/.github/workflows/test_04_malicious_prompts_rag_and_cot_meta_llama_3_2_3b_instruct.yml b/.github/workflows/test_04_malicious_prompts_rag_and_cot_meta_llama_3_2_3b_instruct.yml index ddc7f188f..f8775f4f9 100644 --- a/.github/workflows/test_04_malicious_prompts_rag_and_cot_meta_llama_3_2_3b_instruct.yml +++ b/.github/workflows/test_04_malicious_prompts_rag_and_cot_meta_llama_3_2_3b_instruct.yml @@ -21,8 +21,7 @@ jobs: - name: 'run text generation tests' shell: bash - working-directory: tests/integration - run: pytest test_04_malicious_prompts_rag_and_cot_meta_llama_3_2_3b_instruct.py -s --disable-warnings + run: pytest tests/integration/test_04_malicious_prompts_rag_and_cot_meta_llama_3_2_3b_instruct.py -s --disable-warnings - name: Check for changes id: verify-changed-files diff --git a/.github/workflows/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.01-20.yml b/.github/workflows/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.01-20.yml new file mode 100644 index 000000000..4ed1be9fb --- /dev/null +++ b/.github/workflows/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.01-20.yml @@ -0,0 +1,25 @@ +name: 'Test #4 | RAG + CoT | microsoft/Phi-3-mini-4k-instruct' +on: + workflow_dispatch: + +jobs: + test-all-ranges: + strategy: + matrix: + include: + - offset: 0 + range: "1-20" + - offset: 20 + range: "21-40" + - offset: 40 + range: "41-60" + - offset: 60 + range: "61-80" + - offset: 80 + range: "81-100" + fail-fast: false + uses: ./.github/workflows/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.base.yml + with: + batch_offset: ${{ matrix.offset }} + range_name: ${{ matrix.range }} + batch_size: 2 \ No newline at end of file diff --git a/.github/workflows/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.base.yml b/.github/workflows/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.base.yml new file mode 100644 index 000000000..b513e57fc --- /dev/null +++ b/.github/workflows/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.base.yml @@ -0,0 +1,59 @@ +name: 'Reusable Test #4 | RAG + CoT | microsoft/Phi-3-mini-4k-instruct' +on: + workflow_call: + inputs: + batch_offset: + description: 'Starting prompt index offset' + required: true + type: number + range_name: + description: 'Human readable range name (e.g., "1-20")' + required: true + type: string + batch_size: + description: 'Number of prompts per batch' + required: false + type: number + default: 2 + +jobs: + test: + runs-on: ubuntu-latest + strategy: + # Always 10 batches per workflow + matrix: + batch: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + 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 - range ${{ inputs.range_name }} batch ${{ matrix.batch }}' + shell: bash + env: + PROMPT_BATCH: ${{ matrix.batch }} + BATCH_SIZE: ${{ inputs.batch_size }} + BATCH_OFFSET: ${{ inputs.batch_offset }} + 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: | + if [ -n "$(git status --porcelain)" ]; then + echo "changed=true" >> $GITHUB_OUTPUT + else + echo "changed=false" >> $GITHUB_OUTPUT + fi + - name: Commit and push changes + 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 range ${{ inputs.range_name }} batch ${{ matrix.batch }} [skip ci]" + git push \ No newline at end of file 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 d6ae2e799..e8c07a758 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,29 +1,31 @@ -name: 'Test #4 | RAG + CoT | microsoft/Phi-3-mini-4k-instruct' - +name: 'Test #4 | RAG + CoT | microsoft/Phi-3-mini-4k-instruct | Prompts 1-20' on: workflow_dispatch: - jobs: test: runs-on: ubuntu-latest + strategy: + # Matrix for batches 1-10, each testing 2 prompts (prompts 1-20 total) + matrix: + batch: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + 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 - working-directory: tests/integration - run: pytest test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.py -s --disable-warnings - + env: + PROMPT_BATCH: ${{ matrix.batch }} + BATCH_SIZE: 2 + BATCH_OFFSET: 0 # Prompts 1-20 (offset 0) + 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: | @@ -32,12 +34,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 }} (prompts 1-20) [skip ci]" git push \ No newline at end of file diff --git a/src/text_generation/services/guidelines/base_security_guidelines_service.py b/src/text_generation/services/guidelines/base_security_guidelines_service.py index a95ff9b45..b3bee689a 100644 --- a/src/text_generation/services/guidelines/base_security_guidelines_service.py +++ b/src/text_generation/services/guidelines/base_security_guidelines_service.py @@ -3,7 +3,6 @@ from langchain_core.output_parsers import StrOutputParser from langchain_core.prompts import PromptTemplate, StringPromptTemplate from langchain_core.prompt_values import PromptValue from langchain_core.runnables import RunnablePassthrough -from langchain.prompts import FewShotPromptTemplate from src.text_generation.common.constants import Constants from src.text_generation.domain.abstract_guidelines_processed_completion import AbstractGuidelinesProcessedCompletion @@ -48,7 +47,7 @@ class BaseSecurityGuidelinesService(AbstractSecurityGuidelinesService): | self.response_processing_service.process_text_generation_output ) - def _get_template(self, user_prompt: str) -> StringPromptTemplate: + def _get_template(self, user_prompt: str, template_id: str) -> StringPromptTemplate: """ Get the prompt template for security guidelines. @@ -58,13 +57,14 @@ class BaseSecurityGuidelinesService(AbstractSecurityGuidelinesService): raise NotImplementedError("Subclasses must implement _get_template()") - def apply_guidelines(self, user_prompt: str) -> AbstractGuidelinesProcessedCompletion: + def apply_guidelines(self, user_prompt: str, template_id: str) -> AbstractGuidelinesProcessedCompletion: if not user_prompt: raise ValueError(f"Parameter 'user_prompt' cannot be empty or None") try: - prompt_template: StringPromptTemplate = self._get_template(user_prompt=user_prompt) + prompt_template: StringPromptTemplate = self._get_template(user_prompt=user_prompt, template_id=template_id) + print(f'using prompt template: {template_id}') prompt_value: PromptValue = prompt_template.format_prompt(input=user_prompt) prompt_dict = { "messages": [ diff --git a/src/text_generation/services/guidelines/chain_of_thought_security_guidelines_service.py b/src/text_generation/services/guidelines/chain_of_thought_security_guidelines_service.py index 08e4de4e0..150ad2a90 100644 --- a/src/text_generation/services/guidelines/chain_of_thought_security_guidelines_service.py +++ b/src/text_generation/services/guidelines/chain_of_thought_security_guidelines_service.py @@ -86,6 +86,12 @@ class ChainOfThoughtSecurityGuidelinesService(BaseSecurityGuidelinesService): except Exception: return "" + def _get_template(self, user_prompt: str, template_id: str) -> StringPromptTemplate: + return self.config_builder.get_prompt_template( + template_id=template_id, + user_prompt=user_prompt + ) + def _get_cot_template_id_for_model(self, model_identifier: str) -> str: """ Get the appropriate CoT template ID for the given model. diff --git a/src/text_generation/services/guidelines/guidelines_factory.py b/src/text_generation/services/guidelines/guidelines_factory.py index 5e6900267..fb40eb612 100644 --- a/src/text_generation/services/guidelines/guidelines_factory.py +++ b/src/text_generation/services/guidelines/guidelines_factory.py @@ -27,7 +27,6 @@ class AbstractGuidelinesFactory(ABC): def create_rag_context_guidelines_service( self, foundation_model: AbstractFoundationModel, - prompt_template: StringPromptTemplate, response_processing_service: AbstractResponseProcessingService, prompt_template_service: AbstractPromptTemplateService, llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService, @@ -51,7 +50,6 @@ class GuidelinesFactory(AbstractGuidelinesFactory): def create_rag_context_guidelines_service( self, foundation_model: AbstractFoundationModel, - prompt_template: StringPromptTemplate, response_processing_service: AbstractResponseProcessingService, prompt_template_service: AbstractPromptTemplateService, llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService, @@ -62,8 +60,7 @@ class GuidelinesFactory(AbstractGuidelinesFactory): response_processing_service=response_processing_service, prompt_template_service=prompt_template_service, llm_configuration_introspection_service=llm_configuration_introspection_service, - config_builder=config_builder, - prompt_template=prompt_template + config_builder=config_builder ) def create_cot_guidelines_service( diff --git a/src/text_generation/services/guidelines/rag_context_security_guidelines_service.py b/src/text_generation/services/guidelines/rag_context_security_guidelines_service.py index cb727258f..a55a3bb2a 100644 --- a/src/text_generation/services/guidelines/rag_context_security_guidelines_service.py +++ b/src/text_generation/services/guidelines/rag_context_security_guidelines_service.py @@ -51,6 +51,12 @@ class RagContextSecurityGuidelinesService(BaseSecurityGuidelinesService): "meta-llama/llama-3.2-3b-instruct": self.constants.PromptTemplateIds.LLAMA_1_1B_CHAT__03_FEW_SHOT_EXAMPLES, } + def _get_template(self, user_prompt: str, template_id: str) -> StringPromptTemplate: + return self.config_builder.get_prompt_template( + template_id=template_id, + user_prompt=user_prompt + ) + def _get_model_identifier(self) -> str: """ Get the model identifier from the foundation model. diff --git a/src/text_generation/services/guidelines/rag_plus_cot_security_guidelines_service.py b/src/text_generation/services/guidelines/rag_plus_cot_security_guidelines_service.py index 5fef1cc61..1cd972d66 100644 --- a/src/text_generation/services/guidelines/rag_plus_cot_security_guidelines_service.py +++ b/src/text_generation/services/guidelines/rag_plus_cot_security_guidelines_service.py @@ -54,6 +54,12 @@ class RagPlusCotSecurityGuidelinesService(BaseSecurityGuidelinesService): "meta-llama/llama-3.2-3b-instruct": self.constants.PromptTemplateIds.LLAMA_1_1B_CHAT__04_FEW_SHOT_RAG_PLUS_COT, } + def _get_template(self, user_prompt: str, template_id: str) -> StringPromptTemplate: + return self.config_builder.get_prompt_template( + template_id=template_id, + user_prompt=user_prompt + ) + def _get_model_identifier(self) -> str: """ Get the model identifier from the foundation model. diff --git a/src/text_generation/services/logging/test_run_logging_service.py b/src/text_generation/services/logging/test_run_logging_service.py index 9f3146130..88c66efb9 100644 --- a/src/text_generation/services/logging/test_run_logging_service.py +++ b/src/text_generation/services/logging/test_run_logging_service.py @@ -13,10 +13,16 @@ from src.text_generation.services.logging.abstract_test_run_logging_service impo class TestRunLoggingService(AbstractTestRunLoggingService): - def __init__(self, test_id: int, model_id: ModelId): + def __init__( + self, + test_id: int, + model_id: ModelId, + start: int, + end: int + ): self._lock = threading.Lock() timestamp = calendar.timegm(time.gmtime()) - self.log_file_path = f"./tests/logs/test_{test_id}/{str(model_id.value).replace("/", "_")}/test_{test_id}_logs_{timestamp}.json" + self.log_file_path = f"./tests/logs/test_{test_id}/{start}_{end}/{str(model_id.value).replace("/", "_")}/test_{test_id}_logs_{timestamp}.json" self._ensure_log_file_exists() def _ensure_log_file_exists(self): diff --git a/src/text_generation/services/nlp/text_generation_completion_service.py b/src/text_generation/services/nlp/text_generation_completion_service.py index 2b3ae4dbc..d250a626f 100644 --- a/src/text_generation/services/nlp/text_generation_completion_service.py +++ b/src/text_generation/services/nlp/text_generation_completion_service.py @@ -2,7 +2,7 @@ from enum import Enum from typing import Optional, Dict, Any import logging -from langchain.prompts import StringPromptTemplate +from langchain_core.prompts import StringPromptTemplate from langchain_core.output_parsers import StrOutputParser from langchain_core.runnables import RunnablePassthrough, RunnableConfig from langchain_huggingface import HuggingFacePipeline @@ -177,12 +177,6 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService): logger.info(f"Successfully loaded model: {model_id.value}") - def get_current_model_info(self) -> Optional[Dict[str, Any]]: - """Get information about the currently loaded model""" - if self._current_model and self._current_model.is_loaded(): - return self._current_model.get_model_info() - return None - def _process_prompt_with_guidelines_if_applicable(self, user_prompt: str, target_model_id: ModelId): guidelines_config = ( self._use_zero_shot_chain_of_thought, @@ -267,7 +261,6 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService): """Factory method to create the appropriate guidelines service""" base_params = { 'foundation_model': self._current_model, - 'prompt_template': prompt_template, 'response_processing_service': self.response_processing_service, 'prompt_template_service': self.prompt_template_service, 'llm_configuration_introspection_service': self.llm_configuration_introspection_service @@ -301,7 +294,7 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService): guidelines_service = self._create_guidelines_service(mode, prompt_template) # Apply guidelines and process result - guidelines_result = guidelines_service.apply_guidelines(user_prompt) + guidelines_result = guidelines_service.apply_guidelines(user_prompt, template_id) return self._process_completion_result(guidelines_result) # Simplified handler methods @@ -421,10 +414,10 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService): target_model_id = model_id or self._current_model_id or self.default_model_id if (self._current_model_id != target_model_id or - self._current_model is None or - not self._current_model.is_loaded()): + self._current_model is None + ): self.load_model(target_model_id) print(f'Using model: {target_model_id.value}, guidelines: {self.get_current_config()}') - completion_result = self._process_prompt_with_guidelines_if_applicable(user_prompt=user_prompt, model_id=target_model_id) + completion_result = self._process_prompt_with_guidelines_if_applicable(user_prompt=user_prompt, target_model_id=target_model_id) return completion_result \ No newline at end of file diff --git a/tests/create_dirs.sh b/tests/create_dirs.sh new file mode 100755 index 000000000..3953e8728 --- /dev/null +++ b/tests/create_dirs.sh @@ -0,0 +1,26 @@ +#!/bin/bash + +# Array of model names +models=( + "apple/OpenELM-3B-Instruct" + "meta-llama/Llama-3.2-3B-Instruct" + "microsoft/Phi-3-mini-4k-instruct" +) + +# Create base integration directory if it doesn't exist + +# Iterate through test directories 1-5 +for i in {0..4}; do + # Iterate through each model + for model in "${models[@]}"; do + # Replace / with + dir_name="${model//\//_}" + + # Create the directory structure + mkdir -p "logs/test_${i}/${dir_name}" + + echo "Created: logs/test_${i}/${dir_name}" + done +done + +echo "All directories created successfully!" \ No newline at end of file diff --git a/tests/integration/test_01_malicious_prompts_no_mitigation_apple_openelm_3b_instruct.py b/tests/integration/test_01_malicious_prompts_no_mitigation_apple_openelm_3b_instruct.py index f82ec78a3..66dcdc1fd 100644 --- a/tests/integration/test_01_malicious_prompts_no_mitigation_apple_openelm_3b_instruct.py +++ b/tests/integration/test_01_malicious_prompts_no_mitigation_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 @@ -21,7 +21,7 @@ def test_01_malicious_prompts_no_mitigation_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, + 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_01_malicious_prompts_no_mitigation_meta_llama_3_2_3b_instruct.py b/tests/integration/test_01_malicious_prompts_no_mitigation_meta_llama_3_2_3b_instruct.py index a53e44866..78fdf20ae 100644 --- a/tests/integration/test_01_malicious_prompts_no_mitigation_meta_llama_3_2_3b_instruct.py +++ b/tests/integration/test_01_malicious_prompts_no_mitigation_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 get_prompt_batch, run_prompt_analysis_test import pytest @@ -21,7 +21,7 @@ def test_01_malicious_prompts_no_mitigation_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, + 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_01_malicious_prompts_no_mitigation_microsoft_phi_3_mini4k_instruct.py b/tests/integration/test_01_malicious_prompts_no_mitigation_microsoft_phi_3_mini4k_instruct.py index 6c9dc255d..1e93c5c7d 100644 --- a/tests/integration/test_01_malicious_prompts_no_mitigation_microsoft_phi_3_mini4k_instruct.py +++ b/tests/integration/test_01_malicious_prompts_no_mitigation_microsoft_phi_3_mini4k_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 @@ -21,7 +21,7 @@ def test_01_malicious_prompts_no_mitigation_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_02_malicious_prompts_cot_apple_openelm_3b_instruct.py b/tests/integration/test_02_malicious_prompts_cot_apple_openelm_3b_instruct.py index d092f7ea6..cb60fa7b7 100644 --- a/tests/integration/test_02_malicious_prompts_cot_apple_openelm_3b_instruct.py +++ b/tests/integration/test_02_malicious_prompts_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 @@ -22,7 +22,7 @@ def test_02_malicious_prompts_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, + 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_02_malicious_prompts_cot_meta_llama_3_2_3b_instruct.py b/tests/integration/test_02_malicious_prompts_cot_meta_llama_3_2_3b_instruct.py index 466ffaad4..1f027797a 100644 --- a/tests/integration/test_02_malicious_prompts_cot_meta_llama_3_2_3b_instruct.py +++ b/tests/integration/test_02_malicious_prompts_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 get_prompt_batch, run_prompt_analysis_test import pytest @@ -22,7 +22,7 @@ def test_02_malicious_prompts_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, + 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_02_malicious_prompts_cot_microsoft_phi_3_mini4k_instruct.py b/tests/integration/test_02_malicious_prompts_cot_microsoft_phi_3_mini4k_instruct.py index cd49b7a8f..41672593b 100644 --- a/tests/integration/test_02_malicious_prompts_cot_microsoft_phi_3_mini4k_instruct.py +++ b/tests/integration/test_02_malicious_prompts_cot_microsoft_phi_3_mini4k_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 @@ -22,7 +22,7 @@ def test_02_malicious_prompts_cot_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_03_malicious_prompts_rag_apple_openelm_3b_instruct.py b/tests/integration/test_03_malicious_prompts_rag_apple_openelm_3b_instruct.py index c5912b589..63bcea86a 100644 --- a/tests/integration/test_03_malicious_prompts_rag_apple_openelm_3b_instruct.py +++ b/tests/integration/test_03_malicious_prompts_rag_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 @@ -22,7 +22,7 @@ def test_03_malicious_prompts_rag_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, + 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_03_malicious_prompts_rag_meta_llama_3_2_3b_instruct.py b/tests/integration/test_03_malicious_prompts_rag_meta_llama_3_2_3b_instruct.py index 9b05d3d93..d359e7a4d 100644 --- a/tests/integration/test_03_malicious_prompts_rag_meta_llama_3_2_3b_instruct.py +++ b/tests/integration/test_03_malicious_prompts_rag_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 get_prompt_batch, run_prompt_analysis_test import pytest @@ -22,7 +22,7 @@ def test_03_malicious_prompts_rag_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, + 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_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..4fc7a230b 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,17 @@ def test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct( return (service .with_chain_of_thought_guidelines() .with_rag_context_guidelines()) + + prompt_batch, start, end = get_prompt_batch(malicious_prompts) 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=prompt_batch, + start=start, + end=end, 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..7c8edeaf5 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,12 +10,40 @@ 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: List[str], batch_size=10, env_var='PROMPT_BATCH'): + + batch_size = int(os.getenv('BATCH_SIZE', '2')) + batch_num = int(os.getenv('PROMPT_BATCH', '1')) + + if 'BATCH_OFFSET' in os.environ: + # Option 1: Fixed offset per workflow + offset = int(os.getenv('BATCH_OFFSET', '0')) + else: + # Option 2: Configurable range + prompt_range = int(os.getenv('PROMPT_RANGE', '1')) + offset = (prompt_range - 1) * 20 + + # Calculate start and end indices + start_idx = offset + (batch_num - 1) * batch_size + end_idx = min(start_idx + batch_size, len(prompts)) + + # Get the subset of prompts for this batch + prompt_subset: List[str] = prompts[start_idx:end_idx] + + print(f"Running batch {batch_num} (range offset {offset}): prompts {start_idx+1}-{end_idx} ({len(prompt_subset)} prompts)") + + return prompt_subset, start_idx, end_idx + + def run_prompt_analysis_test( test_id: int, model_id: ModelId, text_generation_completion_service: AbstractTextGenerationCompletionService, semantic_similarity_service: AbstractSemanticSimilarityService, prompts: List, + start: int, + end: int, comparison_texts: List, service_configurator: Callable, max_prompts: int = 100 @@ -53,13 +82,18 @@ def run_prompt_analysis_test( print(f'{i}/{len(prompts)} Max Score: {result.max}') print(f'{i}/{len(prompts)} Avg Score: {result.mean}') - TestRunLoggingService(test_id=test_id, model_id=model_id).log_results( + TestRunLoggingService( + test_id=test_id, + model_id=model_id, + start=start, + end=end + ).log_results( id=inspect.currentframe().f_back.f_code.co_name, text_generation_completion_result=completion_result, final_completion_text_score=result.max, is_rag_few_shot_enabled=text_generation_completion_service.is_rag_context_enabled(), is_cot_enabled=text_generation_completion_service.is_chain_of_thought_enabled(), - is_reflexion_enabled=text_generation_completion_service.is_reflexion_enabled(), + is_reflexion_enabled=False, original_llm_config=completion_result.original_result.llm_config ) diff --git a/tests/integration/test_violation_rate.py b/tests/integration/test_violation_rate.py index 694fd1c00..3b5f206b7 100644 --- a/tests/integration/test_violation_rate.py +++ b/tests/integration/test_violation_rate.py @@ -211,7 +211,7 @@ def test_example_with_fluent_service_call( final_completion=completion_result.original_completion, is_rag_few_shot_enabled=text_generation_completion_service.is_rag_context_enabled(), is_cot_enabled=text_generation_completion_service.is_chain_of_thought_enabled(), - is_reflexion_enabled=text_generation_completion_service.is_reflexion_enabled(), + is_reflexion_enabled=False, scores=result.scores, mean=result.mean, max=result.max) diff --git a/tests/logs/test_0/apple_OpenELM-3B-Instruct/_.txt b/tests/logs/test_0/apple_OpenELM-3B-Instruct/_.txt new file mode 100644 index 000000000..e69de29bb diff --git a/tests/logs/test_0/meta-llama_Llama-3.2-3B-Instruct/_.txt b/tests/logs/test_0/meta-llama_Llama-3.2-3B-Instruct/_.txt new file mode 100644 index 000000000..e69de29bb diff --git a/tests/logs/test_0/microsoft_Phi-3-mini-4k-instruct/_.txt b/tests/logs/test_0/microsoft_Phi-3-mini-4k-instruct/_.txt new file mode 100644 index 000000000..e69de29bb diff --git a/tests/logs/test_1/apple_OpenELM-3B-Instruct/_.txt b/tests/logs/test_1/apple_OpenELM-3B-Instruct/_.txt new file mode 100644 index 000000000..e69de29bb diff --git a/tests/logs/test_1/meta-llama_Llama-3.2-3B-Instruct/_.txt b/tests/logs/test_1/meta-llama_Llama-3.2-3B-Instruct/_.txt new file mode 100644 index 000000000..e69de29bb diff --git a/tests/logs/test_1/microsoft_Phi-3-mini-4k-instruct/_.txt b/tests/logs/test_1/microsoft_Phi-3-mini-4k-instruct/_.txt new file mode 100644 index 000000000..e69de29bb diff --git a/tests/logs/test_2/apple_OpenELM-3B-Instruct/_.txt b/tests/logs/test_2/apple_OpenELM-3B-Instruct/_.txt new file mode 100644 index 000000000..e69de29bb diff --git a/tests/logs/test_2/meta-llama_Llama-3.2-3B-Instruct/_.txt b/tests/logs/test_2/meta-llama_Llama-3.2-3B-Instruct/_.txt new file mode 100644 index 000000000..e69de29bb diff --git a/tests/logs/test_2/microsoft_Phi-3-mini-4k-instruct/_.txt b/tests/logs/test_2/microsoft_Phi-3-mini-4k-instruct/_.txt new file mode 100644 index 000000000..e69de29bb diff --git a/tests/logs/test_3/apple_OpenELM-3B-Instruct/_.txt b/tests/logs/test_3/apple_OpenELM-3B-Instruct/_.txt new file mode 100644 index 000000000..e69de29bb diff --git a/tests/logs/test_3/meta-llama_Llama-3.2-3B-Instruct/_.txt b/tests/logs/test_3/meta-llama_Llama-3.2-3B-Instruct/_.txt new file mode 100644 index 000000000..e69de29bb diff --git a/tests/logs/test_3/microsoft_Phi-3-mini-4k-instruct/_.txt b/tests/logs/test_3/microsoft_Phi-3-mini-4k-instruct/_.txt new file mode 100644 index 000000000..e69de29bb diff --git a/tests/logs/test_4/apple_OpenELM-3B-Instruct/_.txt b/tests/logs/test_4/apple_OpenELM-3B-Instruct/_.txt new file mode 100644 index 000000000..e69de29bb diff --git a/tests/logs/test_4/meta-llama_Llama-3.2-3B-Instruct/_.txt b/tests/logs/test_4/meta-llama_Llama-3.2-3B-Instruct/_.txt new file mode 100644 index 000000000..e69de29bb diff --git a/tests/logs/test_4/microsoft_Phi-3-mini-4k-instruct/_.txt b/tests/logs/test_4/microsoft_Phi-3-mini-4k-instruct/_.txt new file mode 100644 index 000000000..e69de29bb