Fix get_build_dir when custom_build is enabled

backend/chainlit/server.py

@@ -1,66 +1,66 @@
 import asyncio
 import glob
 import json
 import mimetypes
 import os
 import re
 import shutil
 import urllib.parse
 import webbrowser
 from contextlib import asynccontextmanager
 from pathlib import Path
 from typing import Any, Optional, Union
 
 import socketio
 from chainlit.auth import create_jwt, get_configuration, get_current_user
 from chainlit.config import (
     APP_ROOT,
     BACKEND_ROOT,
     DEFAULT_HOST,
     FILES_DIRECTORY,
     PACKAGE_ROOT,
     config,
     load_module,
     reload_config,
 )
 from chainlit.data import get_data_layer
 from chainlit.data.acl import is_thread_author
 from chainlit.logger import logger
 from chainlit.markdown import get_markdown_str
 from chainlit.oauth_providers import get_oauth_provider
 from chainlit.secret import random_secret
 from chainlit.types import (
     DeleteFeedbackRequest,
     DeleteThreadRequest,
     GetThreadsRequest,
     Theme,
     UpdateFeedbackRequest,
 )
 from chainlit.user import PersistedUser, User
 from fastapi import (
     APIRouter,
     Depends,
     FastAPI,
     File,
     Form,
     HTTPException,
     Query,
     Request,
     Response,
     UploadFile,
     status,
 )
 from fastapi.responses import FileResponse, HTMLResponse, JSONResponse, RedirectResponse
 from fastapi.security import OAuth2PasswordRequestForm
 from fastapi.staticfiles import StaticFiles
 from starlette.datastructures import URL
 from starlette.middleware.cors import CORSMiddleware
 from typing_extensions import Annotated
 from watchfiles import awatch
 
 from ._utils import is_path_inside
 
 mimetypes.add_type("application/javascript", ".js")
 mimetypes.add_type("text/css", ".css")
 
 ROOT_PATH = os.environ.get("CHAINLIT_ROOT_PATH", "")
@@ -154,35 +154,42 @@
         os._exit(0)
 
 
-def get_build_dir(local_target: str, packaged_target: str) -> str:
+def get_build_dir(target: str) -> str:
     """
     Get the build directory based on the UI build strategy.
+    Looks in order in the following directories:
+    - Custom build directory, if set in the config.toml (<BUILD_ROOT>/<target>/dist)
+    - Local build directory (<PACKAGE_ROOT>/libs/<target>/dist)
+    - Packaged build directory (<BACKEND_ROOT>/<target>/dist)
 
     Args:
-        local_target (str): The local target directory.
-        packaged_target (str): The packaged target directory.
+        target (str): The target directory.
 
     Returns:
-        str: The build directory
+        str: The path of the target build directory
     """
 
-    local_build_dir = os.path.join(PACKAGE_ROOT, local_target, "dist")
-    packaged_build_dir = os.path.join(BACKEND_ROOT, packaged_target, "dist")
+    local_build_dir = os.path.join(PACKAGE_ROOT, "libs", target, "dist")
+    packaged_build_dir = os.path.join(BACKEND_ROOT, target, "dist")
 
-    if config.ui.custom_build and os.path.exists(
-        os.path.join(APP_ROOT, config.ui.custom_build)
-    ):
-        return os.path.join(APP_ROOT, config.ui.custom_build)
-    elif os.path.exists(local_build_dir):
+    # Check if the custom build directory exists
+    if config.ui.custom_build:
+        custom_build_dir = os.path.join(
+            APP_ROOT, config.ui.custom_build, target, "dist"
+        )
+        if os.path.exists(custom_build_dir):
+            return custom_build_dir
+
+    if os.path.exists(local_build_dir):
         return local_build_dir
     elif os.path.exists(packaged_build_dir):
         return packaged_build_dir
     else:
-        raise FileNotFoundError(f"{local_target} built UI dir not found")
+        raise FileNotFoundError(f"{target} built UI dir not found")
 
 
-build_dir = get_build_dir("frontend", "frontend")
-copilot_build_dir = get_build_dir(os.path.join("libs", "copilot"), "copilot")
+build_dir = get_build_dir("frontend")
+copilot_build_dir = get_build_dir("copilot")
 
 app = FastAPI(lifespan=lifespan)
 
@@ -247,9 +254,9 @@
 # -------------------------------------------------------------------------------
 
 if os.environ.get("TEAMS_APP_ID") and os.environ.get("TEAMS_APP_PASSWORD"):
     from botbuilder.schema import Activity
     from chainlit.teams.app import adapter, bot
 
     @router.post("/teams/events")
     async def teams_endpoint(req: Request):
         body = await req.json()
@@ -895,7 +902,7 @@
             detail="Unauthorized",
         )
 
-     #TODO: Causes 401 error. See https://github.com/Chainlit/chainlit/issues/1472
+    # TODO: Causes 401 error. See https://github.com/Chainlit/chainlit/issues/1472
     # if current_user:
     #     if not session.user or session.user.identifier != current_user.identifier:
     #         raise HTTPException(