chainlit 1.0.506__py3-none-any.whl → 1.1.0__py3-none-any.whl

chainlit/slack/app.py

@@ -0,0 +1,368 @@
+import asyncio
+import os
+import re
+import uuid
+from functools import partial
+from typing import Dict, List, Optional, Union
+
+import httpx
+from chainlit.config import config
+from chainlit.context import ChainlitContext, HTTPSession, context_var
+from chainlit.data import get_data_layer
+from chainlit.element import Element, ElementDict
+from chainlit.emitter import BaseChainlitEmitter
+from chainlit.message import Message, StepDict
+from chainlit.types import Feedback
+from chainlit.user import PersistedUser, User
+from chainlit.user_session import user_session
+from slack_bolt.adapter.fastapi.async_handler import AsyncSlackRequestHandler
+from slack_bolt.async_app import AsyncApp
+
+
+class SlackEmitter(BaseChainlitEmitter):
+    def __init__(
+        self,
+        session: HTTPSession,
+        app: AsyncApp,
+        channel_id: str,
+        say,
+        enabled=False,
+        thread_ts: Optional[str] = None,
+    ):
+        super().__init__(session)
+        self.app = app
+        self.channel_id = channel_id
+        self.say = say
+        self.enabled = enabled
+        self.thread_ts = thread_ts
+
+    async def send_element(self, element_dict: ElementDict):
+        if not self.enabled or element_dict.get("display") != "inline":
+            return
+
+        persisted_file = self.session.files.get(element_dict.get("chainlitKey") or "")
+        file: Optional[Union[bytes, str]] = None
+
+        if persisted_file:
+            file = str(persisted_file["path"])
+        elif file_url := element_dict.get("url"):
+            async with httpx.AsyncClient() as client:
+                response = await client.get(file_url)
+                if response.status_code == 200:
+                    file = response.content
+
+        if not file:
+            return
+
+        await self.app.client.files_upload_v2(
+            channel=self.channel_id,
+            thread_ts=self.thread_ts,
+            file=file,
+            title=element_dict.get("name"),
+        )
+
+    async def send_step(self, step_dict: StepDict):
+        if not self.enabled:
+            return
+
+        is_chain_of_thought = bool(step_dict.get("parentId"))
+        is_empty_output = not step_dict.get("output")
+
+        if is_chain_of_thought or is_empty_output:
+            return
+
+        enable_feedback = not step_dict.get("disableFeedback") and get_data_layer()
+        blocks: List[Dict] = [
+            {
+                "type": "section",
+                "text": {"type": "mrkdwn", "text": step_dict["output"]},
+            }
+        ]
+        if enable_feedback:
+            blocks.append(
+                {
+                    "type": "actions",
+                    "elements": [
+                        {
+                            "action_id": "thumbdown",
+                            "type": "button",
+                            "text": {
+                                "type": "plain_text",
+                                "emoji": True,
+                                "text": ":thumbsdown:",
+                            },
+                            "value": step_dict.get("id"),
+                        },
+                        {
+                            "action_id": "thumbup",
+                            "type": "button",
+                            "text": {
+                                "type": "plain_text",
+                                "emoji": True,
+                                "text": ":thumbsup:",
+                            },
+                            "value": step_dict.get("id"),
+                        },
+                    ],
+                }
+            )
+        await self.say(
+            text=step_dict["output"], blocks=blocks, thread_ts=self.thread_ts
+        )
+
+    async def update_step(self, step_dict: StepDict):
+        if not self.enabled:
+            return
+
+        await self.send_step(step_dict)
+
+
+slack_app = AsyncApp(
+    token=os.environ.get("SLACK_BOT_TOKEN"),
+    signing_secret=os.environ.get("SLACK_SIGNING_SECRET"),
+)
+
+
+def init_slack_context(
+    session: HTTPSession,
+    slack_channel_id: str,
+    event,
+    say,
+    thread_ts: Optional[str] = None,
+) -> ChainlitContext:
+    emitter = SlackEmitter(
+        session=session,
+        app=slack_app,
+        channel_id=slack_channel_id,
+        say=say,
+        thread_ts=thread_ts,
+    )
+    context = ChainlitContext(session=session, emitter=emitter)
+    context_var.set(context)
+    user_session.set("slack_event", event)
+    user_session.set(
+        "fetch_slack_message_history",
+        partial(
+            fetch_message_history, channel_id=slack_channel_id, thread_ts=thread_ts
+        ),
+    )
+    return context
+
+
+slack_app_handler = AsyncSlackRequestHandler(slack_app)
+
+users_by_slack_id: Dict[str, Union[User, PersistedUser]] = {}
+
+USER_PREFIX = "slack_"
+
+
+def clean_content(message: str):
+    cleaned_text = re.sub(r"<@[\w]+>", "", message).strip()
+    return cleaned_text
+
+
+async def get_user(slack_user_id: str):
+    slack_user = await slack_app.client.users_info(user=slack_user_id)
+    slack_user_profile = slack_user["user"]["profile"]
+
+    user_email = slack_user_profile.get("email")
+    user = User(identifier=USER_PREFIX + user_email, metadata=slack_user_profile)
+
+    users_by_slack_id[slack_user_id] = user
+
+    if data_layer := get_data_layer():
+        persisted_user = await data_layer.create_user(user)
+        if persisted_user:
+            users_by_slack_id[slack_user_id] = persisted_user
+
+    return users_by_slack_id[slack_user_id]
+
+
+async def fetch_message_history(
+    channel_id: str, thread_ts: Optional[str] = None, limit=30
+):
+    if not thread_ts:
+        result = await slack_app.client.conversations_history(
+            channel=channel_id, limit=limit
+        )
+    else:
+        result = await slack_app.client.conversations_replies(
+            channel=channel_id, ts=thread_ts, limit=limit
+        )
+    if result["ok"]:
+        messages = result["messages"]
+        return messages
+    else:
+        raise Exception(f"Failed to fetch messages: {result['error']}")
+
+
+async def download_slack_file(url, token):
+    headers = {"Authorization": f"Bearer {token}"}
+    async with httpx.AsyncClient() as client:
+        response = await client.get(url, headers=headers)
+        if response.status_code == 200:
+            return response.content
+        else:
+            return None
+
+
+async def download_slack_files(session: HTTPSession, files, token):
+    download_coros = [
+        download_slack_file(file.get("url_private"), token) for file in files
+    ]
+    file_bytes_list = await asyncio.gather(*download_coros)
+    file_refs = []
+    for idx, file_bytes in enumerate(file_bytes_list):
+        if file_bytes:
+            name = files[idx].get("name")
+            mime_type = files[idx].get("mimetype")
+            file_ref = await session.persist_file(
+                name=name, mime=mime_type, content=file_bytes
+            )
+            file_refs.append(file_ref)
+
+    files_dicts = [
+        session.files[file["id"]] for file in file_refs if file["id"] in session.files
+    ]
+
+    file_elements = [Element.from_dict(file_dict) for file_dict in files_dicts]
+
+    return file_elements
+
+
+async def process_slack_message(
+    event,
+    say,
+    thread_name: Optional[str] = None,
+    bind_thread_to_user=False,
+    thread_ts: Optional[str] = None,
+):
+    user = await get_user(event["user"])
+
+    channel_id = event["channel"]
+    thread_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, thread_ts or channel_id))
+
+    text = event.get("text")
+    slack_files = event.get("files", [])
+
+    session_id = str(uuid.uuid4())
+    session = HTTPSession(
+        id=session_id,
+        thread_id=thread_id,
+        user=user,
+        client_type="slack",
+    )
+
+    ctx = init_slack_context(
+        session=session,
+        slack_channel_id=channel_id,
+        event=event,
+        say=say,
+        thread_ts=thread_ts,
+    )
+
+    file_elements = await download_slack_files(
+        session, slack_files, slack_app.client.token
+    )
+
+    msg = Message(
+        content=clean_content(text),
+        elements=file_elements,
+        type="user_message",
+        author=user.metadata.get("real_name"),
+    )
+
+    await msg.send()
+
+    ctx.emitter.enabled = True
+
+    if on_chat_start := config.code.on_chat_start:
+        await on_chat_start()
+
+    if on_message := config.code.on_message:
+        await on_message(msg)
+
+    if on_chat_end := config.code.on_chat_end:
+        await on_chat_end()
+
+    if data_layer := get_data_layer():
+        user_id = None
+        if isinstance(user, PersistedUser):
+            user_id = user.id if bind_thread_to_user else None
+
+        await data_layer.update_thread(
+            thread_id=thread_id,
+            name=thread_name or msg.content,
+            metadata=ctx.session.to_persistable(),
+            user_id=user_id,
+        )
+
+    ctx.session.delete()
+
+
+@slack_app.event("app_home_opened")
+async def handle_app_home_opened(event, say):
+    pass
+
+
+@slack_app.event("app_mention")
+async def handle_app_mentions(event, say):
+    thread_ts = event.get("thread_ts", event["ts"])
+    await process_slack_message(event, say, thread_ts=thread_ts)
+
+
+@slack_app.event("message")
+async def handle_message(message, say):
+    user = await get_user(message["user"])
+    thread_name = f"{user.identifier} Slack DM"
+    await process_slack_message(message, say, thread_name)
+
+
+@slack_app.block_action("thumbdown")
+async def thumb_down(ack, context, body):
+    await ack()
+    step_id = body["actions"][0]["value"]
+
+    if data_layer := get_data_layer():
+        await data_layer.upsert_feedback(Feedback(forId=step_id, value=0))
+
+    text = body["message"]["text"]
+    blocks = body["message"]["blocks"]
+    updated_blocks = [block for block in blocks if block["type"] != "actions"]
+    updated_blocks.append(
+        {
+            "type": "section",
+            "text": {"type": "mrkdwn", "text": ":thumbsdown: Feedback received."},
+        }
+    )
+    await context.client.chat_update(
+        channel=body["channel"]["id"],
+        ts=body["container"]["message_ts"],
+        text=text,
+        blocks=updated_blocks,
+    )
+
+
+@slack_app.block_action("thumbup")
+async def thumb_up(ack, context, body):
+    await ack()
+    step_id = body["actions"][0]["value"]
+
+    if data_layer := get_data_layer():
+        await data_layer.upsert_feedback(Feedback(forId=step_id, value=1))
+
+    text = body["message"]["text"]
+    blocks = body["message"]["blocks"]
+    updated_blocks = [block for block in blocks if block["type"] != "actions"]
+    updated_blocks.append(
+        {
+            "type": "section",
+            "text": {"type": "mrkdwn", "text": ":thumbsup: Feedback received."},
+        }
+    )
+    await context.client.chat_update(
+        channel=body["channel"]["id"],
+        ts=body["container"]["message_ts"],
+        text=text,
+        blocks=updated_blocks,
+    )

