Skip to content

feat(studio): Support images in chat #364

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
steramae-nvidia merged 8 commits into from
Jun 22, 2026
Merged

feat(studio): Support images in chat #364

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
dd1842b
feat(studio): Support images in chat
steramae-nvidia Jun 16, 2026
7c8dc3d
add dumb guard
steramae-nvidia Jun 16, 2026
1165cd2
fix(studio): Address CodeRabbit review comments on image-chat PR
steramae-nvidia Jun 16, 2026
743119f
remove ff
steramae-nvidia Jun 16, 2026
f53b42e
break up components, fix test
steramae-nvidia Jun 16, 2026
20eb460
remove some claude comments
steramae-nvidia Jun 16, 2026
d29913a
try to fix timeout
steramae-nvidia Jun 16, 2026
ac05d85
Merge branch 'main' into steramae/image-chat
steramae-nvidia Jun 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 33 additions & 0 deletions web/packages/common/src/components/AssistantChat/AssistantChatMessageContent.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

import { MessagePrimitive } from '@assistant-ui/react';
import type { MessageRenderProps } from '@nemo/common/src/components/AssistantChat/types';
import { MessageContent } from '@nemo/common/src/components/Chat/MessageContent';
import { Banner } from '@nvidia/foundations-react-core';

export const AssistantChatMessageContent = ({
messageContentProps,
toolCallPartComponent,
}: MessageRenderProps) => (
<>
<MessagePrimitive.Parts
components={{
Text: ({ text }) => <MessageContent content={text} {...messageContentProps} />,
Image: ({ image, filename }) => (
<img
src={image}
alt={filename ?? 'Attached image'}
className="mt-density-xs max-h-64 w-auto rounded-lg border border-base object-contain"
/>
),
tools: { Fallback: toolCallPartComponent },
}}
/>
<MessagePrimitive.Error>
<Banner kind="inline" status="error" className="mt-density-sm">
There was an error generating a response.
</Banner>
</MessagePrimitive.Error>
</>
);
310 changes: 19 additions & 291 deletions web/packages/common/src/components/AssistantChat/AssistantChatThread.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,84 +1,21 @@
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

import { ThreadPrimitive } from '@assistant-ui/react';
import { AssistantComposer } from '@nemo/common/src/components/AssistantChat/AssistantComposer';
import { AssistantMessage } from '@nemo/common/src/components/AssistantChat/AssistantMessage';
import {
ActionBarPrimitive,
ComposerPrimitive,
MessagePrimitive,
ThreadPrimitive,
type ToolCallMessagePartComponent,
} from '@assistant-ui/react';
import { ComposerMode } from '@nemo/common/src/components/AssistantChat/types';
ComposerMode,
type AssistantChatThreadProps,
} from '@nemo/common/src/components/AssistantChat/types';
import { UserEditComposer } from '@nemo/common/src/components/AssistantChat/UserEditComposer';
import { UserMessage } from '@nemo/common/src/components/AssistantChat/UserMessage';
import { ChatEmptyState } from '@nemo/common/src/components/Chat/ChatEmptyState';
import {
MessageContent,
type MessageContentProps,
} from '@nemo/common/src/components/Chat/MessageContent';
import {
Banner,
Button,
Flex,
Skeleton,
Stack,
Text,
TextArea,
Tooltip,
} from '@nvidia/foundations-react-core';
import { Flex, Stack } from '@nvidia/foundations-react-core';
import cn from 'classnames';
import { ArrowUp, Check, Copy, Pencil, RefreshCw, RotateCcw, Square, X } from 'lucide-react';
import { useCallback, useMemo, type ComponentProps, type ReactNode } from 'react';

export interface AssistantChatThreadAttributes {
ThreadViewport?: ComponentProps<typeof ThreadPrimitive.Viewport>;
}
import { useCallback, useMemo } from 'react';

type AssistantChatMessageContentProps = Pick<MessageContentProps, 'markdownLinkComponent'>;

interface AssistantChatThreadProps {
disabled?: boolean;
placeholder: string;
onReset: () => void;
showRunningIndicator?: boolean;
attributes?: AssistantChatThreadAttributes;
composerMode?: ComposerMode;
slotComposerStart?: ReactNode;
emptyState?: {
slotHeading?: string;
slotSubheading?: string;
};
contentClassName?: string;
composerContainerClassName?: string;
hideAssistantMessageActions?: boolean;
toolCallPartComponent?: ToolCallMessagePartComponent;
viewportClassName?: string;
composerOverride?: ReactNode;
messageContentProps?: AssistantChatMessageContentProps;
}

