chainlit 1.0.401__py3-none-any.whl → 2.0.3__py3-none-any.whl

chainlit/auth/cookie.py

@@ -0,0 +1,122 @@
+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,39 @@
+import os
+from datetime import datetime, timedelta, timezone
+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.now(timezone.utc)
+            + timedelta(seconds=config.project.user_session_timeout),
+            "iat": datetime.now(timezone.utc),  # Add issued at time
+        }
+    )
+
+    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

@@ -0,0 +1,362 @@
+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
+from chainlit.telemetry import trace
+from chainlit.types import ChatProfile, Starter, ThreadDict
+from chainlit.user import User
+from chainlit.utils import wrap_user_function
+
+
+@trace
+def password_auth_callback(
+    func: Callable[[str, str], Awaitable[Optional[User]]],
+) -> Callable:
+    """
+    Framework agnostic decorator to authenticate the user.
+
+    Args:
+        func (Callable[[str, str], Awaitable[Optional[User]]]): The authentication callback to execute. Takes the email and password as parameters.
+
+    Example:
+        @cl.password_auth_callback
+        async def password_auth_callback(username: str, password: str) -> Optional[User]:
+
+    Returns:
+        Callable[[str, str], Awaitable[Optional[User]]]: The decorated authentication callback.
+    """
+
+    config.code.password_auth_callback = wrap_user_function(func)
+    return func
+
+
+@trace
+def header_auth_callback(
+    func: Callable[[Headers], Awaitable[Optional[User]]],
+) -> Callable:
+    """
+    Framework agnostic decorator to authenticate the user via a header
+
+    Args:
+        func (Callable[[Headers], Awaitable[Optional[User]]]): The authentication callback to execute.
+
+    Example:
+        @cl.header_auth_callback
+        async def header_auth_callback(headers: Headers) -> Optional[User]:
+
+    Returns:
+        Callable[[Headers], Awaitable[Optional[User]]]: The decorated authentication callback.
+    """
+
+    config.code.header_auth_callback = wrap_user_function(func)
+    return func
+
+
+@trace
+def oauth_callback(
+    func: Callable[
+        [str, str, Dict[str, str], User, Optional[str]], Awaitable[Optional[User]]
+    ],
+) -> Callable:
+    """
+    Framework agnostic decorator to authenticate the user via oauth
+
+    Args:
+        func (Callable[[str, str, Dict[str, str], User, Optional[str]], Awaitable[Optional[User]]]): The authentication callback to execute.
+
+    Example:
+        @cl.oauth_callback
+        async def oauth_callback(provider_id: str, token: str, raw_user_data: Dict[str, str], default_app_user: User, id_token: Optional[str]) -> Optional[User]:
+
+    Returns:
+        Callable[[str, str, Dict[str, str], User, Optional[str]], Awaitable[Optional[User]]]: The decorated authentication callback.
+    """
+
+    if len(get_configured_oauth_providers()) == 0:
+        raise ValueError(
+            "You must set the environment variable for at least one oauth provider to use oauth authentication."
+        )
+
+    config.code.oauth_callback = wrap_user_function(func)
+    return func
+
+
+@trace
+def on_logout(func: Callable[[Request, Response], Any]) -> Callable:
+    """
+    Function called when the user logs out.
+    Takes the FastAPI request and response as parameters.
+    """
+
+    config.code.on_logout = wrap_user_function(func)
+    return func
+
+
+@trace
+def on_message(func: Callable) -> Callable:
+    """
+    Framework agnostic decorator to react to messages coming from the UI.
+    The decorated function is called every time a new message is received.
+
+    Args:
+        func (Callable[[Message], Any]): The function to be called when a new message is received. Takes a cl.Message.
+
+    Returns:
+        Callable[[str], Any]: The decorated on_message function.
+    """
+
+    async def with_parent_id(message: Message):
+        async with Step(name="on_message", type="run", parent_id=message.id) as s:
+            s.input = message.content
+            if len(inspect.signature(func).parameters) > 0:
+                await func(message)
+            else:
+                await func()
+
+    config.code.on_message = wrap_user_function(with_parent_id)
+    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:
+    """
+    Hook to react to the user websocket connection event.
+
+    Args:
+        func (Callable[], Any]): The connection hook to execute.
+
+    Returns:
+        Callable[], Any]: The decorated hook.
+    """
+
+    config.code.on_chat_start = wrap_user_function(
+        step(func, name="on_chat_start", type="run"), with_task=True
+    )
+    return func
+
+
+@trace
+def on_chat_resume(func: Callable[[ThreadDict], Any]) -> Callable:
+    """
+    Hook to react to resume websocket connection event.
+
+    Args:
+        func (Callable[], Any]): The connection hook to execute.
+
+    Returns:
+        Callable[], Any]: The decorated hook.
+    """
+
+    config.code.on_chat_resume = wrap_user_function(func, with_task=True)
+    return func
+
+
+@trace
+def set_chat_profiles(
+    func: Callable[[Optional["User"]], Awaitable[List["ChatProfile"]]],
+) -> Callable:
+    """
+    Programmatic declaration of the available chat profiles (can depend on the User from the session if authentication is setup).
+
+    Args:
+        func (Callable[[Optional["User"]], Awaitable[List["ChatProfile"]]]): The function declaring the chat profiles.
+
+    Returns:
+        Callable[[Optional["User"]], Awaitable[List["ChatProfile"]]]: The decorated function.
+    """
+
+    config.code.set_chat_profiles = wrap_user_function(func)
+    return func
+
+
+@trace
+def set_starters(
+    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).
+
+    Args:
+        func (Callable[[Optional["User"]], Awaitable[List["Starter"]]]): The function declaring the starters.
+
+    Returns:
+        Callable[[Optional["User"]], Awaitable[List["Starter"]]]: The decorated function.
+    """
+
+    config.code.set_starters = wrap_user_function(func)
+    return func
+
+
+@trace
+def on_chat_end(func: Callable) -> Callable:
+    """
+    Hook to react to the user websocket disconnect event.
+
+    Args:
+        func (Callable[], Any]): The disconnect hook to execute.
+
+    Returns:
+        Callable[], Any]: The decorated hook.
+    """
+
+    config.code.on_chat_end = wrap_user_function(func, with_task=True)
+    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 (InputAudioChunk): The audio chunk being sent.
+
+    Returns:
+        Callable[], Any]: The decorated hook.
+    """
+
+    config.code.on_audio_chunk = wrap_user_function(func, with_task=False)
+    return func
+
+
+@trace
+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.
+
+    Returns:
+        Callable[], Any]: The decorated hook.
+    """
+
+    config.code.on_audio_end = wrap_user_function(
+        step(func, name="on_audio_end", type="run"), with_task=True
+    )
+    return func
+
+
+@trace
+def author_rename(
+    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.
+    Args:
+        func (Callable[[str], Awaitable[str]]): The function to be called to rename an author. Takes the original author name as parameter.
+
+    Returns:
+        Callable[[Any, str], Awaitable[Any]]: The decorated function.
+    """
+
+    config.code.author_rename = wrap_user_function(func)
+    return func
+
+
+@trace
+def on_stop(func: Callable) -> Callable:
+    """
+    Hook to react to the user stopping a thread.
+
+    Args:
+        func (Callable[[], Any]): The stop hook to execute.
+
+    Returns:
+        Callable[[], Any]: The decorated stop hook.
+    """
+
+    config.code.on_stop = wrap_user_function(func)
+    return func
+
+
+def action_callback(name: str) -> Callable:
+    """
+    Callback to call when an action is clicked in the UI.
+
+    Args:
+        func (Callable[[Action], Any]): The action callback to execute. First parameter is the action.
+    """
+
+    def decorator(func: Callable[[Action], Any]):
+        config.code.action_callbacks[name] = wrap_user_function(func, with_task=False)
+        return func
+
+    return decorator
+
+
+def on_settings_update(
+    func: Callable[[Dict[str, Any]], Any],
+) -> Callable[[Dict[str, Any]], Any]:
+    """
+    Hook to react to the user changing any settings.
+
+    Args:
+        func (Callable[], Any]): The hook to execute after settings were changed.
+
+    Returns:
+        Callable[], Any]: The decorated hook.
+    """
+
+    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

