Added ability to have custom icons

Author: Allaoua Benchikh

backend/chainlit/step.py

@@ -63,6 +63,7 @@ class StepDict(TypedDict, total=False):
     showInput: Optional[Union[bool, str]]
     defaultOpen: Optional[bool]
     language: Optional[str]
+    icon: Optional[str]
     feedback: Optional[FeedbackDict]
 
 
@@ -83,6 +84,7 @@ def step(
     tags: Optional[List[str]] = None,
     metadata: Optional[Dict] = None,
     language: Optional[str] = None,
+    icon: Optional[str] = None,
     show_input: Union[bool, str] = "json",
     default_open: bool = False,
 ):
@@ -106,6 +108,7 @@ async def async_wrapper(*args, **kwargs):
                     parent_id=parent_id,
                     tags=tags,
                     language=language,
+                    icon=icon,
                     show_input=show_input,
                     default_open=default_open,
                     metadata=metadata,
@@ -135,6 +138,7 @@ def sync_wrapper(*args, **kwargs):
                     parent_id=parent_id,
                     tags=tags,
                     language=language,
+                    icon=icon,
                     show_input=show_input,
                     default_open=default_open,
                     metadata=metadata,
@@ -182,6 +186,7 @@ class Step:
     end: Union[str, None]
     generation: Optional[BaseGeneration]
     language: Optional[str]
+    icon: Optional[str]
     default_open: Optional[bool]
     elements: Optional[List[Element]]
     fail_on_persist_error: bool
