Adds debug mode to chat handler

chat/src/handlers/chat.py

@@ -30,16 +30,19 @@ def send(self, data):
 
 
 class StreamingSocketCallbackHandler(BaseCallbackHandler):
-    def __init__(self, socket: Websocket):
+    def __init__(self, socket: Websocket, debug_mode: bool):
         self.socket = socket
+        self.debug_mode = debug_mode
 
     def on_llm_new_token(self, token: str, **kwargs):
-        self.socket.send({"token": token})
+        if not self.debug_mode:
+            self.socket.send({"token": token})
 
 
 def handler(event, context):
     try:
         payload = json.loads(event.get("body", "{}"))
+        debug_mode = payload.get("debug", False)
 
         request_context = event.get("requestContext", {})
         connection_id = request_context.get("connectionId")
@@ -71,7 +74,8 @@ def handler(event, context):
         )
 
         client = setup.openai_chat_client(
-            callbacks=[StreamingSocketCallbackHandler(socket)], streaming=True
+            callbacks=[StreamingSocketCallbackHandler(socket, debug_mode)],
+            streaming=True,
         )
 
         prompt_text = (
@@ -101,20 +105,28 @@ def handler(event, context):
 
         try:
             doc_response = [doc.__dict__ for doc in docs]
-            socket.send({"question": question, "source_documents": doc_response})
+            if not debug_mode:
+                socket.send({"question": question, "source_documents": doc_response})
             response = chain({"question": question, "input_documents": docs})
-            response = {
-                "answer": response["output_text"],
-            }
-            socket.send(response)
+            if debug_mode:
+                final_response = {
+                    "answer": response["output_text"],
+                    "source_documents": doc_response,
+                    "ref": ref,
+                    "attributes": attributes,
+                    "prompt": prompt_text,
+                    "isSuperuser": api_token.is_superuser(),
+                }
+            else:
+                final_response = {"answer": response["output_text"], "ref": ref}
         except InvalidRequestError as err:
-            response = {
+            final_response = {
                 "question": question,
                 "answer": str(err),
                 "source_documents": [],
             }
-            socket.send(response)
 
+        socket.send(final_response)
         return {"statusCode": 200}
     except Exception as err:
         print(event)