@@ -0,0 +1,64 @@
+from typing import TYPE_CHECKING, Dict, List
+
+from chainlit.context import context
+
+if TYPE_CHECKING:
+    from chainlit.message import Message
+
+chat_contexts: Dict[str, List["Message"]] = {}
+
+
+class ChatContext:
+    def get(self) -> List["Message"]:
+        if not context.session:
+            return []
+
+        if context.session.id not in chat_contexts:
+            # Create a new chat context
+            chat_contexts[context.session.id] = []
+
+        return chat_contexts[context.session.id].copy()
+
+    def add(self, message: "Message"):
+        if not context.session:
+            return
+
+        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:
+        if not context.session:
+            return False
+
+        if context.session.id not in chat_contexts:
+            return False
+
+        if message in chat_contexts[context.session.id]:
+            chat_contexts[context.session.id].remove(message)
+            return True
+
+        return False
+
+    def clear(self) -> None:
+        if context.session and context.session.id in chat_contexts:
+            chat_contexts[context.session.id] = []
+
+    def to_openai(self):
+        messages = []
+        for message in self.get():
+            if message.type == "assistant_message":
+                messages.append({"role": "assistant", "content": message.content})
+            elif message.type == "user_message":
+                messages.append({"role": "user", "content": message.content})
+            else:
+                messages.append({"role": "system", "content": message.content})
+
+        return messages
+
+
+chat_context = ChatContext()

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