dynamic template and model selection

2026-07-06 21:17:51 +02:00 · 2025-08-18 11:22:50 -06:00
parent a1e07e6a4f
commit 36c11703cb
8 changed files with 271 additions and 141 deletions
@@ -1,12 +1,8 @@
-from src.text_generation.adapters.foundation_models.base.base_model_config import BaseModelConfig
-from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel
-
-
-from transformers import pipeline
-
-
 from abc import abstractmethod
 from typing import Any
+from transformers import pipeline
+from src.text_generation.adapters.foundation_models.base.base_model_config import BaseModelConfig
+from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel


 class BaseFoundationModel(AbstractFoundationModel):
@@ -0,0 +1,8 @@
+from enum import Enum
+
+class GuidelinesMode(Enum):
+    """Enum to define different guidelines processing modes"""
+    RAG_PLUS_COT = "rag_plus_cot"
+    COT_ONLY = "cot_only"
+    RAG_ONLY = "rag_only"
+    NONE = "none"
@@ -4,4 +4,4 @@ from enum import Enum
 class ModelId(Enum):
    APPLE_OPENELM_3B_INSTRUCT = "apple/OpenELM-3B-Instruct"
    META_LLAMA_3_2_3B_INSTRUCT = "meta-llama/Llama-3.2-3B-Instruct"
