From 7f77e82a5c6b6738fb885047e89e46158a777e15 Mon Sep 17 00:00:00 2001 From: Adam Wilson Date: Fri, 18 Jul 2025 13:48:19 -0600 Subject: [PATCH] guidelines close to complete; prior to tests --- src/text_generation/common/constants.py | 1 + .../dependency_injection_container.py | 33 +++++++++++-- .../abstract_security_guidelines_service.py | 6 ++- .../base_security_guidelines_service.py | 23 ++++++---- ..._of_thought_security_guidelines_service.py | 32 ++++++++++++- ...curity_guidelines_configuration_builder.py | 4 +- ...rag_context_security_guidelines_service.py | 33 ++++++++++++- ...ag_plus_cot_security_guidelines_service.py | 46 +++++++++++++++++++ .../services/nlp/prompt_template_service.py | 6 +-- .../nlp/text_generation_completion_service.py | 44 +++++++++++++----- 10 files changed, 194 insertions(+), 34 deletions(-) create mode 100644 src/text_generation/services/guidelines/rag_plus_cot_security_guidelines_service.py diff --git a/src/text_generation/common/constants.py b/src/text_generation/common/constants.py index fa3356903..892c5bd35 100644 --- a/src/text_generation/common/constants.py +++ b/src/text_generation/common/constants.py @@ -11,4 +11,5 @@ class Constants: PHI_3_MINI_4K_INSTRUCT_BASIC = "phi-3-mini-4k-instruct-basic" PHI_3_MINI_4K_INSTRUCT_ZERO_SHOT_CHAIN_OF_THOUGHT = "phi-3-mini-4k-instruct-zero-shot-cot" PHI_3_MINI_4K_INSTRUCT_FEW_SHOT_EXAMPLES = "phi-3-mini-4k-instruct-few-shot" + PHI_3_MINI_4K_INSTRUCT_FEW_SHOT_RAG_PLUS_COT = "phi-3-mini-4k-instruct-few-shot-rag-plus-cot" PHI_3_MINI_4K_INSTRUCT_REFLEXION = "phi-3-mini-4k-instruct-reflexion" \ No newline at end of file diff --git a/src/text_generation/dependency_injection_container.py b/src/text_generation/dependency_injection_container.py index 93f3922e0..039f8e06c 100644 --- a/src/text_generation/dependency_injection_container.py +++ b/src/text_generation/dependency_injection_container.py @@ -1,6 +1,7 @@ from dependency_injector import containers, providers 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.entrypoints.http_api_controller import HttpApiController @@ -11,10 +12,12 @@ from src.text_generation.services.guidelines.rag_context_security_guidelines_con from src.text_generation.services.guidelines.rag_context_security_guidelines_service import RagContextSecurityGuidelinesService, RetrievalAugmentedGenerationContextSecurityGuidelinesService 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.rag_plus_cot_security_guidelines_service import RagPlusCotSecurityGuidelinesService from src.text_generation.services.logging.json_web_traffic_logging_service import JSONWebTrafficLoggingService from src.text_generation.services.nlp.prompt_template_service import PromptTemplateService from src.text_generation.services.nlp.semantic_similarity_service import SemanticSimilarityService from src.text_generation.services.nlp.text_generation_completion_service import TextGenerationCompletionService +from src.text_generation.services.prompt_injection.prompt_injection_example_service import PromptInjectionExampleService from src.text_generation.services.utilities.response_processing_service import ResponseProcessingService @@ -71,6 +74,10 @@ class DependencyInjectionContainer(containers.DeclarativeContainer): semantic_similarity_service=semantic_similarity_service ) + rag_config_builder = providers.Factory( + RetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder + ) + # Register security guideline services chain_of_thought_guidelines = providers.Factory( ChainOfThoughtSecurityGuidelinesService, @@ -83,21 +90,41 @@ class DependencyInjectionContainer(containers.DeclarativeContainer): RagContextSecurityGuidelinesService, foundation_model=foundation_model, response_processing_service=response_processing_service, - prompt_template_service=prompt_template_service + prompt_template_service=prompt_template_service, + config_builder=rag_config_builder ).provides(AbstractSecurityGuidelinesService) reflexion_guardrails = providers.Factory( ReflexionSecurityGuardrailsService ) - # Main service + prompt_injection_example_repository = providers.Factory( + PromptInjectionExampleRepository + ) + + prompt_injection_example_service = providers.Factory( + PromptInjectionExampleService, + repository=prompt_injection_example_repository + ) + + rag_plus_cot_guidelines = providers.Factory( + RagPlusCotSecurityGuidelinesService, + foundation_model=foundation_model, + response_processing_service=response_processing_service, + prompt_template_service=prompt_template_service + ) + text_generation_completion_service = providers.Factory( 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, - reflexion_guardrails=reflexion_guardrails + rag_plus_cot_guidelines=rag_plus_cot_guidelines, + reflexion_guardrails=reflexion_guardrails, + semantic_similarity_service=semantic_similarity_service, + prompt_injection_example_service=prompt_injection_example_service ) api_controller = providers.Factory( diff --git a/src/text_generation/services/guidelines/abstract_security_guidelines_service.py b/src/text_generation/services/guidelines/abstract_security_guidelines_service.py index 18cb0772a..4b49df747 100644 --- a/src/text_generation/services/guidelines/abstract_security_guidelines_service.py +++ b/src/text_generation/services/guidelines/abstract_security_guidelines_service.py @@ -1,5 +1,7 @@ import abc +from langchain_core.prompts import StringPromptTemplate + from src.text_generation.domain.abstract_guidelines_processed_completion import AbstractGuidelinesProcessedCompletion @@ -9,7 +11,7 @@ class AbstractSecurityGuidelinesService(abc.ABC): pass -class AbstractRetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder(abc.ABC): +class AbstractSecurityGuidelinesConfigurationBuilder(abc.ABC): @abc.abstractmethod - def get_prompt_template(self) -> str: + def get_prompt_template(self, template_id: str, user_prompt: str) -> StringPromptTemplate: raise NotImplementedError diff --git a/src/text_generation/services/guidelines/base_security_guidelines_service.py b/src/text_generation/services/guidelines/base_security_guidelines_service.py index c10272f17..afd0bdf5a 100644 --- a/src/text_generation/services/guidelines/base_security_guidelines_service.py +++ b/src/text_generation/services/guidelines/base_security_guidelines_service.py @@ -1,12 +1,12 @@ +from typing import Optional from langchain_core.output_parsers import StrOutputParser -from langchain_core.prompts import PromptTemplate +from langchain_core.prompts import PromptTemplate, StringPromptTemplate from langchain_core.runnables import RunnablePassthrough from src.text_generation.common.constants import Constants from src.text_generation.domain.abstract_guidelines_processed_completion import AbstractGuidelinesProcessedCompletion -from src.text_generation.domain.guidelines_result import GuidelinesResult from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel -from src.text_generation.services.guidelines.abstract_security_guidelines_service import AbstractSecurityGuidelinesService +from src.text_generation.services.guidelines.abstract_security_guidelines_service import AbstractSecurityGuidelinesConfigurationBuilder, AbstractSecurityGuidelinesService from src.text_generation.services.nlp.abstract_prompt_template_service import AbstractPromptTemplateService from src.text_generation.services.utilities.abstract_response_processing_service import AbstractResponseProcessingService @@ -17,12 +17,14 @@ class BaseSecurityGuidelinesService(AbstractSecurityGuidelinesService): self, foundation_model: AbstractFoundationModel, response_processing_service: AbstractResponseProcessingService, - prompt_template_service: AbstractPromptTemplateService): + prompt_template_service: AbstractPromptTemplateService, + config_builder: Optional[AbstractSecurityGuidelinesConfigurationBuilder] = None): 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.config_builder = config_builder def _create_chain(self, prompt_template: PromptTemplate): return ( @@ -33,16 +35,21 @@ class BaseSecurityGuidelinesService(AbstractSecurityGuidelinesService): | self.response_processing_service.process_text_generation_output ) - def _get_template_id(self) -> str: - raise NotImplementedError("Subclasses must implement _get_template_id()") + def _get_template(self) -> StringPromptTemplate: + """ + Get the prompt template for security guidelines. + + Returns: + StringPromptTemplate: Template for processing security guidelines + """ + raise NotImplementedError("Subclasses must implement _get_template()") def apply_guidelines(self, user_prompt: str) -> AbstractGuidelinesProcessedCompletion: if not user_prompt: raise ValueError(f"Parameter 'user_prompt' cannot be empty or None") try: - template_id = self._get_template_id() - prompt_template: PromptTemplate = self.prompt_template_service.get(id=template_id) + prompt_template = self._get_template() chain = self._create_chain(prompt_template) return chain.invoke(user_prompt) except Exception as e: diff --git a/src/text_generation/services/guidelines/chain_of_thought_security_guidelines_service.py b/src/text_generation/services/guidelines/chain_of_thought_security_guidelines_service.py index 2cb40b132..f28534c37 100644 --- a/src/text_generation/services/guidelines/chain_of_thought_security_guidelines_service.py +++ b/src/text_generation/services/guidelines/chain_of_thought_security_guidelines_service.py @@ -1,7 +1,35 @@ +from typing import Optional +from langchain_core.prompts import StringPromptTemplate + +from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel +from src.text_generation.services.guidelines.abstract_security_guidelines_service import AbstractSecurityGuidelinesConfigurationBuilder from src.text_generation.services.guidelines.base_security_guidelines_service import BaseSecurityGuidelinesService +from src.text_generation.services.nlp.abstract_prompt_template_service import AbstractPromptTemplateService +from src.text_generation.services.utilities.abstract_response_processing_service import AbstractResponseProcessingService class ChainOfThoughtSecurityGuidelinesService(BaseSecurityGuidelinesService): """Service for zero-shot chain-of-thought security guidelines.""" - def _get_template_id(self) -> str: - return self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_ZERO_SHOT_CHAIN_OF_THOUGHT \ No newline at end of file + def __init__( + self, + foundation_model: AbstractFoundationModel, + response_processing_service: AbstractResponseProcessingService, + prompt_template_service: AbstractPromptTemplateService, + config_builder: Optional[AbstractSecurityGuidelinesConfigurationBuilder] = None): + super().__init__( + foundation_model=foundation_model, + response_processing_service=response_processing_service, + prompt_template_service=prompt_template_service, + config_builder=config_builder + ) + + def _get_template(self, user_prompt: str) -> StringPromptTemplate: + """ + Get chain of thought security guidelines template. + + Returns: + StringPromptTemplate: Template configured for CoT processing + """ + return self.prompt_template_service.get( + id=self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_ZERO_SHOT_CHAIN_OF_THOUGHT + ) \ No newline at end of file diff --git a/src/text_generation/services/guidelines/rag_context_security_guidelines_configuration_builder.py b/src/text_generation/services/guidelines/rag_context_security_guidelines_configuration_builder.py index 8aece890d..74b7ace53 100644 --- a/src/text_generation/services/guidelines/rag_context_security_guidelines_configuration_builder.py +++ b/src/text_generation/services/guidelines/rag_context_security_guidelines_configuration_builder.py @@ -8,12 +8,12 @@ from src.text_generation.adapters.embedding_model import EmbeddingModel from src.text_generation.common.constants import Constants from src.text_generation.ports.abstract_prompt_injection_example_repository import AbstractPromptInjectionExampleRepository from src.text_generation.ports.abstract_embedding_model import AbstractEmbeddingModel -from src.text_generation.services.guidelines.abstract_security_guidelines_service import AbstractRetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder +from src.text_generation.services.guidelines.abstract_security_guidelines_service import AbstractSecurityGuidelinesConfigurationBuilder from src.text_generation.services.nlp.abstract_prompt_template_service import AbstractPromptTemplateService class RetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder( - AbstractRetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder): + AbstractSecurityGuidelinesConfigurationBuilder): def __init__( self, diff --git a/src/text_generation/services/guidelines/rag_context_security_guidelines_service.py b/src/text_generation/services/guidelines/rag_context_security_guidelines_service.py index 97d3abc35..5c9da8db4 100644 --- a/src/text_generation/services/guidelines/rag_context_security_guidelines_service.py +++ b/src/text_generation/services/guidelines/rag_context_security_guidelines_service.py @@ -1,7 +1,36 @@ +from langchain_core.prompts import StringPromptTemplate + +from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel +from src.text_generation.services.guidelines.abstract_security_guidelines_service import AbstractSecurityGuidelinesConfigurationBuilder from src.text_generation.services.guidelines.base_security_guidelines_service import BaseSecurityGuidelinesService +from src.text_generation.services.nlp.abstract_prompt_template_service import AbstractPromptTemplateService +from src.text_generation.services.utilities.abstract_response_processing_service import AbstractResponseProcessingService class RagContextSecurityGuidelinesService(BaseSecurityGuidelinesService): """Service for RAG context security guidelines.""" - def _get_template_id(self) -> str: - return self.constants.PromptTemplateIds.RAG_CONTEXT_SECURITY_GUIDELINES \ No newline at end of file + def __init__( + self, + foundation_model: AbstractFoundationModel, + response_processing_service: AbstractResponseProcessingService, + prompt_template_service: AbstractPromptTemplateService, + config_builder: AbstractSecurityGuidelinesConfigurationBuilder): + super().__init__( + foundation_model=foundation_model, + response_processing_service=response_processing_service, + prompt_template_service=prompt_template_service, + config_builder=config_builder + ) + + def _get_template(self, user_prompt: str) -> StringPromptTemplate: + """ + Get RAG context security guidelines template. + + Returns: + StringPromptTemplate: Template configured for RAG processing + """ + template_id = self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_FEW_SHOT_EXAMPLES + return self.config_builder.get_prompt_template( + template_id=template_id, + user_prompt=user_prompt + ) diff --git a/src/text_generation/services/guidelines/rag_plus_cot_security_guidelines_service.py b/src/text_generation/services/guidelines/rag_plus_cot_security_guidelines_service.py new file mode 100644 index 000000000..ee564fc34 --- /dev/null +++ b/src/text_generation/services/guidelines/rag_plus_cot_security_guidelines_service.py @@ -0,0 +1,46 @@ +from langchain_core.prompts import StringPromptTemplate + +from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel +from src.text_generation.services.guidelines.abstract_security_guidelines_service import AbstractSecurityGuidelinesConfigurationBuilder +from src.text_generation.services.guidelines.base_security_guidelines_service import BaseSecurityGuidelinesService +from src.text_generation.services.nlp.abstract_prompt_template_service import AbstractPromptTemplateService +from src.text_generation.services.utilities.abstract_response_processing_service import AbstractResponseProcessingService + +class RagPlusCotSecurityGuidelinesService(BaseSecurityGuidelinesService): + """ + Service that combines Retrieval Augmented Generation (RAG) with + Chain of Thought (CoT) security guidelines. + """ + + def __init__( + self, + foundation_model: AbstractFoundationModel, + response_processing_service: AbstractResponseProcessingService, + prompt_template_service: AbstractPromptTemplateService, + config_builder: AbstractSecurityGuidelinesConfigurationBuilder): + super().__init__( + foundation_model=foundation_model, + response_processing_service=response_processing_service, + prompt_template_service=prompt_template_service, + config_builder=config_builder + ) + + def _get_template(self, user_prompt: str) -> StringPromptTemplate: + """ + Get RAG context security guidelines template. + + Returns: + StringPromptTemplate: Template configured for RAG processing + """ + return self.prompt_template_service.get( + id=self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_FEW_SHOT_RAG_PLUS_COT + ) + + def _get_template_id(self) -> str: + """ + Get template ID for combined RAG + CoT processing. + + Returns: + str: Template ID for RAG + CoT security guidelines + """ + return self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_FEW_SHOT_RAG_PLUS_COT \ No newline at end of file diff --git a/src/text_generation/services/nlp/prompt_template_service.py b/src/text_generation/services/nlp/prompt_template_service.py index fcc1bf062..009fa117d 100644 --- a/src/text_generation/services/nlp/prompt_template_service.py +++ b/src/text_generation/services/nlp/prompt_template_service.py @@ -1,4 +1,4 @@ -from langchain_core.prompts import PromptTemplate +from langchain_core.prompts import PromptTemplate, StringPromptTemplate from src.text_generation.ports.abstract_prompt_template_repository import AbstractPromptTemplateRepository from src.text_generation.services.nlp.abstract_prompt_template_service import AbstractPromptTemplateService @@ -10,8 +10,8 @@ class PromptTemplateService(AbstractPromptTemplateService): super().__init__() self.prompt_template_repository = prompt_template_repository - def get(self, id: str) -> PromptTemplate: - prompt_template: PromptTemplate = self.prompt_template_repository.get(id) + def get(self, id: str) -> StringPromptTemplate: + prompt_template: StringPromptTemplate = self.prompt_template_repository.get(id) return prompt_template def add(self, id: str, prompt_template: PromptTemplate) -> None: 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 81cd91cca..3e3352a1f 100644 --- a/src/text_generation/services/nlp/text_generation_completion_service.py +++ b/src/text_generation/services/nlp/text_generation_completion_service.py @@ -29,6 +29,7 @@ class TextGenerationCompletionService( prompt_template_service: AbstractPromptTemplateService, chain_of_thought_guidelines: AbstractSecurityGuidelinesService, rag_context_guidelines: AbstractSecurityGuidelinesService, + rag_plus_cot_guidelines: AbstractSecurityGuidelinesService, reflexion_guardrails: AbstractGeneratedTextGuardrailService, semantic_similarity_service: AbstractSemanticSimilarityService, prompt_injection_example_service: AbstractPromptInjectionExampleService): @@ -46,9 +47,10 @@ class TextGenerationCompletionService( # guidelines services self.chain_of_thought_guidelines: AbstractSecurityGuidelinesService = chain_of_thought_guidelines self.rag_context_guidelines: AbstractSecurityGuidelinesService = rag_context_guidelines + self.rag_plus_cot_guidelines: AbstractSecurityGuidelinesService = rag_plus_cot_guidelines # guardrails services - self.reflexion_guardrails: AbstractSecurityGuidelinesService = reflexion_guardrails + self.reflexion_guardrails: AbstractGeneratedTextGuardrailService = reflexion_guardrails # constants self.COSINE_SIMILARITY_RISK_THRESHOLD = 0.5 @@ -67,7 +69,7 @@ class TextGenerationCompletionService( } # default guardrails settings - self._use_reflexion = False + self._use_reflexion_guardrails = False def _process_prompt_with_guidelines_if_applicable(self, user_prompt: str): guidelines_config = ( @@ -117,28 +119,43 @@ class TextGenerationCompletionService( return completion_result # Handler methods for each guidelines combination - def _handle_cot_and_rag(self, user_prompt: str): + def _handle_cot_and_rag(self, user_prompt: str) -> TextGenerationCompletionResult: """Handle: CoT=True, RAG=True""" - return f"CoT+RAG:" + guidelines_result = self.rag_plus_cot_guidelines.apply_guidelines(user_prompt) + return self._process_guidelines_result(guidelines_result) - def _handle_cot_only(self, user_prompt: str): + def _handle_cot_only(self, user_prompt: str) -> TextGenerationCompletionResult: """Handle: CoT=True, RAG=False""" guidelines_result = self.chain_of_thought_guidelines.apply_guidelines(user_prompt) return self._process_guidelines_result(guidelines_result) - def _handle_rag_only(self, user_prompt: str): + def _handle_rag_only(self, user_prompt: str) -> TextGenerationCompletionResult: """Handle: CoT=False, RAG=True""" guidelines_result = self.rag_context_guidelines.apply_guidelines(user_prompt) return self._process_guidelines_result(guidelines_result) - def _handle_without_guidelines(self, user_prompt: str): + def _handle_without_guidelines(self, user_prompt: str) -> TextGenerationCompletionResult: """Handle: CoT=False, RAG=False""" try: chain = self._create_chain_without_guidelines() - return chain.invoke(user_prompt) + completion = chain.invoke(user_prompt) + return TextGenerationCompletionResult( + original_completion=completion + ) except Exception as e: 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 + # Configuration methods def set_config(self, use_cot=False, use_rag=False): """Set guidelines configuration""" @@ -165,7 +182,7 @@ class TextGenerationCompletionService( return self def with_reflexion_guardrails(self) -> AbstractTextGenerationCompletionService: - self._use_reflexion = True + self._use_reflexion_guardrails = True return self def _create_chain_without_guidelines(self): @@ -180,7 +197,10 @@ class TextGenerationCompletionService( | self.response_processing_service.process_text_generation_output ) - def invoke(self, user_prompt: str) -> str: + def invoke(self, user_prompt: str) -> TextGenerationCompletionResult: if not user_prompt: raise ValueError(f"Parameter 'user_prompt' cannot be empty or None") - self._process_prompt_with_guidelines_if_applicable(user_prompt) \ No newline at end of file + completion_result: TextGenerationCompletionResult = self._process_prompt_with_guidelines_if_applicable(user_prompt) + if not self._use_reflexion_guardrails: + return completion_result + return self._handle_reflexion_guardrails(completion_result)