diff --git a/src/text_generation/domain/guardrails_result.py b/src/text_generation/domain/guardrails_result.py index 23c56cae1..d903539f1 100644 --- a/src/text_generation/domain/guardrails_result.py +++ b/src/text_generation/domain/guardrails_result.py @@ -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 \ No newline at end of file + def is_original_completion_malicious(self) -> bool: + return self.cosine_similarity_score >= self.cosine_similarity_risk_threshold diff --git a/src/text_generation/domain/guidelines_result.py b/src/text_generation/domain/guidelines_result.py index 6a7f1e0c0..341d43479 100644 --- a/src/text_generation/domain/guidelines_result.py +++ b/src/text_generation/domain/guidelines_result.py @@ -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( diff --git a/src/text_generation/services/guardrails/reflexion_security_guidelines_service.py b/src/text_generation/services/guardrails/reflexion_security_guidelines_service.py index 21f9cc1d6..ff79a1f39 100644 --- a/src/text_generation/services/guardrails/reflexion_security_guidelines_service.py +++ b/src/text_generation/services/guardrails/reflexion_security_guidelines_service.py @@ -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 \ No newline at end of file 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 59d1d0966..1f044cded 100644 --- a/src/text_generation/services/nlp/text_generation_completion_service.py +++ b/src/text_generation/services/nlp/text_generation_completion_service.py @@ -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): diff --git a/tests/integration/test_utils.py b/tests/integration/test_utils.py index b171c7703..f1f4b1317 100644 --- a/tests/integration/test_utils.py +++ b/tests/integration/test_utils.py @@ -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)