-    MICROSOFT_PHI_3_MINI4K_INSTRUCT = "microsoft/Phi-3-mini-4k-instruct"
+    MICROSOFT_PHI_3_MINI4K_INSTRUCT = "microsoft/Phi-3-mini-4k-instruct"
@@ -0,0 +1,8 @@
+from enum import Enum
+
+
+class PromptTemplateType(Enum):
+    BASIC = "basic"
+    ZERO_SHOT_COT = "zero_shot_cot"
+    FEW_SHOT = "few_shot"
+    RAG_PLUS_COT = "rag_plus_cot"
@@ -0,0 +1,99 @@
+from abc import ABC, abstractmethod
+from langchain.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.chain_of_thought_security_guidelines_service import ChainOfThoughtSecurityGuidelinesService
+from src.text_generation.services.guidelines.rag_context_security_guidelines_service import RagContextSecurityGuidelinesService
+from src.text_generation.services.guidelines.rag_plus_cot_security_guidelines_service import RagPlusCotSecurityGuidelinesService
+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
+
+
+class AbstractGuidelinesFactory(ABC):
+
+    @abstractmethod
+    def create_cot_guidelines_service(
+        self,
+        foundation_model: AbstractFoundationModel,
+        response_processing_service: AbstractResponseProcessingService,
+        prompt_template_service: AbstractPromptTemplateService,
+        llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService,
+        config_builder: AbstractSecurityGuidelinesConfigurationBuilder
+    ) -> ChainOfThoughtSecurityGuidelinesService:
+        raise NotImplementedError
+
+    @abstractmethod
+    def create_rag_context_guidelines_service(
+        self,
+        foundation_model: AbstractFoundationModel,
+        prompt_template: StringPromptTemplate,
+        response_processing_service: AbstractResponseProcessingService,
+        prompt_template_service: AbstractPromptTemplateService,
+        llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService,
+        config_builder=AbstractSecurityGuidelinesConfigurationBuilder
+    ) -> RagContextSecurityGuidelinesService:
+        raise NotImplementedError
+    
+    @abstractmethod
+    def create_rag_plus_cot_context_guidelines_service(
+        self,
+        foundation_model: AbstractFoundationModel,
+        response_processing_service: AbstractResponseProcessingService,
+        prompt_template_service: AbstractPromptTemplateService,
+        llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService,
+        config_builder: AbstractSecurityGuidelinesConfigurationBuilder
+    ) -> RagPlusCotSecurityGuidelinesService:
+        raise NotImplementedError
+
+class GuidelinesFactory(AbstractGuidelinesFactory):
+
+    def create_rag_context_guidelines_service(
+        self,
+        foundation_model: AbstractFoundationModel,
+        prompt_template: StringPromptTemplate,
+        response_processing_service: AbstractResponseProcessingService,
+        prompt_template_service: AbstractPromptTemplateService,
+        llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService,
+        config_builder=AbstractSecurityGuidelinesConfigurationBuilder
+    ) -> RagContextSecurityGuidelinesService:
+        return RagContextSecurityGuidelinesService(
+            foundation_model=foundation_model,
+            response_processing_service=response_processing_service,
+            prompt_template_service=prompt_template_service,
+            llm_configuration_introspection_service=llm_configuration_introspection_service,
+            config_builder=config_builder,
+            prompt_template=prompt_template
+        )
+    
+    def create_cot_guidelines_service(
+        self,
+        foundation_model: AbstractFoundationModel,
+        response_processing_service: AbstractResponseProcessingService,
+        prompt_template_service: AbstractPromptTemplateService,
+        llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService,
+        config_builder: AbstractSecurityGuidelinesConfigurationBuilder
+    ) -> ChainOfThoughtSecurityGuidelinesService:
+        return ChainOfThoughtSecurityGuidelinesService(
+            foundation_model=foundation_model,
+            response_processing_service=response_processing_service,
+            prompt_template_service=prompt_template_service,
+            llm_configuration_introspection_service=llm_configuration_introspection_service,
+            config_builder=config_builder
+        )
+
+    def create_rag_plus_cot_context_guidelines_service(
+        self,
+        foundation_model: AbstractFoundationModel,
+        response_processing_service: AbstractResponseProcessingService,
+        prompt_template_service: AbstractPromptTemplateService,
+        llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService,
+        config_builder: AbstractSecurityGuidelinesConfigurationBuilder
+    ) -> RagPlusCotSecurityGuidelinesService: 
+        return RagPlusCotSecurityGuidelinesService(
+            foundation_model=foundation_model,
+            response_processing_service=response_processing_service,
+            prompt_template_service=prompt_template_service,
+            llm_configuration_introspection_service=llm_configuration_introspection_service,
+            config_builder=config_builder
+        )
@@ -17,15 +17,15 @@ class RetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder(
        self,
        embedding_model: AbstractEmbeddingModel,
        prompt_template_service: AbstractPromptTemplateService,
-        prompt_injection_example_repository: AbstractPromptInjectionExampleRepository):
-        
+        prompt_injection_example_repository: AbstractPromptInjectionExampleRepository
+    ):
        self.constants = Constants()
        self.embedding_model: EmbeddingModel = embedding_model
        self.prompt_template_service = prompt_template_service
        self.prompt_injection_example_repository = prompt_injection_example_repository
-        self.vectorstore = self._setup_vectorstore()
+        self.vectorstore = self._init_vectorstore()

-    def _setup_vectorstore(self):
+    def _init_vectorstore(self):
        documents = self._load_examples()
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=500,
@@ -11,14 +11,17 @@ from langchain_core.prompt_values import PromptValue
 from src.text_generation.adapters.foundation_models.base.base_model_config import BaseModelConfig
 from src.text_generation.adapters.foundation_models.factories.foundation_model_factory import FoundationModelFactory
 from src.text_generation.common.constants import Constants
+from src.text_generation.common.guidelines_mode import GuidelinesMode
 from src.text_generation.common.model_id import ModelId
+from src.text_generation.common.prompt_template_type import PromptTemplateType
 from src.text_generation.domain.alternate_completion_result import AlternateCompletionResult
 from src.text_generation.domain.guidelines_result import GuidelinesResult
 from src.text_generation.domain.original_completion_result import OriginalCompletionResult
 from src.text_generation.domain.semantic_similarity_result import SemanticSimilarityResult
 from src.text_generation.domain.text_generation_completion_result import TextGenerationCompletionResult
 from src.text_generation.services.guardrails.abstract_generated_text_guardrail_service import AbstractGeneratedTextGuardrailService
