+ model support (Apple OpenELM 270M Instruct, Meta TinyLlama 1.1B Chat)

src/text_generation/adapters/foundation_models/apple_openelm_foundation_model.py

@@ -0,0 +1,33 @@
+from langchain_huggingface import HuggingFacePipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from src.text_generation.adapters.foundation_models.base.base_foundation_model import BaseFoundationModel
+from src.text_generation.adapters.foundation_models.config.apple_openelm_config import AppleOpenELMConfig
+from src.text_generation.common.model_id import ModelId
+
+
+class AppleOpenELMFoundationModel(BaseFoundationModel):
+    """Apple OpenELM 270M implementation"""
+
+    def __init__(self, config: AppleOpenELMConfig = AppleOpenELMConfig()):
+        self.config = config
+        super().__init__()
+        self.MODEL_ID = ModelId.APPLE_OPENELM_270M_INSTRUCT.value
+
+    def _load_model(self) -> None:
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.MODEL_ID,
+            local_files_only=self.config.local_files_only
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.MODEL_ID,
+            local_files_only=self.config.local_files_only
+        )
+
+    def create_pipeline(self) -> HuggingFacePipeline:
+        pipe = self._create_base_pipeline()
+        return HuggingFacePipeline(
+            pipeline=pipe,
+            pipeline_kwargs={
+                "return_full_text": False,
+                "stop_sequence": ["</s>", "[/INST]"]
+            })

src/text_generation/adapters/foundation_models/base/base_foundation_model.py

@@ -0,0 +1,39 @@
+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
+
+
+class BaseFoundationModel(AbstractFoundationModel):
+    """Base class for all foundation models"""
+
+    def __init__(self, config: BaseModelConfig):
+        self.config = config
+        self.model = None
+        self.tokenizer = None
+        self._load_model()
+
+    @abstractmethod
+    def _load_model(self) -> None:
+        """Load model implementation"""
+        pass
+
+    def _create_base_pipeline(self) -> Any:
+        """Create common pipeline configuration"""
+        return pipeline(
+            "text-generation",
+            do_sample=True,
+            max_new_tokens=self.config.max_new_tokens,
+            model=self.model,
+            repetition_penalty=self.config.repetition_penalty,
+            temperature=self.config.temperature,
+            tokenizer=self.tokenizer,
+            use_fast=self.config.use_fast,
+            pad_token_id=self.tokenizer.eos_token_id,
+            eos_token_id=self.tokenizer.eos_token_id
+        )

src/text_generation/adapters/foundation_models/base/base_model_config.py

@@ -0,0 +1,15 @@
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class BaseModelConfig:
+    """Base configuration for foundation models"""
+    max_new_tokens: int = 512
+    temperature: float = 0.3
+    repetition_penalty: float = 1.1
+    use_fast: bool = True
+    local_files_only: bool = False
+
+
+

src/text_generation/adapters/foundation_models/base/base_model_pipeline.py

@@ -0,0 +1,57 @@
+from transformers import pipeline
+from langchain.llms import HuggingFacePipeline
+from typing import Dict, Any, List
+from abc import ABC, abstractmethod
+
+class BaseModelPipeline(ABC):
+    def __init__(self, model, tokenizer):
+        self.model = model
+        self.tokenizer = tokenizer
+    
+    def get_common_config(self) -> Dict[str, Any]:
+        """Common configuration shared across all models"""
+        return {
+            "do_sample": True,
+            "temperature": 0.3,
+            "repetition_penalty": 1.1,
+            "use_fast": True,
+            "pad_token_id": self.tokenizer.eos_token_id,
+            "eos_token_id": self.tokenizer.eos_token_id,
+        }
+    
+    @abstractmethod
+    def get_model_specific_config(self) -> Dict[str, Any]:
+        """Model-specific configuration overrides"""
+        pass
+    
+    @abstractmethod
+    def get_stop_sequences(self) -> List[str]:
+        """Model-specific stop sequences"""
+        pass
+    
+    def _create_base_pipeline(self):
+        """Create the base pipeline with merged configurations"""
+        config = self.get_common_config()
+        config.update(self.get_model_specific_config())
+        
+        return pipeline(
+            "text-generation",
+            model=self.model,
+            tokenizer=self.tokenizer,
+            **config
+        )
+    
+    def create_pipeline(self) -> HuggingFacePipeline:
+        """Create the final HuggingFace pipeline"""
+        pipe = self._create_base_pipeline()
+        
+        return HuggingFacePipeline(
+            pipeline=pipe,
+            pipeline_kwargs={
+                "return_full_text": False,
+                "stop_sequence": self.get_stop_sequences()
+            }
+        )
+
+
+

