From 3044426536f0c858ed77e5bd487f88d23bf05176 Mon Sep 17 00:00:00 2001
From: Adam Wilson <lightbrok3r@gmail.com>
Date: Tue, 19 Aug 2025 10:40:57 -0600
Subject: [PATCH] fix test_utils call

---
 ...ign_prompts_no_mitigation_apple_openelm_3b_instruct.py | 8 ++++++--
 ...gn_prompts_no_mitigation_meta_llama_3_2_3b_instruct.py | 8 ++++++--
 ...ompts_no_mitigation_microsoft_phi_3_mini4k_instruct.py | 8 ++++++--
 ...ous_prompts_no_mitigation_apple_openelm_3b_instruct.py | 6 +++++-
 ...us_prompts_no_mitigation_meta_llama_3_2_3b_instruct.py | 6 +++++-
 ...ompts_no_mitigation_microsoft_phi_3_mini4k_instruct.py | 6 +++++-
 ..._02_malicious_prompts_cot_apple_openelm_3b_instruct.py | 6 +++++-
 ...02_malicious_prompts_cot_meta_llama_3_2_3b_instruct.py | 6 +++++-
 ...licious_prompts_cot_microsoft_phi_3_mini4k_instruct.py | 6 +++++-
 ..._03_malicious_prompts_rag_apple_openelm_3b_instruct.py | 6 +++++-
 ...03_malicious_prompts_rag_meta_llama_3_2_3b_instruct.py | 6 +++++-
 ...licious_prompts_rag_microsoft_phi_3_mini4k_instruct.py | 6 +++++-
 ...cious_prompts_rag_and_cot_apple_openelm_3b_instruct.py | 6 +++++-
 ...ious_prompts_rag_and_cot_meta_llama_3_2_3b_instruct.py | 6 +++++-
 14 files changed, 73 insertions(+), 17 deletions(-)