-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.guidelines.guidelines_factory import AbstractGuidelinesFactory
 from src.text_generation.services.nlp.abstract_prompt_template_service import AbstractPromptTemplateService
 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
@@ -35,13 +38,13 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService):
            self, 
            response_processing_service: AbstractResponseProcessingService,
            prompt_template_service: AbstractPromptTemplateService,
-            chain_of_thought_guidelines: AbstractSecurityGuidelinesService,
-            rag_context_guidelines: AbstractSecurityGuidelinesService,
-            rag_plus_cot_guidelines: AbstractSecurityGuidelinesService,
+            guidelines_factory: AbstractGuidelinesFactory,
+            guidelines_config_builder: AbstractSecurityGuidelinesConfigurationBuilder,
            semantic_similarity_service: AbstractSemanticSimilarityService,
            prompt_injection_example_service: AbstractPromptInjectionExampleService,
            llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService,
-            default_model_type: ModelId = ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT):
+            default_model_type: ModelId = ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT
+    ):
        
        super().__init__()
        self.constants = Constants()
@@ -66,9 +69,8 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService):
        )

        # Guidelines services
-        self.chain_of_thought_guidelines = chain_of_thought_guidelines
-        self.rag_context_guidelines = rag_context_guidelines
-        self.rag_plus_cot_guidelines = rag_plus_cot_guidelines
+        self.guidelines_factory = guidelines_factory
+        self.guidelines_config_builder = guidelines_config_builder

        # Constants and settings
        self.COSINE_SIMILARITY_RISK_THRESHOLD = 0.8
@@ -84,31 +86,40 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService):
            (False, False): self._handle_without_guidelines,
        }

-        # Initialize dynamic template mapping
-        self._basic_template_mapping = self._build_basic_template_mapping()
-
        # Load default model
        self.load_model(default_model_type)