chainlit/socket.py

@@ -9,12 +9,18 @@ from chainlit.auth import get_current_user, require_login
 from chainlit.config import config
 from chainlit.context import init_ws_context
 from chainlit.data import get_data_layer
+from chainlit.element import Element
 from chainlit.logger import logger
 from chainlit.message import ErrorMessage, Message
 from chainlit.server import socket
 from chainlit.session import WebsocketSession
 from chainlit.telemetry import trace_event
-from chainlit.types import UIMessagePayload
+from chainlit.types import (
+    AudioChunk,
+    AudioChunkPayload,
+    AudioEndPayload,
+    UIMessagePayload,
+)
 from chainlit.user_session import user_sessions
 
 
@@ -93,9 +99,13 @@ def build_anon_user_identifier(environ):
 
 @socket.on("connect")
 async def connect(sid, environ, auth):
-    if not config.code.on_chat_start and not config.code.on_message:
+    if (
+        not config.code.on_chat_start
+        and not config.code.on_message
+        and not config.code.on_audio_chunk
+    ):
         logger.warning(
-            "You need to configure at least an on_chat_start or an on_message callback"
+            "You need to configure at least one of on_chat_start, on_message or on_audio_chunk callback"
         )
         return False
     user = None
@@ -113,18 +123,10 @@ async def connect(sid, environ, auth):
 
     # 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)
 
     # 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)
 
     session_id = environ.get("HTTP_X_CHAINLIT_SESSION_ID")
