dependency fixes, test setup

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"""