chainlit 0.7.604rc2__py3-none-any.whl → 1.0.0rc0__py3-none-any.whl

chainlit/element.py

@@ -2,16 +2,14 @@ import json
 import uuid
 from enum import Enum
 from io import BytesIO
-from typing import Any, ClassVar, Dict, List, Optional, TypeVar, Union, cast
+from typing import Any, ClassVar, List, Literal, Optional, TypedDict, TypeVar, Union
 
-import aiofiles
 import filetype
-from chainlit.client.base import ElementDict, ElementDisplay, ElementSize, ElementType
-from chainlit.client.cloud import ChainlitCloudClient
 from chainlit.context import context
-from chainlit.data import chainlit_client
+from chainlit.data import get_data_layer
 from chainlit.logger import logger
 from chainlit.telemetry import trace_event
+from chainlit.types import FileDict
 from pydantic.dataclasses import Field, dataclass
 from syncer import asyncio
 
@@ -21,16 +19,38 @@ mime_types = {
     "plotly": "application/json",
 }
 
+ElementType = Literal[
+    "image", "avatar", "text", "pdf", "tasklist", "audio", "video", "file", "plotly"
+]
+ElementDisplay = Literal["inline", "side", "page"]
+ElementSize = Literal["small", "medium", "large"]
+
+
+class ElementDict(TypedDict):
+    id: str
+    threadId: Optional[str]
+    type: ElementType
+    chainlitKey: Optional[str]
+    url: Optional[str]
+    objectKey: Optional[str]
+    name: str
+    display: ElementDisplay
+    size: Optional[ElementSize]
+    language: Optional[str]
+    forId: Optional[str]
+    mime: Optional[str]
+
 
 @dataclass
 class Element:
     # The type of the element. This will be used to determine how to display the element in the UI.
     type: ClassVar[ElementType]
-
+    # Name of the element, this will be used to reference the element in the UI.
+    name: str
     # The ID of the element. This is set automatically when the element is sent to the UI.
     id: str = Field(default_factory=lambda: str(uuid.uuid4()))
-    # Name of the element, this will be used to reference the element in the UI.
-    name: Optional[str] = None
+    # The key of the element hosted on Chainlit.
+    chainlit_key: Optional[str] = None
     # The URL of the element if already hosted somehwere else.
     url: Optional[str] = None
     # The S3 object key.
@@ -44,7 +64,7 @@ class Element:
     # Controls element size
     size: Optional[ElementSize] = None
     # The ID of the message this element is associated with.
-    for_ids: List[str] = Field(default_factory=list)
+    for_id: Optional[str] = None
     # The language, if relevant
     language: Optional[str] = None
     # Mime type, infered based on content if not provided
@@ -53,6 +73,8 @@ class Element:
     def __post_init__(self) -> None:
         trace_event(f"init {self.__class__.__name__}")
         self.persisted = False
+        self.updatable = False
+        self.thread_id = context.session.thread_id
 
         if not self.url and not self.path and not self.content:
             raise ValueError("Must provide url, path or content to instantiate element")
@@ -61,136 +83,92 @@ class Element:
         _dict = ElementDict(
             {
                 "id": self.id,
+                "threadId": self.thread_id,
                 "type": self.type,
-                "url": self.url or "",
-                "name": self.name or "",
+                "url": self.url,
+                "chainlitKey": self.chainlit_key,
+                "name": self.name,
                 "display": self.display,
                 "objectKey": getattr(self, "object_key", None),
                 "size": getattr(self, "size", None),
                 "language": getattr(self, "language", None),
-                "forIds": getattr(self, "for_ids", None),
+                "forId": getattr(self, "for_id", None),
                 "mime": getattr(self, "mime", None),
-                "conversationId": None,
             }
         )
         return _dict
 
     @classmethod