src/text_generation/adapters/foundation_models/config/apple_openelm_config.py

@@ -0,0 +1,10 @@
+from typing import Optional
+from dataclasses import dataclass
+from src.text_generation.adapters.foundation_models.base.base_model_config import BaseModelConfig
+
+
+@dataclass
+class AppleOpenELMConfig(BaseModelConfig):
+    """OpenELM-specific configuration"""
+    use_cache: bool = True
+    pad_token_id: Optional[int] = None

src/text_generation/adapters/foundation_models/config/meta_tinyllama_config.py

@@ -0,0 +1,10 @@
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+from src.text_generation.adapters.foundation_models.base.base_model_config import BaseModelConfig
+
+
+@dataclass
+class MetaTinyLlamaConfig(BaseModelConfig):
+    """TinyLlama-specific configuration"""
+    use_flash_attention: bool = False
+    rope_scaling: Optional[Dict[str, Any]] = None

src/text_generation/adapters/foundation_models/config/microsoft_phi3mini4k_config.py

@@ -0,0 +1,9 @@
+from dataclasses import dataclass
+from src.text_generation.adapters.foundation_models.base.base_model_config import BaseModelConfig
+
+
+@dataclass
+class MicrosoftPhi3Mini4KConfig(BaseModelConfig):
+    """Phi3-specific configuration"""
+    trust_remote_code: bool = True
+    torch_dtype: str = "auto"

src/text_generation/adapters/foundation_models/factories/foundation_model_factory.py

@@ -0,0 +1,56 @@
+# Factory for creating foundation models
+from typing import Optional
+
+from langchain_huggingface import HuggingFacePipeline
+from src.text_generation.adapters.foundation_models.apple_openelm_foundation_model import AppleOpenELMFoundationModel
+from src.text_generation.adapters.foundation_models.base.base_foundation_model import BaseFoundationModel
+from src.text_generation.adapters.foundation_models.base.base_model_config import BaseModelConfig
+from src.text_generation.adapters.foundation_models.meta_tinyllama_foundation_model import MetaTinyLlamaFoundationModel
+from src.text_generation.adapters.foundation_models.microsoft_phi3_foundation_model import MicrosoftPhi3FoundationModel
+from src.text_generation.adapters.foundation_models.pipelines.apple_openelm_pipeline import AppleOpenELMPipeline
+from src.text_generation.adapters.foundation_models.pipelines.meta_tinyllama_pipeline import MetaTinyLlamaPipeline
+from src.text_generation.adapters.foundation_models.pipelines.microsoft_phi3mini_pipeline import MicrosoftPhi3MiniPipeline
+from src.text_generation.common.model_id import ModelId
+
+
+class FoundationModelFactory:
+    """Factory for creating foundation model instances"""
+
+    @staticmethod
+    def create_model(model_id: ModelId, config: Optional[BaseModelConfig] = None) -> BaseFoundationModel:
+        if config is None:
+            config = BaseModelConfig()
+
+        model_map = {
+            ModelId.APPLE_OPENELM_270M_INSTRUCT.value: AppleOpenELMFoundationModel,
+            ModelId.META_TINYLLAMA_1_1B_CHAT.value: MetaTinyLlamaFoundationModel,
+            ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT.value: MicrosoftPhi3FoundationModel
+        }
+
+        if model_id not in model_map:
+            raise ValueError(f"Unsupported model type: {model_id}")
+
+        return model_map[model_id](config)
+    
+    # Factory function to create the appropriate pipeline
+    def create_model_pipeline(model_id: ModelId, model, tokenizer) -> HuggingFacePipeline:
+        """Factory function to create the appropriate pipeline based on model name"""
+        
+        pipeline_classes = {
+            ModelId.APPLE_OPENELM_270M_INSTRUCT.value: AppleOpenELMPipeline,
+            ModelId.META_TINYLLAMA_1_1B_CHAT.value: MetaTinyLlamaPipeline,
+            ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT.value: MicrosoftPhi3MiniPipeline
+        }
+        
+        # Determine model type from name
+        model_type = None
+        for key in pipeline_classes.keys():
+            if key in model_id:
+                model_type = key
+                break
+        
+        if model_type is None:
+            raise ValueError(f"Unsupported model: {model_id}")
+        
+        pipeline_class = pipeline_classes[model_type]
+        return pipeline_class(model, tokenizer).create_pipeline()