@@ -196,6 +201,7 @@ def __init__(
         metadata: Optional[Dict] = None,
         tags: Optional[List[str]] = None,
         language: Optional[str] = None,
+        icon: Optional[str] = None,
         default_open: Optional[bool] = False,
         show_input: Union[bool, str] = "json",
         thread_id: Optional[str] = None,
@@ -214,6 +220,7 @@ def __init__(
         self.parent_id = parent_id
 
         self.language = language
+        self.icon = icon
         self.default_open = default_open
         self.generation = None
         self.elements = elements or []
@@ -303,6 +310,7 @@ def to_dict(self) -> StepDict:
             "start": self.start,
             "end": self.end,
             "language": self.language,
+            "icon": self.icon,
             "defaultOpen": self.default_open,
             "showInput": self.show_input,
             "generation": self.generation.to_dict() if self.generation else None,

backend/tests/test_emitter.py

@@ -54,6 +54,22 @@ async def test_send_step(
     mock_websocket_session.emit.assert_called_once_with("new_message", step_dict)
 
 
+async def test_send_step_with_icon(
+    emitter: ChainlitEmitter, mock_websocket_session: MagicMock
+) -> None:
+    step_dict: StepDict = {
+        "id": "test_step_with_icon",
+        "type": "tool",
+        "name": "Test Step with Icon",
+        "output": "This is a test step with an icon",
+        "icon": "search",
+    }
+
+    await emitter.send_step(step_dict)
+
+    mock_websocket_session.emit.assert_called_once_with("new_message", step_dict)
+
+
 async def test_update_step(
     emitter: ChainlitEmitter, mock_websocket_session: MagicMock
 ) -> None:
@@ -69,6 +85,22 @@ async def test_update_step(
     mock_websocket_session.emit.assert_called_once_with("update_message", step_dict)
 
 
+async def test_update_step_with_icon(
+    emitter: ChainlitEmitter, mock_websocket_session: MagicMock
+) -> None:
+    step_dict: StepDict = {
+        "id": "test_step_with_icon",
+        "type": "tool",
+        "name": "Updated Test Step with Icon",
+        "output": "This is an updated test step with an icon",
+        "icon": "database",
+    }
+
+    await emitter.update_step(step_dict)
+
+    mock_websocket_session.emit.assert_called_once_with("update_message", step_dict)
+
+
 async def test_delete_step(
     emitter: ChainlitEmitter, mock_websocket_session: MagicMock
 ) -> None:
@@ -139,6 +171,20 @@ async def test_stream_start(
     mock_websocket_session.emit.assert_called_once_with("stream_start", step_dict)
 
 
+async def test_stream_start_with_icon(
+    emitter: ChainlitEmitter, mock_websocket_session: MagicMock
+) -> None:
+    step_dict: StepDict = {
+        "id": "test_stream_with_icon",
+        "type": "tool",
+        "name": "Test Stream with Icon",
+        "output": "This is a test stream with an icon",
+        "icon": "cpu",
+    }
+    await emitter.stream_start(step_dict)
+    mock_websocket_session.emit.assert_called_once_with("stream_start", step_dict)
+
+
 async def test_send_toast(
     emitter: ChainlitEmitter, mock_websocket_session: MagicMock
 ) -> None:

cypress/e2e/step_icon/main.py

@@ -0,0 +1,33 @@
+import chainlit as cl
+
+
+@cl.step(name="search", type="tool", icon="search")
+async def search():
+    await cl.sleep(1)
+    return "Response from search"
+
+
+@cl.step(name="database", type="tool", icon="database")
+async def database():
+    await cl.sleep(1)
+    return "Response from database"
+
+
+@cl.step(name="regular", type="tool")
+async def regular():
+    await cl.sleep(1)
+    return "Response from regular"
+
+
+async def cpu():
+    async with cl.Step(name="cpu", type="tool", icon="cpu") as s:
+        await cl.sleep(1)
+        s.output = "Response from cpu"
+
+
+@cl.on_message
+async def main(message: cl.Message):
+    await search()
+    await database()
+    await regular()
+    await cpu()

cypress/e2e/step_icon/spec.cy.ts

@@ -0,0 +1,17 @@
+import { submitMessage } from '../../support/testUtils';
+
+describe('Step with Icon', () => {
+  it('should be able to use steps with icons', () => {
+    submitMessage('Hello');
+
+    cy.get('#step-search').should('exist').click();
+
+    cy.get('#step-database').should('exist').click();
+
+    cy.get('#step-regular').should('exist').click();
+
+    cy.get('#step-cpu').should('exist');
+
+    cy.get('.step').should('have.length', 5);
+  });
+});

frontend/src/components/chat/Messages/Message/Avatar.tsx

@@ -8,6 +8,7 @@ import {
   useConfig
 } from '@chainlit/react-client';
 
+import Icon from '@/components/Icon';
 import { Avatar, AvatarFallback, AvatarImage } from '@/components/ui/avatar';
 import { Skeleton } from '@/components/ui/skeleton';
 import {
@@ -21,9 +22,10 @@ interface Props {
   author?: string;
   hide?: boolean;
   isError?: boolean;
+  iconName?: string;
 }
 
-const MessageAvatar = ({ author, hide, isError }: Props) => {
+const MessageAvatar = ({ author, hide, isError, iconName }: Props) => {
   const apiClient = useContext(ChainlitContext);
   const { chatProfile } = useChatSession();
   const { config } = useConfig();
@@ -48,22 +50,29 @@ const MessageAvatar = ({ author, hide, isError }: Props) => {
     );
   }
 
+  // Render icon or avatar based on iconName
+  const avatarContent = iconName ? (
+    <span className="inline-flex">
+      <Icon name={iconName} className="h-5 w-5 mt-[3px]" />
+    </span>
+  ) : (
+    <Avatar className="h-5 w-5 mt-[3px]">
+      <AvatarImage
+        src={avatarUrl}
+        alt={`Avatar for ${author || 'default'}`}
+        className="bg-transparent"
+      />
+      <AvatarFallback className="bg-transparent">
+        <Skeleton className="h-full w-full rounded-full" />
+      </AvatarFallback>
+    </Avatar>
+  );
+
   return (
     <span className={cn('inline-block', hide && 'invisible')}>
       <TooltipProvider>
         <Tooltip>
-          <TooltipTrigger asChild>
-            <Avatar className="h-5 w-5 mt-[3px]">
-              <AvatarImage
-                src={avatarUrl}
-                alt={`Avatar for ${author || 'default'}`}
-                className="bg-transparent"
-              />
-              <AvatarFallback className="bg-transparent">
-                <Skeleton className="h-full w-full rounded-full" />
-              </AvatarFallback>
-            </Avatar>
-          </TooltipTrigger>
+          <TooltipTrigger asChild>{avatarContent}</TooltipTrigger>
           <TooltipContent>
             <p>{author}</p>
           </TooltipContent>

frontend/src/components/chat/Messages/Message/index.tsx

@@ -99,6 +99,7 @@ const Message = memo(
                     <MessageAvatar
                       author={message.metadata?.avatarName || message.name}
                       isError={message.isError}
+                      iconName={message.icon}
                     />
                   ) : null}
                   {/* Display the step and its children */}

libs/react-client/src/types/step.ts

@@ -16,6 +16,7 @@ export interface IStep {
   id: string;
   name: string;
   type: StepType;
+  icon?: string;
   threadId?: string;
   parentId?: string;
   isError?: boolean;