chainlit 1.3.1__py3-none-any.whl → 2.0.0__py3-none-any.whl

chainlit/__init__.py

@@ -14,6 +14,9 @@ if env_found:
 import asyncio
 from typing import TYPE_CHECKING, Any, Dict
 
+from literalai import ChatGeneration, CompletionGeneration, GenerationMessage
+from pydantic.dataclasses import dataclass
+
 import chainlit.input_widget as input_widget
 from chainlit.action import Action
 from chainlit.cache import cache
@@ -22,7 +25,8 @@ from chainlit.chat_settings import ChatSettings
 from chainlit.context import context
 from chainlit.element import (
     Audio,
-    Component,
+    CustomElement,
+    Dataframe,
     File,
     Image,
     Pdf,
@@ -43,21 +47,21 @@ from chainlit.message import (
 )
 from chainlit.step import Step, step
 from chainlit.sync import make_async, run_sync
-from chainlit.types import AudioChunk, ChatProfile, Starter
+from chainlit.types import ChatProfile, InputAudioChunk, OutputAudioChunk, Starter
 from chainlit.user import PersistedUser, User
 from chainlit.user_session import user_session
 from chainlit.utils import make_module_getattr
 from chainlit.version import __version__
-from literalai import ChatGeneration, CompletionGeneration, GenerationMessage
-from pydantic.dataclasses import dataclass
 
 from .callbacks import (
     action_callback,
     author_rename,
+    data_layer,
     header_auth_callback,
     oauth_callback,
     on_audio_chunk,
     on_audio_end,
+    on_audio_start,
     on_chat_end,
     on_chat_resume,
     on_chat_start,
@@ -65,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,
 )
@@ -111,77 +117,73 @@ __getattr__ = make_module_getattr(
 )
 
 __all__ = [
-    "__version__",
-    "ChatProfile",
-    "Starter",
-    "user_session",
-    "chat_context",
-    "CopilotFunction",
-    "AudioChunk",
     "Action",
-    "User",
-    "PersistedUser",
+    "AskActionMessage",
+    "AskFileMessage",
+    "AskUserMessage",
+    "AsyncLangchainCallbackHandler",
     "Audio",
+    "ChatGeneration",
+    "ChatProfile",
+    "ChatSettings",
+    "CompletionGeneration",
+    "CopilotFunction",
+    "CustomElement",
+    "Dataframe",
+    "ErrorMessage",
+    "File",
+    "GenerationMessage",
+    "HaystackAgentCallbackHandler",
+    "Image",
+    "InputAudioChunk",
+    "LangchainCallbackHandler",
+    "LlamaIndexCallbackHandler",
+    "Message",
+    "OutputAudioChunk",
     "Pdf",
+    "PersistedUser",
     "Plotly",
-    "Image",
-    "Text",
-    "Component",
     "Pyplot",
-    "File",
+    "Starter",
+    "Step",
     "Task",
     "TaskList",
     "TaskStatus",
+    "Text",
+    "User",
     "Video",
-    "ChatSettings",
-    "input_widget",
-    "Message",
-    "ErrorMessage",
-    "AskUserMessage",
-    "AskActionMessage",
-    "AskFileMessage",
-    "Step",
-    "step",
-    "ChatGeneration",
-    "CompletionGeneration",
-    "GenerationMessage",
-    "on_logout",
-    "on_chat_start",
-    "on_chat_end",
-    "on_chat_resume",
-    "on_stop",
+    "__version__",
     "action_callback",
     "author_rename",
-    "on_settings_update",
-    "password_auth_callback",
-    "header_auth_callback",
-    "sleep",
-    "run_sync",
-    "make_async",
     "cache",
+    "chat_context",
     "context",
-    "LangchainCallbackHandler",
-    "AsyncLangchainCallbackHandler",
-    "LlamaIndexCallbackHandler",
-    "HaystackAgentCallbackHandler",
-    "instrument_openai",
-    "instrument_mistralai",
-    "password_auth_callback",
+    "data_layer",
     "header_auth_callback",
+    "input_widget",
+    "instrument_mistralai",
+    "instrument_openai",
+    "make_async",
     "oauth_callback",
+    "on_audio_chunk",
+    "on_audio_end",
+    "on_audio_start",
+    "on_chat_end",
+    "on_chat_resume",
+    "on_chat_start",
     "on_logout",
     "on_message",
-    "on_chat_start",
-    "on_chat_resume",
+    "on_settings_update",
+    "on_stop",
+    "on_window_message",
+    "password_auth_callback",
+    "run_sync",
+    "send_window_message",
     "set_chat_profiles",
     "set_starters",
-    "on_chat_end",
-    "on_audio_chunk",
-    "on_audio_end",
-    "author_rename",
-    "on_stop",
-    "action_callback",
-    "on_settings_update",
+    "sleep",
+    "step",
+    "user_session",
 ]
 
 

chainlit/action.py

@@ -1,28 +1,30 @@
 import uuid
-from typing import Optional
+from typing import Dict, 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
 class Action(DataClassJsonMixin):
     # Name of the action, this should be used in the action_callback
     name: str
-    # The value associated with the action. This is useful to differentiate between multiple actions with the same name.
-    value: str
-    # The label of the action. This is what the user will see. If not provided the name will be used.
+    # The parameters to call this action with.
+    payload: Dict
+    # The label of the action. This is what the user will see.
     label: str = ""
-    # The description of the action. This is what the user will see when they hover the action.
-    description: str = ""
+    # The tooltip of the action button. This is what the user will see when they hover the action.
+    tooltip: str = ""
+    # The lucid icon name for this action.
+    icon: Optional[str] = None
     # This should not be set manually, only used internally.
     forId: Optional[str] = None
     # The ID of the action
     id: str = Field(default_factory=lambda: str(uuid.uuid4()))
-    # Show the action in a drawer menu
-    collapsed: bool = False
 
     def __post_init__(self) -> None:
         trace_event(f"init {self.__class__.__name__}")

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():
@@ -45,45 +41,36 @@ def get_configuration():
         "oauthProviders": (
             get_configured_oauth_providers() if is_oauth_enabled() else []
         ),
+        "default_theme": config.ui.default_theme,
     }
 
 
