Fix issue with FastAPI mounts.

integration_tests/run_integration_tests.sh

@@ -3,7 +3,6 @@
 set -e # Fail if any of the commands below fail.
 
 ./integration_tests/start_server_test.sh
-./integration_tests/fastapi_mount_test.sh
 
 echo
 echo "CLI integration tests passed."

lilac/data/cluster_titling.py

@@ -14,6 +14,7 @@
 from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_random_exponential
 
 from ..batch_utils import group_by_sorted_key_iter
+from ..env import env
 from ..schema import (
   Item,
 )
@@ -24,7 +25,6 @@
 )
 from ..tasks import TaskInfo
 from ..utils import chunks, log
-from ..env import env
 
 _TOP_K_CENTRAL_DOCS = 7
 _TOP_K_CENTRAL_TITLES = 20
@@ -189,7 +189,6 @@ def _openai_client() -> Any:
   try:
     import openai
 
-
   except ImportError:
     raise ImportError(
       'Could not import the "openai" python package. '

lilac/server.py

@@ -96,9 +96,13 @@ def __init__(self, app: ASGIApp) -> None:
   async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
     """Redirect trailing slashes to non-trailing slashes."""
     url = URL(scope=scope).path
+
     root_path = scope.get('root_path') or ''
     ends_with_slash = (
-      url.endswith('/') and url != '/' and root_path and not url.startswith(root_path + '/api')
+      url.endswith('/')
+      and url != '/'
+      and url != f'{root_path}/'
+      and not url.startswith(root_path + '/api')
     )
 
     if scope['type'] == 'http' and ends_with_slash:

lilac/server_mount_test.py

@@ -0,0 +1,37 @@
+"""Tests the FastAPI server can be mounted."""
+
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from .server import app as lilac_app
+
+app = FastAPI()
+
+MOUNT_POINT = '/lilac_sub'
+
+
+@app.get('/')
+def read_main() -> dict:
+  """The main endpoint."""
+  return {'message': 'hello world'}
+
+
+app.mount(MOUNT_POINT, lilac_app)
+client = TestClient(app)
+
+
+def test_mount_root() -> None:
+  response = client.get('/', allow_redirects=False)
+  assert response.status_code == 200
+  assert response.json() == {'message': 'hello world'}
+
+
+def test_mount_slash_redirect() -> None:
+  response = client.get(f'{MOUNT_POINT}/auth_info/', allow_redirects=False)
+  assert response.status_code == 307
+  # We should redirect to the URL with slash removed.
+  assert response.headers['location'] == f'{MOUNT_POINT}/auth_info'
+
+  # Allow redirects to follow through.
+  response = client.get(f'{MOUNT_POINT}/auth_info/', allow_redirects=True)
+  assert response.status_code == 200

lilac/server_test.py

@@ -174,3 +174,14 @@ def user() -> UserInfo:
     ),
     auth_enabled=True,
   )
+
+
+def test_slash_redirect() -> None:
+  response = client.get('/auth_info/', allow_redirects=False)
+  assert response.status_code == 307
+  # We should redirect to the URL with slash removed.
+  assert response.headers['location'] == '/auth_info'
+
+  # Allow redirects to follow through.
+  response = client.get('/auth_info/', allow_redirects=True)
+  assert response.status_code == 200