Merge branch 'feature/ZC/protect-prod-app'

.vscode/launch.json

@@ -18,8 +18,6 @@
                 "${workspaceFolder}\\log.ini",
                 "--reload-dir",
                 "${workspaceFolder}\\templates"
-
-
             ],
             "envFile": "${workspaceFolder}/.env",
             "jinja": true,

Makefile

@@ -1,6 +1,6 @@
 .PHONY: start
 start:
-	uvicorn main:app --reload --port 9000
+	uvicorn main:app --reload --port 9000 --log-config log.ini
 
 .PHONY: format
 format:

main.py

@@ -1,5 +1,6 @@
 """Main entrypoint for the app."""
 import logging
+import os
 import pickle
 from pathlib import Path
 from typing import Optional
@@ -36,15 +37,19 @@ async def get(request: Request):
 @app.websocket("/chat/{client_id}")
 async def websocket_endpoint(websocket: WebSocket, client_id: str):
     await websocket.accept()
+
+    if client_id not in os.environ.get("AUTHORIZED_CLIENT_IDS", []):
+        return
+
     course_manager = CourseManager(client_id)
+    logging.info(f"[{client_id}] - Connection authorized.")
     question_handler = QuestionGenCallbackHandler(websocket, client_id)
     stream_handler = StreamingLLMCallbackHandler(websocket, client_id)
     chat_history = []
     qa_chain = get_chain(vectorstore, question_handler, stream_handler)
     # Use the below line instead of the above line to enable tracing
     # Ensure `langchain-server` is running
     # qa_chain = get_chain(vectorstore, question_handler, stream_handler, tracing=True)
-    logging.error(client_id)
     while True:
         try:
             # Receive and send back the client message