CalvinBackup/llmsecops-research
1
0
Fork 0
mirror of https://github.com/lightbroker/llmsecops-research.git synced 2026-07-07 05:27:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #22 from lightbroker/model-support-expansion

Browse Source 
Model support expansion
This commit is contained in:
Adam Wilson
2025-08-18 18:00:45 -06:00
committed by GitHub
parent a48ea0c626 8c6230e0dc
commit c142235923
42 changed files with 232 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/test_04_malicious_prompts_rag_and_cot_apple_openelm_3b_instruct.yml
+1 -2
View File
@@ -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
.github/workflows/test_04_malicious_prompts_rag_and_cot_meta_llama_3_2_3b_instruct.yml
+1 -2
View File
@@ -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
.github/workflows/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.01-20.yml
+25
View File
@@ -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
.github/workflows/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.base.yml
+59
View File
@@ -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
.github/workflows/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.yml
+14 -13
View File
@@ -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
src/text_generation/services/guidelines/base_security_guidelines_service.py
+4 -4
View File
@@ -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": [
src/text_generation/services/guidelines/chain_of_thought_security_guidelines_service.py
+6
View File
@@ -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.
src/text_generation/services/guidelines/guidelines_factory.py
+1 -4
View File
@@ -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(
src/text_generation/services/guidelines/rag_context_security_guidelines_service.py
+6
View File
@@ -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.
src/text_generation/services/guidelines/rag_plus_cot_security_guidelines_service.py
+6
View File
@@ -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.
src/text_generation/services/logging/test_run_logging_service.py
+8 -2
View File
@@ -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):
src/text_generation/services/nlp/text_generation_completion_service.py
+5 -12
View File
@@ -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
tests/create_dirs.sh
Executable
+26
View File
@@ -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!"
tests/integration/test_01_malicious_prompts_no_mitigation_apple_openelm_3b_instruct.py
+2 -2
View File
@@ -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
)
tests/integration/test_01_malicious_prompts_no_mitigation_meta_llama_3_2_3b_instruct.py
+2 -2
View File
@@ -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
)
tests/integration/test_01_malicious_prompts_no_mitigation_microsoft_phi_3_mini4k_instruct.py
+2 -2
View File
@@ -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
)
tests/integration/test_02_malicious_prompts_cot_apple_openelm_3b_instruct.py
+2 -2
View File
@@ -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
)
tests/integration/test_02_malicious_prompts_cot_meta_llama_3_2_3b_instruct.py
+2 -2
View File
@@ -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
)
tests/integration/test_02_malicious_prompts_cot_microsoft_phi_3_mini4k_instruct.py
+2 -2
View File
@@ -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
)
tests/integration/test_03_malicious_prompts_rag_apple_openelm_3b_instruct.py
+2 -2
View File
@@ -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
)
tests/integration/test_03_malicious_prompts_rag_meta_llama_3_2_3b_instruct.py
+2 -2
View File
@@ -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
)
tests/integration/test_03_malicious_prompts_rag_microsoft_phi_3_mini4k_instruct.py
+4 -2
View File
@@ -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
)
tests/integration/test_04_malicious_prompts_rag_and_cot_apple_openelm_3b_instruct.py
+2 -2
View File
@@ -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
)
tests/integration/test_04_malicious_prompts_rag_and_cot_meta_llama_3_2_3b_instruct.py
+2 -2
View File
@@ -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
)
tests/integration/test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct.py
+9 -5
View File
@@ -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
)
tests/integration/test_utils.py
+36 -2
View File
@@ -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
)
tests/integration/test_violation_rate.py
+1 -1
View File
@@ -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)
tests/logs/test_0/apple_OpenELM-3B-Instruct/_.txt
View File
tests/logs/test_0/meta-llama_Llama-3.2-3B-Instruct/_.txt
View File
tests/logs/test_0/microsoft_Phi-3-mini-4k-instruct/_.txt
View File
tests/logs/test_1/apple_OpenELM-3B-Instruct/_.txt
View File
tests/logs/test_1/meta-llama_Llama-3.2-3B-Instruct/_.txt
View File
tests/logs/test_1/microsoft_Phi-3-mini-4k-instruct/_.txt
View File
tests/logs/test_2/apple_OpenELM-3B-Instruct/_.txt
View File
tests/logs/test_2/meta-llama_Llama-3.2-3B-Instruct/_.txt
View File
tests/logs/test_2/microsoft_Phi-3-mini-4k-instruct/_.txt
View File
tests/logs/test_3/apple_OpenELM-3B-Instruct/_.txt
View File
tests/logs/test_3/meta-llama_Llama-3.2-3B-Instruct/_.txt
View File
tests/logs/test_3/microsoft_Phi-3-mini-4k-instruct/_.txt
View File
tests/logs/test_4/apple_OpenELM-3B-Instruct/_.txt
View File
tests/logs/test_4/meta-llama_Llama-3.2-3B-Instruct/_.txt
View File
tests/logs/test_4/microsoft_Phi-3-mini-4k-instruct/_.txt
View File