src/text_generation/adapters/foundation_models/factory/foundation_model_factory.py

@@ -0,0 +1,28 @@
+# Factory for creating foundation models
+from typing import Optional
+from src.text_generation.adapters.foundation_models.apple_openelm_foundation_model import AppleOpenELMFoundationModel
+from src.text_generation.adapters.foundation_models.base.base_foundation_model import BaseFoundationModel
+from src.text_generation.adapters.foundation_models.base.base_model_config import BaseModelConfig
+from src.text_generation.adapters.foundation_models.meta_tinyllama_foundation_model import MetaTinyLlamaFoundationModel
+from src.text_generation.adapters.foundation_models.microsoft_phi3_foundation_model import MicrosoftPhi3FoundationModel
+from src.text_generation.common.model_id import ModelId
+
+
+class FoundationModelFactory:
+    """Factory for creating foundation model instances"""
+
+    @staticmethod
+    def create_model(model_id: ModelId, config: Optional[BaseModelConfig] = None) -> BaseFoundationModel:
+        if config is None:
+            config = BaseModelConfig()
+
+        model_map = {
+            ModelId.APPLE_OPENELM_270M_INSTRUCT.value: AppleOpenELMFoundationModel,
+            ModelId.META_TINYLLAMA_1_1B_CHAT.value: MetaTinyLlamaFoundationModel,
+            ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT.value: MicrosoftPhi3FoundationModel
+        }
+
+        if model_id not in model_map:
+            raise ValueError(f"Unsupported model type: {model_id}")
+
+        return model_map[model_id](config)

src/text_generation/adapters/foundation_models/meta_tinyllama_foundation_model.py

@@ -0,0 +1,36 @@
+from src.text_generation.adapters.foundation_models.config.meta_tinyllama_config import MetaTinyLlamaConfig
+from src.text_generation.adapters.foundation_models.base.base_foundation_model import BaseFoundationModel
+
+
+from langchain_huggingface import HuggingFacePipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from src.text_generation.common.model_id import ModelId
+
+
+class MetaTinyLlamaFoundationModel(BaseFoundationModel):
+    """Meta TinyLlama 1.1B implementation"""
+
+    def __init__(self, config: MetaTinyLlamaConfig = MetaTinyLlamaConfig()):
+        self.config = config
+        super().__init__()
+        self.MODEL_ID = ModelId.META_TINYLLAMA_1_1B_CHAT.value
+
+    def _load_model(self) -> None:
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.MODEL_ID,
+            local_files_only=self.config.local_files_only
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.MODEL_ID,
+            local_files_only=self.config.local_files_only
+        )
+
+    def create_pipeline(self) -> HuggingFacePipeline:
+        pipe = self._create_base_pipeline()
+        return HuggingFacePipeline(
+            pipeline=pipe,
+            pipeline_kwargs={
+                "return_full_text": False,
+                "stop_sequence": ["</s>", "[/INST]"]
+            })

src/text_generation/adapters/foundation_models/microsoft_phi3_foundation_model.py

@@ -0,0 +1,33 @@
+from langchain_huggingface import HuggingFacePipeline
+from optimum.onnxruntime import ORTModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from src.text_generation.adapters.foundation_models.config.microsoft_phi3mini4k_config import MicrosoftPhi3Mini4KConfig
+from src.text_generation.adapters.foundation_models.base.base_foundation_model import BaseFoundationModel
+from src.text_generation.common.model_id import ModelId
+
+
+class MicrosoftPhi3FoundationModel(BaseFoundationModel):
+    """Microsoft Phi3 Mini 4K implementation"""
+    def __init__(self, config: MicrosoftPhi3Mini4KConfig = MicrosoftPhi3Mini4KConfig()):
+        self.config = config
+        super().__init__()
+        self.MODEL_ID = ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT.value
+
+    def _load_model(self) -> None:
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.MODEL_ID,
+            local_files_only=self.config.local_files_only
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.MODEL_ID,
+            local_files_only=self.config.local_files_only
+        )
+
+    def create_pipeline(self) -> HuggingFacePipeline:
+        pipe = self._create_base_pipeline()
+        return HuggingFacePipeline(
+            pipeline=pipe,
+            pipeline_kwargs={
+                "return_full_text": False,
+                "stop_sequence": ["<|end|>", "<|user|>", "</s>"]
+            })

