chainlit 2.0rc0__py3-none-any.whl → 2.0rc1__py3-none-any.whl

chainlit/server.py

@@ -1,4 +1,5 @@
 import asyncio
+import fnmatch
 import glob
 import json
 import mimetypes
@@ -9,10 +10,37 @@ import urllib.parse
 import webbrowser
 from contextlib import asynccontextmanager
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import List, Optional, Union
 
 import socketio
+from fastapi import (
+    APIRouter,
+    Depends,
+    FastAPI,
+    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 chainlit.auth import create_jwt, get_configuration, get_current_user
+from chainlit.auth.cookie import (
+    clear_auth_cookie,
+    clear_oauth_state_cookie,
+    set_auth_cookie,
+    set_oauth_state_cookie,
+    validate_oauth_state_cookie,
+)
 from chainlit.config import (
     APP_ROOT,
     BACKEND_ROOT,
@@ -37,26 +65,6 @@ from chainlit.types import (
     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
 
@@ -248,6 +256,7 @@ if os.environ.get("SLACK_BOT_TOKEN") and os.environ.get("SLACK_SIGNING_SECRET"):
 
 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")
@@ -299,7 +308,10 @@ def get_html_template():
     <meta property="og:url" content="{url}">
     <meta property="og:root_path" content="{ROOT_PATH}">"""
 
-    js = f"""<script>{f"window.theme = {json.dumps(config.ui.theme.to_dict())}; " if config.ui.theme else ""}</script>"""
+    js = f"""<script>
+{f"window.theme = {json.dumps(config.ui.theme.to_dict())}; " if config.ui.theme else ""}
+{f"window.transports = {json.dumps(config.project.transports)}; " if config.project.transports else "undefined"}
+</script>"""
 
     css = None
     if config.ui.custom_css:
@@ -316,7 +328,7 @@ def get_html_template():
 
     index_html_file_path = os.path.join(build_dir, "index.html")
 
-    with open(index_html_file_path, "r", encoding="utf-8") as f:
+    with open(index_html_file_path, encoding="utf-8") as f:
         content = f.read()
         content = content.replace(PLACEHOLDER, tags)
         if js:
@@ -361,43 +373,109 @@ async def auth(request: Request):
     return get_configuration()
 
 
-@router.post("/login")
-async def login(form_data: OAuth2PasswordRequestForm = Depends()):
-    """
-    Login a user using the password auth callback.
-    """
-    if not config.code.password_auth_callback:
-        raise HTTPException(
-            status_code=status.HTTP_400_BAD_REQUEST, detail="No auth_callback defined"
+def _get_response_dict(access_token: str) -> dict:
+    """Get the response dictionary for the auth response."""
+
+    if not config.project.cookie_auth:
+        # Legacy auth
+        return {
+            "access_token": access_token,
+            "token_type": "bearer",
+        }
+
+    return {"success": True}
+
+
+def _get_auth_response(access_token: str, redirect_to_callback: bool) -> Response:
+    """Get the redirect params for the OAuth callback."""
+
+    response_dict = _get_response_dict(access_token)
+
+    if redirect_to_callback:
+        root_path = os.environ.get("CHAINLIT_ROOT_PATH", "")
+        redirect_url = (
+            f"{root_path}/login/callback?{urllib.parse.urlencode(response_dict)}"
         )
 
-    user = await config.code.password_auth_callback(
-        form_data.username, form_data.password
+        return RedirectResponse(
+            # FIXME: redirect to the right frontend base url to improve the dev environment
+            url=redirect_url,
+            status_code=302,
+        )
+
+    return JSONResponse(response_dict)
+
+
+def _get_oauth_redirect_error(error: str) -> Response:
+    """Get the redirect response for an OAuth error."""
+    params = urllib.parse.urlencode(
+        {
+            "error": error,
+        }
     )
+    response = RedirectResponse(
+        # FIXME: redirect to the right frontend base url to improve the dev environment
+        url=f"/login?{params}",  # Shouldn't there be {root_path} here?
+    )
+    return response
+
+
+async def _authenticate_user(
+    user: Optional[User], redirect_to_callback: bool = False
+) -> Response:
+    """Authenticate a user and return the response."""
 
     if not user:
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="credentialssignin",
         )
-    access_token = create_jwt(user)
+
+    # If a data layer is defined, attempt to persist user.
     if data_layer := get_data_layer():
         try:
             await data_layer.create_user(user)
         except Exception as e:
+            # Catch and log exceptions during user creation.
+            # TODO: Make this catch only specific errors and allow others to propagate.
             logger.error(f"Error creating user: {e}")
 
-    return {
-        "access_token": access_token,
-        "token_type": "bearer",
-    }
+    access_token = create_jwt(user)
+
+    response = _get_auth_response(access_token, redirect_to_callback)
+
+    if config.project.cookie_auth:
+        set_auth_cookie(response, access_token)
+
+    return response
+
+
+@router.post("/login")
+async def login(response: Response, form_data: OAuth2PasswordRequestForm = Depends()):
+    """
+    Login a user using the password auth callback.
+    """
+    if not config.code.password_auth_callback:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST, detail="No auth_callback defined"
+        )
+
+    user = await config.code.password_auth_callback(
+        form_data.username, form_data.password
+    )
+
+    return await _authenticate_user(user)
 
 
 @router.post("/logout")
 async def logout(request: Request, response: Response):
     """Logout the user by calling the on_logout callback."""
+    if config.project.cookie_auth:
+        clear_auth_cookie(response)
+
     if config.code.on_logout:
         return await config.code.on_logout(request, response)
+
     return {"success": True}
 
 
@@ -412,23 +490,7 @@ async def header_auth(request: Request):
 
     user = await config.code.header_auth_callback(request.headers)
 
-    if not user:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Unauthorized",
-        )
-
-    access_token = create_jwt(user)
-    if data_layer := get_data_layer():
-        try:
-            await data_layer.create_user(user)
-        except Exception as e:
-            logger.error(f"Error creating user: {e}")
-
-    return {
-        "access_token": access_token,
-        "token_type": "bearer",
-    }
+    return await _authenticate_user(user)
 
 
 @router.get("/auth/oauth/{provider_id}")
@@ -460,16 +522,9 @@ async def oauth_login(provider_id: str, request: Request):
     response = RedirectResponse(
         url=f"{provider.authorize_url}?{params}",
     )
-    samesite: Any = os.environ.get("CHAINLIT_COOKIE_SAMESITE", "lax")
-    secure = samesite.lower() == "none"
-    response.set_cookie(
-        "oauth_state",
-        random,
-        httponly=True,
-        samesite=samesite,
-        secure=secure,
-        max_age=3 * 60,
-    )
+
+    set_oauth_state_cookie(response, random)
+
     return response
 
 
@@ -497,16 +552,7 @@ async def oauth_callback(
         )
 
     if error:
-        params = urllib.parse.urlencode(
-            {
-                "error": error,
-            }
-        )
-        response = RedirectResponse(
-            # FIXME: redirect to the right frontend base url to improve the dev environment
-            url=f"/login?{params}",
-        )
-        return response
+        return _get_oauth_redirect_error(error)
 
     if not code or not state:
         raise HTTPException(
@@ -514,9 +560,11 @@ async def oauth_callback(
             detail="Missing code or state",
         )
 
-    # Check the state from the oauth provider against the browser cookie
-    oauth_state = request.cookies.get("oauth_state")
-    if oauth_state != state:
+    try:
+        validate_oauth_state_cookie(request, state)
+    except Exception as e:
+        logger.exception("Unable to validate oauth state: %1", e)
+
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
             detail="Unauthorized",
@@ -531,34 +579,10 @@ async def oauth_callback(
         provider_id, token, raw_user_data, default_user
     )
 
-    if not user:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Unauthorized",
-        )
-
-    access_token = create_jwt(user)
-
-    if data_layer := get_data_layer():
-        try:
-            await data_layer.create_user(user)
-        except Exception as e:
-            logger.error(f"Error creating user: {e}")
-
-    params = urllib.parse.urlencode(
-        {
-            "access_token": access_token,
-            "token_type": "bearer",
-        }
-    )
+    response = await _authenticate_user(user, redirect_to_callback=True)
 
-    root_path = os.environ.get("CHAINLIT_ROOT_PATH", "")
+    clear_oauth_state_cookie(response)
 
-    response = RedirectResponse(
-        # FIXME: redirect to the right frontend base url to improve the dev environment
-        url=f"{root_path}/login/callback?{params}",
-    )
-    response.delete_cookie("oauth_state")
     return response
 
 
@@ -587,16 +611,7 @@ async def oauth_azure_hf_callback(
         )
 
     if error:
-        params = urllib.parse.urlencode(
-            {
-                "error": error,
-            }
-        )
-        response = RedirectResponse(
-            # FIXME: redirect to the right frontend base url to improve the dev environment
-            url=f"/login?{params}",
-        )
-        return response
+        return _get_oauth_redirect_error(error)
 
     if not code:
         raise HTTPException(
@@ -613,36 +628,20 @@ async def oauth_azure_hf_callback(
         provider_id, token, raw_user_data, default_user, id_token
     )
 
-    if not user:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Unauthorized",
-        )
+    response = await _authenticate_user(user, redirect_to_callback=True)
 
-    access_token = create_jwt(user)
+    clear_oauth_state_cookie(response)
 
-    if data_layer := get_data_layer():
-        try:
-            await data_layer.create_user(user)
-        except Exception as e:
-            logger.error(f"Error creating user: {e}")
+    return response
 
-    params = urllib.parse.urlencode(
-        {
-            "access_token": access_token,
-            "token_type": "bearer",
-        }
-    )
 
-    root_path = os.environ.get("CHAINLIT_ROOT_PATH", "")
+GenericUser = Union[User, PersistedUser]
+UserParam = Annotated[GenericUser, Depends(get_current_user)]
 
-    response = RedirectResponse(
-        # FIXME: redirect to the right frontend base url to improve the dev environment
-        url=f"{root_path}/login/callback?{params}",
-        status_code=302,
-    )
-    response.delete_cookie("oauth_state")
-    return response
+
+@router.get("/user")
+async def get_user(current_user: UserParam) -> GenericUser:
+    return current_user
 
 
 _language_pattern = (
@@ -670,7 +669,7 @@ async def project_translations(
 
 @router.get("/project/settings")
 async def project_settings(
-    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)],
+    current_user: UserParam,
     language: str = Query(
         default="en-US", description="Language code", pattern=_language_pattern
     ),
@@ -721,7 +720,7 @@ async def project_settings(
 async def update_feedback(
     request: Request,
     update: UpdateFeedbackRequest,
-    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)],
+    current_user: UserParam,
 ):
     """Update the human feedback for a particular message."""
     data_layer = get_data_layer()
@@ -740,7 +739,7 @@ async def update_feedback(
 async def delete_feedback(
     request: Request,
     payload: DeleteFeedbackRequest,
-    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)],
+    current_user: UserParam,
 ):
     """Delete a feedback."""
 
@@ -759,7 +758,7 @@ async def delete_feedback(
 async def get_user_threads(
     request: Request,
     payload: GetThreadsRequest,
-    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)],
+    current_user: UserParam,
 ):
     """Get the threads page by page."""
 
@@ -784,7 +783,7 @@ async def get_user_threads(
 async def get_thread(
     request: Request,
     thread_id: str,
-    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)],
+    current_user: UserParam,
 ):
     """Get a specific thread."""
     data_layer = get_data_layer()
@@ -803,7 +802,7 @@ async def get_thread_element(
     request: Request,
     thread_id: str,
     element_id: str,
-    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)],
+    current_user: UserParam,
 ):
     """Get a specific thread element."""
     data_layer = get_data_layer()
@@ -821,7 +820,7 @@ async def get_thread_element(
 async def delete_thread(
     request: Request,
     payload: DeleteThreadRequest,
-    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)],
+    current_user: UserParam,
 ):
     """Delete a thread."""
 
@@ -840,7 +839,7 @@ async def delete_thread(
 
 @router.post("/project/file")
 async def upload_file(
-    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)],
+    current_user: UserParam,
     session_id: str,
     file: UploadFile,
 ):
@@ -870,6 +869,11 @@ async def upload_file(
     assert file.filename, "No filename for uploaded file"
     assert file.content_type, "No content type for uploaded file"
 
+    try:
+        validate_file_upload(file)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
     file_response = await session.persist_file(
         name=file.filename, content=content, mime=file.content_type
     )
@@ -877,14 +881,87 @@ async def upload_file(
     return JSONResponse(content=file_response)
 
 
+def validate_file_upload(file: UploadFile):
+    """Validate the file upload as configured in config.features.spontaneous_file_upload.
+    Args:
+        file (UploadFile): The file to validate.
+    Raises:
+        ValueError: If the file is not allowed.
+    """
+    if config.features.spontaneous_file_upload is None:
+        """Default for a missing config is to allow the fileupload without any restrictions"""
+        return
+    if config.features.spontaneous_file_upload.enabled is False:
+        raise ValueError("File upload is not enabled")
+
+    validate_file_mime_type(file)
+    validate_file_size(file)
+
+
+def validate_file_mime_type(file: UploadFile):
+    """Validate the file mime type as configured in config.features.spontaneous_file_upload.
+    Args:
+        file (UploadFile): The file to validate.
+    Raises:
+        ValueError: If the file type is not allowed.
+    """
+    accept = config.features.spontaneous_file_upload.accept
+    if accept is None:
+        "Accept is not configured, allowing all file types"
+        return
+
+    assert (
+        isinstance(accept, List) or isinstance(accept, dict)
+    ), "Invalid configuration for spontaneous_file_upload, accept must be a list or a dict"
+
+    if isinstance(accept, List):
+        for pattern in accept:
+            if fnmatch.fnmatch(file.content_type, pattern):
+                return
+    elif isinstance(accept, dict):
+        for pattern, extensions in accept.items():
+            if fnmatch.fnmatch(file.content_type, pattern):
+                if len(extensions) == 0:
+                    return
+                for extension in extensions:
+                    if file.filename is not None and file.filename.endswith(extension):
+                        return
+    raise ValueError("File type not allowed")
+
+
+def validate_file_size(file: UploadFile):
+    """Validate the file size as configured in config.features.spontaneous_file_upload.
+    Args:
+        file (UploadFile): The file to validate.
+    Raises:
+        ValueError: If the file size is too large.
+    """
+    if config.features.spontaneous_file_upload.max_size_mb is None:
+        return
+
+    if (
+        file.size is not None
+        and file.size
+        > config.features.spontaneous_file_upload.max_size_mb * 1024 * 1024
+    ):
+        raise ValueError("File size too large")
+
+
 @router.get("/project/file/{file_id}")
 async def get_file(
     file_id: str,
     session_id: str,
-    # current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)], #TODO: Causes 401 error. See https://github.com/Chainlit/chainlit/issues/1472
+    current_user: UserParam,
 ):
     """Get a file from the session files directory."""
 
+    if not config.project.cookie_auth:
+        # We cannot make this work safely without cookie auth, so disable it.
+        raise HTTPException(
+            status_code=404,
+            detail="File downloads unavailable.",
+        )
+
     from chainlit.session import WebsocketSession
 
     session = WebsocketSession.get_by_id(session_id) if session_id else None
@@ -895,13 +972,12 @@ async def get_file(
             detail="Unauthorized",
         )
 
-     #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(
-    #             status_code=401,
-    #             detail="You are not authorized to download files from this session",
-    #         )
+    if current_user:
+        if not session.user or session.user.identifier != current_user.identifier:
+            raise HTTPException(
+                status_code=401,
+                detail="You are not authorized to download files from this session",
+            )
 
     if file_id in session.files:
         file = session.files[file_id]
@@ -913,7 +989,7 @@ async def get_file(
 @router.get("/files/{filename:path}")
 async def serve_file(
     filename: str,
-    current_user: Annotated[Union[User, PersistedUser], Depends(get_current_user)],
+    current_user: UserParam,
 ):
     """Serve a file from the local filesystem."""
 

chainlit/session.py

@@ -6,6 +6,7 @@ import uuid
 from typing import TYPE_CHECKING, Any, Callable, Deque, Dict, Literal, Optional, Union
 
 import aiofiles
+
 from chainlit.logger import logger
 from chainlit.types import FileReference
 
@@ -17,9 +18,9 @@ ClientType = Literal["webapp", "copilot", "teams", "slack", "discord"]
 
 
 class JSONEncoderIgnoreNonSerializable(json.JSONEncoder):
-    def default(self, obj):
+    def default(self, o):
         try:
-            return super(JSONEncoderIgnoreNonSerializable, self).default(obj)
+            return super().default(o)
         except TypeError:
             return None
 
@@ -112,9 +113,10 @@ class BaseSession:
 
         if path:
             # Copy the file from the given path
-            async with aiofiles.open(path, "rb") as src, aiofiles.open(
-                file_path, "wb"
-            ) as dst:
+            async with (
+                aiofiles.open(path, "rb") as src,
+                aiofiles.open(file_path, "wb") as dst,
+            ):
                 await dst.write(await src.read())
         elif content:
             # Write the provided content to the file

chainlit/slack/app.py

@@ -7,6 +7,9 @@ from functools import partial
 from typing import Dict, List, Optional, Union
 
 import httpx
+from slack_bolt.adapter.fastapi.async_handler import AsyncSlackRequestHandler
+from slack_bolt.async_app import AsyncApp
+
 from chainlit.config import config
 from chainlit.context import ChainlitContext, HTTPSession, context, context_var
 from chainlit.data import get_data_layer
@@ -18,8 +21,6 @@ from chainlit.telemetry import trace
 from chainlit.types import Feedback
 from chainlit.user import PersistedUser, User
 from chainlit.user_session import user_session
-from slack_bolt.adapter.fastapi.async_handler import AsyncSlackRequestHandler
-from slack_bolt.async_app import AsyncApp
 
 
 class SlackEmitter(BaseChainlitEmitter):