feat: Require login

app/local.env

@@ -31,12 +31,6 @@ LOG_ENABLE_AUDIT=FALSE
 # Change the message length for the human readable formatter
 # LOG_HUMAN_READABLE_FORMATTER__MESSAGE_WIDTH=50
 
-############################
-# Authentication
-############################
-# The auth token used by the local endpoints
-API_AUTH_TOKEN=LOCAL_AUTH_12345678
-
 ############################
 # DB Environment Variables
 ############################
@@ -72,4 +66,8 @@ AWS_SECRET_ACCESS_KEY=DO_NOT_SET_HERE
 
 AWS_DEFAULT_REGION=us-east-1
 
+###########################
+# DST app configuration
+###########################
+
 EMBEDDING_MODEL=/app/models/multi-qa-mpnet-base-dot-v1

app/src/app_config.py

@@ -3,3 +3,4 @@
 
 class AppConfig(PydanticBaseEnvConfig):
     embedding_model: str = "multi-qa-mpnet-base-dot-v1"
+    global_password: str | None = None

app/src/cache.py

@@ -0,0 +1,24 @@
+from sentence_transformers import SentenceTransformer
+
+from src.app_config import AppConfig
+
+_app_config: AppConfig | None = None
+
+
+def get_appconfig() -> AppConfig:
+    global _app_config
+
+    if not _app_config:
+        _app_config = AppConfig()
+    return _app_config
+
+
+_embedding_model: SentenceTransformer | None = None
+
+
+def get_embedding_model() -> SentenceTransformer:
+    global _embedding_model
+
+    if not _embedding_model:
+        _embedding_model = SentenceTransformer(get_appconfig().embedding_model)
+    return _embedding_model

app/src/chainlit.py

@@ -1,17 +1,16 @@
 import logging
 
-from sentence_transformers import SentenceTransformer
-
 import chainlit as cl
 import src.adapters.db as db
-from src.app_config import AppConfig
+from src.cache import get_embedding_model
 from src.format import format_guru_cards
 from src.generate import generate
+from src.login import require_login
 from src.retrieve import retrieve
 
 logger = logging.getLogger(__name__)
 
-embedding_model = SentenceTransformer(AppConfig().embedding_model)
+require_login()
 
 
 @cl.on_message
@@ -21,7 +20,7 @@ async def main(message: cl.Message) -> None:
     with db.PostgresDBClient().get_session() as db_session:
         chunks = retrieve(
             db_session,
-            embedding_model,
+            get_embedding_model(),
             message.content,
         )
 

app/src/ingest_guru_cards.py

@@ -6,7 +6,7 @@
 from smart_open import open
 
 import src.adapters.db as db
-from src.app_config import AppConfig
+from src.cache import get_embedding_model
 from src.db.models.document import Chunk, Document
 from src.util.html import get_text_from_html
 
@@ -59,7 +59,7 @@ def main() -> None:
 
     logger.info(f"Processing Guru cards at {guru_cards_filepath}")
 
-    embedding_model = SentenceTransformer(AppConfig().embedding_model)
+    embedding_model = get_embedding_model()
 
     with db.PostgresDBClient().get_session() as db_session:
         _ingest_cards(db_session, embedding_model, guru_cards_filepath)

app/src/login.py

@@ -0,0 +1,14 @@
+import chainlit as cl
+from src.cache import get_appconfig
+
+
+def login_callback(username: str, password: str) -> cl.User | None:
+    if password == get_appconfig().global_password:
+        return cl.User(identifier=username)
+    else:
+        return None
+
+
+def require_login() -> None:
+    if get_appconfig().global_password:
+        cl.password_auth_callback(login_callback)

app/tests/src/test_login.py

@@ -0,0 +1,32 @@
+import os
+
+import chainlit.config
+import src.cache
+from src.login import login_callback, require_login
+
+
+def test_require_login_no_password(monkeypatch):
+    if "GLOBAL_PASSWORD" in os.environ:
+        monkeypatch.delenv("GLOBAL_PASSWORD")
+    src.cache._app_config = None
+
+    require_login()
+
+    assert not chainlit.config.code.password_auth_callback
+
+
+def test_require_login_with_password(monkeypatch):
+    monkeypatch.setenv("GLOBAL_PASSWORD", "password")
+    src.cache._app_config = None
+
+    require_login()
+
+    assert chainlit.config.code.password_auth_callback
+
+
+def test_login_callback(monkeypatch):
+    monkeypatch.setenv("GLOBAL_PASSWORD", "correct pass")
+    src.cache._app_config = None
+
+    assert login_callback("some user", "wrong pass") is None
+    assert login_callback("some user", "correct pass").identifier == "some user"

infra/app/app-config/env-config/environment-variables.tf

@@ -19,11 +19,15 @@ locals {
   #   }
   # }
   secrets = {
-    # Example generated secret
-    # RANDOM_SECRET = {
-    #   manage_method     = "generated"
-    #   secret_store_name = "/${var.app_name}-${var.environment}/random-secret"
-    # }
+    CHAINLIT_AUTH_SECRET = {
+      manage_method     = "generated"
+      secret_store_name = "/${var.app_name}-${var.environment}/CHAINLIT_AUTH_SECRET"
+    }
+
+    GLOBAL_PASSWORD = {
+      manage_method     = "manual"
+      secret_store_name = "/${var.app_name}-${var.environment}/GLOBAL_PASSWORD"
+    }
 
     OPENAI_API_KEY = {
       manage_method     = "manual"