implement singular guidelines calls in main service

src/text_generation/services/guardrails/abstract_generated_text_guardrail_service.py

@@ -1,8 +1,9 @@
 import abc
 from src.text_generation.domain.abstract_guardrails_processed_completion import AbstractGuardrailsProcessedCompletion
+from src.text_generation.domain.abstract_text_generation_completion_result import AbstractTextGenerationCompletionResult
 
 
 class AbstractGeneratedTextGuardrailService(abc.ABC):
     @abc.abstractmethod
-    def process_generated_text(self, model_generated_text: str) -> AbstractGuardrailsProcessedCompletion:
+    def process_generated_text(self, model_generated_text: AbstractTextGenerationCompletionResult) -> AbstractGuardrailsProcessedCompletion:
         raise NotImplementedError

src/text_generation/services/guardrails/generated_text_guardrail_service.py

@@ -1,10 +1,12 @@
 from src.text_generation.domain.abstract_guardrails_processed_completion import AbstractGuardrailsProcessedCompletion
-from src.text_generation.domain.guardrails_processed_completion import GuardrailsProcessedCompletion
+from src.text_generation.domain.abstract_text_generation_completion_result import AbstractTextGenerationCompletionResult
+from src.text_generation.domain.guardrails_result import GuardrailsResult
 from src.text_generation.services.guardrails.abstract_generated_text_guardrail_service import AbstractGeneratedTextGuardrailService
 from src.text_generation.services.nlp.abstract_semantic_similarity_service import AbstractSemanticSimilarityService
 
 
-class GeneratedTextGuardrailService(AbstractGeneratedTextGuardrailService):
+class GeneratedTextGuardrailService(
+    AbstractGeneratedTextGuardrailService):
     def __init__(
             self,
             semantic_similarity_service: AbstractSemanticSimilarityService):
@@ -15,10 +17,10 @@ class GeneratedTextGuardrailService(AbstractGeneratedTextGuardrailService):
     def use_comparison_texts(self, comparison_texts: list[str]):
         self.semantic_similarity_service.use_comparison_texts(comparison_texts)
 
-    def process_generated_text(self, model_generated_text: str) -> AbstractGuardrailsProcessedCompletion:
-        score: float = self.semantic_similarity_service.analyze(text=model_generated_text)
-        response = GuardrailsProcessedCompletion(
-            score=score,
+    def process_generated_text(self, completion: AbstractTextGenerationCompletionResult) -> AbstractGuardrailsProcessedCompletion:
+        score: float = self.semantic_similarity_service.analyze(text=completion)
+        response = GuardrailsResult(
+            cosine_similarity_score=score,
             cosine_similarity_risk_threshold=self.cosine_similarity_risk_threshold,
-            original_completion=model_generated_text)
+            original_completion=completion)
         return response

src/text_generation/services/guidelines/abstract_security_guidelines_service.py

@@ -1,9 +1,11 @@
 import abc
 
+from src.text_generation.domain.abstract_guidelines_processed_completion import AbstractGuidelinesProcessedCompletion
+
 
 class AbstractSecurityGuidelinesService(abc.ABC):
     @abc.abstractmethod
-    def apply_guidelines(self, user_prompt: str) -> str:
+    def apply_guidelines(self, user_prompt: str) -> AbstractGuidelinesProcessedCompletion:
         pass
 
 

src/text_generation/services/guidelines/base_security_guidelines_service.py

@@ -3,6 +3,8 @@ from langchain_core.prompts import PromptTemplate
 from langchain_core.runnables import RunnablePassthrough
 
 from src.text_generation.common.constants import Constants
+from src.text_generation.domain.abstract_guidelines_processed_completion import AbstractGuidelinesProcessedCompletion
+from src.text_generation.domain.guidelines_result import GuidelinesResult
 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
@@ -32,10 +34,9 @@ class BaseSecurityGuidelinesService(AbstractSecurityGuidelinesService):
         )
 
     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:
+    def apply_guidelines(self, user_prompt: str) -> AbstractGuidelinesProcessedCompletion:
         if not user_prompt:
             raise ValueError(f"Parameter 'user_prompt' cannot be empty or None")
         

src/text_generation/services/nlp/abstract_semantic_similarity_service.py

@@ -1,9 +1,11 @@
 import abc
 