src/text_generation/adapters/foundation_models/pipelines/apple_openelm_pipeline.py

@@ -0,0 +1,16 @@
+from src.text_generation.adapters.foundation_models.base.base_model_pipeline import BaseModelPipeline
+
+
+from typing import Any, Dict, List
+
+
+class AppleOpenELMPipeline(BaseModelPipeline):
+    def get_model_specific_config(self) -> Dict[str, Any]:
+        return {
+            "max_new_tokens": 256,  # Smaller model, might need fewer tokens
+            "temperature": 0.4,     # Override common temperature for this model
+            "top_k": 40,           # Add top-k sampling
+        }
+
+    def get_stop_sequences(self) -> List[str]:
+        return ["</s>", "[/INST]"]

src/text_generation/adapters/foundation_models/pipelines/meta_tinyllama_pipeline.py

@@ -0,0 +1,16 @@
+from src.text_generation.adapters.foundation_models.base.base_model_pipeline import BaseModelPipeline
+
+
+from typing import Any, Dict, List
+
+
+class MetaTinyLlamaPipeline(BaseModelPipeline):
+    def get_model_specific_config(self) -> Dict[str, Any]:
+        return {
+            "max_new_tokens": 512,
+            # TinyLlama might need slightly different settings
+            "top_p": 0.9,  # Add nucleus sampling for better diversity
+        }
+
+    def get_stop_sequences(self) -> List[str]:
+        return ["</s>", "[/INST]"]

src/text_generation/adapters/foundation_models/pipelines/microsoft_phi3mini_pipeline.py

@@ -0,0 +1,15 @@
+from src.text_generation.adapters.foundation_models.base.base_model_pipeline import BaseModelPipeline
+
+
+from typing import Any, Dict, List
+
+
+class MicrosoftPhi3MiniPipeline(BaseModelPipeline):
+    def get_model_specific_config(self) -> Dict[str, Any]:
+        return {
+            "max_new_tokens": 512,
+            # Remove max_length to fix the warning - max_new_tokens takes precedence
+        }
+
+    def get_stop_sequences(self) -> List[str]:
+        return ["<|end|>", "<|user|>", "</s>"]

src/text_generation/adapters/foundation_models/text_generation_foundation_model.py

@@ -6,7 +6,6 @@ from transformers import AutoTokenizer, pipeline
 from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel
 
 
-
 class TextGenerationFoundationModel(AbstractFoundationModel):
 
     def __init__(self):

src/text_generation/common/constants.py

@@ -1,15 +1,43 @@
 class Constants:
+    # Phi-3 tokens
     PHI_3_ASSISTANT_START_TOKEN = "<|assistant|>"
     PHI_3_END_TOKEN =             "<|end|>"
     PHI_3_SYSTEM_START_TOKEN =    "<|system|>"
     PHI_3_USER_START_TOKEN =      "<|user|>"
-    INPUT_VARIABLE_TOKEN =        "input"
 
+    # OpenELM tokens
+    OPENELM_SYSTEM_START_TOKEN =  "[INST]"
+    OPENELM_USER_START_TOKEN =    "[INST]"
+    OPENELM_ASSISTANT_TOKEN =     "[/INST]"
+    OPENELM_END_TOKEN =           "</s>"
+
+    # TinyLlama tokens
+    TINYLLAMA_SYSTEM_START_TOKEN = "<|system|>"
+    TINYLLAMA_USER_START_TOKEN =   "<|user|>"
+    TINYLLAMA_ASSISTANT_TOKEN =    "<|assistant|>"
+    TINYLLAMA_END_TOKEN =          "</s>"
+
+    INPUT_VARIABLE_TOKEN =        "input"
     ALT_COMPLETION_TEXT =         "I can't answer that."
 
     class PromptTemplateIds:
