Added ability to include a message in the trace from the component when it triggers an early exit in the pipeline.

src/architectures.py

@@ -11,7 +11,7 @@ import traceback
 from abc import ABC, abstractmethod
 from enum import Enum
 from time import time
-from typing import List, Optional
+from typing import List
 
 from src.common import config_dir, data_dir, hf_api_token
 from src.models import HFLlamaChatModel
@@ -24,9 +24,10 @@ class ArchitectureRequest:
     is a stack which can be modified through life.
     """
     def __init__(self, query: str):
-        self._request = [query]  # Stack for the request text as it evolves down the pipeline
-        self._response = []  # Stack for the response text as it evolves down the pipeline
-        self.early_exit = False
+        self._request: List[str] = [query]  # Stack for the request text as it evolves down the pipeline
+        self._response: List[str] = []  # Stack for the response text as it evolves down the pipeline
+        self.early_exit: bool = False
+        self.early_exit_message: str = None
 
     @property
     def request(self):
@@ -79,14 +80,15 @@ class ArchitectureTraceStep:
         self.start_ms = int(time() * 1000)
         self.end_ms = None
         self.outcome = ArchitectureTraceOutcome.NONE
-        self._exception = None
+        self._exception: str = None
+        self.early_exit_message: str = None
 
     def end(self, outcome: ArchitectureTraceOutcome):
         self.end_ms = int(time() * 1000)
         self.outcome = outcome
 
     @property
-    def exception(self):
+    def exception(self) -> str:
         return self._exception
 
     @exception.setter
@@ -105,7 +107,7 @@ class ArchitectureTraceStep:
         if self.outcome == ArchitectureTraceOutcome.SUCCESS:
             outcome = "Success"
         elif self.outcome == ArchitectureTraceOutcome.EARLY_EXIT:
-            outcome = "Early Exit"
+            outcome = f"Early Exit ({self.early_exit_message})"
         elif self.outcome == ArchitectureTraceOutcome.EXCEPTION:
             outcome = f"Exception ({self._exception})"
         md += f"  - **Outcome**: {outcome}"
@@ -124,10 +126,12 @@ class ArchitectureTrace:
     def start_trace(self, name: str):
         self.steps.append(ArchitectureTraceStep(name=name))
 
-    def end_trace(self, outcome):
+    def end_trace(self, outcome: ArchitectureTraceOutcome, early_exit_message: str = None):
         assert len(self.steps) > 0
         assert self.steps[-1].outcome == ArchitectureTraceOutcome.NONE
         self.steps[-1].end(outcome=outcome)
+        if early_exit_message is not None:
+            self.steps[-1].early_exit_message = early_exit_message
 
     def as_markdown(self) -> str:
         """
@@ -238,7 +242,8 @@ class Architecture:
             try:
                 component.process_request(request)
                 if request.early_exit:
-                    trace.end_trace(outcome=ArchitectureTraceOutcome.EARLY_EXIT)
+                    trace.end_trace(outcome=ArchitectureTraceOutcome.EARLY_EXIT,
+                                    early_exit_message=request.early_exit_message)
                     break
                 else:
                     trace.end_trace(outcome=ArchitectureTraceOutcome.SUCCESS)
@@ -264,6 +269,7 @@ class InputRequestScreener(ArchitectureComponent):
         if response[0:2].lower() != 'no':  # Lean cautious even if the model fails to return yes/no
             request.response = "Sorry - I cannot answer this question. Please try and rephrase it."
             request.early_exit = True
+            request.early_exit_message = response
 
 
 class OutputResponseScreener(ArchitectureComponent):
@@ -276,7 +282,6 @@ class OutputResponseScreener(ArchitectureComponent):
             raise ValueError(f'Screener model "meta-llama/Llama-2-7b-chat-hf" not set up')
         response = llm(request.response, system_prompt=system_prompt)
         if response[0:2].lower() != 'no':  # Lean cautious even if the model fails to return yes/no
-            print(f'OutputResponseScreener responded {response}')
             request.response = "Sorry - I cannot answer this question. Please try and rephrase it."
             request.early_exit = True