dependency fixes, test setup

src/text_generation/dependency_injection_container.py

@@ -9,7 +9,7 @@ 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.guidelines.rag_context_security_guidelines_service import RagContextSecurityGuidelinesService
 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_plus_cot_security_guidelines_service import RagPlusCotSecurityGuidelinesService
@@ -37,23 +37,6 @@ class DependencyInjectionContainer(containers.DeclarativeContainer):
     embedding_model = providers.Singleton(
         EmbeddingModel
     )
-    
-    rag_guidelines_service = providers.Factory(
-        RetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder,
-        embedding_model=embedding_model
-    )
-
-    response_processing_service = providers.Factory(
-        ResponseProcessingService
-    )
-
-    rag_response_service = providers.Factory(
-        RetrievalAugmentedGenerationCompletionService,
-        foundation_model=foundation_model,
-        embedding_model=embedding_model,
-        rag_guidelines_service=rag_guidelines_service,
-        response_processing_service=response_processing_service
-    )
 
     prompt_template_repository = providers.Factory(
         PromptTemplateRepository
@@ -64,6 +47,15 @@ class DependencyInjectionContainer(containers.DeclarativeContainer):
         prompt_template_repository=prompt_template_repository
     )
 
+    prompt_injection_example_repository = providers.Factory(
+        PromptInjectionExampleRepository
+    )
+
+
+    response_processing_service = providers.Factory(
+        ResponseProcessingService
+    )
+
     semantic_similarity_service = providers.Factory(
         SemanticSimilarityService,
         embedding_model=embedding_model
@@ -75,7 +67,10 @@ class DependencyInjectionContainer(containers.DeclarativeContainer):
     )
 
     rag_config_builder = providers.Factory(
-        RetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder
+        RetrievalAugmentedGenerationSecurityGuidelinesConfigurationBuilder,
+        embedding_model=embedding_model,
+        prompt_template_service=prompt_template_service,
+        prompt_injection_example_repository=prompt_injection_example_repository
     )
 
     # Register security guideline services
@@ -83,8 +78,9 @@ class DependencyInjectionContainer(containers.DeclarativeContainer):
         ChainOfThoughtSecurityGuidelinesService,
         foundation_model=foundation_model,
         response_processing_service=response_processing_service,
-        prompt_template_service=prompt_template_service
-    ).provides(AbstractSecurityGuidelinesService)
+        prompt_template_service=prompt_template_service,
+        config_builder=None
+    )
 
     rag_context_guidelines = providers.Factory(
         RagContextSecurityGuidelinesService,
@@ -92,8 +88,8 @@ class DependencyInjectionContainer(containers.DeclarativeContainer):
         response_processing_service=response_processing_service,
         prompt_template_service=prompt_template_service,
         config_builder=rag_config_builder
-    ).provides(AbstractSecurityGuidelinesService)
-    
+    )
+
     reflexion_guardrails = providers.Factory(
         ReflexionSecurityGuardrailsService
     )
@@ -111,7 +107,8 @@ class DependencyInjectionContainer(containers.DeclarativeContainer):
         RagPlusCotSecurityGuidelinesService,
         foundation_model=foundation_model,
         response_processing_service=response_processing_service,
-        prompt_template_service=prompt_template_service
+        prompt_template_service=prompt_template_service,
+        config_builder=rag_config_builder
     )
 
     text_generation_completion_service = providers.Factory(
@@ -131,7 +128,6 @@ class DependencyInjectionContainer(containers.DeclarativeContainer):
         HttpApiController,
         logging_service=logging_service,
         text_generation_response_service=text_generation_completion_service,
-        rag_response_service=rag_response_service,
         generated_text_guardrail_service=generated_text_guardrail_service
     )
 

src/text_generation/entrypoints/http_api_controller.py

@@ -1,6 +1,8 @@
 import json
 import traceback
+from typing import Callable
 
+from src.text_generation.domain.text_generation_completion_result import TextGenerationCompletionResult
 from src.text_generation.services.logging.abstract_web_traffic_logging_service import AbstractWebTrafficLoggingService
 from src.text_generation.services.nlp.abstract_text_generation_completion_service import AbstractTextGenerationCompletionService
 from src.text_generation.services.guardrails.abstract_generated_text_guardrail_service import AbstractGeneratedTextGuardrailService
@@ -12,12 +14,10 @@ class HttpApiController:
             self, 
             logging_service: AbstractWebTrafficLoggingService,
             text_generation_response_service: AbstractTextGenerationCompletionService,
-            rag_response_service: AbstractTextGenerationCompletionService,
             generated_text_guardrail_service: AbstractGeneratedTextGuardrailService
     ):
         self.logging_service = logging_service        
         self.text_generation_response_service = text_generation_response_service
-        self.rag_response_service = rag_response_service
         self.generated_text_guardrail_service = generated_text_guardrail_service
         self.routes = {}
         self.register_routes()
@@ -30,12 +30,14 @@ class HttpApiController:
             print(f"Args: {args}")
             print(f"Kwargs: {kwargs}")
             raise e
-
-
+        
     def register_routes(self):
         self.routes[('GET', '/')] = self.health_check
         self.routes[('POST', '/api/completions')] = self.handle_conversations
+        self.routes[('POST', '/api/completions/cot-guided')] = self.handle_conversations_with_cot
         self.routes[('POST', '/api/completions/rag-guided')] = self.handle_conversations_with_rag
+        self.routes[('POST', '/api/completions/cot-and-rag-guided')] = self.handle_conversations_with_cot_and_rag
+        # TODO: add guardrails route(s), or add to all of the above?
 
     def format_response(self, data):
         response_data = {'response': data}
