455 lines
13 KiB
Python
455 lines
13 KiB
Python
|
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)
|