fix for template selection

src/text_generation/services/guidelines/base_security_guidelines_service.py

@@ -3,7 +3,6 @@ from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import PromptTemplate, StringPromptTemplate
 from langchain_core.prompt_values import PromptValue
 from langchain_core.runnables import RunnablePassthrough
-from langchain.prompts import FewShotPromptTemplate
 
 from src.text_generation.common.constants import Constants
 from src.text_generation.domain.abstract_guidelines_processed_completion import AbstractGuidelinesProcessedCompletion
@@ -48,7 +47,7 @@ class BaseSecurityGuidelinesService(AbstractSecurityGuidelinesService):
             | self.response_processing_service.process_text_generation_output
         )
 
-    def _get_template(self, user_prompt: str) -> StringPromptTemplate:
+    def _get_template(self, user_prompt: str, template_id: str) -> StringPromptTemplate:
         """
         Get the prompt template for security guidelines.
         
@@ -58,13 +57,14 @@ class BaseSecurityGuidelinesService(AbstractSecurityGuidelinesService):
         raise NotImplementedError("Subclasses must implement _get_template()")
 
 
-    def apply_guidelines(self, user_prompt: str) -> AbstractGuidelinesProcessedCompletion:
+    def apply_guidelines(self, user_prompt: str, template_id: str) -> AbstractGuidelinesProcessedCompletion:
 
         if not user_prompt:
             raise ValueError(f"Parameter 'user_prompt' cannot be empty or None")
         
         try:
-            prompt_template: StringPromptTemplate = self._get_template(user_prompt=user_prompt)
+            prompt_template: StringPromptTemplate = self._get_template(user_prompt=user_prompt, template_id=template_id)
+            print(f'using prompt template: {template_id}')
             prompt_value: PromptValue = prompt_template.format_prompt(input=user_prompt)
             prompt_dict = {
                 "messages": [

src/text_generation/services/guidelines/chain_of_thought_security_guidelines_service.py

@@ -86,6 +86,12 @@ class ChainOfThoughtSecurityGuidelinesService(BaseSecurityGuidelinesService):
         except Exception:
             return ""
     
+    def _get_template(self, user_prompt: str, template_id: str) -> StringPromptTemplate:
+        return self.config_builder.get_prompt_template(
+            template_id=template_id,
+            user_prompt=user_prompt
+        )
+
     def _get_cot_template_id_for_model(self, model_identifier: str) -> str:
         """
         Get the appropriate CoT template ID for the given model.

src/text_generation/services/guidelines/rag_context_security_guidelines_service.py

@@ -51,6 +51,12 @@ class RagContextSecurityGuidelinesService(BaseSecurityGuidelinesService):
             "meta-llama/llama-3.2-3b-instruct": self.constants.PromptTemplateIds.LLAMA_1_1B_CHAT__03_FEW_SHOT_EXAMPLES,
         }
 
+    def _get_template(self, user_prompt: str, template_id: str) -> StringPromptTemplate:
+        return self.config_builder.get_prompt_template(
+            template_id=template_id,
+            user_prompt=user_prompt
+        )
+
     def _get_model_identifier(self) -> str:
         """
         Get the model identifier from the foundation model.

src/text_generation/services/guidelines/rag_plus_cot_security_guidelines_service.py

@@ -54,6 +54,12 @@ class RagPlusCotSecurityGuidelinesService(BaseSecurityGuidelinesService):
             "meta-llama/llama-3.2-3b-instruct": self.constants.PromptTemplateIds.LLAMA_1_1B_CHAT__04_FEW_SHOT_RAG_PLUS_COT,
         }
 
+    def _get_template(self, user_prompt: str, template_id: str) -> StringPromptTemplate:
+        return self.config_builder.get_prompt_template(
+            template_id=template_id,
+            user_prompt=user_prompt
+        )
+
     def _get_model_identifier(self) -> str:
         """
         Get the model identifier from the foundation model.

src/text_generation/services/nlp/text_generation_completion_service.py

@@ -2,7 +2,7 @@ from enum import Enum
 from typing import Optional, Dict, Any
 import logging
 
-from langchain.prompts import StringPromptTemplate
+from langchain_core.prompts import StringPromptTemplate
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnablePassthrough, RunnableConfig
 from langchain_huggingface import HuggingFacePipeline
@@ -294,7 +294,7 @@ class TextGenerationCompletionService(AbstractTextGenerationCompletionService):
         guidelines_service = self._create_guidelines_service(mode, prompt_template)
         
         # Apply guidelines and process result
-        guidelines_result = guidelines_service.apply_guidelines(user_prompt)
+        guidelines_result = guidelines_service.apply_guidelines(user_prompt, template_id)
         return self._process_completion_result(guidelines_result)
 
     # Simplified handler methods