@@ -135,6 +137,7 @@ async def connect(sid, environ, auth):
     user_env = load_user_env(user_env_string)
 
     client_type = environ.get("HTTP_X_CHAINLIT_CLIENT_TYPE")
+    http_referer = environ.get("HTTP_REFERER")
 
     ws_session = WebsocketSession(
         id=session_id,
@@ -148,6 +151,7 @@ async def connect(sid, environ, auth):
         chat_profile=environ.get("HTTP_X_CHAINLIT_CHAT_PROFILE"),
         thread_id=environ.get("HTTP_X_CHAINLIT_THREAD_ID"),
         languages=environ.get("HTTP_ACCEPT_LANGUAGE"),
+        http_referer=http_referer,
     )
 
     trace_event("connection_successful")
@@ -173,46 +177,53 @@ async def connection_successful(sid):
                 "first_interaction",
                 {"interaction": "resume", "thread_id": thread.get("id")},
             )
-            await context.emitter.resume_thread(thread)
             await config.code.on_chat_resume(thread)
+            await context.emitter.resume_thread(thread)
             return
 
     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")
 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):
+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))
 
 
@@ -223,10 +234,11 @@ async def stop(sid):
 
         init_ws_context(session)
         await Message(
-            author="System", content="Task stopped by the user.", disable_feedback=True
+            author="System", content="Task manually stopped.", disable_feedback=True
         ).send()
 
-        session.should_stop = True
+        if session.current_task:
+            session.current_task.cancel()
 
         if config.code.on_stop:
             await config.code.on_stop()