-    def _build_basic_template_mapping(self) -> Dict[str, str]:
+    def _prompt_template_map(self) -> Dict[str, Dict[str, str]]:
        """
-        Build mapping from model identifiers to their corresponding basic template IDs.
+        Build mapping from model identifiers to their corresponding template IDs for all template types.
        
        Returns:
-            Dict[str, str]: Mapping from model name/identifier to basic template ID
+            Dict[str, Dict[str, str]]: Mapping from model name/identifier to all template IDs
        """
        return {
            # Phi-3 models
-            "phi-3-mini-4k-instruct": self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__01_BASIC,
-            "microsoft/phi-3-mini-4k-instruct": self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__01_BASIC,
-            
+            "microsoft/phi-3-mini-4k-instruct": {
+                PromptTemplateType.BASIC.value: self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__01_BASIC,
+                PromptTemplateType.ZERO_SHOT_COT.value: self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__02_ZERO_SHOT_CHAIN_OF_THOUGHT,
+                PromptTemplateType.FEW_SHOT.value: self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__03_FEW_SHOT_EXAMPLES,
+                PromptTemplateType.RAG_PLUS_COT.value: 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__01_BASIC,
-            "apple/openelm-3b-instruct": self.constants.PromptTemplateIds.OPENELM_3B_INSTRUCT__01_BASIC,
-            
+            "apple/openelm-3b-instruct": {
+                PromptTemplateType.BASIC.value: self.constants.PromptTemplateIds.OPENELM_3B_INSTRUCT__01_BASIC,
+                PromptTemplateType.ZERO_SHOT_COT.value: self.constants.PromptTemplateIds.OPENELM_3B_INSTRUCT__02_ZERO_SHOT_CHAIN_OF_THOUGHT,
+                PromptTemplateType.FEW_SHOT.value: self.constants.PromptTemplateIds.OPENELM_3B_INSTRUCT__03_FEW_SHOT_EXAMPLES,
+                PromptTemplateType.RAG_PLUS_COT.value: 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__01_BASIC,
-            "meta-llama/llama-3.2-3b-instruct": self.constants.PromptTemplateIds.LLAMA_1_1B_CHAT__01_BASIC,
+            "meta-llama/llama-3.2-3b-instruct": {
+                PromptTemplateType.BASIC.value: self.constants.PromptTemplateIds.LLAMA_1_1B_CHAT__01_BASIC,
+                PromptTemplateType.ZERO_SHOT_COT.value: self.constants.PromptTemplateIds.LLAMA_1_1B_CHAT__02_ZERO_SHOT_CHAIN_OF_THOUGHT,
+                PromptTemplateType.FEW_SHOT.value: self.constants.PromptTemplateIds.LLAMA_1_1B_CHAT__03_FEW_SHOT_EXAMPLES,
+                PromptTemplateType.RAG_PLUS_COT.value: self.constants.PromptTemplateIds.LLAMA_1_1B_CHAT__04_FEW_SHOT_RAG_PLUS_COT,
+            }
        }

    def _get_model_identifier_from_model_id(self, model_id: ModelId) -> str:
@@ -143,29 +154,6 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService):
        
        return ""

-    def _get_basic_template_id_for_model(self, model_identifier: str) -> str:
-        """
-        Get the appropriate basic template ID for the given model.
-        
-        Args:
-            model_identifier: The model identifier/name
-            
-        Returns:
-            str: The template ID for basic prompting
-        """
-        # Try exact match first
-        if model_identifier in self._basic_template_mapping:
-            return self._basic_template_mapping[model_identifier]
-        
-        # Try partial matches for flexibility
-        for model_key, template_id in self._basic_template_mapping.items():
-            if model_key in model_identifier or model_identifier in model_key:
-                return template_id
-        
-        # Default fallback to Phi-3 if no match found
-        logger.warning(f"No basic template found for model '{model_identifier}', falling back to Phi-3 template")
-        return self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__01_BASIC
-
    def load_model(
        self, 
        model_id: ModelId, 
@@ -175,8 +163,8 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService):
        """Load a specific model"""
        if (not force_reload and 
            self._current_model is not None and 
-            self._current_model_id == model_id and
-            self._current_model.is_loaded()):
+            self._current_model_id == model_id
+        ):
            logger.info(f"Model {model_id.value} already loaded")
            return
        
@@ -184,23 +172,18 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService):
            self._current_model.unload()
        
        self._current_model = self.factory.create_model(model_id, config)
-        self._current_model.load()
        self._current_model_id: ModelId = model_id
        self.foundation_model_pipeline = self._current_model.create_pipeline()
        
        logger.info(f"Successfully loaded model: {model_id.value}")

-    def switch_model(self, model_id: ModelId, config: Optional[BaseModelConfig] = None) -> None:
-        """Switch to a different model"""
-        self.load_model(model_id, config, force_reload=True)
-
    def get_current_model_info(self) -> Optional[Dict[str, Any]]:
        """Get information about the currently loaded model"""
        if self._current_model and self._current_model.is_loaded():
            return self._current_model.get_model_info()
        return None

-    def _process_prompt_with_guidelines_if_applicable(self, user_prompt: str):
+    def _process_prompt_with_guidelines_if_applicable(self, user_prompt: str, target_model_id: ModelId):
        guidelines_config = (
            self._use_zero_shot_chain_of_thought,
            self._use_rag_context
@@ -210,7 +193,7 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService):
            # fall back to unfiltered LLM invocation
            self._handle_without_guidelines
        )
