backend (auth): Remove Session usage, add JWT bearer auth logic, dynamically set router dependencies

.env-template

@@ -40,9 +40,8 @@ USE_EXPERIMENTAL_LANGCHAIN=False
 # Community features
 USE_COMMUNITY_FEATURES='True'
 
-# For setting up authentication, see: docs/auth_guide.md 
-# Authentication session
-SESSION_SECRET_KEY=<GENERATE_A_SECRET_KEY>
+# For setting up authentication, see: docs/auth_guide.md
+JWT_SECRET_KEY=
 
 # Google OAuth 
 GOOGLE_CLIENT_ID=<GOOGLE_CLIENT_ID>

Makefile

@@ -35,4 +35,4 @@ first-run:
 win-first-run:
 	make win-setup
 	make migrate
-	make dev
+	make dev

docs/auth_guide.md

@@ -9,3 +9,13 @@ The list of implemented authentication strategies exist in `src/backend/services
 - GoogleOAuth: requires setting up [Google OAuth 2.0](https://support.google.com/cloud/answer/6158849?hl=en). You will need to retrieve a client ID and client secret and set them as environment variables.
 
 To enable one or more of these strategies, simply add them to the `ENABLED_AUTH_STRATEGIES` list in the configurations.
+
+After enabling one or more strategies, you must create a secret key to be used to encrypt the JWT tokens generated by the backend and store it in the `JWT_SECRET_KEY` environment variable.
+
+For testing use-cases, you can enter any string value.
+For production use-cases, We recommend running the following in a local CLI to generate a random key:
+
+```
+import secrets
+print(secrets.token_hex(32))
+```

pyproject.toml

@@ -36,6 +36,7 @@ xmltodict = "^0.13.0"
 authlib = "^1.3.0"
 itsdangerous = "^2.2.0"
 bcrypt = "^4.1.2"
+pyjwt = "^2.8.0"
 
 [tool.poetry.group.dev]
 optional = true
@@ -45,6 +46,7 @@ pytest = "^7.1.2"
 pytest-env = "^1.1.3"
 pytest-cov = "^5.0.0"
 factory-boy = "^3.3.0"
+freezegun = "^1.5.1"
 pre-commit = "^2.20.0"
 ruff = "^0.0.94"
 isort = "^5.12.0"

src/backend/chat/custom/custom.py

@@ -12,12 +12,12 @@
 from backend.schemas.tool import Category, Tool
 from backend.services.logger import get_logger
 
+logger = get_logger()
+
 
 class CustomChat(BaseChat):
     """Custom chat flow not using integrations for models."""
 
-    logger = get_logger()
-
     def chat(self, chat_request: CohereChatRequest, **kwargs: Any) -> Any:
         """
         Chat flow for custom models.
@@ -31,7 +31,7 @@ def chat(self, chat_request: CohereChatRequest, **kwargs: Any) -> Any:
         """
         # Choose the deployment model - validation already performed by request validator
         deployment_model = get_deployment(kwargs.get("deployment_name"), **kwargs)
-        self.logger.info(f"Using deployment {deployment_model.__class__.__name__}")
+        logger.info(f"Using deployment {deployment_model.__class__.__name__}")
 
         if len(chat_request.tools) > 0 and len(chat_request.documents) > 0:
             raise HTTPException(
@@ -68,13 +68,13 @@ def chat(self, chat_request: CohereChatRequest, **kwargs: Any) -> Any:
             queries = deployment_model.invoke_search_queries(
                 chat_request.message, chat_history
             )
-            self.logger.info(f"Search queries generated: {queries}")
+            logger.info(f"Search queries generated: {queries}")
 
             # Fetch Documents
             retrievers = self.get_retrievers(
                 kwargs.get("file_paths", []), [tool.name for tool in chat_request.tools]
             )
-            self.logger.info(
+            logger.info(
                 f"Using retrievers: {[retriever.__class__.__name__ for retriever in retrievers]}"
             )
 

src/backend/config/auth.py

@@ -1,9 +1,23 @@
 from backend.services.auth import BasicAuthentication, GoogleOAuth
 
-# Modify this to enable auth strategies.
+# Add Auth strategy classes here to enable them
+# Ex: [BasicAuthentication]
 ENABLED_AUTH_STRATEGIES = []
 
 # Define the mapping from Auth strategy name to class obj - does not need to be manually modified.
 # During runtime, this will create an instance of each enabled strategy class.
 # Ex: {"Basic": BasicAuthentication()}
 ENABLED_AUTH_STRATEGY_MAPPING = {cls.NAME: cls() for cls in ENABLED_AUTH_STRATEGIES}
+
+
+def is_authentication_enabled() -> bool:
+    """
+    Check whether any form of authentication was enabled.
+
+    Returns:
+        bool: Whether authentication is enabled.
+    """
+    if ENABLED_AUTH_STRATEGIES:
+        return True
+
+    return False

src/backend/config/routers.py

@@ -0,0 +1,90 @@
+from enum import StrEnum
+
+from fastapi import Depends
+
+from backend.database_models import get_session
+from backend.services.auth.request_validators import validate_authorization
+from backend.services.request_validators import (
+    validate_chat_request,
+    validate_user_header,
+)
+
+
+# Important! Any new routers must have a corresponding RouterName entry and Router dependencies
+# mapping below. Make sure they use the correct ones depending on whether authentication is enabled or not.
+class RouterName(StrEnum):
+    AUTH = "auth"
+    CHAT = "chat"
+    CONVERSATION = "conversation"
+    DEPLOYMENT = "deployment"
+    EXPERIMENTAL_FEATURES = "experimental_features"
+    TOOL = "tool"
+    USER = "user"
+
+
+# Router dependency mappings
+ROUTER_DEPENDENCIES = {
+    RouterName.AUTH: {
+        "default": [
+            Depends(get_session),
+        ],
+        "auth": [
+            Depends(get_session),
+        ],
+    },
+    RouterName.CHAT: {
+        "default": [
+            Depends(get_session),
+            Depends(validate_chat_request),
+            Depends(validate_user_header),
+        ],
+        "auth": [
+            Depends(get_session),
+            Depends(validate_chat_request),
+            Depends(validate_authorization),
+        ],
+    },
+    RouterName.CONVERSATION: {
+        "default": [
+            Depends(get_session),
+            Depends(validate_user_header),
+        ],
+        "auth": [
+            Depends(get_session),
+            Depends(validate_authorization),
+        ],
+    },
+    RouterName.DEPLOYMENT: {
+        "default": [
+            Depends(get_session),
+        ],
+        "auth": [
+            Depends(get_session),
+            Depends(validate_authorization),
+        ],
+    },
+    RouterName.EXPERIMENTAL_FEATURES: {
+        "default": [
+            Depends(get_session),
+        ],
+        "auth": [
+            Depends(get_session),
+            Depends(validate_authorization),
+        ],
+    },
+    RouterName.TOOL: {
+        "default": [],
+        "auth": [
+            Depends(validate_authorization),
+        ],
+    },
+    RouterName.USER: {
+        "default": [
+            Depends(get_session),
+        ],
+        "auth": [
+            Depends(get_session),
+            Depends(validate_authorization),
+        ],
+    },
+}

src/backend/main.py

@@ -1,15 +1,13 @@
-import logging
-import os
 from contextlib import asynccontextmanager
 
 from alembic.command import upgrade
 from alembic.config import Config
 from dotenv import load_dotenv
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from starlette.middleware.sessions import SessionMiddleware
 
-from backend.config.auth import ENABLED_AUTH_STRATEGY_MAPPING
+from backend.config.auth import is_authentication_enabled
+from backend.config.routers import ROUTER_DEPENDENCIES
 from backend.routers.auth import router as auth_router
 from backend.routers.chat import router as chat_router
 from backend.routers.conversation import router as conversation_router
@@ -35,14 +33,28 @@ async def lifespan(app: FastAPI):
 def create_app():
     app = FastAPI(lifespan=lifespan)
 
-    # Add routers
-    app.include_router(auth_router)
-    app.include_router(chat_router)
-    app.include_router(user_router)
-    app.include_router(conversation_router)
-    app.include_router(tool_router)
-    app.include_router(deployment_router)
-    app.include_router(experimental_feature_router)
+    routers = [
+        auth_router,
+        chat_router,
+        user_router,
+        conversation_router,
+        tool_router,
+        deployment_router,
+        experimental_feature_router,
+    ]
+
+    # Dynamically set router dependencies
+    # These values must be set in config/routers.py
+    dependencies_type = "default"
+    if is_authentication_enabled():
+        dependencies_type = "auth"
+    for router in routers:
+        if getattr(router, "name", "") in ROUTER_DEPENDENCIES.keys():
+            router_name = router.name
+            dependencies = ROUTER_DEPENDENCIES[router_name][dependencies_type]
+            app.include_router(router, dependencies=dependencies)
+        else:
+            app.include_router(router)
 
     # Add middleware
     app.add_middleware(
@@ -54,22 +66,6 @@ def create_app():
     )
     app.add_middleware(LoggingMiddleware)
 
-    # Handle Authentication enabled
-    if ENABLED_AUTH_STRATEGY_MAPPING:
-        secret_key = os.environ.get("SESSION_SECRET_KEY", None)
-
-        if not secret_key:
-            raise ValueError(
-                "Missing SESSION_SECRET_KEY environment variable to enable Authentication."
-            )
-
-        # Handle User sessions and Auth
-        app.add_middleware(
-            SessionMiddleware,
-            secret_key=secret_key,
-            max_age=SESSION_EXPIRY,
-        )
-
     return app