chainlit 1.1.404__py3-none-any.whl → 1.2.0rc0__py3-none-any.whl

chainlit/callbacks.py

@@ -0,0 +1,340 @@
+import functools
+import inspect
+from typing import Any, Awaitable, Callable, Dict, List, Optional
+
+from chainlit.action import Action
+from chainlit.assistant import Assistant
+from chainlit.assistant_settings import AssistantSettings
+from chainlit.config import config
+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
+from fastapi import Request, Response
+from starlette.datastructures import Headers
+
+
+def experimental(func):
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        print(f"\033[1;33mexperimental feature: {func.__name__}\033[0m")
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+@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
+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_chunk(func: Callable) -> Callable:
+    """
+    Hook to react to the audio chunks being sent.
+
+    Args:
+        chunk (AudioChunk): 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.
+
+    Args:
+    elements ([List[Element]): The files that were uploaded before starting the audio stream (if any).
+
+    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=True)
+        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
+
+
+# Experimental
+@trace
+@experimental
+def on_create_assistant(
+    func: Callable[[Optional[User], AssistantSettings], Any]
+) -> Callable[[Optional[User], AssistantSettings], Any]:
+    config.code.on_create_assistant = wrap_user_function(func)
+    return func
+
+
+# Experimental
+@trace
+@experimental
+def on_list_assistants(
+    func: Callable[[Optional[User]], List[Assistant]]
+) -> Callable[[Optional[User]], List[Assistant]]:
+    config.code.on_list_assistants = wrap_user_function(func)
+    return func

chainlit/cli/__init__.py

@@ -183,7 +183,7 @@ def chainlit_run(
         no_cache = True
         # This is required to have OpenAI LLM providers available for the CI run
         os.environ["OPENAI_API_KEY"] = "sk-FAKE-OPENAI-API-KEY"
-        # This is required for authenticationt tests
+        # This is required for authentication tests
         os.environ["CHAINLIT_AUTH_SECRET"] = "SUPER_SECRET"
     else:
         trace_event("chainlit run")

chainlit/config.py

@@ -4,7 +4,17 @@ import site
 import sys
 from importlib import util
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Union,
+)
 
 import tomli
 from chainlit.logger import logger
@@ -14,15 +24,17 @@ from dataclasses_json import DataClassJsonMixin
 from pydantic.dataclasses import Field, dataclass
 from starlette.datastructures import Headers
 
+from ._utils import is_path_inside
+
 if TYPE_CHECKING:
     from chainlit.action import Action
+    from chainlit.assistant import Assistant
     from chainlit.element import ElementBased
     from chainlit.message import Message
     from chainlit.types import AudioChunk, ChatProfile, Starter, ThreadDict
     from chainlit.user import User
     from fastapi import Request, Response
 
-
 BACKEND_ROOT = os.path.dirname(__file__)
 PACKAGE_ROOT = os.path.dirname(os.path.dirname(BACKEND_ROOT))
 TRANSLATIONS_DIR = os.path.join(BACKEND_ROOT, "translations")
@@ -270,10 +282,14 @@ class CodeSettings:
     # Module object loaded from the module_name
     module: Any = None
     # Bunch of callbacks defined by the developer
-    password_auth_callback: Optional[Callable[[str, str], Optional["User"]]] = None
-    header_auth_callback: Optional[Callable[[Headers], Optional["User"]]] = None
+    password_auth_callback: Optional[
+        Callable[[str, str], Awaitable[Optional["User"]]]
+    ] = None
+    header_auth_callback: Optional[Callable[[Headers], Awaitable[Optional["User"]]]] = (
+        None
+    )
     oauth_callback: Optional[
-        Callable[[str, str, Dict[str, str], "User"], Optional["User"]]
+        Callable[[str, str, Dict[str, str], "User"], Awaitable[Optional["User"]]]
     ] = None
     on_logout: Optional[Callable[["Request", "Response"], Any]] = None
     on_stop: Optional[Callable[[], Any]] = None
@@ -284,12 +300,18 @@ class CodeSettings:
     on_audio_chunk: Optional[Callable[["AudioChunk"], Any]] = None
     on_audio_end: Optional[Callable[[List["ElementBased"]], Any]] = None
 
-    author_rename: Optional[Callable[[str], str]] = None
+    author_rename: Optional[Callable[[str], Awaitable[str]]] = None
     on_settings_update: Optional[Callable[[Dict[str, Any]], Any]] = None
-    set_chat_profiles: Optional[Callable[[Optional["User"]], List["ChatProfile"]]] = (
+    set_chat_profiles: Optional[
+        Callable[[Optional["User"]], Awaitable[List["ChatProfile"]]]
+    ] = None
+    set_starters: Optional[Callable[[Optional["User"]], Awaitable[List["Starter"]]]] = (
         None
     )
-    set_starters: Optional[Callable[[Optional["User"]], List["Starter"]]] = None
+
+    # assistant-related callback function
+    on_create_assistant: Optional[Callable[[Optional["User"], Any], Any]] = None
+    on_list_assistants: Optional[Callable[[Optional["User"]], List["Assistant"]]] = None
 
 
 @dataclass()
@@ -327,33 +349,41 @@ class ChainlitConfig:
         # fallback to root language (ex: `de` when `de-DE` is not found)
         parent_language = language.split("-")[0]
 
-        translation_lib_file_path = os.path.join(
-            config_translation_dir, f"{language}.json"
-        )
-        translation_lib_parent_language_file_path = os.path.join(
-            config_translation_dir, f"{parent_language}.json"
-        )
-        default_translation_lib_file_path = os.path.join(
-            config_translation_dir, f"{default_language}.json"
-        )
+        translation_dir = Path(config_translation_dir)
 
-        if os.path.exists(translation_lib_file_path):
-            with open(translation_lib_file_path, "r", encoding="utf-8") as f:
-                translation = json.load(f)
-        elif os.path.exists(translation_lib_parent_language_file_path):
+        translation_lib_file_path = translation_dir / f"{language}.json"
+        translation_lib_parent_language_file_path = (
+            translation_dir / f"{parent_language}.json"
+        )
+        default_translation_lib_file_path = translation_dir / f"{default_language}.json"
+
+        if (
+            is_path_inside(translation_lib_file_path, translation_dir)
+            and translation_lib_file_path.is_file()
+        ):
+            translation = json.loads(
+                translation_lib_file_path.read_text(encoding="utf-8")
+            )
+        elif (
+            is_path_inside(translation_lib_parent_language_file_path, translation_dir)
+            and translation_lib_parent_language_file_path.is_file()
+        ):
             logger.warning(
                 f"Translation file for {language} not found. Using parent translation {parent_language}."
             )
-            with open(
-                translation_lib_parent_language_file_path, "r", encoding="utf-8"
-            ) as f:
-                translation = json.load(f)
-        elif os.path.exists(default_translation_lib_file_path):
+            translation = json.loads(
+                translation_lib_parent_language_file_path.read_text(encoding="utf-8")
+            )
+        elif (
+            is_path_inside(default_translation_lib_file_path, translation_dir)
+            and default_translation_lib_file_path.is_file()
+        ):
             logger.warning(
                 f"Translation file for {language} not found. Using default translation {default_language}."
             )
-            with open(default_translation_lib_file_path, "r", encoding="utf-8") as f:
-                translation = json.load(f)
+            translation = json.loads(
+                default_translation_lib_file_path.read_text(encoding="utf-8")
+            )
 
         return translation