-        return guidelines_handler(user_prompt)
+        return guidelines_handler(user_prompt, target_model_id)

    def _process_completion_result(self, completion_result: TextGenerationCompletionResult) -> TextGenerationCompletionResult:
        """
@@ -251,40 +234,103 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService):
        completion_result.finalize_completion_text()
        return completion_result

-    # Handler methods for each guidelines combination
-    def _handle_cot_and_rag(self, user_prompt: str) -> TextGenerationCompletionResult:
-        """Handle: CoT=True, RAG=True"""
-        guidelines_result = self.rag_plus_cot_guidelines.apply_guidelines(user_prompt)
+    def _get_template_for_mode(self, mode: GuidelinesMode, target_model_id: Optional[ModelId] = None) -> str:
+        """Get the appropriate template ID based on the guidelines mode and model"""
+        if target_model_id:
+            model_identifier = self._get_model_identifier_from_model_id(target_model_id)
+        else:
+            model_identifier = self._get_current_model_identifier()
+        
+        template_map = {
+            GuidelinesMode.RAG_PLUS_COT: PromptTemplateType.RAG_PLUS_COT.value,
+            GuidelinesMode.COT_ONLY: PromptTemplateType.ZERO_SHOT_COT.value,
+            GuidelinesMode.RAG_ONLY: PromptTemplateType.FEW_SHOT.value,
+            GuidelinesMode.NONE: PromptTemplateType.BASIC.value
+        }
+        
+        return self._prompt_template_map()[model_identifier][template_map[mode]]
+
+    def _ensure_model_loaded(self, target_model_id: ModelId) -> None:
+        """Ensure the correct model is loaded"""
+        if (self._current_model_id != target_model_id or 
+            self._current_model is None):
+            self.load_model(target_model_id)
+
+    def _get_prompt_template(self, template_id: str) -> StringPromptTemplate:
+        """Get and validate prompt template"""
+        prompt_template = self.prompt_template_service.get(id=template_id)
+        if prompt_template is None:
+            raise ValueError(f"Prompt template not found for ID: {template_id}")
+        return prompt_template
+
+    def _create_guidelines_service(self, mode: GuidelinesMode, prompt_template: StringPromptTemplate) -> AbstractSecurityGuidelinesService:
+        """Factory method to create the appropriate guidelines service"""
+        base_params = {
+            'foundation_model': self._current_model,
+            'prompt_template': prompt_template,
+            'response_processing_service': self.response_processing_service,
+            'prompt_template_service': self.prompt_template_service,
+            'llm_configuration_introspection_service': self.llm_configuration_introspection_service
+        }
+        
+        if mode == GuidelinesMode.RAG_PLUS_COT:
+            return self.guidelines_factory.create_rag_plus_cot_context_guidelines_service(
+                **base_params,
+                config_builder=self.guidelines_config_builder
+            )
+        elif mode == GuidelinesMode.COT_ONLY:
+            return self.guidelines_factory.create_cot_guidelines_service(
+                **base_params,
+                config_builder=self.guidelines_config_builder
+            )
+        elif mode == GuidelinesMode.RAG_ONLY:
+            return self.guidelines_factory.create_rag_context_guidelines_service(**base_params)
+        else:
+            raise ValueError(f"Unsupported guidelines mode: {mode}")
+
+    def _handle_with_guidelines(self, user_prompt: str, target_model_id: ModelId, mode: GuidelinesMode) -> TextGenerationCompletionResult:
+        """Generic handler for guidelines-based processing"""
+        # Get template ID and load template
+        template_id = self._get_template_for_mode(mode, target_model_id)
+        prompt_template = self._get_prompt_template(template_id)
+        
+        # Ensure correct model is loaded
+        self._ensure_model_loaded(target_model_id)
+        
+        # Create appropriate guidelines service
+        guidelines_service = self._create_guidelines_service(mode, prompt_template)
+        
+        # Apply guidelines and process result
+        guidelines_result = guidelines_service.apply_guidelines(user_prompt)
        return self._process_completion_result(guidelines_result)

-    def _handle_cot_only(self, user_prompt: str) -> TextGenerationCompletionResult:
+    # Simplified handler methods
+    def _handle_cot_and_rag(self, user_prompt: str, target_model_id: ModelId) -> TextGenerationCompletionResult:
+        """Handle: CoT=True, RAG=True"""
+        return self._handle_with_guidelines(user_prompt, target_model_id, GuidelinesMode.RAG_PLUS_COT)
+
+    def _handle_cot_only(self, user_prompt: str, target_model_id: ModelId) -> TextGenerationCompletionResult:
        """Handle: CoT=True, RAG=False"""
-        guidelines_result = self.chain_of_thought_guidelines.apply_guidelines(user_prompt)
-        return self._process_completion_result(guidelines_result)
+        return self._handle_with_guidelines(user_prompt, target_model_id, GuidelinesMode.COT_ONLY)
    
-    def _handle_rag_only(self, user_prompt: str) -> TextGenerationCompletionResult:
+    def _handle_rag_only(self, user_prompt: str, target_model_id: ModelId) -> TextGenerationCompletionResult:
        """Handle: CoT=False, RAG=True"""
-        guidelines_result = self.rag_context_guidelines.apply_guidelines(user_prompt)
-        return self._process_completion_result(guidelines_result)
+        return self._handle_with_guidelines(user_prompt, target_model_id, GuidelinesMode.RAG_ONLY)
    
    def _handle_without_guidelines(self, user_prompt: str) -> TextGenerationCompletionResult:
        """Handle: CoT=False, RAG=False - now with dynamic template selection"""
        try:
-            # Get the current model identifier
-            model_identifier = self._get_current_model_identifier()
+            # Get template ID and load template
+            template_id = self._get_template_for_mode(GuidelinesMode.NONE)
+            prompt_template = self._get_prompt_template(template_id)
            
-            # Get the appropriate basic template ID for this model
-            template_id = self._get_basic_template_id_for_model(model_identifier)
-            
-            # Get the template from the service
-            prompt_template: StringPromptTemplate = self.prompt_template_service.get(id=template_id)
-
-            if prompt_template is None:
-                raise ValueError(f"Prompt template not found for ID: {template_id}")
+            print(f'using template: {template_id}')
            
+            # Create chain and get config
            chain = self._create_chain_without_guidelines(prompt_template)
            llm_config = self.llm_configuration_introspection_service.get_config(chain)

+            # Format prompt
            prompt_value: PromptValue = prompt_template.format_prompt(input=user_prompt)
            prompt_dict = {
                "messages": [
@@ -294,14 +340,17 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService):
                "string_representation": prompt_value.to_string(),
            }

+            # Create and return result
            result = TextGenerationCompletionResult(
                original_result=OriginalCompletionResult(
                    user_prompt=user_prompt,
-                    completion_text=chain.invoke({ self.constants.INPUT_VARIABLE_TOKEN: user_prompt }),
+                    completion_text=chain.invoke({self.constants.INPUT_VARIABLE_TOKEN: user_prompt}),
                    llm_config=llm_config,
                    full_prompt=prompt_dict
-            ))
+                )
+            )
            return self._process_completion_result(result)
+            
        except Exception as e:
            logger.error(f"Error in _handle_without_guidelines: {str(e)}")
            raise e
@@ -354,7 +403,7 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService):
            model_identifier: The model identifier/name
            basic_template_id: The corresponding basic template ID
        """