const AssistantChatMessageContent = ({
messageContentProps,
toolCallPartComponent,
}: {
messageContentProps?: AssistantChatMessageContentProps;
toolCallPartComponent?: ToolCallMessagePartComponent;
}) => (
<>
<MessagePrimitive.Parts
components={{
Text: ({ text }) => <MessageContent content={text} {...messageContentProps} />,
tools: { Fallback: toolCallPartComponent },
}}
/>
<MessagePrimitive.Error>
<Banner kind="inline" status="error" className="mt-density-sm">
There was an error generating a response.
</Banner>
</MessagePrimitive.Error>
</>
);

const ACTION_BUTTON_CLASS =
'flex cursor-pointer size-8 items-center justify-center rounded text-base bg-surface-raised hover:bg-surface-sunken disabled:cursor-not-allowed disabled:opacity-50';
export type { AssistantChatThreadAttributes } from '@nemo/common/src/components/AssistantChat/types';

const CHAT_VIEWPORT_SCROLLBAR_CLASS = [
'[scrollbar-width:thin]',
Expand All @@ -91,221 +28,6 @@ const CHAT_VIEWPORT_SCROLLBAR_CLASS = [
'[&::-webkit-scrollbar-thumb:hover]:bg-[var(--border-color-interaction-strong)]',
].join(' ');

const CopyAction = () => (
<Tooltip slotContent="Copy message">
<ActionBarPrimitive.Copy aria-label="Copy message" className={ACTION_BUTTON_CLASS}>
<MessagePrimitive.If copied>
<Check size={16} />
</MessagePrimitive.If>
<MessagePrimitive.If copied={false}>
<Copy size={16} />
</MessagePrimitive.If>
</ActionBarPrimitive.Copy>
</Tooltip>
);

const AssistantMessage = ({
hideAssistantMessageActions,
messageContentProps,
showRunningIndicator = true,
toolCallPartComponent,
}: {
hideAssistantMessageActions?: boolean;
messageContentProps?: AssistantChatMessageContentProps;
showRunningIndicator?: boolean;
toolCallPartComponent?: ToolCallMessagePartComponent;
}) => (
<MessagePrimitive.Root
data-testid="assistant-chat-message"
data-testspeaker="assistant"
className="group/message self-stretch whitespace-normal"
>
<AssistantChatMessageContent
messageContentProps={messageContentProps}
toolCallPartComponent={toolCallPartComponent}
/>
{showRunningIndicator ? (
<MessagePrimitive.If last>
<ThreadPrimitive.If running>
<div
className="mt-density-xs flex h-6 items-center"
data-testid="assistant-chat-running-indicator"
>
<Skeleton className="h-density-4 w-full" data-testid="assistant-chat-skeleton" />
</div>
</ThreadPrimitive.If>
</MessagePrimitive.If>
) : null}
{!hideAssistantMessageActions ? (
<div
className="mt-density-xs flex h-8 items-center"
data-testid="assistant-chat-message-actions"
>
<ActionBarPrimitive.Root
hideWhenRunning
className="flex gap-density-xs opacity-0 transition-opacity group-hover/message:opacity-100 group-focus-within/message:opacity-100 [@media(hover:none)]:opacity-100"
>
<Tooltip slotContent="Regenerate response">
<ActionBarPrimitive.Reload
aria-label="Regenerate response"
className={ACTION_BUTTON_CLASS}
>
<RefreshCw size={16} />
</ActionBarPrimitive.Reload>
</Tooltip>
<CopyAction />
</ActionBarPrimitive.Root>
</div>
) : null}
</MessagePrimitive.Root>
);

const UserMessage = ({
messageContentProps,
toolCallPartComponent,
}: {
messageContentProps?: AssistantChatMessageContentProps;
toolCallPartComponent?: ToolCallMessagePartComponent;
}) => (
<MessagePrimitive.Root
data-testid="assistant-chat-message"
data-testspeaker="user"
className="group/message flex w-full flex-col items-end gap-density-xs whitespace-normal"
>
<div className="max-w-[80%] rounded-xl rounded-br-none bg-surface-overlay px-3 py-2">
<AssistantChatMessageContent
messageContentProps={messageContentProps}
toolCallPartComponent={toolCallPartComponent}
/>
</div>
<div className="flex h-8 shrink-0 items-center">
<ActionBarPrimitive.Root
hideWhenRunning
className="flex gap-density-xs opacity-0 transition-opacity group-hover/message:opacity-100 group-focus-within/message:opacity-100 [@media(hover:none)]:opacity-100"
>
<Tooltip slotContent="Edit message">
<ActionBarPrimitive.Edit aria-label="Edit message" className={ACTION_BUTTON_CLASS}>
<Pencil size={16} />
</ActionBarPrimitive.Edit>
</Tooltip>
<CopyAction />
</ActionBarPrimitive.Root>
</div>
</MessagePrimitive.Root>
);

const UserEditComposer = () => (
<MessagePrimitive.Root
data-testid="assistant-chat-edit-composer"
className="w-full max-w-[80%] self-end rounded-xl rounded-br-none bg-surface-overlay px-4 py-2"
>
<ComposerPrimitive.Root className="w-full">
<ComposerPrimitive.Input
aria-label="Edit message"
addAttachmentOnPaste={false}
autoFocus
submitMode="enter"
rows={3}
render={
<TextArea
resizeable="auto"
size="large"
className="w-full max-h-64"
slotEnd={
<Flex
gap="density-sm"
align="center"
justify="end"
className="mt-density-sm self-end"
>
<Tooltip slotContent="Cancel edit">
<ComposerPrimitive.Cancel
aria-label="Cancel edit"
className="cursor-pointer flex size-8 items-center justify-center rounded border border-base bg-surface-raised hover:bg-surface-sunken disabled:cursor-not-allowed disabled:opacity-50"
>
<X />
</ComposerPrimitive.Cancel>
</Tooltip>
<ComposerPrimitive.Send asChild>
<Button aria-label="Save edit" color="brand" size="small" className="h-full">
<Text kind="label/regular/sm">Send</Text>
</Button>
</ComposerPrimitive.Send>
</Flex>
}
/>
}
/>
</ComposerPrimitive.Root>
</MessagePrimitive.Root>
);

type AssistantComposerProps = Pick<
AssistantChatThreadProps,
'disabled' | 'placeholder' | 'onReset' | 'slotComposerStart'
> & {
className?: string;
};

const AssistantComposer = ({
disabled,
placeholder,
onReset,
slotComposerStart,
className,
}: AssistantComposerProps) => (
<div className="flex w-full flex-col gap-2">
{slotComposerStart && <div className="shrink-0">{slotComposerStart}</div>}
<ComposerPrimitive.Root
data-testid="assistant-chat-composer"
className={cn(
'flex w-full items-end gap-density-xs rounded-lg border border-base bg-surface-base p-1',
className
)}
>
<ComposerPrimitive.Input
aria-label="Task prompt"
addAttachmentOnPaste={false}
disabled={disabled}
placeholder={placeholder}
submitMode="enter"
className="max-h-32 min-h-[24px] flex-1 resize-none border-0 bg-transparent px-density-md py-density-md text-sm leading-6 outline-none disabled:cursor-not-allowed disabled:text-fg-disabled"
/>
<Tooltip slotContent="Clear chat thread">
<Button
aria-label="Reset"
kind="tertiary"
size="small"
onClick={onReset}
type="button"
disabled={disabled}
>
<RotateCcw />
</Button>
</Tooltip>
<ThreadPrimitive.If running>
<ComposerPrimitive.Cancel asChild>
<Button aria-label="Stop" color="danger" size="small" className="size-8 rounded-full p-0">
<Square size={14} />
</Button>
</ComposerPrimitive.Cancel>
</ThreadPrimitive.If>
<ThreadPrimitive.If running={false}>
<ComposerPrimitive.Send asChild>
<Button
aria-label="Submit"
color="brand"
size="small"
className="size-8 rounded-full p-0"
>
<ArrowUp size={16} />
</Button>
</ComposerPrimitive.Send>
</ThreadPrimitive.If>
</ComposerPrimitive.Root>
</div>
);

export const AssistantChatThread = ({
disabled,
placeholder,
Expand All @@ -322,6 +44,7 @@ export const AssistantChatThread = ({
viewportClassName,
composerOverride,
messageContentProps,
enableImageAttachments,
}: AssistantChatThreadProps) => {
const { className: threadViewportClassName, ...threadViewportAttributes } =
attributes?.ThreadViewport ?? {};
Expand All @@ -345,14 +68,18 @@ export const AssistantChatThread = ({
),
[messageContentProps, toolCallPartComponent]
);
const UserEditComposerWithAttachments = useCallback(
() => <UserEditComposer enableImageAttachments={enableImageAttachments} />,
[enableImageAttachments]
);
const messageComponents = useMemo(
() => ({
AssistantMessage: AssistantMessageWithToolCallPart,
UserMessage: UserMessageWithToolCallPart,
UserEditComposer,
UserEditComposer: UserEditComposerWithAttachments,
SystemMessage: AssistantMessageWithToolCallPart,
}),
[AssistantMessageWithToolCallPart, UserMessageWithToolCallPart]
[AssistantMessageWithToolCallPart, UserMessageWithToolCallPart, UserEditComposerWithAttachments]
);

return (
Expand Down Expand Up @@ -394,6 +121,7 @@ export const AssistantChatThread = ({
placeholder={placeholder}
onReset={onReset}
slotComposerStart={slotComposerStart}
enableImageAttachments={enableImageAttachments}
/>
)}
</Flex>
Expand Down
Loading
Loading