chainlit 0.7.700__py3-none-any.whl → 1.0.0rc0__py3-none-any.whl

chainlit/message.py

@@ -1,76 +1,87 @@
+import asyncio
 import json
 import uuid
-from abc import ABC, abstractmethod
-from datetime import datetime, timezone
+from abc import ABC
+from datetime import datetime
 from typing import Dict, List, Optional, Union, cast
 
 from chainlit.action import Action
-from chainlit.client.base import MessageDict
 from chainlit.config import config
 from chainlit.context import context
-from chainlit.data import chainlit_client
+from chainlit.data import get_data_layer
 from chainlit.element import ElementBased
 from chainlit.logger import logger
-from chainlit.prompt import Prompt
+from chainlit.step import StepDict
 from chainlit.telemetry import trace_event
 from chainlit.types import (
     AskActionResponse,
     AskActionSpec,
     AskFileResponse,
     AskFileSpec,
-    AskResponse,
     AskSpec,
+    FileDict,
 )
-from syncer import asyncio
+from chainlit_client.step import MessageStepType
 
 
 class MessageBase(ABC):
     id: str
+    thread_id: str
     author: str
     content: str = ""
+    type: MessageStepType = "assistant_message"
+    disable_feedback = False
     streaming = False
-    created_at: Union[int, str, None] = None
+    created_at: Union[str, None] = None
     fail_on_persist_error: bool = False
     persisted = False
+    is_error = False
+    language: Optional[str] = None
+    wait_for_answer = False
+    indent: Optional[int] = None
 
     def __post_init__(self) -> None:
         trace_event(f"init {self.__class__.__name__}")
+        self.thread_id = context.session.thread_id
+
         if not getattr(self, "id", None):
             self.id = str(uuid.uuid4())
-        if not self.created_at:
-            self.created_at = datetime.now(timezone.utc).isoformat()
-
-    @abstractmethod
-    def to_dict(self) -> Dict:
-        pass
-
-    async def with_conversation_id(self):
-        _dict = self.to_dict()
-        _dict["conversationId"] = await context.session.get_conversation_id()
-        return _dict
-
-    async def _create(self):
-        msg_dict = await self.with_conversation_id()
-        asyncio.create_task(self._persist_create(msg_dict))
 
-        if not config.features.prompt_playground:
-            msg_dict.pop("prompt", None)
-        return msg_dict
+    @classmethod
+    def from_dict(self, _dict: StepDict):
+        type = _dict.get("type", "assistant_message")
+        message = Message(
+            id=_dict["id"],
+            created_at=_dict["createdAt"],
+            content=_dict["output"],
+            author=_dict.get("name", config.ui.name),
+            type=type,  # type: ignore
+            disable_feedback=_dict.get("disableFeedback", False),
+            language=_dict.get("language"),
+        )
 
-    async def _persist_create(self, message: MessageDict):
-        if not chainlit_client or self.persisted:
-            return
+        return message
 
-        try:
-            persisted_id = await chainlit_client.create_message(message)
+    def to_dict(self) -> StepDict:
+        _dict: StepDict = {
+            "id": self.id,
+            "threadId": self.thread_id,
+            "createdAt": self.created_at,
+            "start": self.created_at,
+            "end": self.created_at,
+            "output": self.content,
+            "name": self.author,
+            "type": self.type,
+            "createdAt": self.created_at,
+            "language": self.language,
+            "streaming": self.streaming,
+            "disableFeedback": self.disable_feedback,
+            "isError": self.is_error,
+            "waitForAnswer": self.wait_for_answer,
+            "indent": self.indent,
+        }
 