diff --git a/tests/integration/test_00_benign_prompts_no_mitigation_apple_openelm_3b_instruct.py b/tests/integration/test_00_benign_prompts_no_mitigation_apple_openelm_3b_instruct.py
index d1c0989da..5b9f3abc4 100644
--- a/tests/integration/test_00_benign_prompts_no_mitigation_apple_openelm_3b_instruct.py
+++ b/tests/integration/test_00_benign_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
 
 
 def test_00_benign_prompts_no_mitigation_apple_openelm_3b_instruct(
@@ -15,12 +15,16 @@ def test_00_benign_prompts_no_mitigation_apple_openelm_3b_instruct(
     if len(prompt_injection_all_example_responses) == 0:
         raise ValueError('Unable to load prompt injection examples')
 
+    prompt_batch, start, end = get_prompt_batch(math_prompts)
+    
     results = run_prompt_analysis_test(
         test_id=0,
         model_id=ModelId.APPLE_OPENELM_3B_INSTRUCT,
         text_generation_completion_service=text_generation_completion_service,
         semantic_similarity_service=semantic_similarity_service,
-        prompts=math_prompts,
+        prompts=prompt_batch,
+        start=start,
+        end=end,
         comparison_texts=prompt_injection_all_example_responses,
         service_configurator=configure_service
     )
diff --git a/tests/integration/test_00_benign_prompts_no_mitigation_meta_llama_3_2_3b_instruct.py b/tests/integration/test_00_benign_prompts_no_mitigation_meta_llama_3_2_3b_instruct.py
index 26e63dbce..a72b17280 100644
--- a/tests/integration/test_00_benign_prompts_no_mitigation_meta_llama_3_2_3b_instruct.py
+++ b/tests/integration/test_00_benign_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
 
 
 def test_00_benign_prompts_no_mitigation_meta_llama_3_2_3b_instruct(
@@ -15,12 +15,16 @@ def test_00_benign_prompts_no_mitigation_meta_llama_3_2_3b_instruct(
     if len(prompt_injection_all_example_responses) == 0:
         raise ValueError('Unable to load prompt injection examples')
 
+    prompt_batch, start, end = get_prompt_batch(math_prompts)
+
     results = run_prompt_analysis_test(
         test_id=0,
         model_id=ModelId.META_LLAMA_3_2_3B_INSTRUCT,
         text_generation_completion_service=text_generation_completion_service,
         semantic_similarity_service=semantic_similarity_service,
-        prompts=math_prompts,
+        prompts=prompt_batch,
+        start=start,
+        end=end,
         comparison_texts=prompt_injection_all_example_responses,
         service_configurator=configure_service
     )
diff --git a/tests/integration/test_00_benign_prompts_no_mitigation_microsoft_phi_3_mini4k_instruct.py b/tests/integration/test_00_benign_prompts_no_mitigation_microsoft_phi_3_mini4k_instruct.py
index 0eeff1990..1790b425a 100644
--- a/tests/integration/test_00_benign_prompts_no_mitigation_microsoft_phi_3_mini4k_instruct.py
+++ b/tests/integration/test_00_benign_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
 
 
 def test_00_benign_prompts_no_mitigation_microsoft_phi_3_mini4k_instruct(
@@ -15,12 +15,16 @@ def test_00_benign_prompts_no_mitigation_microsoft_phi_3_mini4k_instruct(
     if len(prompt_injection_all_example_responses) == 0:
         raise ValueError('Unable to load prompt injection examples')
 
+    prompt_batch, start, end = get_prompt_batch(math_prompts)
+    
     results = run_prompt_analysis_test(
         test_id=0,
         model_id=ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT,
         text_generation_completion_service=text_generation_completion_service,
         semantic_similarity_service=semantic_similarity_service,
-        prompts=math_prompts,
+        prompts=prompt_batch,
+        start=start,
+        end=end,
         comparison_texts=prompt_injection_all_example_responses,
         service_configurator=configure_service
     )
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 66dcdc1fd..52aa48c91 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
@@ -16,12 +16,16 @@ def test_01_malicious_prompts_no_mitigation_apple_openelm_3b_instruct(
     def configure_service(service):
         return service.without_guidelines()
 
+    prompt_batch, start, end = get_prompt_batch(malicious_prompts)
+    
     results = run_prompt_analysis_test(
         test_id=1,
         model_id=ModelId.APPLE_OPENELM_3B_INSTRUCT,
         text_generation_completion_service=text_generation_completion_service,
         semantic_similarity_service=semantic_similarity_service,
-        prompts=get_prompt_batch(malicious_prompts),
+        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_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 78fdf20ae..98197a71b 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
@@ -16,12 +16,16 @@ def test_01_malicious_prompts_no_mitigation_meta_llama_3_2_3b_instruct(
     def configure_service(service):
         return service.without_guidelines()
 
+    prompt_batch, start, end = get_prompt_batch(malicious_prompts)
+    
     results = run_prompt_analysis_test(
         test_id=1,
         model_id=ModelId.META_LLAMA_3_2_3B_INSTRUCT,
         text_generation_completion_service=text_generation_completion_service,
         semantic_similarity_service=semantic_similarity_service,
-        prompts=get_prompt_batch(malicious_prompts),
+        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_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 1e93c5c7d..138697f91 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
@@ -16,12 +16,16 @@ def test_01_malicious_prompts_no_mitigation_microsoft_phi_3_mini4k_instruct(
     def configure_service(service):
         return service.without_guidelines()
 
+    prompt_batch, start, end = get_prompt_batch(malicious_prompts)
+    
     results = run_prompt_analysis_test(
         test_id=1,
         model_id=ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT,
         text_generation_completion_service=text_generation_completion_service,
         semantic_similarity_service=semantic_similarity_service,
-        prompts=get_prompt_batch(malicious_prompts),
+        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_02_malicious_prompts_cot_apple_openelm_3b_instruct.py b/tests/integration/test_02_malicious_prompts_cot_apple_openelm_3b_instruct.py
index cb60fa7b7..e10912b7c 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
@@ -17,12 +17,16 @@ def test_02_malicious_prompts_cot_apple_openelm_3b_instruct(
         return (service
                 .with_chain_of_thought_guidelines())
 
+    prompt_batch, start, end = get_prompt_batch(malicious_prompts)
+    
     results = run_prompt_analysis_test(
         test_id=2,
         model_id=ModelId.APPLE_OPENELM_3B_INSTRUCT,
         text_generation_completion_service=text_generation_completion_service,
         semantic_similarity_service=semantic_similarity_service,
-        prompts=get_prompt_batch(malicious_prompts),
+        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_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 1f027797a..567c2f27b 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
@@ -17,12 +17,16 @@ def test_02_malicious_prompts_cot_meta_llama_3_2_3b_instruct(
         return (service
                 .with_chain_of_thought_guidelines())
 
+    prompt_batch, start, end = get_prompt_batch(malicious_prompts)
+    
     results = run_prompt_analysis_test(
         test_id=2,
         model_id=ModelId.META_LLAMA_3_2_3B_INSTRUCT,
         text_generation_completion_service=text_generation_completion_service,
         semantic_similarity_service=semantic_similarity_service,
-        prompts=get_prompt_batch(malicious_prompts),
+        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_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 41672593b..8e0f52acf 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
@@ -17,12 +17,16 @@ def test_02_malicious_prompts_cot_microsoft_phi_3_mini4k_instruct(
         return (service
                 .with_chain_of_thought_guidelines())
 
+    prompt_batch, start, end = get_prompt_batch(malicious_prompts)
+    
     results = run_prompt_analysis_test(
         test_id=2,
         model_id=ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT,
         text_generation_completion_service=text_generation_completion_service,
         semantic_similarity_service=semantic_similarity_service,
-        prompts=get_prompt_batch(malicious_prompts),
+        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_03_malicious_prompts_rag_apple_openelm_3b_instruct.py b/tests/integration/test_03_malicious_prompts_rag_apple_openelm_3b_instruct.py
index 63bcea86a..82cf3b423 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
@@ -17,12 +17,16 @@ def test_03_malicious_prompts_rag_apple_openelm_3b_instruct(
         return (service
                 .with_rag_context_guidelines())
 
+    prompt_batch, start, end = get_prompt_batch(malicious_prompts)
+    
     results = run_prompt_analysis_test(
         test_id=3,
         model_id=ModelId.APPLE_OPENELM_3B_INSTRUCT,
         text_generation_completion_service=text_generation_completion_service,
         semantic_similarity_service=semantic_similarity_service,
-        prompts=get_prompt_batch(malicious_prompts),
+        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_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 d359e7a4d..7305b85b9 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
@@ -17,12 +17,16 @@ def test_03_malicious_prompts_rag_meta_llama_3_2_3b_instruct(
         return (service
                 .with_rag_context_guidelines())
 
+    prompt_batch, start, end = get_prompt_batch(malicious_prompts)
+    
     results = run_prompt_analysis_test(
         test_id=3,
         model_id=ModelId.META_LLAMA_3_2_3B_INSTRUCT,
         text_generation_completion_service=text_generation_completion_service,
         semantic_similarity_service=semantic_similarity_service,
-        prompts=get_prompt_batch(malicious_prompts),
+        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_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 64539eb24..ca5490c37 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
@@ -15,12 +15,16 @@ def test_03_malicious_prompts_rag_microsoft_phi_3_mini4k_instruct(
         return (service
                 .with_rag_context_guidelines())
 
+    prompt_batch, start, end = get_prompt_batch(malicious_prompts)
+    
     results = run_prompt_analysis_test(
         test_id=3,
         model_id=ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT,
         text_generation_completion_service=text_generation_completion_service,
         semantic_similarity_service=semantic_similarity_service,
-        prompts=get_prompt_batch(malicious_prompts),
+        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_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 99a670efc..180310c00 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
@@ -18,12 +18,16 @@ def test_04_malicious_prompts_rag_and_cot_apple_openelm_3b_instruct(
             .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.APPLE_OPENELM_3B_INSTRUCT,
         text_generation_completion_service=text_generation_completion_service,
         semantic_similarity_service=semantic_similarity_service,
-        prompts=get_prompt_batch(malicious_prompts),
+        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_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 13e119e0b..07c9575e7 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
@@ -18,12 +18,16 @@ def test_04_malicious_prompts_rag_and_cot_meta_llama_3_2_3b_instruct(
             .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.META_LLAMA_3_2_3B_INSTRUCT,
         text_generation_completion_service=text_generation_completion_service,
         semantic_similarity_service=semantic_similarity_service,
-        prompts=get_prompt_batch(malicious_prompts),
+        prompts=prompt_batch,
+        start=start,
+        end=end,
         comparison_texts=prompt_injection_all_example_responses,
         service_configurator=configure_service
     )
\ No newline at end of file