-        self._basic_template_mapping[model_identifier.lower()] = basic_template_id
+        self._prompt_template_map()[model_identifier.lower()] = basic_template_id

    def get_supported_models(self) -> list[str]:
        """
@@ -363,7 +412,7 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService):
        Returns:
            list[str]: List of supported model identifiers
        """
-        return list(self._basic_template_mapping.keys())
+        return list(self._prompt_template_map().keys())

    def invoke(self, user_prompt: str, model_id: Optional[ModelId] = None) -> TextGenerationCompletionResult:
        """Generate text using specified or current model"""
@@ -377,5 +426,5 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService):
            self.load_model(target_model_id)

        print(f'Using model: {target_model_id.value}, guidelines: {self.get_current_config()}')
-        completion_result = self._process_prompt_with_guidelines_if_applicable(user_prompt)        
+        completion_result = self._process_prompt_with_guidelines_if_applicable(user_prompt=user_prompt, model_id=target_model_id)        
        return completion_result
@@ -23,6 +23,7 @@ from src.text_generation.common.constants import Constants
 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.chain_of_thought_security_guidelines_service import ChainOfThoughtSecurityGuidelinesService
+from src.text_generation.services.guidelines.guidelines_factory import GuidelinesFactory
 from src.text_generation.services.guidelines.rag_context_security_guidelines_configuration_builder import RetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder
 from src.text_generation.services.guidelines.rag_context_security_guidelines_service import RagContextSecurityGuidelinesService
 from src.text_generation.services.guidelines.rag_plus_cot_security_guidelines_service import RagPlusCotSecurityGuidelinesService