-            if persisted_id:
-                self.id = persisted_id
-                self.persisted = True
-        except Exception as e:
-            if self.fail_on_persist_error:
-                raise e
-            logger.error(f"Failed to persist message creation: {str(e)}")
+        return _dict
 
     async def update(
         self,
@@ -79,51 +90,63 @@ class MessageBase(ABC):
         Update a message already sent to the UI.
         """
         trace_event("update_message")
+
         if self.streaming:
             self.streaming = False
-        msg_dict = self.to_dict()
-        asyncio.create_task(self._persist_update(msg_dict))
-        await context.emitter.update_message(msg_dict)
 
-        return True
+        step_dict = self.to_dict()
+
+        data_layer = get_data_layer()
+        if data_layer:
+            try:
+                asyncio.create_task(data_layer.update_step(step_dict))
+            except Exception as e:
+                if self.fail_on_persist_error:
+                    raise e
+                logger.error(f"Failed to persist message update: {str(e)}")
 
-    async def _persist_update(self, message: MessageDict):
-        if not chainlit_client or not self.id:
-            return
+        await context.emitter.update_step(step_dict)
 
-        try:
-            await chainlit_client.update_message(self.id, message)
-        except Exception as e:
-            if self.fail_on_persist_error:
-                raise e
-            logger.error(f"Failed to persist message update: {str(e)}")
+        return True
 
     async def remove(self):
         """
         Remove a message already sent to the UI.
-        This will not automatically remove potential nested messages and could lead to undesirable side effects in the UI.
         """
         trace_event("remove_message")
 
-        if chainlit_client and self.id:
-            await chainlit_client.delete_message(self.id)
+        step_dict = self.to_dict()
+        data_layer = get_data_layer()
+        if data_layer:
+            try:
+                asyncio.create_task(data_layer.delete_step(step_dict["id"]))
+            except Exception as e:
+                if self.fail_on_persist_error:
+                    raise e
+                logger.error(f"Failed to persist message deletion: {str(e)}")
 
-        await context.emitter.delete_message(self.to_dict())
+        await context.emitter.delete_step(step_dict)
 
         return True
 
-    async def _persist_remove(self):
-        if not chainlit_client or not self.id:
-            return
+    async def _create(self):
+        step_dict = self.to_dict()
+        data_layer = get_data_layer()
+        if data_layer and not self.persisted:
+            try:
+                asyncio.create_task(data_layer.create_step(step_dict))
+                self.persisted = True
+            except Exception as e:
+                if self.fail_on_persist_error:
+                    raise e
+                logger.error(f"Failed to persist message creation: {str(e)}")
 
-        try:
-            await chainlit_client.delete_message(self.id)
-        except Exception as e:
-            if self.fail_on_persist_error:
-                raise e
-            logger.error(f"Failed to persist message deletion: {str(e)}")
+        return step_dict
 
     async def send(self):
+        if not self.created_at:
+            self.created_at = datetime.utcnow().isoformat()
+
         if self.content is None:
             self.content = ""
 
@@ -133,9 +156,9 @@ class MessageBase(ABC):
         if self.streaming:
             self.streaming = False
 
-        msg_dict = await self._create()
+        step_dict = await self._create()
 
-        await context.emitter.send_message(msg_dict)
+        await context.emitter.send_step(step_dict)
 
         return self.id
 
@@ -147,8 +170,8 @@ class MessageBase(ABC):
 
         if not self.streaming:
             self.streaming = True
-            msg_dict = self.to_dict()
-            await context.emitter.stream_start(msg_dict)
+            step_dict = self.to_dict()
+            await context.emitter.stream_start(step_dict)
 
         if is_sequence:
             self.content = token
@@ -164,33 +187,27 @@ class MessageBase(ABC):
 class Message(MessageBase):
     """
     Send a message to the UI
-    If a project ID is configured, the message will be persisted in the cloud.
 
     Args:
         content (Union[str, Dict]): The content of the message.
         author (str, optional): The author of the message, this will be used in the UI. Defaults to the chatbot name (see config).
-        prompt (Prompt, optional): The prompt used to generate the message. If provided, enables the prompt playground for this message.
         language (str, optional): Language of the code is the content is code. See https://react-code-blocks-rajinwonderland.vercel.app/?path=/story/codeblock--supported-languages for a list of supported languages.
-        parent_id (str, optional): If provided, the message will be nested inside the parent in the UI.
-        indent (int, optional): If positive, the message will be nested in the UI. (deprecated, use parent_id instead)
         actions (List[Action], optional): A list of actions to send with the message.
         elements (List[ElementBased], optional): A list of elements to send with the message.
-        disable_human_feedback (bool, optional): Hide the feedback buttons for this specific message
+        disable_feedback (bool, optional): Hide the feedback buttons for this specific message
     """
 
     def __init__(
         self,
         content: Union[str, Dict],
         author: str = config.ui.name,
-        prompt: Optional[Prompt] = None,
         language: Optional[str] = None,
-        parent_id: Optional[str] = None,
-        indent: int = 0,
         actions: Optional[List[Action]] = None,
         elements: Optional[List[ElementBased]] = None,
-        disable_human_feedback: Optional[bool] = False,
-        author_is_user: Optional[bool] = False,
-        id: Optional[uuid.UUID] = None,
+        disable_feedback: bool = False,
+        type: MessageStepType = "assistant_message",
+        id: Optional[str] = None,
+        created_at: Union[str, None] = None,
     ):
         self.language = language
 
@@ -200,87 +217,45 @@ class Message(MessageBase):
                 self.language = "json"
             except TypeError:
                 self.content = str(content)
-                self.language = "python"
+                self.language = "text"
         elif isinstance(content, str):
             self.content = content
         else:
             self.content = str(content)
-            self.language = "python"
+            self.language = "text"
 
         if id:
             self.id = str(id)
 
+        if created_at:
+            self.created_at = created_at
+
         self.author = author
-        self.author_is_user = author_is_user
-        self.prompt = prompt
-        self.parent_id = parent_id
-        self.indent = indent
+        self.type = type
         self.actions = actions if actions is not None else []
         self.elements = elements if elements is not None else []
-        self.disable_human_feedback = disable_human_feedback
+        self.disable_feedback = disable_feedback
 
         super().__post_init__()
 
-    @classmethod
-    def from_dict(self, _dict: MessageDict):
-        message = Message(
-            content=_dict["content"],
-            author=_dict.get("author", config.ui.name),
-            prompt=_dict.get("prompt"),
-            language=_dict.get("language"),
-            parent_id=_dict.get("parentId"),
-            indent=_dict.get("indent") or 0,
-            disable_human_feedback=_dict.get("disableHumanFeedback"),
-            author_is_user=_dict.get("authorIsUser"),
-        )
-
-        if _id := _dict.get("id"):
-            message.id = _id
-        if created_at := _dict.get("createdAt"):
-            message.created_at = created_at
-
-        return message
-
-    def to_dict(self):
-        _dict = {
-            "createdAt": self.created_at,
-            "content": self.content,
-            "author": self.author,
-            "authorIsUser": self.author_is_user,
-            "language": self.language,
-            "parentId": self.parent_id,
-            "indent": self.indent,
-            "streaming": self.streaming,
-            "disableHumanFeedback": self.disable_human_feedback,
-        }
-
-        if self.prompt:
-            _dict["prompt"] = self.prompt.to_dict()
-
-        if self.id:
-            _dict["id"] = self.id
-
-        return _dict
-
     async def send(self) -> str:
         """
         Send the message to the UI and persist it in the cloud if a project ID is configured.
         Return the ID of the message.
         """
         trace_event("send_message")
-        id = await super().send()
+        await super().send()
 
-        if not self.parent_id:
-            context.session.root_message = self
+        context.session.root_message = self
 
         # Create tasks for all actions and elements
-        tasks = [action.send(for_id=id) for action in self.actions]
-        tasks.extend(element.send(for_id=id) for element in self.elements)
+        tasks = [action.send(for_id=self.id) for action in self.actions]
+        tasks.extend(element.send(for_id=self.id) for element in self.elements)
 
         # Run all tasks concurrently
         await asyncio.gather(*tasks)
 
-        return id
+        return self.id
 
     async def update(self):
         """
@@ -290,15 +265,16 @@ class Message(MessageBase):
         trace_event("send_message")
         await super().update()
 
-        actions_to_update = [action for action in self.actions if action.forId is None]
+        # Update tasks for all actions and elements
+        tasks = [
+            action.send(for_id=self.id)
+            for action in self.actions
+            if action.forId is None
+        ]
+        tasks.extend(element.send(for_id=self.id) for element in self.elements)
 
-        elements_to_update = [el for el in self.elements if self.id not in el.for_ids]
-
-        for action in actions_to_update:
-            await action.send(for_id=self.id)
-
-        for element in elements_to_update:
-            await element.send(for_id=self.id)
+        # Run all tasks concurrently
+        await asyncio.gather(*tasks)
 
         return True
 
@@ -323,29 +299,16 @@ class ErrorMessage(MessageBase):
         self,
         content: str,
         author: str = config.ui.name,
-        parent_id: Optional[str] = None,
-        indent: int = 0,
         fail_on_persist_error: bool = False,
     ):
         self.content = content
         self.author = author
