chainlit 2.7.0__py3-none-any.whl → 2.7.1__py3-none-any.whl

chainlit/socket.py

@@ -1,381 +0,0 @@
-import asyncio
-import json
-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.auth import (
-    get_current_user,
-    get_token_from_cookies,
-    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 sio
-from chainlit.session import WebsocketSession
-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."""
-    if session := WebsocketSession.get_by_id(session_id):
-        session.restore(new_socket_id=sid)
-        session.emit = emit_fn
-        session.emit_call = emit_call_fn
-        return True
-    return False
-
-
-async def persist_user_session(thread_id: str, metadata: Dict):
-    if data_layer := get_data_layer():
-        await data_layer.update_thread(thread_id=thread_id, metadata=metadata)
-
-
-async def resume_thread(session: WebsocketSession):
-    data_layer = get_data_layer()
-    if not data_layer or not session.user or not session.thread_id_to_resume:
-        return
-    thread = await data_layer.get_thread(thread_id=session.thread_id_to_resume)
-    if not thread:
-        return
-
-    author = thread.get("userIdentifier")
-    user_is_author = author == session.user.identifier
-
-    if user_is_author:
-        metadata = thread.get("metadata") or {}
-        if isinstance(metadata, str):
-            metadata = json.loads(metadata)
-        user_sessions[session.id] = metadata.copy()
-        if chat_profile := metadata.get("chat_profile"):
-            session.chat_profile = chat_profile
-        if chat_settings := metadata.get("chat_settings"):
-            session.chat_settings = chat_settings
-
-        return thread
-
-
-def load_user_env(user_env):
-    if user_env:
-        user_env_dict = json.loads(user_env)
-    # Check user env
-    if config.project.user_env:
-        if not user_env_dict:
-            raise ConnectionRefusedError("Missing user environment variables")
-        # Check if requested user environment variables are provided
-        for key in config.project.user_env:
-            if key not in user_env_dict:
-                raise ConnectionRefusedError(
-                    "Missing user environment variable: " + key
-                )
-    return user_env_dict
-
-
-def _get_token_from_cookie(environ: WSGIEnvironment) -> Optional[str]:
-    if cookie_header := environ.get("HTTP_COOKIE", None):
-        cookies = cookie_parser(cookie_header)
-        return get_token_from_cookies(cookies)
-
-    return None
-
-
-def _get_token(environ: WSGIEnvironment, auth: dict) -> Optional[str]:
-    """Take WSGI environ, return access token."""
-    return _get_token_from_cookie(environ)
-
-
-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
-
-    return None, None
-
-
-@sio.on("connect")  # pyright: ignore [reportOptionalCall]
-async def connect(sid, environ, auth):
-    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):
-        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):
-        return sio.call(event, data, timeout=timeout, to=sid)
-
-    session_id = auth.get("sessionId")
-    if restore_existing_session(sid, session_id, emit_fn, emit_call_fn):
-        return True
-
-    user_env_string = auth.get("userEnv")
-    user_env = load_user_env(user_env_string)
-
-    client_type = auth.get("clientType")
-    url_encoded_chat_profile = auth.get("chatProfile")
-    chat_profile = (
-        unquote(url_encoded_chat_profile) if url_encoded_chat_profile else None
-    )
-
-    WebsocketSession(
-        id=session_id,
-        socket_id=sid,
-        emit=emit_fn,
-        emit_call=emit_call_fn,
-        client_type=client_type,
-        user_env=user_env,
-        user=user,
-        token=token,
-        chat_profile=chat_profile,
-        thread_id=auth.get("threadId"),
-        environ=environ,
-    )
-
-    return True
-
-
-@sio.on("connection_successful")  # pyright: ignore [reportOptionalCall]
-async def connection_successful(sid):
-    context = init_ws_context(sid)
-
-    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",
-                {"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:
-        task = asyncio.create_task(config.code.on_chat_start())
-        context.session.current_task = task
-
-
-@sio.on("clear_session")  # pyright: ignore [reportOptionalCall]
-async def clean_session(sid):
-    session = WebsocketSession.get(sid)
-    if session:
-        session.to_clear = True
-
-
-@sio.on("disconnect")  # pyright: ignore [reportOptionalCall]
-async def disconnect(sid):
-    session = WebsocketSession.get(sid)
-
-    if not session:
-        return
-
-    init_ws_context(session)
-
-    if config.code.on_chat_end:
-        await config.code.on_chat_end()
-
-    if session.thread_id and session.has_first_interaction:
-        await persist_user_session(session.thread_id, session.to_persistable())
-
-    async 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
-            await session.delete()
-
-    if session.to_clear:
-        await clear(sid)
-    else:
-
-        async def clear_on_timeout(_sid):
-            await asyncio.sleep(config.project.session_timeout)
-            await clear(_sid)
-
-        asyncio.ensure_future(clear_on_timeout(sid))
-
-
-@sio.on("stop")  # pyright: ignore [reportOptionalCall]
-async def stop(sid):
-    if session := WebsocketSession.get(sid):
-        init_ws_context(session)
-        await Message(content="Task manually stopped.").send()
-
-        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: MessagePayload):
-    """Process a message from the user."""
-    try:
-        context = init_ws_context(session)
-        await context.emitter.task_start()
-        message = await context.emitter.process_message(payload)
-
-        if config.code.on_message:
-            await asyncio.sleep(0.001)
-            await config.code.on_message(message)
-    except asyncio.CancelledError:
-        pass
-    except Exception as e:
-        logger.exception(e)
-        await ErrorMessage(
-            author="Error", content=str(e) or e.__class__.__name__
-        ).send()
-    finally:
-        await context.emitter.task_end()
-
-
-@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)
-    context = init_ws_context(session)
-
-    messages = chat_context.get()
-
-    orig_message = None
-
-    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
-
-    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()
-
-
-@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.features.audio.enabled:
-        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.features.audio.enabled:
-        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:
-        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.features.audio.enabled:
-            await config.code.on_audio_end()
-
-    except asyncio.CancelledError:
-        pass
-    except Exception as e:
-        logger.exception(e)
-        await ErrorMessage(
-            author="Error", content=str(e) or e.__class__.__name__
-        ).send()
-    finally:
-        await context.emitter.task_end()
-
-
-@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)
-
-    for key, value in settings.items():
-        context.session.chat_settings[key] = value
-
-    if config.code.on_settings_update:
-        await config.code.on_settings_update(settings)