@@ -51,59 +53,66 @@ class HttpApiController:
         start_response('200 OK', response_headers)    
         return [response_body]
     
-    def handle_conversations(self, env, start_response):
-        """POST /api/completions"""
+    def _handle_completion_request(self, env, start_response, service_configurator: Callable[[AbstractTextGenerationCompletionService], AbstractTextGenerationCompletionService]):
+        """Helper method to handle common completion request logic"""
         try:
             request_body_size = int(env.get('CONTENT_LENGTH', 0))
         except ValueError:
             request_body_size = 0
-
+        
         request_body = env['wsgi.input'].read(request_body_size)
         request_json = json.loads(request_body.decode('utf-8'))
         prompt = request_json.get('prompt')
-
+        
         if not prompt:
             response_body = json.dumps({'error': 'Missing prompt in request body'}).encode('utf-8')
             response_headers = [('Content-Type', 'application/json'), ('Content-Length', str(len(response_body)))]
             start_response('400 Bad Request', response_headers)
             return [response_body]
-
-        response_text = self.text_generation_response_service.invoke(user_prompt=prompt)
-        score = self.generated_text_guardrail_service.process_generated_text(response_text)
-        response_body = self.format_response(response_text)
         
-        http_status_code = 200 # make enum
+        # Apply the service configuration (with or without guidelines)
+        configured_service = service_configurator(self.text_generation_response_service)
+        result: TextGenerationCompletionResult = configured_service.invoke(user_prompt=prompt)
+        
+        response_body = self.format_response(result.final)
+        http_status_code = 200
         response_headers = [('Content-Type', 'application/json'), ('Content-Length', str(len(response_body)))]
         start_response(f'{http_status_code} OK', response_headers)
-        self.logging_service.log_request_response(request=prompt, response=response_text)
+        
+        self.logging_service.log_request_response(request=prompt, response=result.final)
         return [response_body]
 
+    def handle_conversations(self, env, start_response):
+        """POST /api/completions"""
+        return self._handle_completion_request(
+            env, 
+            start_response, 
+            lambda service: service.without_guidelines()
+        )
+
     def handle_conversations_with_rag(self, env, start_response):
         """POST /api/completions/rag-guided"""
-        try:
-            request_body_size = int(env.get('CONTENT_LENGTH', 0))
-        except ValueError:
-            request_body_size = 0
+        return self._handle_completion_request(
+            env, 
+            start_response, 
+            lambda service: service.with_rag_context_guidelines()
+        )
 
-        request_body = env['wsgi.input'].read(request_body_size)
-        request_json = json.loads(request_body.decode('utf-8'))
-        prompt = request_json.get('prompt')
+    def handle_conversations_with_cot(self, env, start_response):
+        """POST /api/completions/cot-guided"""
+        return self._handle_completion_request(
+            env, 
+            start_response, 
+            lambda service: service.with_chain_of_thought_guidelines()
+        )
 
-        if not prompt:
-            response_body = json.dumps({'error': 'Missing prompt in request body'}).encode('utf-8')
-            response_headers = [('Content-Type', 'application/json'), ('Content-Length', str(len(response_body)))]
-            start_response('400 Bad Request', response_headers)
-            return [response_body]
-
-        response_text = self.rag_response_service.invoke(user_prompt=prompt)
-        score = self.generated_text_guardrail_service.process_generated_text(response_text)
-        response_body = self.format_response(response_text)
-        
-        http_status_code = 200 # make enum
-        response_headers = [('Content-Type', 'application/json'), ('Content-Length', str(len(response_body)))]
-        start_response(f'{http_status_code} OK', response_headers)
-        self.logging_service.log_request_response(request=prompt, response=response_text)
-        return [response_body]
+    def handle_conversations_with_cot_and_rag(self, env, start_response):
+        """POST /api/completions/cot-and-rag-guided"""
+        return self._handle_completion_request(
+            env, 
+            start_response, 
+            lambda service: service.with_rag_context_guidelines().with_chain_of_thought_guidelines()
+        )
 
     def _http_200_ok(self, env, start_response):
         """Default handler for other routes"""

src/text_generation/services/nlp/abstract_text_generation_completion_service.py

@@ -20,11 +20,6 @@ class AbstractTextGenerationCompletionService(abc.ABC):
         """Enable RAG context security guidelines"""
         raise NotImplementedError
 
-    @abc.abstractmethod
-    def with_prompt_injection_guidelines(self) -> 'AbstractTextGenerationCompletionService':
-        """Apply security guidelines using few-shot malicious prompt examples"""
-        raise NotImplementedError
-
     @abc.abstractmethod
     def with_reflexion_guardrails(self) -> 'AbstractTextGenerationCompletionService':
         """Apply security guardrails using the reflexion technique"""

src/text_generation/services/nlp/semantic_similarity_service.py

@@ -23,6 +23,7 @@ class SemanticSimilarityService(AbstractSemanticSimilarityService):
         """
             Perfect alignment (similarity) results in a score of 1; opposite is 0
         """
+        print(f'===== Using {len(self.comparison_texts)} comparison texts')
         query_embedding = array(self.embeddings.embed_query(text)).reshape(1, -1)
         doc_embeddings = array(self.embeddings.embed_documents(self.comparison_texts))
 

src/text_generation/services/nlp/text_generation_completion_service.py

@@ -9,9 +9,6 @@ from src.text_generation.domain.semantic_similarity_result import SemanticSimila
 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
@@ -200,6 +197,7 @@ class TextGenerationCompletionService(
     def invoke(self, user_prompt: str) -> TextGenerationCompletionResult:
         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)
         if not self._use_reflexion_guardrails:
             return completion_result