@@ -124,48 +125,6 @@ def rag_config_builder(
 def llm_configuration_introspection_service():
    return LLMConfigurationIntrospectionService()

-@pytest.fixture(scope="session")
-def rag_context_guidelines(
-        foundation_model,
-        response_processing_service,
-        prompt_template_service,
-        llm_configuration_introspection_service,
-        rag_config_builder):
-    return RagContextSecurityGuidelinesService(
-        foundation_model=foundation_model,
-        response_processing_service=response_processing_service,    
-        prompt_template_service=prompt_template_service,
-        llm_configuration_introspection_service=llm_configuration_introspection_service,
-        config_builder=rag_config_builder
-    )
-
-@pytest.fixture(scope="session")
-def chain_of_thought_guidelines(
-        foundation_model,
-        response_processing_service,
-        llm_configuration_introspection_service,
-        prompt_template_service):
-    return ChainOfThoughtSecurityGuidelinesService(
-        foundation_model=foundation_model,
-        response_processing_service=response_processing_service,
-        llm_configuration_introspection_service=llm_configuration_introspection_service,
-        prompt_template_service=prompt_template_service
-    )
-
-@pytest.fixture(scope="session")
-def rag_plus_cot_guidelines(
-        foundation_model,
-        response_processing_service,
-        prompt_template_service,
-        llm_configuration_introspection_service,
-        rag_config_builder):
-    return RagPlusCotSecurityGuidelinesService(
-        foundation_model=foundation_model,
-        response_processing_service=response_processing_service,    
-        prompt_template_service=prompt_template_service,
-        llm_configuration_introspection_service=llm_configuration_introspection_service,
-        config_builder=rag_config_builder
-    )

@pytest.fixture(scope="session")
 def prompt_injection_example_service(prompt_injection_example_repository):
@@ -196,24 +155,35 @@ def response_processing_service():
 def llm_configuration_introspection_service():
    return LLMConfigurationIntrospectionService()

+@pytest.fixture(scope="session")
+def guidelines_factory():
+    return GuidelinesFactory()
+
+@pytest.fixture(scope="session")
+def guidelines_config_builder(
+        embedding_model,
+        prompt_template_service,
+        prompt_injection_example_repository):
+    return RetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder(
+        embedding_model=embedding_model,
+        prompt_template_service=prompt_template_service,
+        prompt_injection_example_repository=prompt_injection_example_repository
+    )
+
@pytest.fixture(scope="session")
 def text_generation_completion_service(
        response_processing_service,
        prompt_template_service,
-        chain_of_thought_guidelines,
-        rag_context_guidelines,
-        rag_plus_cot_guidelines,
-        reflexion_guardrails,
+        guidelines_factory,
+        guidelines_config_builder,
        semantic_similarity_service,
        prompt_injection_example_service,
        llm_configuration_introspection_service):
    return TextGenerationCompletionService(
        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,
-        rag_plus_cot_guidelines=rag_plus_cot_guidelines,
-        reflexion_guardrails=reflexion_guardrails,
+        guidelines_factory=guidelines_factory,
+        guidelines_config_builder=guidelines_config_builder,
        semantic_similarity_service=semantic_similarity_service,
        prompt_injection_example_service=prompt_injection_example_service,
        llm_configuration_introspection_service=llm_configuration_introspection_service