@@ -243,7 +255,7 @@ async def process_message(session: WebsocketSession, payload: UIMessagePayload):
             # Sleep 1ms to make sure any children step starts after the message step start
             time.sleep(0.001)
             await config.code.on_message(message)
-    except InterruptedError:
+    except asyncio.CancelledError:
         pass
     except Exception as e:
         logger.exception(e)
@@ -258,9 +270,55 @@ async def process_message(session: WebsocketSession, payload: UIMessagePayload):
 async def message(sid, payload: UIMessagePayload):
     """Handle a message sent by the User."""
     session = WebsocketSession.require(sid)
-    session.should_stop = False
 
-    await process_message(session, payload)
+    task = asyncio.create_task(process_message(session, payload))
+    session.current_task = task
+
+
+@socket.on("audio_chunk")
+async def audio_chunk(sid, payload: AudioChunkPayload):
+    """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(AudioChunk(**payload)))
+
+
+@socket.on("audio_end")
+async def audio_end(sid, payload: AudioEndPayload):
+    """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"))
+
+        file_elements = []
+        if config.code.on_audio_end:
+            file_refs = payload.get("fileReferences")
+            if file_refs:
+                files = [
+                    session.files[file["id"]]
+                    for file in file_refs
+                    if file["id"] in session.files
+                ]
+                file_elements = [Element.from_dict(file) for file in files]
+
+            await config.code.on_audio_end(file_elements)
+    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()
 
 
 async def process_action(action: Action):
@@ -288,7 +346,7 @@ async def call_action(sid, action):
             id=action.id, status=True, response=res if isinstance(res, str) else None
         )
 
-    except InterruptedError:
+    except asyncio.CancelledError:
         await context.emitter.send_action_response(
             id=action.id, status=False, response="Action interrupted by the user"
         )

chainlit/step.py

@@ -193,10 +193,34 @@ class Step:
         self.persisted = False
         self.fail_on_persist_error = False
 
+    def _clean_content(self, content):
+        """
+        Recursively checks and converts bytes objects in content.
+        """
+
+        def handle_bytes(item):
+            if isinstance(item, bytes):
+                return "STRIPPED_BINARY_DATA"
+            elif isinstance(item, dict):
+                return {k: handle_bytes(v) for k, v in item.items()}
+            elif isinstance(item, list):
+                return [handle_bytes(i) for i in item]
+            elif isinstance(item, tuple):
+                return tuple(handle_bytes(i) for i in item)
+            return item
+
+        return handle_bytes(content)
+
     def _process_content(self, content, set_language=False):
         if content is None:
             return ""
-        if isinstance(content, dict):
+        content = self._clean_content(content)
+
+        if (
+            isinstance(content, dict)
+            or isinstance(content, list)
+            or isinstance(content, tuple)
+        ):
             try:
                 processed_content = json.dumps(content, indent=4, ensure_ascii=False)
                 if set_language:

chainlit/types.py

@@ -1,4 +1,5 @@
 from enum import Enum
+from pathlib import Path
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -144,7 +145,7 @@ class FileReference(TypedDict):
 class FileDict(TypedDict):
     id: str
     name: str
-    path: str
+    path: Path
     size: int
     type: str
 
@@ -154,6 +155,25 @@ class UIMessagePayload(TypedDict):
     fileReferences: Optional[List[FileReference]]
 
 
+class AudioChunkPayload(TypedDict):
+    isStart: bool
+    mimeType: str
+    elapsedTime: float
+    data: bytes
+
+
+@dataclass
+class AudioChunk:
+    isStart: bool
+    mimeType: str
+    elapsedTime: float
+    data: bytes
+
+
+class AudioEndPayload(TypedDict):
+    fileReferences: Optional[List[FileReference]]
+
+
 @dataclass
 class AskFileResponse:
     id: str

chainlit/user_session.py

@@ -1,6 +1,6 @@
 from typing import Dict
 
-from chainlit.context import context
+from chainlit.context import WebsocketSession, context
 
 user_sessions: Dict[str, Dict] = {}
 
@@ -27,7 +27,11 @@ class UserSession:
         user_session["chat_settings"] = context.session.chat_settings
         user_session["user"] = context.session.user
         user_session["chat_profile"] = context.session.chat_profile
-        user_session["languages"] = context.session.languages
+        user_session["http_referer"] = context.session.http_referer
+        user_session["client_type"] = context.session.client_type
+
+        if isinstance(context.session, WebsocketSession):
+            user_session["languages"] = context.session.languages
 
         if context.session.root_message:
             user_session["root_message"] = context.session.root_message