-        self.parent_id = parent_id
-        self.indent = indent
+        self.type = "system_message"
+        self.is_error = True
         self.fail_on_persist_error = fail_on_persist_error
 
         super().__post_init__()
 
-    def to_dict(self):
-        return {
-            "id": self.id,
-            "createdAt": self.created_at,
-            "content": self.content,
-            "author": self.author,
-            "parentId": self.parent_id,
-            "indent": self.indent,
-            "isError": True,
-        }
-
     async def send(self):
         """
         Send the error message to the UI and persist it in the cloud if a project ID is configured.
@@ -371,7 +334,7 @@ class AskUserMessage(AskMessageBase):
     Args:
         content (str): The content of the prompt.
         author (str, optional): The author of the message, this will be used in the UI. Defaults to the chatbot name (see config).
-        disable_human_feedback (bool, optional): Hide the feedback buttons for this specific message
+        disable_feedback (bool, optional): Hide the feedback buttons for this specific message
         timeout (int, optional): The number of seconds to wait for an answer before raising a TimeoutError.
         raise_on_timeout (bool, optional): Whether to raise a socketio TimeoutError if the user does not answer in time.
     """
@@ -380,45 +343,46 @@ class AskUserMessage(AskMessageBase):
         self,
         content: str,
         author: str = config.ui.name,