+        # Phi-3 templates
         PHI_3_MINI_4K_INSTRUCT__01_BASIC =                      "phi-3-mini-4k-instruct.01-basic"
         PHI_3_MINI_4K_INSTRUCT__02_ZERO_SHOT_CHAIN_OF_THOUGHT = "phi-3-mini-4k-instruct.02-zero-shot-cot"
         PHI_3_MINI_4K_INSTRUCT__03_FEW_SHOT_EXAMPLES =          "phi-3-mini-4k-instruct.03-few-shot"
         PHI_3_MINI_4K_INSTRUCT__04_FEW_SHOT_RAG_PLUS_COT =      "phi-3-mini-4k-instruct.04-few-shot-rag-plus-cot"
-        PHI_3_MINI_4K_INSTRUCT__05_REFLEXION =                  "phi-3-mini-4k-instruct.05-reflexion"
+        PHI_3_MINI_4K_INSTRUCT__05_REFLEXION =                  "phi-3-mini-4k-instruct.05-reflexion"
+
+        # OpenELM templates
+        OPENELM_270M_INSTRUCT__01_BASIC =                      "openelm-270m-instruct.01-basic"
+        OPENELM_270M_INSTRUCT__02_ZERO_SHOT_CHAIN_OF_THOUGHT = "openelm-270m-instruct.02-zero-shot-cot"
+        OPENELM_270M_INSTRUCT__03_FEW_SHOT_EXAMPLES =          "openelm-270m-instruct.03-few-shot"
+        OPENELM_270M_INSTRUCT__04_FEW_SHOT_RAG_PLUS_COT =      "openelm-270m-instruct.04-few-shot-rag-plus-cot"
+        OPENELM_270M_INSTRUCT__05_REFLEXION =                  "openelm-270m-instruct.05-reflexion"
+
+        # TinyLlama templates
+        TINYLLAMA_1_1B_CHAT__01_BASIC =                      "tinyllama-1.1b-chat.01-basic"
+        TINYLLAMA_1_1B_CHAT__02_ZERO_SHOT_CHAIN_OF_THOUGHT = "tinyllama-1.1b-chat.02-zero-shot-cot"
+        TINYLLAMA_1_1B_CHAT__03_FEW_SHOT_EXAMPLES =          "tinyllama-1.1b-chat.03-few-shot"
+        TINYLLAMA_1_1B_CHAT__04_FEW_SHOT_RAG_PLUS_COT =      "tinyllama-1.1b-chat.04-few-shot-rag-plus-cot"
+        TINYLLAMA_1_1B_CHAT__05_REFLEXION =                  "tinyllama-1.1b-chat.05-reflexion"

src/text_generation/common/model_id.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class ModelId(Enum):
+    APPLE_OPENELM_270M_INSTRUCT = "apple/openelm-270m-instruct"
+    META_TINYLLAMA_1_1B_CHAT = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+    MICROSOFT_PHI_3_MINI4K_INSTRUCT = "microsoft/Phi-3-mini-4k-instruct-onnx"

src/text_generation/dependency_injection_container.py

@@ -3,7 +3,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.adapters.foundation_models.text_generation_foundation_model import TextGenerationFoundationModel
 from src.text_generation.entrypoints.http_api_controller import HttpApiController
 from src.text_generation.entrypoints.server import RestApiServer
 from src.text_generation.services.guidelines.abstract_security_guidelines_service import AbstractSecurityGuidelinesService

src/text_generation/services/nlp/text_generation_completion_service.py

@@ -1,10 +1,17 @@
+from enum import Enum
+from typing import Optional, Dict, Any
+import logging
+
 from langchain.prompts import StringPromptTemplate
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnablePassthrough, RunnableConfig
 from langchain_huggingface import HuggingFacePipeline
 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.model_id import ModelId
 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
@@ -21,11 +28,12 @@ from src.text_generation.services.utilities.abstract_llm_configuration_introspec
 from src.text_generation.services.utilities.abstract_response_processing_service import AbstractResponseProcessingService
 
 
