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

chainlit/socket.py

@@ -1,22 +1,27 @@
 import asyncio
 import json
-import time
-import uuid
-from typing import Any, Dict, Literal
+from typing import Any, Dict, Literal, Optional, Tuple, Union
+from urllib.parse import unquote
+
+from starlette.requests import cookie_parser
+from typing_extensions import TypeAlias
 
-from chainlit.action import Action
 from chainlit.auth import get_current_user, require_login
+from chainlit.chat_context import chat_context
 from chainlit.config import config
 from chainlit.context import init_ws_context
 from chainlit.data import get_data_layer
 from chainlit.logger import logger
 from chainlit.message import ErrorMessage, Message
-from chainlit.server import socket
+from chainlit.server import sio
 from chainlit.session import WebsocketSession
 from chainlit.telemetry import trace_event
-from chainlit.types import UIMessagePayload
+from chainlit.types import InputAudioChunk, InputAudioChunkPayload, MessagePayload
+from chainlit.user import PersistedUser, User
 from chainlit.user_session import user_sessions
 
+WSGIEnvironment: TypeAlias = dict[str, Any]
+
 
 def restore_existing_session(sid, session_id, emit_fn, emit_call_fn):
     """Restore a session from the sessionId provided by the client."""
@@ -42,11 +47,11 @@ async def resume_thread(session: WebsocketSession):
     if not thread:
         return
 
-    author = thread.get("user").get("identifier") if thread["user"] else None
+    author = thread.get("userIdentifier")
     user_is_author = author == session.user.identifier
 
     if user_is_author:
-        metadata = thread.get("metadata", {})
+        metadata = thread.get("metadata") or {}
         user_sessions[session.id] = metadata.copy()
         if chat_profile := metadata.get("chat_profile"):
             session.chat_profile = chat_profile
@@ -75,68 +80,69 @@ def load_user_env(user_env):
     return user_env
 
 
-def build_anon_user_identifier(environ):
-    scope = environ.get("asgi.scope", {})
-    client_ip, _ = scope.get("client")
-    ip = environ.get("HTTP_X_FORWARDED_FOR", client_ip)
+def _get_token_from_cookie(environ: WSGIEnvironment) -> Optional[str]:
+    if cookie_header := environ.get("HTTP_COOKIE", None):
+        cookies = cookie_parser(cookie_header)
+        return cookies.get("access_token", None)
 
-    try:
-        headers = scope.get("headers", {})
-        user_agent = next(
-            (v.decode("utf-8") for k, v in headers if k.decode("utf-8") == "user-agent")
-        )
-        return str(uuid.uuid5(uuid.NAMESPACE_DNS, user_agent + ip))
+    return None
+
+
+def _get_token(environ: WSGIEnvironment, auth: dict) -> Optional[str]:
+    """Take WSGI environ, return access token."""
+    return _get_token_from_cookie(environ)
 
-    except StopIteration:
-        return str(uuid.uuid5(uuid.NAMESPACE_DNS, ip))
 
+async def _authenticate_connection(
+    environ,
+    auth,
+) -> Union[Tuple[Union[User, PersistedUser], str], Tuple[None, None]]:
+    if token := _get_token(environ, auth):
+        user = await get_current_user(token=token)
+        if user:
+            return user, token
 
-@socket.on("connect")
+    return None, None
+
+
+@sio.on("connect")  # pyright: ignore [reportOptionalCall]
 async def connect(sid, environ, auth):
-    if not config.code.on_chat_start and not config.code.on_message:
-        logger.warning(
-            "You need to configure at least an on_chat_start or an on_message callback"
-        )
-        return False
-    user = None
-    token = None
-    login_required = require_login()
-    try:
-        # Check if the authentication is required
-        if login_required:
-            authorization_header = environ.get("HTTP_AUTHORIZATION")
-            token = authorization_header.split(" ")[1] if authorization_header else None
-            user = await get_current_user(token=token)
-    except Exception as e:
-        logger.info("Authentication failed")
-        return False
+    user = token = None
+
+    if require_login():
+        try:
+            user, token = await _authenticate_connection(environ, auth)
+        except Exception as e:
+            logger.exception("Exception authenticating connection: %s", e)
+
+        if not user:
+            logger.error("Authentication failed in websocket connect.")
+            raise ConnectionRefusedError("authentication failed")
 
     # Session scoped function to emit to the client
     def emit_fn(event, data):
-        if session := WebsocketSession.get(sid):
-            if session.should_stop:
-                session.should_stop = False
-                raise InterruptedError("Task stopped by user")
-        return socket.emit(event, data, to=sid)
+        return sio.emit(event, data, to=sid)
 
     # Session scoped function to emit to the client and wait for a response
     def emit_call_fn(event: Literal["ask", "call_fn"], data, timeout):
-        if session := WebsocketSession.get(sid):
-            if session.should_stop:
-                session.should_stop = False
-                raise InterruptedError("Task stopped by user")
-        return socket.call(event, data, timeout=timeout, to=sid)
+        return sio.call(event, data, timeout=timeout, to=sid)
 