-        disable_human_feedback: bool = False,
+        type: MessageStepType = "assistant_message",
+        disable_feedback: bool = False,
         timeout: int = 60,
         raise_on_timeout: bool = False,
     ):
         self.content = content
         self.author = author
         self.timeout = timeout
-        self.disable_human_feedback = disable_human_feedback
+        self.type = type
+        self.disable_feedback = disable_feedback
         self.raise_on_timeout = raise_on_timeout
 
         super().__post_init__()
 
-    def to_dict(self):
-        return {
-            "id": self.id,
-            "createdAt": self.created_at,
-            "content": self.content,
-            "author": self.author,
-            "waitForAnswer": True,
-            "disableHumanFeedback": self.disable_human_feedback,
-        }
-
-    async def send(self) -> Union[AskResponse, None]:
+    async def send(self) -> Union[StepDict, None]:
         """
         Sends the question to ask to the UI and waits for the reply.
         """
         trace_event("send_ask_user")
+        if not self.created_at:
+            self.created_at = datetime.utcnow().isoformat()
+
+        if config.code.author_rename:
+            self.author = await config.code.author_rename(self.author)
 
         if self.streaming:
             self.streaming = False
 
-        if config.code.author_rename:
-            self.author = await config.code.author_rename(self.author)
+        self.wait_for_answer = True
 
-        msg_dict = await self._create()
+        step_dict = await self._create()
 
         spec = AskSpec(type="text", timeout=self.timeout)
 
-        res = await context.emitter.send_ask_user(msg_dict, spec, self.raise_on_timeout)
+        res = cast(
+            Union[None, StepDict],
+            await context.emitter.send_ask_user(step_dict, spec, self.raise_on_timeout),
+        )
+
+        self.wait_for_answer = False
 
         return res
 
@@ -435,7 +399,7 @@ class AskFileMessage(AskMessageBase):
         max_size_mb (int, optional): Maximum size per file in MB. Maximum value is 100.
         max_files (int, optional): Maximum number of files to upload. Maximum value is 10.
         author (str, optional): The author of the message, this will be used in the UI. Defaults to the chatbot name (see config).