-def create_jwt(data: User) -> str:
-    to_encode: Dict[str, Any] = data.to_dict()
-    to_encode.update(
-        {
-            "exp": datetime.utcnow() + timedelta(minutes=60 * 24 * 15),  # 15 days
-        }
-    )
-    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)
-    except Exception:
-        raise HTTPException(status_code=401, detail="Invalid authentication token")
+        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)):
@@ -91,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,123 @@
+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,
+    )
+
+
+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/cache.py

@@ -1,5 +1,7 @@
+import importlib.util
 import os
 import threading
+from typing import Any
 
 from chainlit.config import config
 from chainlit.logger import logger
@@ -8,11 +10,7 @@ from chainlit.logger import logger
 def init_lc_cache():
     use_cache = config.project.cache is True and config.run.no_cache is False
 
-    if use_cache:
-        try:
-            import langchain
-        except ImportError:
-            return
+    if use_cache and importlib.util.find_spec("langchain") is not None:
         from langchain.cache import SQLiteCache
         from langchain.globals import set_llm_cache
 
@@ -25,7 +23,7 @@ def init_lc_cache():
                 )
 
 
-_cache = {}
+_cache: dict[tuple, Any] = {}
 _cache_lock = threading.Lock()
 
 

chainlit/callbacks.py

@@ -1,8 +1,13 @@
 import inspect
 from typing import Any, Awaitable, Callable, Dict, List, Optional
 
+from fastapi import Request, Response
+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
 from chainlit.step import Step, step
@@ -10,13 +15,11 @@ from chainlit.telemetry import trace
 from chainlit.types import ChatProfile, Starter, ThreadDict
 from chainlit.user import User
 from chainlit.utils import wrap_user_function
-from fastapi import Request, Response
-from starlette.datastructures import Headers
 
 
 @trace
 def password_auth_callback(
-    func: Callable[[str, str], Awaitable[Optional[User]]]
+    func: Callable[[str, str], Awaitable[Optional[User]]],
 ) -> Callable:
     """
     Framework agnostic decorator to authenticate the user.
@@ -38,7 +41,7 @@ def password_auth_callback(
 
 @trace
 def header_auth_callback(
-    func: Callable[[Headers], Awaitable[Optional[User]]]
+    func: Callable[[Headers], Awaitable[Optional[User]]],
 ) -> Callable:
     """
     Framework agnostic decorator to authenticate the user via a header
@@ -123,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:
     """
@@ -177,7 +207,7 @@ def set_chat_profiles(
 
 @trace
 def set_starters(
-    func: Callable[[Optional["User"]], Awaitable[List["Starter"]]]
+    func: Callable[[Optional["User"]], Awaitable[List["Starter"]]],
 ) -> Callable:
     """
     Programmatic declaration of the available starter (can depend on the User from the session if authentication is setup).
@@ -209,13 +239,26 @@ def on_chat_end(func: Callable) -> Callable:
     return func
 
 
+@trace
+def on_audio_start(func: Callable) -> Callable:
+    """
+    Hook to react to the user initiating audio.
+
+    Returns:
+        Callable[], Any]: The decorated hook.
+    """
+
+    config.code.on_audio_start = wrap_user_function(func, with_task=False)
+    return func
+
+
 @trace
 def on_audio_chunk(func: Callable) -> Callable:
     """
     Hook to react to the audio chunks being sent.
 
     Args:
-        chunk (AudioChunk): The audio chunk being sent.
+        chunk (InputAudioChunk): The audio chunk being sent.
 
     Returns:
         Callable[], Any]: The decorated hook.
@@ -230,9 +273,6 @@ def on_audio_end(func: Callable) -> Callable:
     """
     Hook to react to the audio stream ending. This is called after the last audio chunk is sent.
 
-    Args:
-    elements ([List[Element]): The files that were uploaded before starting the audio stream (if any).
-
     Returns:
         Callable[], Any]: The decorated hook.
     """
@@ -245,7 +285,7 @@ def on_audio_end(func: Callable) -> Callable:
 
 @trace
 def author_rename(
-    func: Callable[[str], Awaitable[str]]
+    func: Callable[[str], Awaitable[str]],
 ) -> Callable[[str], Awaitable[str]]:
     """
     Useful to rename the author of message to display more friendly author names in the UI.
@@ -285,7 +325,7 @@ def action_callback(name: str) -> Callable:
     """
 
     def decorator(func: Callable[[Action], Any]):
-        config.code.action_callbacks[name] = wrap_user_function(func, with_task=True)
+        config.code.action_callbacks[name] = wrap_user_function(func, with_task=False)
         return func
 
     return decorator
@@ -306,3 +346,17 @@ def on_settings_update(
 
     config.code.on_settings_update = wrap_user_function(func, with_task=True)
     return func
+
+
+def data_layer(
+    func: Callable[[], BaseDataLayer],
+) -> Callable[[], BaseDataLayer]:
+    """
+    Hook to configure custom data layer.
+    """
+
+    # We don't use wrap_user_function here because:
+    # 1. We don't need to support async here and;
+    # 2. We don't want to change the API for get_data_layer() to be async, everywhere (at this point).
+    config.code.data_layer = func
+    return func

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