From 12db6f52fbcbd2ce525412ece1cd829a1b2edc0c Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Thu, 15 Jun 2023 20:48:01 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(main.py):=20add=20a=20custom?=
 =?UTF-8?q?=20404=20handler=20to=20return=20the=20index.html=20file=20when?=
 =?UTF-8?q?=20a=20404=20error=20occurs=20The=20create=5Fapp=20function=20n?=
 =?UTF-8?q?ow=20accepts=20a=20static=5Fpath=20parameter=20that=20defaults?=
 =?UTF-8?q?=20to=20"static".=20The=20setup=5Fstatic=5Ffiles=20function=20i?=
 =?UTF-8?q?s=20created=20to=20mount=20the=20static=20files=20directory=20t?=
 =?UTF-8?q?o=20the=20app.=20A=20custom=20404=20handler=20is=20added=20to?=
 =?UTF-8?q?=20the=20app=20to=20return=20the=20index.html=20file=20when=20a?=
 =?UTF-8?q?=20404=20error=20occurs.=20This=20allows=20the=20app=20to=20ser?=
 =?UTF-8?q?ve=20static=20files=20such=20as=20HTML,=20CSS,=20and=20JavaScri?=
 =?UTF-8?q?pt=20files.=20=F0=9F=9A=80=20feat(=5F=5Fmain=5F=5F.py,=20main.p?=
 =?UTF-8?q?y):=20add=20support=20for=20serving=20static=20files=20?=
 =?UTF-8?q?=E2=9C=A8=20feat(=5F=5Fmain=5F=5F.py):=20create=20a=20setup=5Fs?=
 =?UTF-8?q?tatic=5Ffiles=20function=20to=20mount=20the=20static=20files=20?=
 =?UTF-8?q?directory=20to=20the=20app?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/backend/langflow/__main__.py | 26 ++++++++++++++++++++------
 src/backend/langflow/main.py     |  7 ++++++-
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/src/backend/langflow/__main__.py b/src/backend/langflow/__main__.py
index 91f68d711..c0acbe30a 100644
--- a/src/backend/langflow/__main__.py
+++ b/src/backend/langflow/__main__.py
@@ -1,5 +1,6 @@
 import sys
 import time
+from fastapi import FastAPI
 import httpx
 from multiprocess import Process, cpu_count  # type: ignore
 import platform
@@ -145,18 +146,16 @@ def serve(
     update_settings(
         config, dev=dev, database_url=database_url, remove_api_keys=remove_api_keys
     )
-    app = create_app()
     # get the directory of the current file
     if not path:
         frontend_path = Path(__file__).parent
         static_files_dir = frontend_path / "frontend"
     else:
         static_files_dir = Path(path)
-    app.mount(
-        "/",
-        StaticFiles(directory=static_files_dir, html=True),
-        name="static",
-    )
+
+    app = create_app(static_files_dir)
+    setup_static_files(app, static_files_dir)
+
     options = {
         "bind": f"{host}:{port}",
         "workers": get_number_of_workers(workers),
@@ -180,6 +179,21 @@ def serve(
         webbrowser.open(f"http://{host}:{port}")
 
 
+def setup_static_files(app: FastAPI, static_files_dir: str):
+    """
+    Setup the static files directory.
+
+    Args:
+        app (FastAPI): FastAPI app.
+        path (str): Path to the static files directory.
+    """
+    app.mount(
+        "/",
+        StaticFiles(directory=static_files_dir, html=True),
+        name="static",
+    )
+
+
 def print_banner(host, port):
     # console = Console()
 
diff --git a/src/backend/langflow/main.py b/src/backend/langflow/main.py
index e25a2257e..a2a74d668 100644
--- a/src/backend/langflow/main.py
+++ b/src/backend/langflow/main.py
@@ -1,11 +1,12 @@
 from fastapi import FastAPI
+from fastapi.responses import FileResponse
 from fastapi.middleware.cors import CORSMiddleware
 
 from langflow.api import router
 from langflow.database.base import create_db_and_tables
 
 
-def create_app():
+def create_app(static_path: str = "static"):
     """Create the FastAPI app and include the router."""
     app = FastAPI()
 
@@ -17,6 +18,10 @@ def create_app():
     def get_health():
         return {"status": "OK"}
 
+    @app.exception_handler(404)
+    async def custom_404_handler(request, __):
+        return FileResponse(f"{static_path}/index.html")
+
     app.add_middleware(
         CORSMiddleware,
         allow_origins=origins,