-    session_id = environ.get("HTTP_X_CHAINLIT_SESSION_ID")
+    session_id = auth.get("sessionId")
     if restore_existing_session(sid, session_id, emit_fn, emit_call_fn):
         return True
 
-    user_env_string = environ.get("HTTP_USER_ENV")
+    user_env_string = auth.get("userEnv")
     user_env = load_user_env(user_env_string)
 
-    client_type = environ.get("HTTP_X_CHAINLIT_CLIENT_TYPE")
+    client_type = auth.get("clientType")
+    http_referer = environ.get("HTTP_REFERER")
+    http_cookie = environ.get("HTTP_COOKIE")
+    url_encoded_chat_profile = auth.get("chatProfile")
+    chat_profile = (
+        unquote(url_encoded_chat_profile) if url_encoded_chat_profile else None
+    )
 
-    ws_session = WebsocketSession(
+    WebsocketSession(
         id=session_id,
         socket_id=sid,
         emit=emit_fn,
@@ -145,101 +151,119 @@ async def connect(sid, environ, auth):
         user_env=user_env,
         user=user,
         token=token,
-        chat_profile=environ.get("HTTP_X_CHAINLIT_CHAT_PROFILE"),
-        thread_id=environ.get("HTTP_X_CHAINLIT_THREAD_ID"),
+        chat_profile=chat_profile,
+        thread_id=auth.get("threadId"),
+        languages=environ.get("HTTP_ACCEPT_LANGUAGE"),
+        http_referer=http_referer,
+        http_cookie=http_cookie,
     )
 
     trace_event("connection_successful")
     return True
 
 
-@socket.on("connection_successful")
+@sio.on("connection_successful")  # pyright: ignore [reportOptionalCall]
 async def connection_successful(sid):
     context = init_ws_context(sid)
 
-    if context.session.restored:
-        return
-
     await context.emitter.task_end()
     await context.emitter.clear("clear_ask")
     await context.emitter.clear("clear_call_fn")
 
+    if context.session.restored:
+        return
+
     if context.session.thread_id_to_resume and config.code.on_chat_resume:
         thread = await resume_thread(context.session)
         if thread:
             context.session.has_first_interaction = True
-            await context.emitter.emit("first_interaction", "resume")
-            await context.emitter.resume_thread(thread)
+            await context.emitter.emit(
+                "first_interaction",
+                {"interaction": "resume", "thread_id": thread.get("id")},
+            )
             await config.code.on_chat_resume(thread)
+
+            for step in thread.get("steps", []):
+                if "message" in step["type"]:
+                    chat_context.add(Message.from_dict(step))
+
+            await context.emitter.resume_thread(thread)
             return
+        else:
+            await context.emitter.send_resume_thread_error("Thread not found.")
 
     if config.code.on_chat_start:
-        await config.code.on_chat_start()
+        task = asyncio.create_task(config.code.on_chat_start())
+        context.session.current_task = task
 
 
-@socket.on("clear_session")
+@sio.on("clear_session")  # pyright: ignore [reportOptionalCall]
 async def clean_session(sid):
-    await disconnect(sid, force_clear=True)
+    session = WebsocketSession.get(sid)
+    if session:
+        session.to_clear = True
 
 
-@socket.on("disconnect")
-async def disconnect(sid, force_clear=False):
+@sio.on("disconnect")  # pyright: ignore [reportOptionalCall]
+async def disconnect(sid):
     session = WebsocketSession.get(sid)
-    if session:
-        init_ws_context(session)
 
-    if config.code.on_chat_end and session:
+    if not session:
+        return
+
+    init_ws_context(session)
+
+    if config.code.on_chat_end:
         await config.code.on_chat_end()
 
-    if session and session.thread_id and session.has_first_interaction:
+    if session.thread_id and session.has_first_interaction:
         await persist_user_session(session.thread_id, session.to_persistable())
 
-    def clear():
-        if session := WebsocketSession.get(sid):
+    def clear(_sid):
+        if session := WebsocketSession.get(_sid):
             # Clean up the user session
             if session.id in user_sessions:
                 user_sessions.pop(session.id)
             # Clean up the session
             session.delete()
 
-    async def clear_on_timeout(sid):
-        await asyncio.sleep(config.project.session_timeout)
-        clear()
-
-    if force_clear:
-        clear()
+    if session.to_clear:
+        clear(sid)
     else:
+
+        async def clear_on_timeout(_sid):
+            await asyncio.sleep(config.project.session_timeout)
+            clear(_sid)
+
         asyncio.ensure_future(clear_on_timeout(sid))
 
 
-@socket.on("stop")
+@sio.on("stop")  # pyright: ignore [reportOptionalCall]
 async def stop(sid):
     if session := WebsocketSession.get(sid):
         trace_event("stop_task")
 
         init_ws_context(session)
-        await Message(
-            author="System", content="Task stopped by the user.", disable_feedback=True
-        ).send()
+        await Message(content="Task manually stopped.").send()
 
-        session.should_stop = True
+        if session.current_task:
+            session.current_task.cancel()
 
         if config.code.on_stop:
             await config.code.on_stop()
 
 
-async def process_message(session: WebsocketSession, payload: UIMessagePayload):
+async def process_message(session: WebsocketSession, payload: MessagePayload):
     """Process a message from the user."""
     try:
         context = init_ws_context(session)
         await context.emitter.task_start()
-        message = await context.emitter.process_user_message(payload)
+        message = await context.emitter.process_message(payload)
 
         if config.code.on_message:
-            # Sleep 1ms to make sure any children step starts after the message step start
-            time.sleep(0.001)
+            await asyncio.sleep(0.001)
             await config.code.on_message(message)
-    except InterruptedError:
+    except asyncio.CancelledError:
         pass
     except Exception as e:
         logger.exception(e)
@@ -250,49 +274,108 @@ async def process_message(session: WebsocketSession, payload: UIMessagePayload):
         await context.emitter.task_end()
 
 
-@socket.on("ui_message")
-async def message(sid, payload: UIMessagePayload):
+@sio.on("edit_message")  # pyright: ignore [reportOptionalCall]
+async def edit_message(sid, payload: MessagePayload):
     """Handle a message sent by the User."""
     session = WebsocketSession.require(sid)
-    session.should_stop = False
+    context = init_ws_context(session)
 
-    await process_message(session, payload)
+    messages = chat_context.get()
 
+    orig_message = None
 
-async def process_action(action: Action):
-    callback = config.code.action_callbacks.get(action.name)
-    if callback:
-        res = await callback(action)
-        return res
-    else:
-        logger.warning("No callback found for action %s", action.name)
+    for message in messages:
+        if orig_message:
+            await message.remove()
 
+        if message.id == payload["message"]["id"]:
+            message.content = payload["message"]["output"]
+            await message.update()
+            orig_message = message
 
-@socket.on("action_call")
-async def call_action(sid, action):
-    """Handle an action call from the UI."""
-    context = init_ws_context(sid)
+    await context.emitter.task_start()
+
+    if config.code.on_message:
+        try:
+            await config.code.on_message(orig_message)
+        except asyncio.CancelledError:
+            pass
+        finally:
+            await context.emitter.task_end()
 
-    action = Action(**action)
 
+@sio.on("client_message")  # pyright: ignore [reportOptionalCall]
+async def message(sid, payload: MessagePayload):
+    """Handle a message sent by the User."""
+    session = WebsocketSession.require(sid)
+
+    task = asyncio.create_task(process_message(session, payload))
+    session.current_task = task
+
+
+@sio.on("window_message")  # pyright: ignore [reportOptionalCall]
+async def window_message(sid, data):
+    """Handle a message send by the host window."""
+    session = WebsocketSession.require(sid)
+    init_ws_context(session)
+
+    if config.code.on_window_message:
+        try:
+            await config.code.on_window_message(data)
+        except asyncio.CancelledError:
+            pass
+
+
+@sio.on("audio_start")  # pyright: ignore [reportOptionalCall]
+async def audio_start(sid):
+    """Handle audio init."""
+    session = WebsocketSession.require(sid)
+
+    context = init_ws_context(session)
+    if config.code.on_audio_start:
+        connected = bool(await config.code.on_audio_start())
+        connection_state = "on" if connected else "off"
+        await context.emitter.update_audio_connection(connection_state)
+
+
+@sio.on("audio_chunk")
+async def audio_chunk(sid, payload: InputAudioChunkPayload):
+    """Handle an audio chunk sent by the user."""
+    session = WebsocketSession.require(sid)
+
+    init_ws_context(session)
+
+    if config.code.on_audio_chunk:
+        asyncio.create_task(config.code.on_audio_chunk(InputAudioChunk(**payload)))
+
+
+@sio.on("audio_end")
+async def audio_end(sid):
+    """Handle the end of the audio stream."""
+    session = WebsocketSession.require(sid)
     try:
-        res = await process_action(action)
-        await context.emitter.send_action_response(
-            id=action.id, status=True, response=res if isinstance(res, str) else None
-        )
-
-    except InterruptedError:
-        await context.emitter.send_action_response(
-            id=action.id, status=False, response="Action interrupted by the user"
-        )
+        context = init_ws_context(session)
+        await context.emitter.task_start()
+
+        if not session.has_first_interaction:
+            session.has_first_interaction = True
+            asyncio.create_task(context.emitter.init_thread("audio"))
+
+        if config.code.on_audio_end:
+            await config.code.on_audio_end()
+
+    except asyncio.CancelledError:
+        pass
     except Exception as e:
         logger.exception(e)
-        await context.emitter.send_action_response(
-            id=action.id, status=False, response="An error occured"
-        )
+        await ErrorMessage(
+            author="Error", content=str(e) or e.__class__.__name__
+        ).send()
+    finally:
+        await context.emitter.task_end()
 
 
-@socket.on("chat_settings_change")
+@sio.on("chat_settings_change")
 async def change_settings(sid, settings: Dict[str, Any]):
     """Handle change settings submit from the UI."""
     context = init_ws_context(sid)