diff --git a/src/text_generation/dependency_injection_container.py b/src/text_generation/dependency_injection_container.py index 57114ca98..93f3922e0 100644 --- a/src/text_generation/dependency_injection_container.py +++ b/src/text_generation/dependency_injection_container.py @@ -5,16 +5,16 @@ from src.text_generation.adapters.prompt_template_repository import PromptTempla from src.text_generation.adapters.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 from src.text_generation.services.guidelines.chain_of_thought_security_guidelines_service import ChainOfThoughtSecurityGuidelinesService +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, RetrievalAugmentedGenerationContextSecurityGuidelinesService +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.rag_context_security_guidelines_service import RetrievalAugmentedGenerationContextSecurityGuidelinesService from src.text_generation.services.logging.json_web_traffic_logging_service import JSONWebTrafficLoggingService from src.text_generation.services.nlp.prompt_template_service import PromptTemplateService from src.text_generation.services.nlp.semantic_similarity_service import SemanticSimilarityService from src.text_generation.services.nlp.text_generation_completion_service import TextGenerationCompletionService -from src.text_generation.services.nlp.retrieval_augmented_generation_completion_service import RetrievalAugmentedGenerationCompletionService -from src.text_generation.services.guardrails.generated_text_guardrail_service import GeneratedTextGuardrailService -from src.text_generation.services.guidelines.rag_context_security_guidelines_configuration_builder import RetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder from src.text_generation.services.utilities.response_processing_service import ResponseProcessingService @@ -71,14 +71,20 @@ class DependencyInjectionContainer(containers.DeclarativeContainer): semantic_similarity_service=semantic_similarity_service ) + # Register security guideline services chain_of_thought_guidelines = providers.Factory( - ChainOfThoughtSecurityGuidelinesService - ) - + ChainOfThoughtSecurityGuidelinesService, + foundation_model=foundation_model, + response_processing_service=response_processing_service, + prompt_template_service=prompt_template_service + ).provides(AbstractSecurityGuidelinesService) + rag_context_guidelines = providers.Factory( - RetrievalAugmentedGenerationContextSecurityGuidelinesService, - embedding_model=embedding_model - ) + RagContextSecurityGuidelinesService, + foundation_model=foundation_model, + response_processing_service=response_processing_service, + prompt_template_service=prompt_template_service + ).provides(AbstractSecurityGuidelinesService) reflexion_guardrails = providers.Factory( ReflexionSecurityGuardrailsService diff --git a/src/text_generation/services/guidelines/base_security_guidelines_service.py b/src/text_generation/services/guidelines/base_security_guidelines_service.py new file mode 100644 index 000000000..f6c73d4b2 --- /dev/null +++ b/src/text_generation/services/guidelines/base_security_guidelines_service.py @@ -0,0 +1,48 @@ +from langchain_core.output_parsers import StrOutputParser +from langchain_core.prompts import PromptTemplate +from langchain_core.runnables import RunnablePassthrough + +from src.text_generation.common.constants import Constants +from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel +from src.text_generation.services.guidelines.abstract_security_guidelines_service import AbstractSecurityGuidelinesService +from src.text_generation.services.nlp.abstract_prompt_template_service import AbstractPromptTemplateService +from src.text_generation.services.utilities.abstract_response_processing_service import AbstractResponseProcessingService + +class BaseSecurityGuidelinesService(AbstractSecurityGuidelinesService): + """Base service for security guidelines implementations.""" + + def __init__( + self, + foundation_model: AbstractFoundationModel, + response_processing_service: AbstractResponseProcessingService, + prompt_template_service: AbstractPromptTemplateService): + super().__init__() + self.constants = Constants() + self.foundation_model_pipeline = foundation_model.create_pipeline() + self.response_processing_service = response_processing_service + self.prompt_template_service = prompt_template_service + + def _create_chain(self, prompt_template: PromptTemplate): + return ( + { f"{self.constants.INPUT_VARIABLE_TOKEN}": RunnablePassthrough() } + | prompt_template + | self.foundation_model_pipeline + | StrOutputParser() + | self.response_processing_service.process_text_generation_output + ) + + def _get_template_id(self) -> str: + """Get the template ID for the specific implementation.""" + raise NotImplementedError("Subclasses must implement _get_template_id()") + + def apply_guidelines(self, user_prompt: str) -> str: + if not user_prompt: + raise ValueError(f"Parameter 'user_prompt' cannot be empty or None") + + try: + template_id = self._get_template_id() + prompt_template: PromptTemplate = self.prompt_template_service.get(id=template_id) + chain = self._create_chain(prompt_template) + return chain.invoke(user_prompt) + except Exception as e: + raise e \ No newline at end of file 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 7be2b86f4..2cb40b132 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,46 +1,7 @@ -from langchain_core.output_parsers import StrOutputParser -from langchain_core.prompts import PromptTemplate -from langchain_core.runnables import RunnablePassthrough +from src.text_generation.services.guidelines.base_security_guidelines_service import BaseSecurityGuidelinesService -from src.text_generation.common.constants import Constants -from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel -from src.text_generation.services.guidelines.abstract_security_guidelines_service import AbstractSecurityGuidelinesService -from src.text_generation.services.nlp.abstract_prompt_template_service import AbstractPromptTemplateService -from src.text_generation.services.nlp.prompt_template_service import PromptTemplateService -from src.text_generation.services.utilities.abstract_response_processing_service import AbstractResponseProcessingService - - -class ChainOfThoughtSecurityGuidelinesService( - AbstractSecurityGuidelinesService): +class ChainOfThoughtSecurityGuidelinesService(BaseSecurityGuidelinesService): """Service for zero-shot chain-of-thought security guidelines.""" - def __init__( - self, - foundation_model: AbstractFoundationModel, - response_processing_service: AbstractResponseProcessingService, - prompt_template_service: AbstractPromptTemplateService): - super().__init__() - self.constants = Constants() - self.foundation_model_pipeline = foundation_model.create_pipeline() - self.response_processing_service = response_processing_service - self.prompt_template_service: PromptTemplateService = prompt_template_service - def _create_chain(self, prompt_template: PromptTemplate): - return ( - { f"{self.constants.INPUT_VARIABLE_TOKEN}": RunnablePassthrough() } - | prompt_template - | self.foundation_model_pipeline - | StrOutputParser() - | self.response_processing_service.process_text_generation_output - ) - - def apply_guidelines(self, user_prompt: str) -> str: - if not user_prompt: - raise ValueError(f"Parameter 'user_prompt' cannot be empty or None") - - try: - template_id = self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_ZERO_SHOT_CHAIN_OF_THOUGHT - prompt_template: PromptTemplate = self.prompt_template_service.get(id=template_id) - chain = self._create_chain(prompt_template) - return chain.invoke(user_prompt) - except Exception as e: - raise e \ No newline at end of file + def _get_template_id(self) -> str: + return self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_ZERO_SHOT_CHAIN_OF_THOUGHT \ 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 dfaa2729d..97d3abc35 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,46 +1,7 @@ -from langchain_core.output_parsers import StrOutputParser -from langchain_core.prompts import PromptTemplate -from langchain_core.runnables import RunnablePassthrough +from src.text_generation.services.guidelines.base_security_guidelines_service import BaseSecurityGuidelinesService -from src.text_generation.common.constants import Constants -from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel -from src.text_generation.services.guidelines.abstract_security_guidelines_service import AbstractSecurityGuidelinesService -from src.text_generation.services.nlp.abstract_prompt_template_service import AbstractPromptTemplateService -from src.text_generation.services.nlp.prompt_template_service import PromptTemplateService -from src.text_generation.services.utilities.abstract_response_processing_service import AbstractResponseProcessingService +class RagContextSecurityGuidelinesService(BaseSecurityGuidelinesService): + """Service for RAG context security guidelines.""" - -class RetrievalAugmentedGenerationContextSecurityGuidelinesService( - AbstractSecurityGuidelinesService): - """Implementation of RAG context security guidelines service.""" - def __init__( - self, - foundation_model: AbstractFoundationModel, - response_processing_service: AbstractResponseProcessingService, - prompt_template_service: AbstractPromptTemplateService): - super().__init__() - self.constants = Constants() - self.foundation_model_pipeline = foundation_model.create_pipeline() - self.response_processing_service = response_processing_service - self.prompt_template_service: PromptTemplateService = prompt_template_service - - def _create_chain(self, prompt_template: PromptTemplate): - return ( - { f"{self.constants.INPUT_VARIABLE_TOKEN}": RunnablePassthrough() } - | prompt_template - | self.foundation_model_pipeline - | StrOutputParser() - | self.response_processing_service.process_text_generation_output - ) - - def apply_guidelines(self, user_prompt: str) -> str: - if not user_prompt: - raise ValueError(f"Parameter 'user_prompt' cannot be empty or None") - - try: - template_id = self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_FEW_SHOT_EXAMPLES - prompt_template: PromptTemplate = self.prompt_template_service.get(id=template_id) - chain = self._create_chain(prompt_template) - return chain.invoke(user_prompt) - except Exception as e: - raise e \ No newline at end of file + def _get_template_id(self) -> str: + return self.constants.PromptTemplateIds.RAG_CONTEXT_SECURITY_GUIDELINES \ No newline at end of file