CalvinBackup/llmsecops-research
1
0
Fork 0
mirror of https://github.com/lightbroker/llmsecops-research.git synced 2026-08-17 00:30:25 +02:00
Code Issues Packages Projects Releases Wiki Activity

support batch tests

Browse Source
This commit is contained in:
Adam Wilson
2025-08-19 20:09:34 -06:00
parent 3585f80414
commit cc124a91a3
12 changed files with 277 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/conftest.py
+23 -59
View File
@@ -18,11 +18,12 @@ from src.text_generation import config
from src.text_generation.adapters.embedding_model import EmbeddingModel
from src.text_generation.adapters.prompt_injection_example_repository import PromptInjectionExampleRepository
from src.text_generation.adapters.prompt_template_repository import PromptTemplateRepository
from src.text_generation.adapters.text_generation_foundation_model import TextGenerationFoundationModel
from src.text_generation.adapters.foundation_models.text_generation_foundation_model import TextGenerationFoundationModel
from src.text_generation.common.constants import Constants
from src.text_generation.services.guardrails.generated_text_guardrail_service import GeneratedTextGuardrailService
from src.text_generation.services.guardrails.reflexion_security_guidelines_service import ReflexionSecurityGuardrailsService
from src.text_generation.services.guidelines.chain_of_thought_security_guidelines_service import ChainOfThoughtSecurityGuidelinesService
from src.text_generation.services.guidelines.guidelines_factory import GuidelinesFactory
from src.text_generation.services.guidelines.rag_context_security_guidelines_configuration_builder import RetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder
from src.text_generation.services.guidelines.rag_context_security_guidelines_service import RagContextSecurityGuidelinesService
from src.text_generation.services.guidelines.rag_plus_cot_security_guidelines_service import RagPlusCotSecurityGuidelinesService
@@ -69,11 +70,14 @@ def pytest_deselected(items):
@pytest.fixture(scope="session", autouse=True)
def setup_test_environment():
"""Setup run before every test automatically."""
# Set test environment variables
os.environ["TESTING"] = "true"
os.environ["LOG_LEVEL"] = "DEBUG"
os.environ["PROMPT_TEMPLATES_DIR"] = "./infrastructure/prompt_templates"
os.environ["INJECTION_DATA_DIR"] = "./tests/security/tests/results/01_garak_no_guidelines"
os.environ["TEST_RUNS"] = "./tests/logs"
os.environ["MODEL_BASE_DIR"] = "./infrastructure/foundation_model"
os.environ["MODEL_CPU_DIR"] = "cpu_and_mobile/cpu-int4-rtn-block-32-acc-level-4"
os.environ["MODEL_DATA_FILENAME"] = "phi3-mini-4k-instruct-cpu-int4-rtn-block-32-acc-level-4.onnx.data"
@@ -93,10 +97,6 @@ def setup_test_environment():
def constants():
return Constants()
@pytest.fixture(scope="session")
def foundation_model():
return TextGenerationFoundationModel()
@pytest.fixture(scope="session")
def embedding_model():
return EmbeddingModel()
@@ -128,48 +128,6 @@ def rag_config_builder(
def llm_configuration_introspection_service():
return LLMConfigurationIntrospectionService()
@pytest.fixture(scope="session")
def rag_context_guidelines(
foundation_model,
response_processing_service,
prompt_template_service,
llm_configuration_introspection_service,
rag_config_builder):
return RagContextSecurityGuidelinesService(
foundation_model=foundation_model,
response_processing_service=response_processing_service,
prompt_template_service=prompt_template_service,
llm_configuration_introspection_service=llm_configuration_introspection_service,
config_builder=rag_config_builder
)
@pytest.fixture(scope="session")
def chain_of_thought_guidelines(
foundation_model,
response_processing_service,
llm_configuration_introspection_service,
prompt_template_service):
return ChainOfThoughtSecurityGuidelinesService(
foundation_model=foundation_model,
response_processing_service=response_processing_service,
llm_configuration_introspection_service=llm_configuration_introspection_service,
prompt_template_service=prompt_template_service
)
@pytest.fixture(scope="session")
def rag_plus_cot_guidelines(
foundation_model,
response_processing_service,
prompt_template_service,
llm_configuration_introspection_service,
rag_config_builder):
return RagPlusCotSecurityGuidelinesService(
foundation_model=foundation_model,
response_processing_service=response_processing_service,
prompt_template_service=prompt_template_service,
llm_configuration_introspection_service=llm_configuration_introspection_service,
config_builder=rag_config_builder
)
@pytest.fixture(scope="session")
def prompt_injection_example_service(prompt_injection_example_repository):
@@ -196,30 +154,36 @@ def reflexion_guardrails(
def response_processing_service():
return ResponseProcessingService()
@pytest.fixture(scope="session")
def llm_configuration_introspection_service():
return LLMConfigurationIntrospectionService()
def guidelines_factory():
return GuidelinesFactory()
@pytest.fixture(scope="session")
def guidelines_config_builder(
embedding_model,
prompt_template_service,
prompt_injection_example_repository):
return RetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder(
embedding_model=embedding_model,
prompt_template_service=prompt_template_service,
prompt_injection_example_repository=prompt_injection_example_repository
)
@pytest.fixture(scope="session")
def text_generation_completion_service(
foundation_model,
response_processing_service,
prompt_template_service,
chain_of_thought_guidelines,
rag_context_guidelines,
rag_plus_cot_guidelines,
reflexion_guardrails,
guidelines_factory,
guidelines_config_builder,
semantic_similarity_service,
prompt_injection_example_service,
llm_configuration_introspection_service):
return TextGenerationCompletionService(
foundation_model=foundation_model,
response_processing_service=response_processing_service,
prompt_template_service=prompt_template_service,
chain_of_thought_guidelines=chain_of_thought_guidelines,
rag_context_guidelines=rag_context_guidelines,
rag_plus_cot_guidelines=rag_plus_cot_guidelines,
reflexion_guardrails=reflexion_guardrails,
guidelines_factory=guidelines_factory,
guidelines_config_builder=guidelines_config_builder,
semantic_similarity_service=semantic_similarity_service,
prompt_injection_example_service=prompt_injection_example_service,
llm_configuration_introspection_service=llm_configuration_introspection_service
tests/integration/test_04_malicious_prompts_rag_and_cot.py
+13 -10
View File
@@ -1,15 +1,13 @@
"""
Usage:
$ pytest -k test_04_malicious_prompts_rag_and_cot -s --disable-warnings
"""
import os
import pytest
from test_utils import run_prompt_analysis_test
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, get_prompt_batch
@pytest.mark.mitigation
def test_04_malicious_prompts_rag_and_cot(
def test_04_malicious_prompts_rag_and_cot_microsoft_phi_3_mini4k_instruct(
text_generation_completion_service: AbstractTextGenerationCompletionService,
semantic_similarity_service,
malicious_prompts,
@@ -17,14 +15,19 @@ def test_04_malicious_prompts_rag_and_cot(
):
def configure_service(service):
return (service
.with_chain_of_thought_guidelines()
.with_rag_context_guidelines())
.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,
prompts=prompt_batch,
start=start,
end=end,
comparison_texts=prompt_injection_all_example_responses,
service_configurator=configure_service
)
tests/integration/test_utils.py
+39 -3
View File
@@ -1,5 +1,7 @@
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
from src.text_generation.domain.text_generation_completion_result import TextGenerationCompletionResult
from src.text_generation.services.logging.test_run_logging_service import TestRunLoggingService
@@ -8,11 +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+1), 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
@@ -42,7 +73,7 @@ def run_prompt_analysis_test(
configured_service: TextGenerationCompletionService = service_configurator(text_generation_completion_service)
print(f'sending prompt {i} to LLM')
completion_result: TextGenerationCompletionResult = configured_service.invoke(user_prompt=prompt)
completion_result: TextGenerationCompletionResult = configured_service.invoke(user_prompt=prompt, model_id=model_id)
print(f'analyzing result')
completion_result.finalize_completion_text()
@@ -51,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).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
)