From a1e07e6a4f7a31038a3f91bef660e9b44468fd34 Mon Sep 17 00:00:00 2001 From: Adam Wilson Date: Sat, 16 Aug 2025 19:47:26 -0600 Subject: [PATCH] support multiple templates --- ..._of_thought_security_guidelines_service.py | 150 ++++++++++++++++-- ...curity_guidelines_configuration_builder.py | 6 +- ...rag_context_security_guidelines_service.py | 144 +++++++++++++++-- ...ag_plus_cot_security_guidelines_service.py | 148 +++++++++++++++-- 4 files changed, 405 insertions(+), 43 deletions(-) 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 17d09b55a..08e4de4e0 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,6 +1,5 @@ -from typing import Optional +from typing import Optional, Dict 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 @@ -8,16 +7,18 @@ from src.text_generation.services.nlp.abstract_prompt_template_service import Ab 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 -class ChainOfThoughtSecurityGuidelinesService(BaseSecurityGuidelinesService): - """Service for zero-shot chain-of-thought security guidelines.""" +class ChainOfThoughtSecurityGuidelinesService(BaseSecurityGuidelinesService): + """Service for zero-shot chain-of-thought security guidelines with dynamic template selection.""" + def __init__( - self, - foundation_model: AbstractFoundationModel, - response_processing_service: AbstractResponseProcessingService, - prompt_template_service: AbstractPromptTemplateService, - llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService, - config_builder: Optional[AbstractSecurityGuidelinesConfigurationBuilder] = None): + self, + foundation_model: AbstractFoundationModel, + response_processing_service: AbstractResponseProcessingService, + prompt_template_service: AbstractPromptTemplateService, + llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService, + config_builder: Optional[AbstractSecurityGuidelinesConfigurationBuilder] = None + ): super().__init__( foundation_model=foundation_model, response_processing_service=response_processing_service, @@ -25,14 +26,129 @@ class ChainOfThoughtSecurityGuidelinesService(BaseSecurityGuidelinesService): llm_configuration_introspection_service=llm_configuration_introspection_service, config_builder=config_builder ) - - def _get_template(self, user_prompt: str) -> StringPromptTemplate: - """ - Get chain of thought security guidelines template. + # Initialize the model-to-template mapping + self._cot_template_mapping = self._build_cot_template_mapping() + + def _build_cot_template_mapping(self) -> Dict[str, str]: + """ + Build mapping from model identifiers to their corresponding CoT template IDs. + + Returns: + Dict[str, str]: Mapping from model name/identifier to CoT template ID + """ + return { + # Phi-3 models + "phi-3-mini-4k-instruct": self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__02_ZERO_SHOT_CHAIN_OF_THOUGHT, + "microsoft/Phi-3-mini-4K-Instruct": self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__02_ZERO_SHOT_CHAIN_OF_THOUGHT, + + # OpenELM models + "openelm-3b-instruct": self.constants.PromptTemplateIds.OPENELM_3B_INSTRUCT__02_ZERO_SHOT_CHAIN_OF_THOUGHT, + "apple/OpenELM-3B-Instruct": self.constants.PromptTemplateIds.OPENELM_3B_INSTRUCT__02_ZERO_SHOT_CHAIN_OF_THOUGHT, + + # Llama models + "llama-3.2-3b-instruct": self.constants.PromptTemplateIds.LLAMA_1_1B_CHAT__02_ZERO_SHOT_CHAIN_OF_THOUGHT, + "meta-llama/Llama-3.2-3B-Instruct": self.constants.PromptTemplateIds.LLAMA_1_1B_CHAT__02_ZERO_SHOT_CHAIN_OF_THOUGHT, + } + + def _get_model_identifier(self) -> str: + """ + Get the model identifier from the foundation model. + + Returns: + str: Model identifier/name + """ + # First try to get from foundation model if available + if hasattr(self, 'foundation_model') and self.foundation_model: + model_info = self.foundation_model.get_model_info() + if model_info: + model_id = ( + model_info.get('model_name') or + model_info.get('model_id') or + model_info.get('name') or + str(model_info) + ) + return model_id.lower() if model_id else "" + + # Fallback to introspection service + try: + model_info = self.llm_configuration_introspection_service.get_model_configuration() + + # Try different possible attribute names for the model identifier + model_id = ( + getattr(model_info, 'model_name', None) or + getattr(model_info, 'model_id', None) or + getattr(model_info, 'name', None) or + str(model_info) + ) + + return model_id.lower() if model_id else "" + except Exception: + return "" + + def _get_cot_template_id_for_model(self, model_identifier: str) -> str: + """ + Get the appropriate CoT template ID for the given model. + + Args: + model_identifier: The model identifier/name + + Returns: + str: The template ID for chain of thought prompting + + Raises: + ValueError: If no CoT template is found for the model + """ + # Try exact match first + if model_identifier in self._cot_template_mapping: + return self._cot_template_mapping[model_identifier] + + # Try partial matches for flexibility + for model_key, template_id in self._cot_template_mapping.items(): + if model_key in model_identifier or model_identifier in model_key: + return template_id + + # If no match found, raise an informative error + available_models = list(self._cot_template_mapping.keys()) + raise ValueError( + f"No chain of thought template found for model '{model_identifier}'. " + f"Available models: {available_models}" + ) + + def get_template(self, user_prompt: str) -> StringPromptTemplate: + """ + Get chain of thought security guidelines template dynamically based on the current model. + + Args: + user_prompt: The user's input prompt + Returns: StringPromptTemplate: Template configured for CoT processing """ - return self.prompt_template_service.get( - id=self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__02_ZERO_SHOT_CHAIN_OF_THOUGHT - ) \ No newline at end of file + # Get the current model identifier + model_identifier = self._get_model_identifier() + + # Get the appropriate CoT template ID for this model + template_id = self._get_cot_template_id_for_model(model_identifier) + + # Return the template from the service + return self.prompt_template_service.get(id=template_id) + + def add_model_template_mapping(self, model_identifier: str, template_id: str) -> None: + """ + Add or update a model-to-template mapping. + + Args: + model_identifier: The model identifier/name + template_id: The corresponding CoT template ID + """ + self._cot_template_mapping[model_identifier.lower()] = template_id + + def get_supported_models(self) -> list[str]: + """ + Get list of supported model identifiers. + + Returns: + list[str]: List of supported model identifiers + """ + return list(self._cot_template_mapping.keys()) \ 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 cbf10c830..da3ce34eb 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 @@ -37,7 +37,6 @@ class RetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder( def _load_examples(self): data = self.prompt_injection_example_repository.get_all() - documents = [] for item in data: content = f"Prompt: {item['prompt_injection_prompt']}\nCompletion: {item['prompt_injection_completion']}" @@ -49,7 +48,6 @@ class RetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder( } ) documents.append(doc) - return documents def _create_context(self, user_prompt: str, top_k: int = 3) -> str: @@ -60,10 +58,10 @@ class RetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder( context_parts = [] for i, doc in enumerate(context_docs, 1): context_parts.append(f"Example {i}:\n{doc.page_content}") - return "\n\n".join(context_parts) def get_prompt_template(self, template_id: str, user_prompt: str) -> PromptTemplate: + """Get the base template from the template service and fill in RAG context""" # Get the base template from the template service base_template = self.prompt_template_service.get(id=template_id) @@ -75,9 +73,9 @@ class RetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder( input_variables=[self.constants.INPUT_VARIABLE_TOKEN], template=base_template.template.replace("{context}", context) ) - return filled_template def get_formatted_prompt(self, template_id: str, user_prompt: str) -> str: + """Get formatted prompt with RAG context""" prompt_template = self.get_prompt_template(template_id, user_prompt) return prompt_template.format(**{self.constants.INPUT_VARIABLE_TOKEN: user_prompt}) \ No newline at end of file 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 4e2f0d2ac..cb727258f 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,3 +1,4 @@ +from typing import Dict from langchain_core.prompts import StringPromptTemplate from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel @@ -8,15 +9,16 @@ from src.text_generation.services.utilities.abstract_llm_configuration_introspec from src.text_generation.services.utilities.abstract_response_processing_service import AbstractResponseProcessingService class RagContextSecurityGuidelinesService(BaseSecurityGuidelinesService): - """Service for RAG context security guidelines.""" - + """Service for RAG context security guidelines with dynamic template selection.""" + def __init__( - self, - foundation_model: AbstractFoundationModel, - response_processing_service: AbstractResponseProcessingService, - prompt_template_service: AbstractPromptTemplateService, - llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService, - config_builder: AbstractSecurityGuidelinesConfigurationBuilder): + self, + foundation_model: AbstractFoundationModel, + response_processing_service: AbstractResponseProcessingService, + prompt_template_service: AbstractPromptTemplateService, + llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService, + config_builder: AbstractSecurityGuidelinesConfigurationBuilder): + super().__init__( foundation_model=foundation_model, response_processing_service=response_processing_service, @@ -24,10 +26,132 @@ class RagContextSecurityGuidelinesService(BaseSecurityGuidelinesService): llm_configuration_introspection_service=llm_configuration_introspection_service, config_builder=config_builder ) + + # Initialize the model-to-few-shot-template mapping + self._few_shot_template_mapping = self._build_few_shot_template_mapping() - def _get_template(self, user_prompt: str) -> StringPromptTemplate: - template_id = self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__03_FEW_SHOT_EXAMPLES + def _build_few_shot_template_mapping(self) -> Dict[str, str]: + """ + Build mapping from model identifiers to their corresponding few-shot template IDs. + + Returns: + Dict[str, str]: Mapping from model name/identifier to few-shot template ID + """ + return { + # Phi-3 models + "phi-3-mini-4k-instruct": self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__03_FEW_SHOT_EXAMPLES, + "microsoft/phi-3-mini-4k-instruct": self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__03_FEW_SHOT_EXAMPLES, + + # OpenELM models + "openelm-3b-instruct": self.constants.PromptTemplateIds.OPENELM_3B_INSTRUCT__03_FEW_SHOT_EXAMPLES, + "apple/openelm-3b-instruct": self.constants.PromptTemplateIds.OPENELM_3B_INSTRUCT__03_FEW_SHOT_EXAMPLES, + + # Llama models + "llama-3.2-3b-instruct": self.constants.PromptTemplateIds.LLAMA_1_1B_CHAT__03_FEW_SHOT_EXAMPLES, + "meta-llama/llama-3.2-3b-instruct": self.constants.PromptTemplateIds.LLAMA_1_1B_CHAT__03_FEW_SHOT_EXAMPLES, + } + + def _get_model_identifier(self) -> str: + """ + Get the model identifier from the foundation model. + + Returns: + str: Model identifier/name + """ + # First try to get from foundation model if available + if hasattr(self, 'foundation_model') and self.foundation_model: + model_info = self.foundation_model.get_model_info() + if model_info: + model_id = ( + model_info.get('model_name') or + model_info.get('model_id') or + model_info.get('name') or + str(model_info) + ) + return model_id.lower() if model_id else "" + + # Fallback to introspection service + try: + model_info = self.llm_configuration_introspection_service.get_model_configuration() + + # Try different possible attribute names for the model identifier + model_id = ( + getattr(model_info, 'model_name', None) or + getattr(model_info, 'model_id', None) or + getattr(model_info, 'name', None) or + str(model_info) + ) + + return model_id.lower() if model_id else "" + except Exception: + return "" + + def _get_few_shot_template_id_for_model(self, model_identifier: str) -> str: + """ + Get the appropriate few-shot template ID for the given model. + + Args: + model_identifier: The model identifier/name + + Returns: + str: The template ID for few-shot prompting + + Raises: + ValueError: If no few-shot template is found for the model + """ + # Try exact match first + if model_identifier in self._few_shot_template_mapping: + return self._few_shot_template_mapping[model_identifier] + + # Try partial matches for flexibility + for model_key, template_id in self._few_shot_template_mapping.items(): + if model_key in model_identifier or model_identifier in model_key: + return template_id + + # If no match found, raise an informative error + available_models = list(self._few_shot_template_mapping.keys()) + raise ValueError( + f"No few-shot template found for model '{model_identifier}'. " + f"Available models: {available_models}" + ) + + def get_template(self, user_prompt: str) -> StringPromptTemplate: + """ + Get RAG context security guidelines template dynamically based on the current model. + + Args: + user_prompt: The user's input prompt + + Returns: + StringPromptTemplate: Template configured for RAG processing + """ + # Get the current model identifier + model_identifier = self._get_model_identifier() + + # Get the appropriate few-shot template ID for this model + template_id = self._get_few_shot_template_id_for_model(model_identifier) + + # Use the config builder to get the template with RAG context return self.config_builder.get_prompt_template( template_id=template_id, user_prompt=user_prompt ) + + def add_model_template_mapping(self, model_identifier: str, template_id: str) -> None: + """ + Add or update a model-to-few-shot-template mapping. + + Args: + model_identifier: The model identifier/name + template_id: The corresponding few-shot template ID + """ + self._few_shot_template_mapping[model_identifier.lower()] = template_id + + def get_supported_models(self) -> list[str]: + """ + Get list of supported model identifiers. + + Returns: + list[str]: List of supported model identifiers + """ + return list(self._few_shot_template_mapping.keys()) 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 index a24695527..5fef1cc61 100644 --- 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 @@ -1,3 +1,4 @@ +from typing import Dict from langchain_core.prompts import StringPromptTemplate from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel @@ -9,17 +10,18 @@ from src.text_generation.services.utilities.abstract_response_processing_service class RagPlusCotSecurityGuidelinesService(BaseSecurityGuidelinesService): """ - Service that combines Retrieval Augmented Generation (RAG) with - Chain of Thought (CoT) security guidelines. + Service that combines Retrieval Augmented Generation (RAG) with + Chain of Thought (CoT) security guidelines with dynamic template selection. """ - + def __init__( - self, - foundation_model: AbstractFoundationModel, - response_processing_service: AbstractResponseProcessingService, - prompt_template_service: AbstractPromptTemplateService, - llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService, - config_builder: AbstractSecurityGuidelinesConfigurationBuilder): + self, + foundation_model: AbstractFoundationModel, + response_processing_service: AbstractResponseProcessingService, + prompt_template_service: AbstractPromptTemplateService, + llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService, + config_builder: AbstractSecurityGuidelinesConfigurationBuilder): + super().__init__( foundation_model=foundation_model, response_processing_service=response_processing_service, @@ -27,10 +29,132 @@ class RagPlusCotSecurityGuidelinesService(BaseSecurityGuidelinesService): llm_configuration_introspection_service=llm_configuration_introspection_service, config_builder=config_builder ) + + # Initialize the model-to-rag-plus-cot-template mapping + self._rag_plus_cot_template_mapping = self._build_rag_plus_cot_template_mapping() - def _get_template(self, user_prompt: str) -> StringPromptTemplate: - template_id = self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__04_FEW_SHOT_RAG_PLUS_COT + def _build_rag_plus_cot_template_mapping(self) -> Dict[str, str]: + """ + Build mapping from model identifiers to their corresponding RAG+CoT template IDs. + + Returns: + Dict[str, str]: Mapping from model name/identifier to RAG+CoT template ID + """ + return { + # Phi-3 models + "phi-3-mini-4k-instruct": self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__04_FEW_SHOT_RAG_PLUS_COT, + "microsoft/phi-3-mini-4k-instruct": self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__04_FEW_SHOT_RAG_PLUS_COT, + + # OpenELM models + "openelm-3b-instruct": self.constants.PromptTemplateIds.OPENELM_3B_INSTRUCT__04_FEW_SHOT_RAG_PLUS_COT, + "apple/openelm-3b-instruct": self.constants.PromptTemplateIds.OPENELM_3B_INSTRUCT__04_FEW_SHOT_RAG_PLUS_COT, + + # Llama models + "llama-3.2-3b-instruct": self.constants.PromptTemplateIds.LLAMA_1_1B_CHAT__04_FEW_SHOT_RAG_PLUS_COT, + "meta-llama/llama-3.2-3b-instruct": self.constants.PromptTemplateIds.LLAMA_1_1B_CHAT__04_FEW_SHOT_RAG_PLUS_COT, + } + + def _get_model_identifier(self) -> str: + """ + Get the model identifier from the foundation model. + + Returns: + str: Model identifier/name + """ + # First try to get from foundation model if available + if hasattr(self, 'foundation_model') and self.foundation_model: + model_info = self.foundation_model.get_model_info() + if model_info: + model_id = ( + model_info.get('model_name') or + model_info.get('model_id') or + model_info.get('name') or + str(model_info) + ) + return model_id.lower() if model_id else "" + + # Fallback to introspection service + try: + model_info = self.llm_configuration_introspection_service.get_model_configuration() + + # Try different possible attribute names for the model identifier + model_id = ( + getattr(model_info, 'model_name', None) or + getattr(model_info, 'model_id', None) or + getattr(model_info, 'name', None) or + str(model_info) + ) + + return model_id.lower() if model_id else "" + except Exception: + return "" + + def _get_rag_plus_cot_template_id_for_model(self, model_identifier: str) -> str: + """ + Get the appropriate RAG+CoT template ID for the given model. + + Args: + model_identifier: The model identifier/name + + Returns: + str: The template ID for RAG+CoT prompting + + Raises: + ValueError: If no RAG+CoT template is found for the model + """ + # Try exact match first + if model_identifier in self._rag_plus_cot_template_mapping: + return self._rag_plus_cot_template_mapping[model_identifier] + + # Try partial matches for flexibility + for model_key, template_id in self._rag_plus_cot_template_mapping.items(): + if model_key in model_identifier or model_identifier in model_key: + return template_id + + # If no match found, raise an informative error + available_models = list(self._rag_plus_cot_template_mapping.keys()) + raise ValueError( + f"No RAG+CoT template found for model '{model_identifier}'. " + f"Available models: {available_models}" + ) + + def get_template(self, user_prompt: str) -> StringPromptTemplate: + """ + Get RAG+CoT security guidelines template dynamically based on the current model. + + Args: + user_prompt: The user's input prompt + + Returns: + StringPromptTemplate: Template configured for RAG+CoT processing + """ + # Get the current model identifier + model_identifier = self._get_model_identifier() + + # Get the appropriate RAG+CoT template ID for this model + template_id = self._get_rag_plus_cot_template_id_for_model(model_identifier) + + # Use the config builder to get the template with RAG context return self.config_builder.get_prompt_template( template_id=template_id, user_prompt=user_prompt - ) \ No newline at end of file + ) + + def add_model_template_mapping(self, model_identifier: str, template_id: str) -> None: + """ + Add or update a model-to-RAG+CoT-template mapping. + + Args: + model_identifier: The model identifier/name + template_id: The corresponding RAG+CoT template ID + """ + self._rag_plus_cot_template_mapping[model_identifier.lower()] = template_id + + def get_supported_models(self) -> list[str]: + """ + Get list of supported model identifiers. + + Returns: + list[str]: List of supported model identifiers + """ + return list(self._rag_plus_cot_template_mapping.keys()) \ No newline at end of file