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

Reflexion guardrails updates

Browse Source
This commit is contained in:
Adam Wilson
2025-07-27 16:39:06 -06:00
parent 99ec0ddf98
commit a621ad82a9
5 changed files with 53 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/text_generation/domain/guardrails_result.py
+17 -10
View File
@@ -1,17 +1,24 @@
from typing import Any
from src.text_generation.domain.abstract_guardrails_processed_completion import AbstractGuardrailsProcessedCompletion
class GuardrailsResult(
AbstractGuardrailsProcessedCompletion):
def __init__(
self,
cosine_similarity_score: float,
cosine_similarity_risk_threshold: float,
original_completion: str,
guardrails_processed_completion_text: str):
is_original_completion_malicious = cosine_similarity_score >= cosine_similarity_risk_threshold
self,
user_prompt: str,
completion_text: str,
full_prompt: dict[str, Any],
llm_config: dict,
cosine_similarity_score: float = 0.0,
cosine_similarity_risk_threshold: float = 0.0):
self.user_prompt = user_prompt
self.completion_text = completion_text
self.full_prompt = full_prompt
self.llm_config = llm_config
self.cosine_similarity_score = cosine_similarity_score
self.cosine_similarity_risk_threshold = cosine_similarity_risk_threshold
self.score = cosine_similarity_score
self.original_completion = original_completion
self.is_original_completion_malicious = is_original_completion_malicious
self.guardrails_processed_completion_text = guardrails_processed_completion_text
def is_original_completion_malicious(self) -> bool:
return self.cosine_similarity_score >= self.cosine_similarity_risk_threshold
src/text_generation/domain/guidelines_result.py
-1
View File
@@ -1,6 +1,5 @@
from typing import Any, List
from src.text_generation.domain.abstract_guidelines_processed_completion import AbstractGuidelinesProcessedCompletion
from src.text_generation.domain.abstract_text_generation_completion_result import AbstractTextGenerationCompletionResult
class GuidelinesResult(
src/text_generation/services/guardrails/reflexion_security_guidelines_service.py
+32 -27
View File
@@ -2,21 +2,17 @@ from typing import Optional
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
from src.text_generation.domain.guardrails_result import GuardrailsResult
from src.text_generation.domain.guidelines_result import GuidelinesResult
from src.text_generation.domain.text_generation_completion_result import TextGenerationCompletionResult
from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel
from src.text_generation.services.guidelines.abstract_security_guidelines_service import AbstractSecurityGuidelinesConfigurationBuilder, AbstractSecurityGuidelinesService
from src.text_generation.services.guidelines.abstract_security_guidelines_service import AbstractSecurityGuidelinesConfigurationBuilder
from src.text_generation.services.nlp.abstract_prompt_template_service import AbstractPromptTemplateService
from src.text_generation.services.utilities.abstract_llm_configuration_introspection_service import AbstractLLMConfigurationIntrospectionService
from src.text_generation.services.utilities.abstract_response_processing_service import AbstractResponseProcessingService
from src.text_generation.domain.abstract_guardrails_processed_completion import AbstractGuardrailsProcessedCompletion
from src.text_generation.domain.abstract_text_generation_completion_result import AbstractTextGenerationCompletionResult
from src.text_generation.services.guardrails.abstract_generated_text_guardrail_service import AbstractGeneratedTextGuardrailService
@@ -24,42 +20,42 @@ from src.text_generation.services.guardrails.abstract_generated_text_guardrail_s
class ReflexionSecurityGuardrailsService(
AbstractGeneratedTextGuardrailService):
"""Basic implementation of reflexion security guidelines service."""
"""Basic implementation of reflexion security guardrails service."""
def __init__(
self,
foundation_model: AbstractFoundationModel,
response_processing_service: AbstractResponseProcessingService,
prompt_template_service: AbstractPromptTemplateService,
llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService,
config_builder: Optional[AbstractSecurityGuidelinesConfigurationBuilder] = None):
llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService):
super().__init__()
self.constants = Constants()
self.foundation_model_pipeline = foundation_model.create_pipeline()
self.response_processing_service = response_processing_service
self.prompt_template_service = prompt_template_service
self.llm_configuration_introspection_service = llm_configuration_introspection_service
self.config_builder = config_builder
def _create_context_from_rag(self, text_generation_completion_result: AbstractTextGenerationCompletionResult) -> str:
if not self.vectorstore:
raise ValueError("Vector store not initialized")
result: TextGenerationCompletionResult = text_generation_completion_result
original_user_prompt = result.original_user_prompt
original_completion = result.original_completion
# context_docs = self.vectorstore.as_retriever().invoke(?)
# context_parts = []
# for i, doc in enumerate(context_docs, 1):
# context_parts.append(f"Example {i}:\n{doc.page_content}")
# Assemble the context showing the original prompt injection attack
context_parts = [
"Original User Prompt:",
f"'{original_user_prompt}'",
"",
"Generated Completion:",
f"'{original_completion}'"
]
# return "\n\n".join(context_parts)
return "\n".join(context_parts)
def _get_template(self, text_generation_completion_result: AbstractTextGenerationCompletionResult) -> StringPromptTemplate:
prompt_template: StringPromptTemplate = self.prompt_template_service.get(id=self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__05_REFLEXION)
prompt_template: StringPromptTemplate = self.prompt_template_service.get(
id=self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__05_REFLEXION
)
context = self._create_context_from_rag(text_generation_completion_result)
# Create a new template with the context filled in
@@ -70,15 +66,23 @@ class ReflexionSecurityGuardrailsService(
return filled_template
def _create_chain(self, prompt_template: StringPromptTemplate):
return prompt_template | self.foundation_model_pipeline | StrOutputParser()
def apply_guardrails(self, text_generation_completion_result: AbstractTextGenerationCompletionResult) -> AbstractTextGenerationCompletionResult:
"""
Apply reflexion-based guardrails to mitigate prompt injection attacks
"""
if not text_generation_completion_result:
raise ValueError(f"Parameter 'user_prompt' cannot be empty or None")
raise ValueError(f"Parameter 'text_generation_completion_result' cannot be empty or None")
try:
prompt_template: StringPromptTemplate = self._get_template(user_prompt=)
prompt_value: PromptValue = prompt_template.format_prompt(input=user_prompt)
result: TextGenerationCompletionResult = text_generation_completion_result
original_user_prompt = result.original_user_prompt
prompt_template: StringPromptTemplate = self._get_template(text_generation_completion_result)
prompt_value: PromptValue = prompt_template.format_prompt(**{self.constants.INPUT_VARIABLE_TOKEN: original_user_prompt})
prompt_dict = {
"messages": [
{"role": msg.type, "content": msg.content, "additional_kwargs": msg.additional_kwargs}
@@ -88,14 +92,15 @@ class ReflexionSecurityGuardrailsService(
}
chain = self._create_chain(prompt_template)
completion_text=chain.invoke({self.constants.INPUT_VARIABLE_TOKEN: user_prompt})
completion_text = chain.invoke({self.constants.INPUT_VARIABLE_TOKEN: original_user_prompt})
llm_config = self.llm_configuration_introspection_service.get_config(chain)
result = GuidelinesResult(
result.guardrails_processed_completion = GuardrailsResult(
completion_text=completion_text,
llm_config=llm_config,
full_prompt=prompt_dict
)
return result
except Exception as e:
raise e
src/text_generation/services/nlp/text_generation_completion_service.py
+2 -9
View File
@@ -153,15 +153,8 @@ class TextGenerationCompletionService(
raise e
def _handle_reflexion_guardrails(self, text_generation_completion_result: TextGenerationCompletionResult) -> TextGenerationCompletionResult:
raise NotImplementedError
try:
chain = self._create_chain_without_guidelines()
completion = chain.invoke(user_prompt)
return TextGenerationCompletionResult(
original_completion=completion
)
except Exception as e:
raise e
result_with_guardrails_applied = self.reflexion_guardrails.apply_guardrails(text_generation_completion_result)
return result_with_guardrails_applied
# Configuration methods
def set_config(self, use_cot=False, use_rag=False):
tests/integration/test_utils.py
+2 -1
View File
@@ -5,6 +5,7 @@ from src.text_generation.domain.text_generation_completion_result import TextGen
from src.text_generation.services.logging.test_run_logging_service import TestRunLoggingService
from src.text_generation.services.nlp.abstract_semantic_similarity_service import AbstractSemanticSimilarityService
from src.text_generation.services.nlp.abstract_text_generation_completion_service import AbstractTextGenerationCompletionService
from src.text_generation.services.nlp.text_generation_completion_service import TextGenerationCompletionService
def run_prompt_analysis_test(
@@ -38,7 +39,7 @@ def run_prompt_analysis_test(
for i, prompt in enumerate(prompts[:max_prompts], 1):
# Configure the service using the provided configurator function
configured_service = service_configurator(text_generation_completion_service)
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)