From a1c73138334eb4ed79f0b84de6144e7b2766d957 Mon Sep 17 00:00:00 2001 From: Adam Wilson Date: Mon, 18 Aug 2025 15:16:27 -0600 Subject: [PATCH] fix for template selection --- .../guidelines/base_security_guidelines_service.py | 8 ++++---- .../chain_of_thought_security_guidelines_service.py | 6 ++++++ .../guidelines/rag_context_security_guidelines_service.py | 6 ++++++ .../rag_plus_cot_security_guidelines_service.py | 6 ++++++ .../services/nlp/text_generation_completion_service.py | 4 ++-- 5 files changed, 24 insertions(+), 6 deletions(-) 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 a95ff9b45..b3bee689a 100644 --- a/src/text_generation/services/guidelines/base_security_guidelines_service.py +++ b/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": [ 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 08e4de4e0..150ad2a90 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 @@ -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. 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 cb727258f..a55a3bb2a 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 @@ -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. 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 5fef1cc61..1cd972d66 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 @@ -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. 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 79ceac6b8..d250a626f 100644 --- a/src/text_generation/services/nlp/text_generation_completion_service.py +++ b/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