mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-08-16 14:38:00 +00:00
feat(UI): adding MVP playground UI (#2828)
# What does this PR do? I've been tinkering a little with a simple chat playground in the UI, so I'm opening the PR with what's kind of a WIP. If you look at the first commit, that includes the big part of the changes. The rest of the files changed come from adding installing the `shadcn` components. Note this is missing a lot; e.g., - sessions - document upload - audio (the shadcn components install these by default from https://shadcn-chatbot-kit.vercel.app/docs/components/chat) I still need to wire up a lot more to make it actually fully functional but it does basic chat using the LS Typescript Client. Basic demo: <img width="1329" height="1430" alt="Image" src="https://github.com/user-attachments/assets/917a2096-36d4-4925-b83b-f1f2cda98698" /> <img width="1319" height="1424" alt="Image" src="https://github.com/user-attachments/assets/fab1583b-1c72-4bf3-baf2-405aee13c6bb" /> <!-- If resolving an issue, uncomment and update the line below --> <!-- Closes #[issue-number] --> ## Test Plan <!-- Describe the tests you ran to verify your changes with result summaries. *Provide clear instructions so the plan can be easily re-executed.* --> --------- Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
This commit is contained in:
Francisco Arceo
GitHub
parent
d6ae2b0f47
commit
f3d5459647
32 changed files with 4876 additions and 31 deletions
349
llama_stack/ui/components/chat-playground/chat.tsx
Normal file
349
llama_stack/ui/components/chat-playground/chat.tsx
Normal file
|
|
@ -0,0 +1,349 @@
|
|||
"use client"
|
||||
|
||||
import {
|
||||
forwardRef,
|
||||
useCallback,
|
||||
useRef,
|
||||
useState,
|
||||
type ReactElement,
|
||||
} from "react"
|
||||
import { ArrowDown, ThumbsDown, ThumbsUp } from "lucide-react"
|
||||
|
||||
import { cn } from "@/lib/utils"
|
||||
import { useAutoScroll } from "@/hooks/use-auto-scroll"
|
||||
import { Button } from "@/components/ui/button"
|
||||
import { type Message } from "@/components/chat-playground/chat-message"
|
||||
import { CopyButton } from "@/components/ui/copy-button"
|
||||
import { MessageInput } from "@/components/chat-playground/message-input"
|
||||
import { MessageList } from "@/components/chat-playground/message-list"
|
||||
import { PromptSuggestions } from "@/components/chat-playground/prompt-suggestions"
|
||||
|
||||
interface ChatPropsBase {
|
||||
handleSubmit: (
|
||||
event?: { preventDefault?: () => void },
|
||||
options?: { experimental_attachments?: FileList }
|
||||
) => void
|
||||
messages: Array<Message>
|
||||
input: string
|
||||
className?: string
|
||||
handleInputChange: React.ChangeEventHandler<HTMLTextAreaElement>
|
||||
isGenerating: boolean
|
||||
stop?: () => void
|
||||
onRateResponse?: (
|
||||
messageId: string,
|
||||
rating: "thumbs-up" | "thumbs-down"
|
||||
) => void
|
||||
setMessages?: (messages: any[]) => void
|
||||
transcribeAudio?: (blob: Blob) => Promise<string>
|
||||
}
|
||||
|
||||
interface ChatPropsWithoutSuggestions extends ChatPropsBase {
|
||||
append?: never
|
||||
suggestions?: never
|
||||
}
|
||||
|
||||
interface ChatPropsWithSuggestions extends ChatPropsBase {
|
||||
append: (message: { role: "user"; content: string }) => void
|
||||
suggestions: string[]
|
||||
}
|
||||
|
||||
type ChatProps = ChatPropsWithoutSuggestions | ChatPropsWithSuggestions
|
||||
|
||||
export function Chat({
|
||||
messages,
|
||||
handleSubmit,
|
||||
input,
|
||||
handleInputChange,
|
||||
stop,
|
||||
isGenerating,
|
||||
append,
|
||||
suggestions,
|
||||
className,
|
||||
onRateResponse,
|
||||
setMessages,
|
||||
transcribeAudio,
|
||||
}: ChatProps) {
|
||||
const lastMessage = messages.at(-1)
|
||||
const isEmpty = messages.length === 0
|
||||
const isTyping = lastMessage?.role === "user"
|
||||
|
||||
const messagesRef = useRef(messages)
|
||||
messagesRef.current = messages
|
||||
|
||||
// Enhanced stop function that marks pending tool calls as cancelled
|
||||
const handleStop = useCallback(() => {
|
||||
stop?.()
|
||||
|
||||
if (!setMessages) return
|
||||
|
||||
const latestMessages = [...messagesRef.current]
|
||||
const lastAssistantMessage = latestMessages.findLast(
|
||||
(m) => m.role === "assistant"
|
||||
)
|
||||
|
||||
if (!lastAssistantMessage) return
|
||||
|
||||
let needsUpdate = false
|
||||
let updatedMessage = { ...lastAssistantMessage }
|
||||
|
||||
if (lastAssistantMessage.toolInvocations) {
|
||||
const updatedToolInvocations = lastAssistantMessage.toolInvocations.map(
|
||||
(toolInvocation) => {
|
||||
if (toolInvocation.state === "call") {
|
||||
needsUpdate = true
|
||||
return {
|
||||
...toolInvocation,
|
||||
state: "result",
|
||||
result: {
|
||||
content: "Tool execution was cancelled",
|
||||
__cancelled: true, // Special marker to indicate cancellation
|
||||
},
|
||||
} as const
|
||||
}
|
||||
return toolInvocation
|
||||
}
|
||||
)
|
||||
|
||||
if (needsUpdate) {
|
||||
updatedMessage = {
|
||||
...updatedMessage,
|
||||
toolInvocations: updatedToolInvocations,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (lastAssistantMessage.parts && lastAssistantMessage.parts.length > 0) {
|
||||
const updatedParts = lastAssistantMessage.parts.map((part: any) => {
|
||||
if (
|
||||
part.type === "tool-invocation" &&
|
||||
part.toolInvocation &&
|
||||
part.toolInvocation.state === "call"
|
||||
) {
|
||||
needsUpdate = true
|
||||
return {
|
||||
...part,
|
||||
toolInvocation: {
|
||||
...part.toolInvocation,
|
||||
state: "result",
|
||||
result: {
|
||||
content: "Tool execution was cancelled",
|
||||
__cancelled: true,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
return part
|
||||
})
|
||||
|
||||
if (needsUpdate) {
|
||||
updatedMessage = {
|
||||
...updatedMessage,
|
||||
parts: updatedParts,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (needsUpdate) {
|
||||
const messageIndex = latestMessages.findIndex(
|
||||
(m) => m.id === lastAssistantMessage.id
|
||||
)
|
||||
if (messageIndex !== -1) {
|
||||
latestMessages[messageIndex] = updatedMessage
|
||||
setMessages(latestMessages)
|
||||
}
|
||||
}
|
||||
}, [stop, setMessages, messagesRef])
|
||||
|
||||
const messageOptions = useCallback(
|
||||
(message: Message) => ({
|
||||
actions: onRateResponse ? (
|
||||
<>
|
||||
<div className="border-r pr-1">
|
||||
<CopyButton
|
||||
content={message.content}
|
||||
copyMessage="Copied response to clipboard!"
|
||||
/>
|
||||
</div>
|
||||
<Button
|
||||
size="icon"
|
||||
variant="ghost"
|
||||
className="h-6 w-6"
|
||||
onClick={() => onRateResponse(message.id, "thumbs-up")}
|
||||
>
|
||||
<ThumbsUp className="h-4 w-4" />
|
||||
</Button>
|
||||
<Button
|
||||
size="icon"
|
||||
variant="ghost"
|
||||
className="h-6 w-6"
|
||||
onClick={() => onRateResponse(message.id, "thumbs-down")}
|
||||
>
|
||||
<ThumbsDown className="h-4 w-4" />
|
||||
</Button>
|
||||
</>
|
||||
) : (
|
||||
<CopyButton
|
||||
content={message.content}
|
||||
copyMessage="Copied response to clipboard!"
|
||||
/>
|
||||
),
|
||||
}),
|
||||
[onRateResponse]
|
||||
)
|
||||
|
||||
return (
|
||||
<ChatContainer className={className}>
|
||||
<div className="flex-1 flex flex-col">
|
||||
{isEmpty && append && suggestions ? (
|
||||
<div className="flex-1 flex items-center justify-center">
|
||||
<PromptSuggestions
|
||||
label="Try these prompts ✨"
|
||||
append={append}
|
||||
suggestions={suggestions}
|
||||
/>
|
||||
</div>
|
||||
) : null}
|
||||
|
||||
{messages.length > 0 ? (
|
||||
<ChatMessages messages={messages}>
|
||||
<MessageList
|
||||
messages={messages}
|
||||
isTyping={isTyping}
|
||||
messageOptions={messageOptions}
|
||||
/>
|
||||
</ChatMessages>
|
||||
) : null}
|
||||
</div>
|
||||
|
||||
<div className="mt-auto border-t bg-background/95 backdrop-blur supports-[backdrop-filter]:bg-background/60">
|
||||
<div className="container max-w-4xl py-4">
|
||||
<ChatForm
|
||||
isPending={isGenerating || isTyping}
|
||||
handleSubmit={handleSubmit}
|
||||
>
|
||||
{({ files, setFiles }) => (
|
||||
<MessageInput
|
||||
value={input}
|
||||
onChange={handleInputChange}
|
||||
allowAttachments
|
||||
files={files}
|
||||
setFiles={setFiles}
|
||||
stop={handleStop}
|
||||
isGenerating={isGenerating}
|
||||
transcribeAudio={transcribeAudio}
|
||||
/>
|
||||
)}
|
||||
</ChatForm>
|
||||
</div>
|
||||
</div>
|
||||
</ChatContainer>
|
||||
)
|
||||
}
|
||||
Chat.displayName = "Chat"
|
||||
|
||||
export function ChatMessages({
|
||||
messages,
|
||||
children,
|
||||
}: React.PropsWithChildren<{
|
||||
messages: Message[]
|
||||
}>) {
|
||||
const {
|
||||
containerRef,
|
||||
scrollToBottom,
|
||||
handleScroll,
|
||||
shouldAutoScroll,
|
||||
handleTouchStart,
|
||||
} = useAutoScroll([messages])
|
||||
|
||||
return (
|
||||
<div
|
||||
className="grid grid-cols-1 overflow-y-auto pb-4"
|
||||
ref={containerRef}
|
||||
onScroll={handleScroll}
|
||||
onTouchStart={handleTouchStart}
|
||||
>
|
||||
<div className="max-w-full [grid-column:1/1] [grid-row:1/1]">
|
||||
{children}
|
||||
</div>
|
||||
|
||||
{!shouldAutoScroll && (
|
||||
<div className="pointer-events-none flex flex-1 items-end justify-end [grid-column:1/1] [grid-row:1/1]">
|
||||
<div className="sticky bottom-0 left-0 flex w-full justify-end">
|
||||
<Button
|
||||
onClick={scrollToBottom}
|
||||
className="pointer-events-auto h-8 w-8 rounded-full ease-in-out animate-in fade-in-0 slide-in-from-bottom-1"
|
||||
size="icon"
|
||||
variant="ghost"
|
||||
>
|
||||
<ArrowDown className="h-4 w-4" />
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
export const ChatContainer = forwardRef<
|
||||
HTMLDivElement,
|
||||
React.HTMLAttributes<HTMLDivElement>
|
||||
>(({ className, ...props }, ref) => {
|
||||
return (
|
||||
<div
|
||||
ref={ref}
|
||||
className={cn("flex flex-col max-h-full w-full", className)}
|
||||
{...props}
|
||||
/>
|
||||
)
|
||||
})
|
||||
ChatContainer.displayName = "ChatContainer"
|
||||
|
||||
interface ChatFormProps {
|
||||
className?: string
|
||||
isPending: boolean
|
||||
handleSubmit: (
|
||||
event?: { preventDefault?: () => void },
|
||||
options?: { experimental_attachments?: FileList }
|
||||
) => void
|
||||
children: (props: {
|
||||
files: File[] | null
|
||||
setFiles: React.Dispatch<React.SetStateAction<File[] | null>>
|
||||
}) => ReactElement
|
||||
}
|
||||
|
||||
export const ChatForm = forwardRef<HTMLFormElement, ChatFormProps>(
|
||||
({ children, handleSubmit, isPending, className }, ref) => {
|
||||
const [files, setFiles] = useState<File[] | null>(null)
|
||||
|
||||
const onSubmit = (event: React.FormEvent) => {
|
||||
// if (isPending) {
|
||||
// event.preventDefault()
|
||||
// return
|
||||
// }
|
||||
|
||||
if (!files) {
|
||||
handleSubmit(event)
|
||||
return
|
||||
}
|
||||
|
||||
const fileList = createFileList(files)
|
||||
handleSubmit(event, { experimental_attachments: fileList })
|
||||
setFiles(null)
|
||||
}
|
||||
|
||||
return (
|
||||
<form ref={ref} onSubmit={onSubmit} className={className}>
|
||||
{children({ files, setFiles })}
|
||||
</form>
|
||||
)
|
||||
}
|
||||
)
|
||||
ChatForm.displayName = "ChatForm"
|
||||
|
||||
function createFileList(files: File[] | FileList): FileList {
|
||||
const dataTransfer = new DataTransfer()
|
||||
for (const file of Array.from(files)) {
|
||||
dataTransfer.items.add(file)
|
||||
}
|
||||
return dataTransfer.files
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue