add size option for Image Elements and language option for Text Element

Author: willydouhard

cypress/e2e/global_elements/spec.cy.ts

@@ -9,7 +9,7 @@ describe("Global Elements", () => {
     cy.get(".message").should("have.length", 1);
 
     // Inlined
-    cy.get(".inlined-image").should("have.length", 1);
+    cy.get(".inline-image").should("have.length", 1);
     cy.get(".element-link").eq(0).should("contain", "text1");
     cy.get(".element-link").eq(0).click();
 

cypress/e2e/scoped_elements/spec.cy.ts

@@ -8,10 +8,10 @@ describe("Scoped Elements", () => {
   it("should be able to display inlined, side and page elements", () => {
     cy.get(".message").should("have.length", 2);
 
-    cy.get(".message").eq(0).find(".inlined-image").should("have.length", 0);
+    cy.get(".message").eq(0).find(".inline-image").should("have.length", 0);
     cy.get(".message").eq(0).find(".element-link").should("have.length", 0);
 
-    cy.get(".message").eq(1).find(".inlined-image").should("have.length", 1);
+    cy.get(".message").eq(1).find(".inline-image").should("have.length", 1);
     cy.get(".message").eq(1).find(".element-link").should("have.length", 2);
   });
 });

src/chainlit/client.py

@@ -3,7 +3,7 @@
 from abc import ABC, abstractmethod
 import uuid
 import requests
-from chainlit.types import ElementType
+from chainlit.types import ElementType, ElementSize
 from chainlit.logger import logger
 
 
@@ -30,6 +30,8 @@ def create_element(
         url: str,
         name: str,
         display: str,
+        size: ElementSize = None,
+        language: str = None,
         for_id: str = None,
     ) -> Dict[str, Any]:
         pass
@@ -120,7 +122,14 @@ def create_message(self, variables: Dict[str, Any]) -> int:
         return int(res["data"]["createMessage"]["id"])
 
     def create_element(
-        self, type: ElementType, url: str, name: str, display: str, for_id: str = None
+        self,
+        type: ElementType,
+        url: str,
+        name: str,
+        display: str,
+        size: ElementSize = None,
+        language: str = None,
+        for_id: str = None,
     ) -> Dict[str, Any]:
         c_id = self.get_conversation_id()
 
@@ -129,13 +138,15 @@ def create_element(
             return None
 
         mutation = """
-        mutation ($conversationId: ID!, $type: String!, $url: String!, $name: String!, $display: String!, $forId: String) {
-            createElement(conversationId: $conversationId, type: $type, url: $url, name: $name, display: $display, forId: $forId) {
+        mutation ($conversationId: ID!, $type: String!, $url: String!, $name: String!, $display: String!, $size: String, $language: String, $forId: String) {
+            createElement(conversationId: $conversationId, type: $type, url: $url, name: $name, display: $display, size: $size, language: $language, forId: $forId) {
                 id,
                 type,
                 url,
                 name,
                 display,
+                size,
+                language,
                 forId
             }
         }
@@ -146,6 +157,8 @@ def create_element(
             "url": url,
             "name": name,
             "display": display,
+            "size": size,
+            "language": language,
             "forId": for_id,
         }
         res = self.mutation(mutation, variables)

src/chainlit/element.py

@@ -4,7 +4,7 @@
 from abc import ABC, abstractmethod
 from chainlit.sdk import get_sdk, BaseClient
 from chainlit.telemetry import trace_event
-from chainlit.types import ElementType, ElementDisplay
+from chainlit.types import ElementType, ElementDisplay, ElementSize
 
 
 @dataclass_json
@@ -43,20 +43,23 @@ def send(self, for_id: str = None):
 
 
 @dataclass
-class LocalElementBase:
-    content: bytes
+class LocalElement(Element):
+    content: bytes = None
 
-
-@dataclass
-class LocalElement(Element, LocalElementBase):
     def persist(self, client: BaseClient, for_id: str = None):
+        if not self.content:
+            raise ValueError("Must provide content")
         url = client.upload_element(content=self.content)
         if url:
+            size = getattr(self, "size", None)
+            language = getattr(self, "language", None)
             element = client.create_element(
                 name=self.name,
                 url=url,
                 type=self.type,
                 display=self.display,
+                size=size,
+                language=language,
                 for_id=for_id,
             )
             return element
@@ -67,54 +70,67 @@ class RemoteElementBase:
     url: str
 
 
+@dataclass
+class ImageBase:
+    type: ElementType = "image"
+    size: ElementSize = "medium"
+
+
 @dataclass
 class RemoteElement(Element, RemoteElementBase):
     def persist(self, client: BaseClient, for_id: str = None):
+        size = getattr(self, "size", None)
+        language = getattr(self, "language", None)
         element = client.create_element(
             name=self.name,
             url=self.url,
             type=self.type,
             display=self.display,
+            size=size,
+            language=language,
             for_id=for_id,
         )
         return element
 
 
-class LocalImage(LocalElement):
-    def __init__(
-        self,
-        name: str,
-        display: ElementDisplay = "side",
-        path: str = None,
-        content: bytes = None,
-    ):
-        if path:
-            with open(path, "rb") as f:
+@dataclass
+class LocalImage(ImageBase, LocalElement):
+    """Useful to send an image living on the local filesystem to the UI."""
+
+    path: str = None
+
+    def __post_init__(self):
+        if self.path:
+            with open(self.path, "rb") as f:
                 self.content = f.read()
-        elif content:
-            self.content = content
+        elif self.content:
+            self.content = self.content
         else:
             raise ValueError("Must provide either path or content")
 
-        self.name = name
-        self.display = display
-        self.type = "image"
 
+@dataclass
+class RemoteImage(ImageBase, RemoteElement):
+    """Useful to send an image based on an URL to the UI."""
+
+    pass
 
-class RemoteImage(RemoteElement):
-    def __init__(self, name: str, url: str, display: ElementDisplay = "side"):
-        self.name = name
-        self.display = display
-        self.type = "image"
-        self.url = url
 
+@dataclass
+class TextBase:
+    text: str
+
+
+@dataclass
+class Text(LocalElement, TextBase):
+    """Useful to send a text (not a message) to the UI."""
+
+    type: ElementType = "text"
+    content = bytes("", "utf-8")
+    language: str = None
 
-class Text(LocalElement):
-    def __init__(self, name: str, text: str, display: ElementDisplay = "side"):
-        self.name = name
-        self.display = display
-        self.type = "text"
-        self.content = bytes(text, "utf-8")
+    def __post_init__(self):
+        self.content = bytes(self.text, "utf-8")
 
     def before_emit(self, text_element):
         if "content" in text_element and isinstance(text_element["content"], bytes):

src/chainlit/frontend/src/components/element/frame.tsx

@@ -0,0 +1,18 @@
+import { Box, BoxProps } from '@mui/material/';
+
+export default function ElementFrame(props: BoxProps) {
+  return (
+    <Box
+      sx={{
+        p: 1,
+        boxSizing: 'border-box',
+        bgcolor: (theme) =>
+          theme.palette.mode === 'light' ? '#EEEEEE' : '#212121',
+        borderRadius: '4px',
+        display: 'flex'
+      }}
+    >
+      {props.children}
+    </Box>
+  );
+}

src/chainlit/frontend/src/components/element/image.tsx

@@ -1,27 +1,36 @@
-import { Box } from '@mui/material';
-import { IElement } from 'state/element';
+import { IImageElement } from 'state/element';
+import ImageFrame from './frame';
 
 interface Props {
-  element: IElement;
+  element: IImageElement;
 }
 
 export default function ImageElement({ element }: Props) {
   const src = element.url || URL.createObjectURL(new Blob([element.content!]));
+  const className = `${element.display}-image`;
   return (
-    <Box
-      sx={{
-        p: 1,
-        boxSizing: 'border-box',
-        bgcolor: (theme) =>
-          theme.palette.mode === 'light' ? '#EEEEEE' : '#212121',
-        borderRadius: '4px'
-      }}
-    >
+    <ImageFrame>
       <img
+        className={className}
         src={src}
-        style={{ objectFit: 'cover', width: '100%' }}
+        onClick={(e) => {
+          if (element.display === 'inline') {
+            const w = window.open('');
+            const target = e.target as HTMLImageElement;
+            w?.document.write(`<img src="${target.src}" />`);
+          }
+        }}
+        style={{
+          objectFit: 'cover',
+          maxWidth: '100%',
+          margin: 'auto',
+          height: 'auto',
+          display: 'block',
+          cursor: element.display === 'inline' ? 'pointer' : 'default'
+        }}
         alt={element.name}
+        loading="lazy"
       />
-    </Box>
+    </ImageFrame>
   );
 }

src/chainlit/frontend/src/components/element/inlined/image.tsx

@@ -1,65 +1,46 @@
 import ImageList from '@mui/material/ImageList';
 import ImageListItem from '@mui/material/ImageListItem';
+import { IImageElement } from 'state/element';
+import ImageElement from '../image';
 
 interface Props {
-  items: {
-    url?: string;
-    src: string;
-    title: string;
-  }[];
+  images: IImageElement[];
 }
 
-export default function InlinedImageList({ items }: Props) {
+function sizeToUnit(image: IImageElement) {
+  if (image.size === 'small') {
+    return 1;
+  } else if (image.size === 'medium') {
+    return 2;
+  } else if (image.size === 'large') {
+    return 4;
+  } else {
+    return 2;
+  }
+}
+
+export default function InlinedImageList({ images }: Props) {
   return (
     <ImageList
       sx={{
         margin: 0,
-        width: '100%',
-        maxWidth: '600px',
-        height: 200,
         // Promote the list into its own layer in Chrome. This costs memory, but helps keeping high FPS.
-        transform: 'translateZ(0)'
+        transform: 'translateZ(0)',
+        width: '100%',
+        maxWidth: 600,
+        maxHeight: 400
       }}
-      rowHeight={200}
-      gap={5}
+      variant="quilted"
+      cols={4}
+      gap={8}
     >
-      {items.map((item) => {
-        const cols = 1;
-        const rows = 1;
+      {images.map((image, i) => {
+        const cols = sizeToUnit(image);
+        const rows = sizeToUnit(image);
 
         return (
-          <ImageListItem
-            key={item.src}
-            cols={cols}
-            rows={rows}
-            sx={{
-              '.MuiImageListItem-img': {
-                height: '100%',
-                width: 'auto',
-                p: 1,
-                boxSizing: 'border-box',
-                bgcolor: (theme) =>
-                  theme.palette.mode === 'light' ? '#EEEEEE' : '#212121',
-                borderRadius: '4px'
-              }
-            }}
-          >
-            <img
-              className="inlined-image"
-              src={item.src}
-              alt={item.title}
-              style={{
-                objectFit: 'contain',
-                cursor: item.url ? 'pointer' : 'default'
-              }}
-              onClick={() => {
-                if (item.url) {
-                  window.open(item.url, '_blank')?.focus();
-                }
-              }}
-              loading="lazy"
-            />
-            {/* <ImageListItemBar title={item.title} position="top" /> */}
+          <ImageListItem key={i} cols={cols} rows={rows}>
+            <ImageElement element={image} />
           </ImageListItem>
         );
       })}