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

chainlit/__init__.py

@@ -69,7 +69,9 @@ from .callbacks import (
     on_message,
     on_settings_update,
     on_stop,
+    on_window_message,
     password_auth_callback,
+    send_window_message,
     set_chat_profiles,
     set_starters,
 )
@@ -132,6 +134,7 @@ __all__ = [
     "Image",
     "Text",
     "Component",
+    "Dataframe",
     "Pyplot",
     "File",
     "Task",
@@ -151,6 +154,8 @@ __all__ = [
     "CompletionGeneration",
     "GenerationMessage",
     "on_logout",
+    "on_window_message",
+    "send_window_message",
     "on_chat_start",
     "on_chat_end",
     "on_chat_resume",

chainlit/action.py

@@ -1,10 +1,12 @@
 import uuid
 from typing import Optional
 
+from dataclasses_json import DataClassJsonMixin
+from pydantic import Field
+from pydantic.dataclasses import dataclass
+
 from chainlit.context import context
 from chainlit.telemetry import trace_event
-from dataclasses_json import DataClassJsonMixin
-from pydantic.dataclasses import Field, dataclass
 
 
 @dataclass

chainlit/{auth.py → auth/__init__.py}

@@ -1,20 +1,16 @@
 import os
-from datetime import datetime, timedelta
-from typing import Any, Dict
 
-import jwt
+from fastapi import Depends, HTTPException
+
 from chainlit.config import config
 from chainlit.data import get_data_layer
+from chainlit.logger import logger
 from chainlit.oauth_providers import get_configured_oauth_providers
-from chainlit.user import User
-from fastapi import Depends, HTTPException
-from fastapi.security import OAuth2PasswordBearer
-
-reuseable_oauth = OAuth2PasswordBearer(tokenUrl="/login", auto_error=False)
 
+from .cookie import OAuth2PasswordBearerWithCookie
+from .jwt import create_jwt, decode_jwt, get_jwt_secret
 
-def get_jwt_secret():
-    return os.environ.get("CHAINLIT_AUTH_SECRET")
+reuseable_oauth = OAuth2PasswordBearerWithCookie(tokenUrl="/login", auto_error=False)
 
 
 def ensure_jwt_secret():
@@ -42,52 +38,39 @@ def get_configuration():
         "requireLogin": require_login(),
         "passwordAuth": config.code.password_auth_callback is not None,
         "headerAuth": config.code.header_auth_callback is not None,
+        "cookieAuth": config.project.cookie_auth,
         "oauthProviders": (
             get_configured_oauth_providers() if is_oauth_enabled() else []
         ),
     }
 
 
-def create_jwt(data: User) -> str:
-    to_encode: Dict[str, Any] = data.to_dict()
-    to_encode.update(
-        {
-            "exp": datetime.utcnow() + timedelta(
-                seconds=config.project.user_session_timeout
-            ),
-        }
-    )
-    encoded_jwt = jwt.encode(to_encode, get_jwt_secret(), algorithm="HS256")
-    return encoded_jwt
-
-
 async def authenticate_user(token: str = Depends(reuseable_oauth)):
     try:
-        dict = jwt.decode(
-            token,
-            get_jwt_secret(),
-            algorithms=["HS256"],
-            options={"verify_signature": True},
-        )
-        del dict["exp"]
-        user = User(**dict)
+        user = decode_jwt(token)
     except Exception as e:
         raise HTTPException(
             status_code=401, detail="Invalid authentication token"
         ) from e
+
     if data_layer := get_data_layer():
+        # Get or create persistent user if we've a data layer available.
         try:
             persisted_user = await data_layer.get_user(user.identifier)
             if persisted_user is None:
                 persisted_user = await data_layer.create_user(user)
-        except Exception:
+                assert persisted_user
+        except Exception as e:
+            logger.exception("Unable to get persisted_user from data layer: %s", e)
             return user
 
         if user and user.display_name:
+            # Copy ephemeral display_name from authenticated user to persistent user.
             persisted_user.display_name = user.display_name
+
         return persisted_user
-    else:
-        return user
+
+    return user
 
 
 async def get_current_user(token: str = Depends(reuseable_oauth)):
@@ -95,3 +78,6 @@ async def get_current_user(token: str = Depends(reuseable_oauth)):
         return None
 
     return await authenticate_user(token)
+
+
+__all__ = ["create_jwt", "get_configuration", "get_current_user"]

chainlit/auth/cookie.py

@@ -0,0 +1,124 @@
+import os
+from typing import Literal, Optional, cast
+
+from fastapi import Request, Response
+from fastapi.exceptions import HTTPException
+from fastapi.security.base import SecurityBase
+from fastapi.security.utils import get_authorization_scheme_param
+from starlette.status import HTTP_401_UNAUTHORIZED
+
+""" Module level cookie settings. """
+_cookie_samesite = cast(
+    Literal["lax", "strict", "none"],
+    os.environ.get("CHAINLIT_COOKIE_SAMESITE", "lax"),
+)
+
+assert (
+    _cookie_samesite
+    in [
+        "lax",
+        "strict",
+        "none",
+    ]
+), "Invalid value for CHAINLIT_COOKIE_SAMESITE. Must be one of 'lax', 'strict' or 'none'."
+_cookie_secure = _cookie_samesite == "none"
+
+_auth_cookie_lifetime = 60 * 60  # 1 hour
+_state_cookie_lifetime = 3 * 60  # 3m
+_auth_cookie_name = "access_token"
+_state_cookie_name = "oauth_state"
+
+
+class OAuth2PasswordBearerWithCookie(SecurityBase):
+    """
+    OAuth2 password flow with cookie support with fallback to bearer token.
+    """
+
+    def __init__(
+        self,
+        tokenUrl: str,
+        scheme_name: Optional[str] = None,
+        auto_error: bool = True,
+    ):
+        self.tokenUrl = tokenUrl
+        self.scheme_name = scheme_name or self.__class__.__name__
+        self.auto_error = auto_error
+
+    async def __call__(self, request: Request) -> Optional[str]:
+        # First try to get the token from the cookie
+        token = request.cookies.get(_auth_cookie_name)
+
+        # If no cookie, try the Authorization header as fallback
+        if not token:
+            # TODO: Only bother to check if cookie auth is explicitly disabled.
+            authorization = request.headers.get("Authorization")
+            if authorization:
+                scheme, token = get_authorization_scheme_param(authorization)
+                if scheme.lower() != "bearer":
+                    if self.auto_error:
+                        raise HTTPException(
+                            status_code=HTTP_401_UNAUTHORIZED,
+                            detail="Invalid authentication credentials",
+                            headers={"WWW-Authenticate": "Bearer"},
+                        )
+                    else:
+                        return None
+            else:
+                if self.auto_error:
+                    raise HTTPException(
+                        status_code=HTTP_401_UNAUTHORIZED,
+                        detail="Not authenticated",
+                        headers={"WWW-Authenticate": "Bearer"},
+                    )
+                else:
+                    return None
+
+        return token
+
+
+def set_auth_cookie(response: Response, token: str):
+    """
+    Helper function to set the authentication cookie with secure parameters
+    """
+
+    response.set_cookie(
+        key=_auth_cookie_name,
+        value=token,
+        httponly=True,
+        secure=_cookie_secure,
+        samesite=_cookie_samesite,
+        max_age=_auth_cookie_lifetime,
+        path="/",  # Why is path set here and not below?
+    )
+
+
+def clear_auth_cookie(response: Response):
+    """
+    Helper function to clear the authentication cookie
+    """
+    response.delete_cookie(key=_auth_cookie_name, path="/")
+
+
+def set_oauth_state_cookie(response: Response, token: str):
+    response.set_cookie(
+        _state_cookie_name,
+        token,
+        httponly=True,
+        samesite=_cookie_samesite,
+        secure=_cookie_secure,
+        max_age=_state_cookie_lifetime,
+    )
+
+
+def validate_oauth_state_cookie(request: Request, state: str):
+    """Check the state from the oauth provider against the browser cookie."""
+
+    oauth_state = request.cookies.get(_state_cookie_name)
+
+    if oauth_state != state:
+        raise Exception("oauth state does not correspond")
+
+
+def clear_oauth_state_cookie(response: Response):
+    """Oauth complete, delete state token."""
+    response.delete_cookie(_state_cookie_name)  # Do we set path here?

chainlit/auth/jwt.py

@@ -0,0 +1,37 @@
+import datetime
+import os
+from typing import Any, Dict, Optional
+
+import jwt as pyjwt
+
+from chainlit.config import config
+from chainlit.user import User
+
+
+def get_jwt_secret() -> Optional[str]:
+    return os.environ.get("CHAINLIT_AUTH_SECRET")
+
+
+def create_jwt(data: User) -> str:
+    to_encode: Dict[str, Any] = data.to_dict()
+    to_encode.update(
+        {
+            "exp": datetime.datetime.utcnow()
+            + datetime.timedelta(seconds=config.project.user_session_timeout),
+        }
+    )
+    secret = get_jwt_secret()
+    assert secret
+    encoded_jwt = pyjwt.encode(to_encode, secret, algorithm="HS256")
+    return encoded_jwt
+
+
+def decode_jwt(token: str) -> User:
+    dict = pyjwt.decode(
+        token,
+        get_jwt_secret(),
+        algorithms=["HS256"],
+        options={"verify_signature": True},
+    )
+    del dict["exp"]
+    return User(**dict)

chainlit/callbacks.py

@@ -6,6 +6,7 @@ from starlette.datastructures import Headers
 
 from chainlit.action import Action
 from chainlit.config import config
+from chainlit.context import context
 from chainlit.data.base import BaseDataLayer
 from chainlit.message import Message
 from chainlit.oauth_providers import get_configured_oauth_providers
@@ -125,6 +126,33 @@ def on_message(func: Callable) -> Callable:
     return func
 
 
+@trace
+async def send_window_message(data: Any):
+    """
+    Send custom data to the host window via a window.postMessage event.
+
+    Args:
+        data (Any): The data to send with the event.
+    """
+    await context.emitter.send_window_message(data)
+
+
+@trace
+def on_window_message(func: Callable[[str], Any]) -> Callable:
+    """
+    Hook to react to javascript postMessage events coming from the UI.
+
+    Args:
+        func (Callable[[str], Any]): The function to be called when a window message is received.
+                                     Takes the message content as a string parameter.
+
+    Returns:
+        Callable[[str], Any]: The decorated on_window_message function.
+    """
+    config.code.on_window_message = wrap_user_function(func)
+    return func
+
+
 @trace
 def on_chat_start(func: Callable) -> Callable:
     """

chainlit/chat_context.py

@@ -25,10 +25,10 @@ class ChatContext:
 
         if context.session.id not in chat_contexts:
             chat_contexts[context.session.id] = []
-            
+
         if message not in chat_contexts[context.session.id]:
             chat_contexts[context.session.id].append(message)
-        
+
         return message
 
     def remove(self, message: "Message") -> bool:

chainlit/chat_settings.py

@@ -1,8 +1,10 @@
 from typing import List
 
+from pydantic import Field
+from pydantic.dataclasses import dataclass
+
 from chainlit.context import context
 from chainlit.input_widget import InputWidget
-from pydantic.dataclasses import Field, dataclass
 
 
 @dataclass

chainlit/cli/__init__.py

@@ -9,6 +9,7 @@ import uvicorn
 nest_asyncio.apply()
 
 # ruff: noqa: E402
+from chainlit.auth import ensure_jwt_secret
 from chainlit.cache import init_lc_cache
 from chainlit.config import (
     BACKEND_ROOT,
@@ -24,7 +25,18 @@ from chainlit.logger import logger
 from chainlit.markdown import init_markdown
 from chainlit.secret import random_secret
 from chainlit.telemetry import trace_event
-from chainlit.utils import check_file, ensure_jwt_secret
+from chainlit.utils import check_file
+
+
+def assert_app():
+    if (
+        not config.code.on_chat_start
+        and not config.code.on_message
+        and not config.code.on_audio_chunk
+    ):
+        raise Exception(
+            "You need to configure at least one of on_chat_start, on_message or on_audio_chunk callback"
+        )
 
 
 # Create the main command group for Chainlit CLI
@@ -66,6 +78,7 @@ def run_chainlit(target: str):
     load_module(config.run.module_name)
 
     ensure_jwt_secret()
+    assert_app()
 
     # Create the chainlit.md file if it doesn't exist
     init_markdown(config.root)

chainlit/config.py

@@ -18,7 +18,8 @@ from typing import (
 
 import tomli
 from dataclasses_json import DataClassJsonMixin
-from pydantic.dataclasses import Field, dataclass
+from pydantic import Field
+from pydantic.dataclasses import dataclass
 from starlette.datastructures import Headers
 
 from chainlit.data.base import BaseDataLayer
@@ -35,7 +36,11 @@ if TYPE_CHECKING:
     from chainlit.message import Message
     from chainlit.types import ChatProfile, InputAudioChunk, Starter, ThreadDict
     from chainlit.user import User
-
+else:
+    # Pydantic needs to resolve forward annotations. Because all of these are used
+    # within `typing.Callable`, alias to `Any` as Pydantic does not perform validation
+    # of callable argument/return types anyway.
+    Request = Response = Action = Message = ChatProfile = InputAudioChunk = Starter = ThreadDict = User = Any  # fmt: off
 
 BACKEND_ROOT = os.path.dirname(__file__)
 PACKAGE_ROOT = os.path.dirname(os.path.dirname(BACKEND_ROOT))
@@ -90,6 +95,9 @@ auto_tag_thread = true
 # Allow users to edit their own messages
 edit_message = true
 
+# Use httponly cookie for client->server authentication, required to be able to use file upload and elements.
+cookie_auth = true
+
 # Authorize users to spontaneously upload files with messages
 [features.spontaneous_file_upload]
     enabled = true
@@ -287,6 +295,7 @@ class CodeSettings:
     on_chat_end: Optional[Callable[[], Any]] = None
     on_chat_resume: Optional[Callable[["ThreadDict"], Any]] = None
     on_message: Optional[Callable[["Message"], Any]] = None
+    on_window_message: Optional[Callable[[str], Any]] = None
     on_audio_start: Optional[Callable[[], Any]] = None
     on_audio_chunk: Optional[Callable[["InputAudioChunk"], Any]] = None
     on_audio_end: Optional[Callable[[], Any]] = None
@@ -305,6 +314,8 @@ class CodeSettings:
 @dataclass()
 class ProjectSettings(DataClassJsonMixin):
     allow_origins: List[str] = Field(default_factory=lambda: ["*"])
+    # Socket.io client transports option
+    transports: Optional[List[str]] = None
     enable_telemetry: bool = True
     # List of environment variables to be provided by each user to use the app. If empty, no environment variables will be asked to the user.
     user_env: Optional[List[str]] = None
@@ -319,6 +330,8 @@ class ProjectSettings(DataClassJsonMixin):
     cache: bool = False
     # Follow symlink for asset mount (see https://github.com/Chainlit/chainlit/issues/317)
     follow_symlink: bool = False
+    # Use httponly cookie for client->server authentication, required to be able to use file upload and elements.
+    cookie_auth: bool = True
 
 
 @dataclass()
@@ -399,7 +412,7 @@ def init_config(log=False):
             dst = os.path.join(config_translation_dir, file)
             if not os.path.exists(dst):
                 src = os.path.join(TRANSLATIONS_DIR, file)
-                with open(src, "r", encoding="utf-8") as f:
+                with open(src, encoding="utf-8") as f:
                     translation = json.load(f)
                     with open(dst, "w", encoding="utf-8") as f:
                         json.dump(translation, f, indent=4)
@@ -514,7 +527,7 @@ def load_config():
 def lint_translations():
     # Load the ground truth (en-US.json file from chainlit source code)
     src = os.path.join(TRANSLATIONS_DIR, "en-US.json")
-    with open(src, "r", encoding="utf-8") as f:
+    with open(src, encoding="utf-8") as f:
         truth = json.load(f)
 
         # Find the local app translations
@@ -522,7 +535,7 @@ def lint_translations():
             if file.endswith(".json"):
                 # Load the translation file
                 to_lint = os.path.join(config_translation_dir, file)
-                with open(to_lint, "r", encoding="utf-8") as f:
+                with open(to_lint, encoding="utf-8") as f:
                     translation = json.load(f)
 
                     # Lint the translation file

chainlit/context.py

@@ -3,9 +3,10 @@ import uuid
 from contextvars import ContextVar
 from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
-from chainlit.session import ClientType, HTTPSession, WebsocketSession
 from lazify import LazyProxy
 
+from chainlit.session import ClientType, HTTPSession, WebsocketSession
+
 if TYPE_CHECKING:
     from chainlit.emitter import BaseChainlitEmitter
     from chainlit.step import Step
@@ -104,7 +105,7 @@ def get_context() -> ChainlitContext:
     try:
         return context_var.get()
     except LookupError as e:
-        raise ChainlitContextException() from e
+        raise ChainlitContextException from e
 
 
 context: ChainlitContext = LazyProxy(get_context, enable_cache=False)