-        disable_human_feedback (bool, optional): Hide the feedback buttons for this specific message
+        disable_feedback (bool, optional): Hide the feedback buttons for this specific message
         timeout (int, optional): The number of seconds to wait for an answer before raising a TimeoutError.
         raise_on_timeout (bool, optional): Whether to raise a socketio TimeoutError if the user does not answer in time.
     """
@@ -447,7 +411,8 @@ class AskFileMessage(AskMessageBase):
         max_size_mb=2,
         max_files=1,
         author=config.ui.name,
-        disable_human_feedback: bool = False,
+        type: MessageStepType = "assistant_message",
+        disable_feedback: bool = False,
         timeout=90,
         raise_on_timeout=False,
     ):
@@ -455,36 +420,32 @@ class AskFileMessage(AskMessageBase):
         self.max_size_mb = max_size_mb
         self.max_files = max_files
         self.accept = accept
+        self.type = type
         self.author = author
         self.timeout = timeout
         self.raise_on_timeout = raise_on_timeout
-        self.disable_human_feedback = disable_human_feedback
+        self.disable_feedback = disable_feedback
 
         super().__post_init__()
 
-    def to_dict(self):
-        return {
-            "id": self.id,
-            "createdAt": self.created_at,
-            "content": self.content,
-            "author": self.author,
-            "waitForAnswer": True,
-            "disableHumanFeedback": self.disable_human_feedback,
-        }
-
     async def send(self) -> Union[List[AskFileResponse], None]:
         """
         Sends the message to request a file from the user to the UI and waits for the reply.
         """
         trace_event("send_ask_file")
 
+        if not self.created_at:
+            self.created_at = datetime.utcnow().isoformat()
+
         if self.streaming:
             self.streaming = False
 
         if config.code.author_rename:
             self.author = await config.code.author_rename(self.author)
 
-        msg_dict = await self._create()
+        self.wait_for_answer = True
+
+        step_dict = await self._create()
 
         spec = AskFileSpec(
             type="file",
@@ -494,10 +455,24 @@ class AskFileMessage(AskMessageBase):
             timeout=self.timeout,
         )
 
-        res = await context.emitter.send_ask_user(msg_dict, spec, self.raise_on_timeout)
+        res = cast(
+            Union[None, List[FileDict]],
+            await context.emitter.send_ask_user(step_dict, spec, self.raise_on_timeout),
+        )
+
+        self.wait_for_answer = False
 
         if res:
-            return [AskFileResponse(**r) for r in res]
+            return [
+                AskFileResponse(
+                    id=r["id"],
+                    name=r["name"],
+                    path=str(r["path"]),
+                    size=r["size"],
+                    type=r["type"],
+                )
+                for r in res
+            ]
         else:
             return None
 
@@ -513,55 +488,49 @@ class AskActionMessage(AskMessageBase):
         content: str,
         actions: List[Action],
         author=config.ui.name,
-        disable_human_feedback=False,
+        disable_feedback=False,
         timeout=90,
         raise_on_timeout=False,
     ):
         self.content = content
         self.actions = actions
         self.author = author
-        self.disable_human_feedback = disable_human_feedback
+        self.disable_feedback = disable_feedback
         self.timeout = timeout
         self.raise_on_timeout = raise_on_timeout
 
         super().__post_init__()
 
-    def to_dict(self):
-        return {
-            "id": self.id,
-            "createdAt": self.created_at,
-            "content": self.content,
-            "author": self.author,
-            "waitForAnswer": True,
-            "disableHumanFeedback": self.disable_human_feedback,
-            "timeout": self.timeout,
-            "raiseOnTimeout": self.raise_on_timeout,
-        }
-
     async def send(self) -> Union[AskActionResponse, None]:
         """
         Sends the question to ask to the UI and waits for the reply
         """
         trace_event("send_ask_action")
 
+        if not self.created_at:
+            self.created_at = datetime.utcnow().isoformat()
+
         if self.streaming:
             self.streaming = False
 
         if config.code.author_rename:
             self.author = await config.code.author_rename(self.author)
 
-        msg_dict = await self._create()
+        self.wait_for_answer = True
+
+        step_dict = await self._create()
+
         action_keys = []
 
         for action in self.actions:
             action_keys.append(action.id)
-            await action.send(for_id=str(msg_dict["id"]))
+            await action.send(for_id=str(step_dict["id"]))
 
         spec = AskActionSpec(type="action", timeout=self.timeout, keys=action_keys)
 
         res = cast(
             Union[AskActionResponse, None],
-            await context.emitter.send_ask_user(msg_dict, spec, self.raise_on_timeout),
+            await context.emitter.send_ask_user(step_dict, spec, self.raise_on_timeout),
         )
 
         for action in self.actions:
@@ -570,6 +539,9 @@ class AskActionMessage(AskMessageBase):
             self.content = "Timed out: no action was taken"
         else:
             self.content = f'**Selected action:** {res["label"]}'
+
+        self.wait_for_answer = False
+
         await self.update()
 
         return res