+from src.text_generation.domain.abstract_semantic_similarity_result import AbstractSemanticSimilarityResult
+
 
 class AbstractSemanticSimilarityService(abc.ABC):
     @abc.abstractmethod
-    def analyze(self, text: str) -> float:
+    def analyze(self, text: str) -> AbstractSemanticSimilarityResult:
         raise NotImplementedError
     
     @abc.abstractmethod

src/text_generation/services/nlp/text_generation_completion_service.py

@@ -3,14 +3,20 @@ from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnablePassthrough
 
 from src.text_generation.common.constants import Constants
+from src.text_generation.domain.alternate_completion_result import AlternateCompletionResult
+from src.text_generation.domain.guidelines_result import GuidelinesResult
+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.chain_of_thought_security_guidelines_service import ChainOfThoughtSecurityGuidelinesService
 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.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
 from src.text_generation.ports.abstract_foundation_model import AbstractFoundationModel
+from src.text_generation.services.prompt_injection.abstract_prompt_injection_example_service import AbstractPromptInjectionExampleService
 from src.text_generation.services.utilities.abstract_response_processing_service import AbstractResponseProcessingService
 
 
@@ -23,13 +29,20 @@ class TextGenerationCompletionService(
             prompt_template_service: AbstractPromptTemplateService,
             chain_of_thought_guidelines: AbstractSecurityGuidelinesService,
             rag_context_guidelines: AbstractSecurityGuidelinesService,
-            reflexion_guardrails: AbstractGeneratedTextGuardrailService):
+            reflexion_guardrails: AbstractGeneratedTextGuardrailService,
+            semantic_similarity_service: AbstractSemanticSimilarityService,
+            prompt_injection_example_service: AbstractPromptInjectionExampleService):
         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
 
+        # set up semantic similarity service and supporting texts
+        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
+
         # guidelines services
         self.chain_of_thought_guidelines: AbstractSecurityGuidelinesService = chain_of_thought_guidelines
         self.rag_context_guidelines: AbstractSecurityGuidelinesService = rag_context_guidelines