+logger = logging.getLogger(__name__)
+
 class TextGenerationCompletionService(
         AbstractTextGenerationCompletionService):
     def __init__(
             self, 
-            foundation_model: AbstractFoundationModel,
             response_processing_service: AbstractResponseProcessingService,
             prompt_template_service: AbstractPromptTemplateService,
             chain_of_thought_guidelines: AbstractSecurityGuidelinesService,
@@ -34,38 +42,47 @@ class TextGenerationCompletionService(
             reflexion_guardrails: AbstractGeneratedTextGuardrailService,
             semantic_similarity_service: AbstractSemanticSimilarityService,
             prompt_injection_example_service: AbstractPromptInjectionExampleService,
-            llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService):
+            llm_configuration_introspection_service: AbstractLLMConfigurationIntrospectionService,
+            default_model_type: ModelId = ModelId.MICROSOFT_PHI_3_MINI4K_INSTRUCT.value):
+        
         super().__init__()
         self.constants = Constants()
-        self.foundation_model_pipeline = foundation_model.create_pipeline()
+        
+        # Model management
+        self._current_model = None
+        self._current_model_id = None
+        self.default_model_id = default_model_type
+        self.factory = FoundationModelFactory()
+        
+        # Services
         self.response_processing_service = response_processing_service
         self.prompt_template_service = prompt_template_service
+        self.semantic_similarity_service = semantic_similarity_service
+        self.llm_configuration_introspection_service = llm_configuration_introspection_service
 
-        # set up semantic similarity service and supporting texts
+        # Set up semantic similarity service
         self.example_prompt_injection_completions = prompt_injection_example_service.get_all_completions()
         self.example_prompt_injection_prompts = prompt_injection_example_service.get_all_prompts()
-        self.semantic_similarity_service = semantic_similarity_service
         self.semantic_similarity_service.use_comparison_texts(
             self.example_prompt_injection_completions
         )
 
-        # 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
+        # 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
         
-        # guardrails services
-        self.reflexion_guardrails: AbstractGeneratedTextGuardrailService = reflexion_guardrails
+        # Guardrails service
+        self.reflexion_guardrails = reflexion_guardrails
 
-        # constants
+        # Constants and settings
         self.COSINE_SIMILARITY_RISK_THRESHOLD = 0.8
-
-        # default guidelines settings
         self._use_guidelines = False
         self._use_zero_shot_chain_of_thought = False
         self._use_rag_context = False
+        self._use_reflexion_guardrails = False
 
-        # dictionary dispatch for handling guidelines combinations
+        # Strategy map for guidelines
         self.guidelines_strategy_map = {
             (True, True):   self._handle_cot_and_rag,
             (True, False):  self._handle_cot_only,
@@ -73,11 +90,43 @@ class TextGenerationCompletionService(
             (False, False): self._handle_without_guidelines,
         }
 
-        # default guardrails settings
-        self._use_reflexion_guardrails = False
+        # Load default model
+        self.load_model(default_model_type)
 
-        # introspection for logging
-        self.llm_configuration_introspection_service = llm_configuration_introspection_service
+
+    def load_model(
+        self, 
+        model_id: ModelId, 
+        config: Optional[BaseModelConfig] = None,
+        force_reload: bool = False
+    ) -> None:
+        """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()):
+            logger.info(f"Model {model_id.value} already loaded")
+            return
+        
+        if self._current_model is not None:
+            self._current_model.unload()
+        
+        self._current_model = self.factory.create_model(model_id, config)
+        self._current_model.load()
+        self._current_model_id = 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):
@@ -236,11 +285,19 @@ class TextGenerationCompletionService(
         return self._use_reflexion_guardrails
 
 
-    def invoke(self, user_prompt: str) -> TextGenerationCompletionResult:
+    def invoke(self, user_prompt: str, model_id: Optional[ModelId] = None) -> TextGenerationCompletionResult:
+        """Generate text using specified or current model"""
         if not user_prompt:
             raise ValueError(f"Parameter 'user_prompt' cannot be empty or None")
-        print(f'Using guidelines: {self.get_current_config()}')
-        completion_result: TextGenerationCompletionResult = self._process_prompt_with_guidelines_if_applicable(user_prompt)
+            
+        target_model_id = model_id or self._current_model_id or self.default_model_id
+        if (self._current_model_id != target_model_id or 
+            self._current_model is None or 
+            not self._current_model.is_loaded()):
+            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)
         
         if not self._use_reflexion_guardrails:
             return completion_result

tests/conftest.py

@@ -18,7 +18,7 @@ from src.text_generation import config
 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.adapters.foundation_models.text_generation_foundation_model import TextGenerationFoundationModel
 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