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

chainlit/element.py

@@ -1,454 +0,0 @@
-import json
-import mimetypes
-import uuid
-from enum import Enum
-from io import BytesIO
-from typing import (
-    Any,
-    ClassVar,
-    Dict,
-    List,
-    Literal,
-    Optional,
-    TypedDict,
-    TypeVar,
-    Union,
-)
-
-import filetype
-from pydantic import Field
-from pydantic.dataclasses import dataclass
-from syncer import asyncio
-
-from chainlit.context import context
-from chainlit.data import get_data_layer
-from chainlit.logger import logger
-
-mime_types = {
-    "text": "text/plain",
-    "tasklist": "application/json",
-    "plotly": "application/json",
-}
-
-ElementType = Literal[
-    "image",
-    "text",
-    "pdf",
-    "tasklist",
-    "audio",
-    "video",
-    "file",
-    "plotly",
-    "dataframe",
-    "custom",
-]
-ElementDisplay = Literal["inline", "side", "page"]
-ElementSize = Literal["small", "medium", "large"]
-
-
-class ElementDict(TypedDict, total=False):
-    id: str
-    threadId: Optional[str]
-    type: ElementType
-    chainlitKey: Optional[str]
-    path: Optional[str]
-    url: Optional[str]
-    objectKey: Optional[str]
-    name: str
-    display: ElementDisplay
-    size: Optional[ElementSize]
-    language: Optional[str]
-    page: Optional[int]
-    props: Optional[Dict]
-    autoPlay: Optional[bool]
-    playerConfig: Optional[dict]
-    forId: Optional[str]
-    mime: Optional[str]
-
-
-@dataclass
-class Element:
-    # Thread id
-    thread_id: str = Field(default_factory=lambda: context.session.thread_id)
-    # 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()))
-    # The key of the element hosted on Chainlit.
-    chainlit_key: Optional[str] = None
-    # The URL of the element if already hosted somewhere else.
-    url: Optional[str] = None
-    # The S3 object key.
-    object_key: Optional[str] = None
-    # The local path of the element.
-    path: Optional[str] = None
-    # The byte content of the element.
-    content: Optional[Union[bytes, str]] = None
-    # Controls how the image element should be displayed in the UI. Choices are “side” (default), “inline”, or “page”.
-    display: ElementDisplay = Field(default="inline")
-    # Controls element size
-    size: Optional[ElementSize] = None
-    # The ID of the message this element is associated with.
-    for_id: Optional[str] = None
-    # The language, if relevant
-    language: Optional[str] = None
-    # Mime type, inferred based on content if not provided
-    mime: Optional[str] = None
-
-    def __post_init__(self) -> None:
-        self.persisted = False
-        self.updatable = False
-
-        if not self.url and not self.path and not self.content:
-            raise ValueError("Must provide url, path or content to instantiate element")
-
-    def to_dict(self) -> ElementDict:
-        _dict = ElementDict(
-            {
-                "id": self.id,
-                "threadId": self.thread_id,
-                "type": self.type,
-                "url": self.url,
-                "chainlitKey": self.chainlit_key,
-                "name": self.name,
-                "display": self.display,
-                "objectKey": getattr(self, "object_key", None),
-                "size": getattr(self, "size", None),
-                "props": getattr(self, "props", None),
-                "page": getattr(self, "page", None),
-                "autoPlay": getattr(self, "auto_play", None),
-                "playerConfig": getattr(self, "player_config", None),
-                "language": getattr(self, "language", None),
-                "forId": getattr(self, "for_id", None),
-                "mime": getattr(self, "mime", None),
-            }
-        )
-        return _dict
-
-    @classmethod
-    def from_dict(cls, e_dict: ElementDict):
-        """
-        Create an Element instance from a dictionary representation.
-
-        Args:
-            _dict (ElementDict): Dictionary containing element data
-
-        Returns:
-            Element: An instance of the appropriate Element subclass
-        """
-        element_id = e_dict.get("id", str(uuid.uuid4()))
-        for_id = e_dict.get("forId")
-        name = e_dict.get("name", "")
-        type = e_dict.get("type", "file")
-        path = str(e_dict.get("path")) if e_dict.get("path") else None
-        url = str(e_dict.get("url")) if e_dict.get("url") else None
-        content = str(e_dict.get("content")) if e_dict.get("content") else None
-        object_key = e_dict.get("objectKey")
-        chainlit_key = e_dict.get("chainlitKey")
-        display = e_dict.get("display", "inline")
-        mime_type = e_dict.get("mime", "")
-
-        # Common parameters for all element types
-        common_params = {
-            "id": element_id,
-            "for_id": for_id,
-            "name": name,
-            "content": content,
-            "path": path,
-            "url": url,
-            "object_key": object_key,
-            "chainlit_key": chainlit_key,
-            "display": display,
-            "mime": mime_type,
-        }
-
-        if type == "image":
-            return Image(size="medium", **common_params)  # type: ignore[arg-type]
-
-        elif type == "audio":
-            return Audio(auto_play=e_dict.get("autoPlay", False), **common_params)  # type: ignore[arg-type]
-
-        elif type == "video":
-            return Video(
-                player_config=e_dict.get("playerConfig"),
-                **common_params,  # type: ignore[arg-type]
-            )
-
-        elif type == "plotly":
-            return Plotly(size=e_dict.get("size", "medium"), **common_params)  # type: ignore[arg-type]
-
-        elif type == "custom":
-            return CustomElement(props=e_dict.get("props", {}), **common_params)  # type: ignore[arg-type]
-        else:
-            # Default to File for any other type
-            return File(**common_params)  # type: ignore[arg-type]
-
-    @classmethod
-    def infer_type_from_mime(cls, mime_type: str):
-        """Infer the element type from a mime type. Useful to know which element to instantiate from a file upload."""
-        if "image" in mime_type:
-            return "image"
-
-        elif mime_type == "application/pdf":
-            return "pdf"
-
-        elif "audio" in mime_type:
-            return "audio"
-
-        elif "video" in mime_type:
-            return "video"
-
-        else:
-            return "file"
-
-    async def _create(self, persist=True) -> bool:
-        if self.persisted and not self.updatable:
-            return True
-
-        if (data_layer := get_data_layer()) and persist:
-            try:
-                asyncio.create_task(data_layer.create_element(self))
-            except Exception as e:
-                logger.error(f"Failed to create element: {e!s}")
-        if not self.url and (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 "",
-            )
-            self.chainlit_key = file_dict["id"]
-
-        self.persisted = True
-
-        return True
-
-    async def remove(self):
-        data_layer = get_data_layer()
-        if data_layer:
-            await data_layer.delete_element(self.id, self.thread_id)
-        await context.emitter.emit("remove_element", {"id": self.id})
-
-    async def send(self, for_id: str, persist=True):
-        self.for_id = for_id
-
-        if not self.mime:
-            if self.type in mime_types:
-                self.mime = mime_types[self.type]
-            elif self.path or isinstance(self.content, (bytes, bytearray)):
-                file_type = filetype.guess(self.path or self.content)
-                if file_type:
-                    self.mime = file_type.mime
-            elif self.url:
-                self.mime = mimetypes.guess_type(self.url)[0]
-
-        await self._create(persist=persist)
-
-        if not self.url and not self.chainlit_key:
-            raise ValueError("Must provide url or chainlit key to send element")
-
-        await context.emitter.send_element(self.to_dict())
-
-
-ElementBased = TypeVar("ElementBased", bound=Element)
-
-
-@dataclass
-class Image(Element):
-    type: ClassVar[ElementType] = "image"
-
-    size: ElementSize = "medium"
-
-
-@dataclass
-class Text(Element):
-    """Useful to send a text (not a message) to the UI."""
-
-    type: ClassVar[ElementType] = "text"
-    language: Optional[str] = None
-
-
-@dataclass
-class Pdf(Element):
-    """Useful to send a pdf to the UI."""
-
-    mime: str = "application/pdf"
-    page: Optional[int] = None
-    type: ClassVar[ElementType] = "pdf"
-
-
-@dataclass
-class Pyplot(Element):
-    """Useful to send a pyplot to the UI."""
-
-    # We reuse the frontend image element to display the chart
-    type: ClassVar[ElementType] = "image"
-
-    size: ElementSize = "medium"
-    # The type is set to Any because the figure is not serializable
-    # and its actual type is checked in __post_init__.
-    figure: Any = None
-
-    def __post_init__(self) -> None:
-        from matplotlib.figure import Figure
-
-        if not isinstance(self.figure, Figure):
-            raise TypeError("figure must be a matplotlib.figure.Figure")
-
-        image = BytesIO()
-        self.figure.savefig(
-            image, dpi=200, bbox_inches="tight", backend="Agg", format="png"
-        )
-        self.content = image.getvalue()
-
-        super().__post_init__()
-
-
-class TaskStatus(Enum):
-    READY = "ready"
-    RUNNING = "running"
-    FAILED = "failed"
-    DONE = "done"
-
-
-@dataclass
-class Task:
-    title: str
-    status: TaskStatus = TaskStatus.READY
-    forId: Optional[str] = None
-
-    def __init__(
-        self,
-        title: str,
-        status: TaskStatus = TaskStatus.READY,
-        forId: Optional[str] = None,
-    ):
-        self.title = title
-        self.status = status
-        self.forId = forId
-
-
-@dataclass
-class TaskList(Element):
-    type: ClassVar[ElementType] = "tasklist"
-    tasks: List[Task] = Field(default_factory=list, exclude=True)
-    status: str = "Ready"
-    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 = [
-            {"title": task.title, "status": task.status.value, "forId": task.forId}
-            for task in self.tasks
-        ]
-
-        # store stringified json in content so that it's correctly stored in the database
-        self.content = json.dumps(
-            {
-                "status": self.status,
-                "tasks": tasks,
-            },
-            indent=4,
-            ensure_ascii=False,
-        )
-
-
-@dataclass
-class Audio(Element):
-    type: ClassVar[ElementType] = "audio"
-    auto_play: bool = False
-
-
-@dataclass
-class Video(Element):
-    type: ClassVar[ElementType] = "video"
-
-    size: ElementSize = "medium"
-    # Override settings for each type of player in ReactPlayer
-    # https://github.com/cookpete/react-player?tab=readme-ov-file#config-prop
-    player_config: Optional[dict] = None
-
-
-@dataclass
-class File(Element):
-    type: ClassVar[ElementType] = "file"
-
-
-@dataclass
-class Plotly(Element):
-    """Useful to send a plotly to the UI."""
-
-    type: ClassVar[ElementType] = "plotly"
-
-    size: ElementSize = "medium"
-    # The type is set to Any because the figure is not serializable
-    # and its actual type is checked in __post_init__.
-    figure: Any = None
-    content: str = ""
-
-    def __post_init__(self) -> None:
-        from plotly import graph_objects as go, io as pio
-
-        if not isinstance(self.figure, go.Figure):
-            raise TypeError("figure must be a plotly.graph_objects.Figure")
-
-        self.figure.layout.autosize = True
-        self.figure.layout.width = None
-        self.figure.layout.height = None
-        self.content = pio.to_json(self.figure, validate=True)
-        self.mime = "application/json"
-
-        super().__post_init__()
-
-
-@dataclass
-class Dataframe(Element):
-    """Useful to send a pandas DataFrame to the UI."""
-
-    type: ClassVar[ElementType] = "dataframe"
-    size: ElementSize = "large"
-    data: Any = None  # The type is Any because it is checked in __post_init__.
-
-    def __post_init__(self) -> None:
-        """Ensures the data is a pandas DataFrame and converts it to JSON."""
-        from pandas import DataFrame
-
-        if not isinstance(self.data, DataFrame):
-            raise TypeError("data must be a pandas.DataFrame")
-
-        self.content = self.data.to_json(orient="split", date_format="iso")
-        super().__post_init__()
-
-
-@dataclass
-class CustomElement(Element):
-    """Useful to send a custom element to the UI."""
-
-    type: ClassVar[ElementType] = "custom"
-    mime: str = "application/json"
-    props: Dict = Field(default_factory=dict)
-
-    def __post_init__(self) -> None:
-        self.content = json.dumps(self.props)
-        super().__post_init__()
-        self.updatable = True
-
-    async def update(self):
-        await super().send(self.for_id)