@@ -37,6 +50,9 @@ class TextGenerationCompletionService(
         # guardrails services
         self.reflexion_guardrails: AbstractSecurityGuidelinesService = reflexion_guardrails
 
+        # constants
+        self.COSINE_SIMILARITY_RISK_THRESHOLD = 0.5
+
         # default guidelines settings
         self._use_guidelines = False
         self._use_zero_shot_chain_of_thought = False
@@ -63,49 +79,66 @@ class TextGenerationCompletionService(
             self._handle_without_guidelines
         )
         return guidelines_handler(user_prompt)
+    
+    def _process_guidelines_result(self, guidelines_result: GuidelinesResult) -> TextGenerationCompletionResult:
+        """
+        Process guidelines result and create completion result with semantic similarity check.
+        
+        Args:
+            guidelines_result: Result from applying security guidelines
+            
+        Returns:
+            TextGenerationCompletionResult with appropriate completion text
+        """
+        self.semantic_similarity_service.use_comparison_texts(
+            self.example_prompt_injection_completions
+        )
+        similarity_result: SemanticSimilarityResult = self.semantic_similarity_service.analyze(
+            text=guidelines_result.original_completion
+        )
+        
+        processed_guidelines_result = GuidelinesResult(
+            original_completion=guidelines_result.original_completion,
+            cosine_similarity_score=similarity_result.mean,
+            cosine_similarity_risk_threshold=self.COSINE_SIMILARITY_RISK_THRESHOLD
+        )
+        
+        completion_result = TextGenerationCompletionResult(
+            original_completion=guidelines_result.original_completion,
+            guidelines_result=processed_guidelines_result
+        )
+        
+        if not processed_guidelines_result.is_original_completion_malicious:
+            return completion_result
+    
+        completion_result.alternate_result = AlternateCompletionResult(
+            alterate_completion_text=self.constants.ALT_COMPLETION_TEXT
+        )
+        return completion_result
 
-    # Handler methods for each combination
+    # Handler methods for each guidelines combination
     def _handle_cot_and_rag(self, user_prompt: str):
         """Handle: CoT=True, RAG=True"""
-        context = self._retrieve_rag_context(query)
-        thought_process = self._apply_chain_of_thought(query, context)
-        return f"CoT+RAG: {thought_process}"
-    
+        return f"CoT+RAG:"
+
     def _handle_cot_only(self, user_prompt: str):
         """Handle: CoT=True, RAG=False"""
-        print("🧠 Using Chain of Thought only")
-        thought_process = self._apply_chain_of_thought(query)
-        return f"CoT: {thought_process}"
+        guidelines_result = self.chain_of_thought_guidelines.apply_guidelines(user_prompt)
+        return self._process_guidelines_result(guidelines_result)
     
     def _handle_rag_only(self, user_prompt: str):
         """Handle: CoT=False, RAG=True"""
-        self.rag_context_guidelines.apply_guidelines(user_prompt)
-        return f"RAG: {response}"
+        guidelines_result = self.rag_context_guidelines.apply_guidelines(user_prompt)
+        return self._process_guidelines_result(guidelines_result)
     
     def _handle_without_guidelines(self, user_prompt: str):
         """Handle: CoT=False, RAG=False"""
-        response = self._basic_generate(query)
-        return f"Basic: {response}"
-    
-    # Helper methods (example implementations)
-    def _retrieve_rag_context(self, user_prompt: str):
-        """Retrieve relevant context from knowledge base"""
-        return f"Context for '{query}'"
-    
-    def _apply_chain_of_thought(self, user_prompt: str):
-        """Apply chain of thought reasoning"""
-        if context:
-            return f"Step-by-step reasoning for '{query}' with {context}"
-        return f"Step-by-step reasoning for '{query}'"
-    
-    def _generate_with_context(self, user_prompt: str):
-        """Generate response using context"""
-        return f"Response to '{query}' using {context}"
-    
-    def _basic_generate(self, user_prompt: str):
-        """Basic generation without special features"""
-        return f"Basic response to '{query}'"
-    
+        try:
+            chain = self._create_chain_without_guidelines()
+            return chain.invoke(user_prompt)
+        except Exception as e:
+            raise e
+        
     # Configuration methods
     def set_config(self, use_cot=False, use_rag=False):
         """Set guidelines configuration"""
@@ -135,7 +168,7 @@ class TextGenerationCompletionService(
         self._use_reflexion = True
         return self
 
-    def create_chain(self):
+    def _create_chain_without_guidelines(self):
         prompt_template = self.prompt_template_service.get(
             id=self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_BASIC
         )
@@ -150,12 +183,4 @@ class TextGenerationCompletionService(
     def invoke(self, user_prompt: str) -> str:
         if not user_prompt:
             raise ValueError(f"Parameter 'user_prompt' cannot be empty or None")
-        
-        if self._use_guidelines == False:
-            try:
-                chain = self.create_chain()
-                return chain.invoke(user_prompt)
-            except Exception as e:
-                raise e
-
         self._process_prompt_with_guidelines_if_applicable(user_prompt)

src/text_generation/services/prompt_injection/abstract_prompt_injection_example_service.py

@@ -0,0 +1,15 @@
+from abc import ABC, abstractmethod
+from typing import List
+
+class AbstractPromptInjectionExampleService(ABC):
+    """Abstract base class for prompt injection example services."""
+    
+    @abstractmethod
+    def get_all_prompts(self) -> List[str]:
+        """Get all prompt injection prompts."""
+        pass
+        
+    @abstractmethod
+    def get_all_completions(self) -> List[str]:
+        """Get all prompt injection completions."""
+        pass

src/text_generation/services/prompt_injection/prompt_injection_example_service.py

@@ -0,0 +1,19 @@
+from typing import List
+from src.text_generation.ports.abstract_prompt_injection_example_repository import AbstractPromptInjectionExampleRepository
+from src.text_generation.services.prompt_injection.abstract_prompt_injection_example_service import AbstractPromptInjectionExampleService
+
+class PromptInjectionExampleService(AbstractPromptInjectionExampleService):
+    """Service for handling prompt injection examples."""
+
+    def __init__(self, repository: AbstractPromptInjectionExampleRepository):
+        self.repository = repository
+
+    def get_all_prompts(self) -> List[str]:
+        """Get all prompt injection prompts."""
+        examples = self.repository.get_all()
+        return [example["prompt_injection_prompt"] for example in examples]
+
+    def get_all_completions(self) -> List[str]:
+        """Get all prompt injection completions."""
+        examples = self.repository.get_all()
+        return [example["prompt_injection_completion"] for example in examples]