-    def from_dict(self, _dict: Dict):
-        if "image" in _dict.get("mime", ""):
+    def from_dict(self, _dict: FileDict):
+        type = _dict.get("type", "")
+        if "image" in type and "svg" not in type:
             return Image(
                 id=_dict.get("id", str(uuid.uuid4())),
-                content=_dict.get("content"),
-                name=_dict.get("name"),
-                url=_dict.get("url"),
-                display=_dict.get("display", "inline"),
-                mime=_dict.get("mime"),
+                name=_dict.get("name", ""),
+                path=str(_dict.get("path")),
+                chainlit_key=_dict.get("id"),
+                display="inline",
+                mime=type,
             )
         else:
             return File(
                 id=_dict.get("id", str(uuid.uuid4())),
-                content=_dict.get("content"),
-                name=_dict.get("name"),
-                url=_dict.get("url"),
-                language=_dict.get("language"),
-                display=_dict.get("display", "inline"),
-                size=_dict.get("size"),
-                mime=_dict.get("mime"),
+                name=_dict.get("name", ""),
+                path=str(_dict.get("path")),
+                chainlit_key=_dict.get("id"),
+                display="inline",
+                mime=type,
             )
 
-    async def with_conversation_id(self):
-        _dict = self.to_dict()
-        _dict["conversationId"] = await context.session.get_conversation_id()
-        return _dict
-
-    async def preprocess_content(self):
-        pass
-
-    async def load(self):
-        if self.path:
-            async with aiofiles.open(self.path, "rb") as f:
-                self.content = await f.read()
-        else:
-            raise ValueError("Must provide path or content to load element")
-
-    async def persist(self, client: ChainlitCloudClient) -> Optional[ElementDict]:
-        if not self.url and self.content and not self.persisted:
-            conversation_id = await context.session.get_conversation_id()
-            upload_res = await client.upload_element(
+    async def _create(self) -> bool:
+        if (self.persisted or self.url) and not self.updatable:
+            return True
+        if data_layer := get_data_layer():
+            try:
+                asyncio.create_task(data_layer.create_element(self))
+            except Exception as e:
+                logger.error(f"Failed to create element: {str(e)}")
+        if not self.chainlit_key or self.updatable:
+            file_dict = await context.session.persist_file(
+                name=self.name,
+                path=self.path,
                 content=self.content,
                 mime=self.mime or "",
-                conversation_id=conversation_id,
             )
-            self.url = upload_res["url"]
-            self.object_key = upload_res["object_key"]
-        element_dict = await self.with_conversation_id()
+            self.chainlit_key = file_dict["id"]
 
-        asyncio.create_task(self._persist(element_dict))
+        self.persisted = True
 
-        return element_dict
-
-    async def _persist(self, element: ElementDict):
-        if not chainlit_client:
-            return
-
-        try:
-            if self.persisted:
-                await chainlit_client.update_element(element)
-            else:
-                await chainlit_client.create_element(element)
-                self.persisted = True
-        except Exception as e:
-            logger.error(f"Failed to persist element: {str(e)}")
-
-    async def before_emit(self, element: Dict) -> Dict:
-        return element
+        return True
 
     async def remove(self):
         trace_event(f"remove {self.__class__.__name__}")
+        data_layer = get_data_layer()
+        if data_layer and self.persisted:
+            await data_layer.delete_element(self.id)
         await context.emitter.emit("remove_element", {"id": self.id})
 
-    async def send(self, for_id: Optional[str] = None):
-        if not self.content and not self.url and self.path:
-            await self.load()
+    async def send(self, for_id: str):
+        if self.persisted and not self.updatable:
+            return
 
-        await self.preprocess_content()
+        self.for_id = for_id
 
         if not self.mime:
             # Only guess the mime type when the content is binary
             self.mime = (
                 mime_types[self.type]
                 if self.type in mime_types
-                else filetype.guess_mime(self.content)
+                else filetype.guess_mime(self.path or self.content)
             )
 
-        if for_id and for_id not in self.for_ids:
-            self.for_ids.append(for_id)
+        await self._create()
 
-        # We have a client, persist the element
-        if chainlit_client:
-            element_dict = await self.persist(chainlit_client)
-            if element_dict:
-                self.id = element_dict["id"]
+        if not self.url and not self.chainlit_key:
+            raise ValueError("Must provide url or chainlit key to send element")
 
-        elif not self.url and not self.content:
-            raise ValueError("Must provide url or content to send element")
-
-        emit_dict = cast(Dict, self.to_dict())
-
-        # Adding this out of to_dict since the dict will be persisted in the DB
-        emit_dict["content"] = self.content
-
-        # Element was already sent
-        if len(self.for_ids) > 1:
-            trace_event(f"update {self.__class__.__name__}")
-            await context.emitter.emit(
-                "update_element",
-                {"id": self.id, "forIds": self.for_ids},
-            )
-        else:
-            trace_event(f"send {self.__class__.__name__}")
-            emit_dict = await self.before_emit(emit_dict)
-            await context.emitter.emit("element", emit_dict)
+        trace_event(f"send {self.__class__.__name__}")
+        await context.emitter.emit("element", self.to_dict())
 
 
 ElementBased = TypeVar("ElementBased", bound=Element)
@@ -208,23 +186,7 @@ class Avatar(Element):
     type: ClassVar[ElementType] = "avatar"
 
     async def send(self):
-        element = None
-
-        if not self.content and not self.url and self.path:
-            await self.load()
-
-        if not self.url and not self.content:
-            raise ValueError("Must provide url or content to send element")
-
-        element = self.to_dict()
-
-        # Adding this out of to_dict since the dict will be persisted in the DB
-        element["content"] = self.content
-
-        if element:
-            trace_event(f"send {self.__class__.__name__}")
-            element = await self.before_emit(element)
-            await context.emitter.emit("element", element)
+        await super().send(for_id="")
 
 
 @dataclass
@@ -232,15 +194,8 @@ class Text(Element):
     """Useful to send a text (not a message) to the UI."""
 
     type: ClassVar[ElementType] = "text"
-
-    content: bytes = b""
     language: Optional[str] = None
 
-    async def before_emit(self, text_element):
-        if "content" in text_element and isinstance(text_element["content"], bytes):
-            text_element["content"] = text_element["content"].decode("utf-8")
-        return text_element
-
 
 @dataclass
 class Pdf(Element):
@@ -312,12 +267,20 @@ class TaskList(Element):
     name: str = "tasklist"
     content: str = "dummy content to pass validation"
 
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.updatable = True
+
     async def add_task(self, task: Task):
         self.tasks.append(task)
 
     async def update(self):
         await self.send()
 
+    async def send(self):
+        await self.preprocess_content()
+        await super().send(for_id="")
+
     async def preprocess_content(self):
         # serialize enum
         tasks = [

chainlit/emitter.py

@@ -1,12 +1,22 @@
 import asyncio
 import uuid
-from typing import Any, Dict, Optional
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Union, cast
 
-from chainlit.client.base import ConversationDict, MessageDict
-from chainlit.element import Element
+from chainlit.data import get_data_layer
+from chainlit.element import Element, File
 from chainlit.message import Message
 from chainlit.session import BaseSession, WebsocketSession
-from chainlit.types import AskSpec, UIMessagePayload
+from chainlit.step import StepDict
+from chainlit.types import (
+    AskActionResponse,
+    AskSpec,
+    FileDict,
+    FileReference,
+    ThreadDict,
+    UIMessagePayload,
+)
+from chainlit.user import PersistedUser
 from socketio.exceptions import TimeoutError
 
 
@@ -30,19 +40,19 @@ class BaseChainlitEmitter:
         """Stub method to get the 'ask_user' property from the session."""
         pass
 
-    async def resume_conversation(self, conv_dict: ConversationDict):
-        """Stub method to resume a conversation."""
+    async def resume_thread(self, thread_dict: ThreadDict):
+        """Stub method to resume a thread."""
         pass
 
-    async def send_message(self, msg_dict: dict):
+    async def send_step(self, step_dict: StepDict):
         """Stub method to send a message to the UI."""
         pass
 
-    async def update_message(self, msg_dict: dict):
+    async def update_step(self, step_dict: StepDict):
         """Stub method to update a message in the UI."""
         pass
 
-    async def delete_message(self, msg_dict: dict):
+    async def delete_step(self, step_dict: StepDict):
         """Stub method to delete a message in the UI."""
         pass
 
@@ -54,15 +64,17 @@ class BaseChainlitEmitter:
         """Stub method to clear the prompt from the UI."""
         pass
 
-    async def init_conversation(self, msg_dict: MessageDict):
-        """Signal the UI that a new conversation (with a user message) exists"""
+    async def init_thread(self, step_dict: StepDict):
+        """Signal the UI that a new thread (with a user message) exists"""
         pass
 
     async def process_user_message(self, payload: UIMessagePayload) -> Message:
         """Stub method to process user message."""
         return Message(content="")
 
-    async def send_ask_user(self, msg_dict: dict, spec, raise_on_timeout=False):
+    async def send_ask_user(
+        self, step_dict: StepDict, spec: AskSpec, raise_on_timeout=False
+    ) -> Optional[Union["StepDict", "AskActionResponse", List["FileDict"]]]:
         """Stub method to send a prompt to the UI and wait for a response."""
         pass
 
@@ -78,7 +90,7 @@ class BaseChainlitEmitter:
         """Stub method to send a task end signal to the UI."""
         pass
 
-    async def stream_start(self, msg_dict: dict):
+    async def stream_start(self, step_dict: StepDict):
         """Stub method to send a stream start signal to the UI."""
         pass
 
@@ -90,6 +102,12 @@ class BaseChainlitEmitter:
         """Stub method to set chat settings."""
         pass
 
+    async def send_action_response(
+        self, id: str, status: bool, response: Optional[str] = None
+    ):
+        """Send an action response to the UI."""
+        pass
+
 
 class ChainlitEmitter(BaseChainlitEmitter):
     """
@@ -123,23 +141,21 @@ class ChainlitEmitter(BaseChainlitEmitter):
         """Get the 'ask_user' property from the session."""
         return self._get_session_property("ask_user")
 
-    def resume_conversation(self, conv_dict: ConversationDict):
-        """Send a conversation to the UI to resume it"""
-        return self.emit("resume_conversation", conv_dict)
+    def resume_thread(self, thread_dict: ThreadDict):
+        """Send a thread to the UI to resume it"""
+        return self.emit("resume_thread", thread_dict)
 
-    def send_message(self, msg_dict: Dict):
+    def send_step(self, step_dict: StepDict):
         """Send a message to the UI."""
-        return self.emit("new_message", msg_dict)
+        return self.emit("new_message", step_dict)
 
-    def update_message(self, msg_dict: Dict):
+    def update_step(self, step_dict: StepDict):
         """Update a message in the UI."""
+        return self.emit("update_message", step_dict)
 
-        return self.emit("update_message", msg_dict)
-
-    def delete_message(self, msg_dict):
+    def delete_step(self, step_dict: StepDict):
         """Delete a message in the UI."""
-
-        return self.emit("delete_message", msg_dict)
+        return self.emit("delete_message", step_dict)
 
     def send_ask_timeout(self):
         """Send a prompt timeout message to the UI."""
@@ -151,22 +167,44 @@ class ChainlitEmitter(BaseChainlitEmitter):
 
         return self.emit("clear_ask", {})
 
-    def init_conversation(self, message: MessageDict):
-        """Signal the UI that a new conversation (with a user message) exists"""
+    async def init_thread(self, step: StepDict):
+        """Signal the UI that a new thread (with a user message) exists"""
+        if data_layer := get_data_layer():
+            if isinstance(self.session.user, PersistedUser):
+                user_id = self.session.user.id
+            else:
+                user_id = None
+            await data_layer.update_thread(
+                thread_id=self.session.thread_id,
+                user_id=user_id,
+                metadata={"name": step["output"]},
+            )
+            await self.session.flush_method_queue()
 
-        return self.emit("init_conversation", message)
+        await self.emit("init_thread", step)
 
     async def process_user_message(self, payload: UIMessagePayload):
-        message_dict = payload["message"]
-        files = payload["files"]
-        # Temporary UUID generated by the frontend should use v4
-        assert uuid.UUID(message_dict["id"]).version == 4
+        step_dict = payload["message"]
+        file_refs = payload["fileReferences"]
+        # UUID generated by the frontend should use v4
+        assert uuid.UUID(step_dict["id"]).version == 4
 
-        message = Message.from_dict(message_dict)
+        message = Message.from_dict(step_dict)
+        # Overwrite the created_at timestamp with the current time
+        message.created_at = datetime.utcnow().isoformat()
 
         asyncio.create_task(message._create())
 
-        if files:
+        if not self.session.has_user_message:
+            self.session.has_user_message = True
+            asyncio.create_task(self.init_thread(message.to_dict()))
+
+        if file_refs:
+            files = [
+                self.session.files[file["id"]]
+                for file in file_refs
+                if file["id"] in self.session.files
+            ]
             file_elements = [Element.from_dict(file) for file in files]
             message.elements = file_elements
 
@@ -176,38 +214,60 @@ class ChainlitEmitter(BaseChainlitEmitter):
 
             asyncio.create_task(send_elements())
 
-        if not self.session.has_user_message:
-            self.session.has_user_message = True
-            await self.init_conversation(await message.with_conversation_id())
-
         self.session.root_message = message
 
         return message
 
     async def send_ask_user(
-        self, msg_dict: Dict, spec: AskSpec, raise_on_timeout=False
+        self, step_dict: StepDict, spec: AskSpec, raise_on_timeout=False
     ):
         """Send a prompt to the UI and wait for a response."""
 
         try:
             # Send the prompt to the UI
-            res = await self.ask_user(
-                {"msg": msg_dict, "spec": spec.to_dict()}, spec.timeout
-            )  # type: Optional["MessageDict"]
+            user_res = await self.ask_user(
+                {"msg": step_dict, "spec": spec.to_dict()}, spec.timeout
+            )  # type: Optional[Union["StepDict", "AskActionResponse", List["FileReference"]]]
 
             # End the task temporarily so that the User can answer the prompt
             await self.task_end()
 
-            if res:
-                # If cloud is enabled, store the response in the database/S3
+            final_res: Optional[
+                Union["StepDict", "AskActionResponse", List["FileDict"]]
+            ] = None
+
+            if user_res:
                 if spec.type == "text":
-                    await self.process_user_message({"message": res, "files": None})
+                    message_dict_res = cast(StepDict, user_res)
+                    await self.process_user_message(
+                        {"message": message_dict_res, "fileReferences": None}
+                    )
+                    final_res = message_dict_res
                 elif spec.type == "file":
-                    # TODO: upload file to S3
-                    pass
-
+                    file_refs = cast(List[FileReference], user_res)
+                    files = [
+                        self.session.files[file["id"]]
+                        for file in file_refs
+                        if file["id"] in self.session.files
+                    ]
+                    final_res = files
+                    if get_data_layer():
+                        coros = [
+                            File(
+                                name=file["name"],
+                                path=str(file["path"]),
+                                mime=file["type"],
+                                chainlit_key=file["id"],
+                                for_id=step_dict["id"],
+                            )._create()
+                            for file in files
+                        ]
+                        await asyncio.gather(*coros)
+                elif spec.type == "action":
+                    action_res = cast(AskActionResponse, user_res)
+                    final_res = action_res
             await self.clear_ask()
-            return res
+            return final_res
         except TimeoutError as e:
             await self.send_ask_timeout()
 
@@ -231,11 +291,11 @@ class ChainlitEmitter(BaseChainlitEmitter):
         """Send a task end signal to the UI."""
         return self.emit("task_end", {})
 
-    def stream_start(self, msg_dict: Dict):
+    def stream_start(self, step_dict: StepDict):
         """Send a stream start signal to the UI."""
         return self.emit(
             "stream_start",
-            msg_dict,
+            step_dict,
         )
 
     def send_token(self, id: str, token: str, is_sequence=False):
@@ -246,3 +306,10 @@ class ChainlitEmitter(BaseChainlitEmitter):
 
     def set_chat_settings(self, settings: Dict[str, Any]):
         self.session.chat_settings = settings
+
+    def send_action_response(
+        self, id: str, status: bool, response: Optional[str] = None
+    ):
+        return self.emit(
+            "action_response", {"id": id, "status": status, "response": response}
+        )