mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-04 18:13:44 +00:00
chore: move src/llama_stack/ui to src/llama_stack_ui (#4068)
# What does this PR do? This better separates UI from backend code, which was a point of confusion often for our beloved AI friends. ## Test Plan CI
This commit is contained in:
ehhuang
GitHub
parent
5850e3473f
commit
95b0493fae
156 changed files with 20 additions and 20 deletions
357
src/llama_stack_ui/components/chat-playground/chat.tsx
Normal file
357
src/llama_stack_ui/components/chat-playground/chat.tsx
Normal file
|
|
@ -0,0 +1,357 @@
|
|||
"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: Message[]) => void;
|
||||
transcribeAudio?: (blob: Blob) => Promise<string>;
|
||||
onRAGFileUpload?: (file: File) => Promise<void>;
|
||||
}
|
||||
|
||||
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,
|
||||
onRAGFileUpload,
|
||||
}: 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: {
|
||||
type: string;
|
||||
toolInvocation?: { state: string; toolName: string };
|
||||
}) => {
|
||||
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}
|
||||
>
|
||||
{() => (
|
||||
<MessageInput
|
||||
value={input}
|
||||
onChange={handleInputChange}
|
||||
allowAttachments={true}
|
||||
files={null}
|
||||
setFiles={() => {}}
|
||||
stop={handleStop}
|
||||
isGenerating={isGenerating}
|
||||
transcribeAudio={transcribeAudio}
|
||||
onRAGFileUpload={onRAGFileUpload}
|
||||
/>
|
||||
)}
|
||||
</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