From 4f6ab674175c195544ee2850297940246d55fab0 Mon Sep 17 00:00:00 2001 From: Allaoua Benchikh Date: Fri, 9 Jan 2026 16:04:37 +0100 Subject: [PATCH 1/2] Added ability to customize step icon --- backend/chainlit/step.py | 14 +++++- backend/tests/test_emitter.py | 46 +++++++++++++++++++ cypress/e2e/step_icon/main.py | 33 +++++++++++++ cypress/e2e/step_icon/spec.cy.ts | 42 +++++++++++++++++ .../chat/Messages/Message/Avatar.tsx | 35 ++++++++------ .../chat/Messages/Message/index.tsx | 1 + libs/react-client/src/types/step.ts | 1 + 7 files changed, 158 insertions(+), 14 deletions(-) create mode 100644 cypress/e2e/step_icon/main.py create mode 100644 cypress/e2e/step_icon/spec.cy.ts diff --git a/backend/chainlit/step.py b/backend/chainlit/step.py index 233f51eabe..da567785e4 100644 --- a/backend/chainlit/step.py +++ b/backend/chainlit/step.py @@ -64,6 +64,7 @@ class StepDict(TypedDict, total=False): showInput: Optional[Union[bool, str]] defaultOpen: Optional[bool] language: Optional[str] + icon: Optional[str] feedback: Optional[FeedbackDict] @@ -84,6 +85,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, ): @@ -107,6 +109,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, @@ -136,6 +139,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, @@ -183,6 +187,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 @@ -197,6 +202,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, @@ -215,6 +221,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 [] @@ -288,6 +295,11 @@ def output(self, content: Union[Dict, str]): self._output = self._process_content(content, set_language=True) def to_dict(self) -> StepDict: + # Move icon into metadata for storage + metadata = dict(self.metadata) + if self.icon: + metadata["icon"] = self.icon + _dict: StepDict = { "name": self.name, "type": self.type, @@ -295,7 +307,7 @@ def to_dict(self) -> StepDict: "threadId": self.thread_id, "parentId": self.parent_id, "streaming": self.streaming, - "metadata": self.metadata, + "metadata": metadata, "tags": self.tags, "input": self.input, "isError": self.is_error, diff --git a/backend/tests/test_emitter.py b/backend/tests/test_emitter.py index 9a8290583c..d9589fae74 100644 --- a/backend/tests/test_emitter.py +++ b/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", + "metadata": {"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", + "metadata": {"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", + "metadata": {"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: diff --git a/cypress/e2e/step_icon/main.py b/cypress/e2e/step_icon/main.py new file mode 100644 index 0000000000..a5587605d8 --- /dev/null +++ b/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() diff --git a/cypress/e2e/step_icon/spec.cy.ts b/cypress/e2e/step_icon/spec.cy.ts new file mode 100644 index 0000000000..a0240af594 --- /dev/null +++ b/cypress/e2e/step_icon/spec.cy.ts @@ -0,0 +1,42 @@ +import { submitMessage } from '../../support/testUtils'; + +describe('Step with Icon', () => { + it('should display icons for steps with icon property', () => { + submitMessage('Hello'); + + cy.get('.step').should('have.length', 5); + + // Check that steps with icons have SVG icons (not avatar images) + // The avatar is a sibling of the step content in the .ai-message container + cy.get('#step-search') + .closest('.ai-message') + .within(() => { + // Should have an svg icon (Lucide icons are SVGs) + cy.get('svg').should('exist'); + // Should NOT have an avatar image + cy.get('img').should('not.exist'); + }); + + cy.get('#step-database') + .closest('.ai-message') + .within(() => { + cy.get('svg').should('exist'); + cy.get('img').should('not.exist'); + }); + + // Check that step without icon has avatar (image) + cy.get('#step-regular') + .closest('.ai-message') + .within(() => { + // Should have an avatar image + cy.get('img').should('exist'); + }); + + cy.get('#step-cpu') + .closest('.ai-message') + .within(() => { + cy.get('svg').should('exist'); + cy.get('img').should('not.exist'); + }); + }); +}); diff --git a/frontend/src/components/chat/Messages/Message/Avatar.tsx b/frontend/src/components/chat/Messages/Message/Avatar.tsx index e9e2626d84..1ea41acf5d 100644 --- a/frontend/src/components/chat/Messages/Message/Avatar.tsx +++ b/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(); @@ -50,22 +52,29 @@ const MessageAvatar = ({ author, hide, isError }: Props) => { ); } + // Render icon or avatar based on iconName + const avatarContent = iconName ? ( + + + + ) : ( + + + + + + + ); + return ( - - - - - - - - + {avatarContent}

{author}

 diff --git a/frontend/src/components/chat/Messages/Message/index.tsx b/frontend/src/components/chat/Messages/Message/index.tsx index 455b6f71c8..455af884b8 100644 --- a/frontend/src/components/chat/Messages/Message/index.tsx +++ b/frontend/src/components/chat/Messages/Message/index.tsx @@ -110,6 +110,7 @@ const Message = memo( ) : null} {/* Display the step and its children */} diff --git a/libs/react-client/src/types/step.ts b/libs/react-client/src/types/step.ts index 44c210d8eb..95bf0fb498 100644 --- a/libs/react-client/src/types/step.ts +++ b/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; From 279e22b9850a7701af25696cbd0f73ac3f5138f2 Mon Sep 17 00:00:00 2001 From: Allaoua Benchikh Date: Fri, 9 Jan 2026 16:53:56 +0100 Subject: [PATCH 2/2] Removed icon from IStep --- libs/react-client/src/types/step.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/libs/react-client/src/types/step.ts b/libs/react-client/src/types/step.ts index 95bf0fb498..44c210d8eb 100644 --- a/libs/react-client/src/types/step.ts +++ b/libs/react-client/src/types/step.ts @@ -16,7 +16,6 @@ export interface IStep { id: string; name: string; type: StepType; - icon?: string; threadId?: string; parentId?: string; isError?: boolean;