Merge branch 'master' into josiah/expose-playground

pyproject.toml

@@ -34,7 +34,7 @@ classifiers = [
 requires-python = ">=3.9"
 dynamic = ["version"]
 dependencies = [
-  "pydantic>=2.0.0",
+  "pydantic>=2.0.0,<2.10.0", # Pinning to resolve issues caused in 2.10.0 release
   "wandb>=0.17.1",
   "packaging>=21.0",         # For version parsing in integrations
   "tenacity>=8.3.0,!=8.4.0", # Excluding 8.4.0 because it had a bug on import of AsyncRetrying

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/ChatView.tsx

@@ -34,7 +34,7 @@ export const ChatView = ({chat}: ChatViewProps) => {
         messages={chat.request?.messages || []}
         scrollLastMessage={scrollLastMessage}
       />
-      {chatResult && chatResult.choices && (
+      {chatResult?.choices && chatResult.choices.length > 0 && (
         <>
           <span className="mt-16 text-sm font-semibold text-moon-800">
             Response

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/ChoiceView.tsx

@@ -15,6 +15,7 @@ export const ChoiceView = ({choice, isStructuredOutput}: ChoiceViewProps) => {
       index={choice.index}
       message={message}
       isStructuredOutput={isStructuredOutput}
+      isChoice
     />
   );
 };

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/MessageList.tsx

@@ -38,17 +38,17 @@ export const MessageList = ({
 const processToolCallMessages = (messages: Messages): Messages => {
   const processedMessages: Message[] = [];
   for (let i = 0; i < messages.length; i++) {
-    const message = messages[i];
+    const message = {
+      ...messages[i],
+      // Store the original index of the message in the message object
+      // so that we can use it to sort the messages later.
+      original_index: i,
+    };
 
     // If there are no tool calls, just add the message to the processed messages
     // and continue to the next iteration.
     if (!message.tool_calls) {
-      processedMessages.push({
-        ...message,
-        // Store the original index of the message in the message object
-        // so that we can use it to sort the messages later.
-        original_index: i,
-      });
+      processedMessages.push(message);
       continue;
     }
 
@@ -58,7 +58,7 @@ const processToolCallMessages = (messages: Messages): Messages => {
     while (i + 1 < messages.length && messages[i + 1].role === 'tool') {
       toolMessages.push({
         ...messages[i + 1],
-        original_index: (messages[i + 1] as any).original_index ?? i + 1,
+        original_index: i + 1,
       });
       i++;
     }

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/MessagePanel.tsx

@@ -3,27 +3,43 @@ import classNames from 'classnames';
 import _ from 'lodash';
 import React, {useEffect, useRef, useState} from 'react';
 
+import {usePlaygroundContext} from '../PlaygroundPage/PlaygroundContext';
 import {MessagePanelPart} from './MessagePanelPart';
+import {PlaygroundMessagePanelButtons} from './PlaygroundMessagePanelButtons';
+import {PlaygroundMessagePanelEditor} from './PlaygroundMessagePanelEditor';
 import {ShowMoreButton} from './ShowMoreButton';
 import {ToolCalls} from './ToolCalls';
-import {Message} from './types';
+import {Message, ToolCall} from './types';
 
 type MessagePanelProps = {
   index: number;
   message: Message;
   isStructuredOutput?: boolean;
+  isChoice?: boolean;
   isNested?: boolean;
+  pendingToolResponseId?: string;
 };
 
 export const MessagePanel = ({
+  index,
   message,
   isStructuredOutput,
+  isChoice,
   isNested,
+  // The id of the tool call response that is pending
+  // If the tool call response is pending, the editor will be shown automatically
+  // and on save the tool call response will be updated and sent to the LLM
+  pendingToolResponseId,
 }: MessagePanelProps) => {
   const [isShowingMore, setIsShowingMore] = useState(false);
   const [isOverflowing, setIsOverflowing] = useState(false);
+  const [isHovering, setIsHovering] = useState(false);
+  const [editorHeight, setEditorHeight] = useState<number | null>(
+    pendingToolResponseId ? 100 : null
+  );
   const contentRef = useRef<HTMLDivElement>(null);
 
+  const {isPlayground} = usePlaygroundContext();
   useEffect(() => {
     if (contentRef.current) {
       setIsOverflowing(contentRef.current.scrollHeight > 400);
@@ -37,10 +53,18 @@ export const MessagePanel = ({
     message.tool_calls != null && message.tool_calls.length > 0;
   const hasContent = message.content != null && message.content.length > 0;
 
+  const responseIndexes: number[] | undefined = hasToolCalls
+    ? message
+        .tool_calls!.map(
+          (toolCall: ToolCall) => toolCall.response?.original_index
+        )
+        .filter((idx): idx is number => idx !== undefined)
+    : undefined;
+
   return (
     <div className={classNames('flex gap-8', {'mt-24': !isTool})}>
       {!isNested && !isSystemPrompt && (
-        <div className="w-32">
+        <div className="w-32 flex-shrink-0">
           {!isUser && !isTool && (
             <Callout
               size="small"
@@ -53,16 +77,19 @@ export const MessagePanel = ({
       )}
 
       <div
-        className={classNames('relative overflow-visible rounded-lg', {
+        className={classNames('relative w-full overflow-visible', {
+          'rounded-lg': !isNested,
           'border-t border-moon-250': isTool,
-          'bg-moon-100': isSystemPrompt,
+          'bg-moon-100': isSystemPrompt || hasToolCalls,
           'bg-cactus-300/[0.24]': isUser,
-          'w-full': !isUser,
-          'max-w-3xl': isUser,
+          'max-w-full': !isUser,
+          'max-w-[768px]': isUser,
           'ml-auto': isUser,
           'mr-auto': !isUser,
           'py-8': hasContent,
-        })}>
+        })}
+        onMouseEnter={() => setIsHovering(true)}
+        onMouseLeave={() => setIsHovering(false)}>
         <div>
           {isSystemPrompt && (
             <div className="flex justify-between px-16">
@@ -73,7 +100,7 @@ export const MessagePanel = ({
           )}
 
           {isTool && (
-            <div className={classNames({'px-16': isNested}, 'pb-8')}>
+            <div className={classNames('pb-8')}>
               <div className="text-sm font-semibold text-moon-500">
                 Response
               </div>
@@ -86,39 +113,75 @@ export const MessagePanel = ({
               'max-h-[400px]': !isShowingMore,
               'max-h-full': isShowingMore,
             })}>
-            {hasContent && (
-              <div
-                className={classNames(hasToolCalls ? 'pb-8' : '', ' text-sm', {
-                  'px-16': isSystemPrompt || isUser,
-                })}>
-                {_.isString(message.content) ? (
-                  <MessagePanelPart
-                    value={message.content}
-                    isStructuredOutput={isStructuredOutput}
-                  />
-                ) : (
-                  message.content!.map((p, i) => (
-                    <MessagePanelPart key={i} value={p} />
-                  ))
+            {isPlayground && editorHeight ? (
+              <PlaygroundMessagePanelEditor
+                message={message}
+                index={index}
+                isChoice={isChoice ?? false}
+                editorHeight={editorHeight}
+                isNested={isNested ?? false}
+                pendingToolResponseId={pendingToolResponseId}
+                setEditorHeight={setEditorHeight}
+              />
+            ) : (
+              <>
+                {hasContent && (
+                  <div
+                    className={classNames(
+                      hasToolCalls ? 'pb-8' : '',
+                      ' text-sm',
+                      {'px-16': isSystemPrompt || isUser}
+                    )}>
+                    {_.isString(message.content) ? (
+                      <MessagePanelPart
+                        value={message.content}
+                        isStructuredOutput={isStructuredOutput}
+                      />
+                    ) : (
+                      message.content!.map((p, i) => (
+                        <MessagePanelPart key={i} value={p} />
+                      ))
+                    )}
+                  </div>
                 )}
-              </div>
-            )}
-            {hasToolCalls && (
-              <div
-                className={classNames({
-                  'border-t border-moon-250 pt-8': hasContent,
-                })}>
-                <ToolCalls toolCalls={message.tool_calls!} />
-              </div>
+                {hasToolCalls && (
+                  <div
+                    className={classNames({
+                      'border-t border-moon-250 pt-8': hasContent,
+                    })}>
+                    <ToolCalls toolCalls={message.tool_calls!} />
+                  </div>
+                )}
+              </>
             )}
           </div>
-          {isOverflowing && (
+
+          {isOverflowing && !editorHeight && (
             <ShowMoreButton
               isUser={isUser}
               isShowingMore={isShowingMore}
               setIsShowingMore={setIsShowingMore}
             />
           )}
+
+          {/* Playground buttons (retry, edit, delete) */}
+          {isPlayground && isHovering && !editorHeight && (
+            <div
+              className={classNames(
+                'absolute flex w-full items-center justify-start pt-20',
+                isNested ? 'bottom-0' : 'bottom-[-32px]'
+              )}>
+              <PlaygroundMessagePanelButtons
+                index={message.original_index ?? index}
+                isChoice={isChoice ?? false}
+                isTool={isTool}
+                hasContent={hasContent}
+                contentRef={contentRef}
+                setEditorHeight={setEditorHeight}
+                responseIndexes={responseIndexes}
+              />
+            </div>
+          )}
         </div>
       </div>
     </div>

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/PlaygroundMessagePanelButtons.tsx

@@ -0,0 +1,79 @@
+import {Button} from '@wandb/weave/components/Button';
+import React from 'react';
+
+import {usePlaygroundContext} from '../PlaygroundPage/PlaygroundContext';
+
+type PlaygroundMessagePanelButtonsProps = {
+  index: number;
+  isChoice: boolean;
+  isTool: boolean;
+  hasContent: boolean;
+  contentRef: React.RefObject<HTMLDivElement>;
+  setEditorHeight: (height: number | null) => void;
+  responseIndexes?: number[];
+};
+
+export const PlaygroundMessagePanelButtons: React.FC<
+  PlaygroundMessagePanelButtonsProps
+> = ({
+  index,
+  isChoice,
+  isTool,
+  hasContent,
+  contentRef,
+  setEditorHeight,
+  responseIndexes,
+}) => {
+  const {deleteMessage, deleteChoice, retry} = usePlaygroundContext();
+
+  return (
+    <div className="z-10 flex gap-4 rounded-lg border border-moon-250 bg-white p-4">
+      <Button
+        variant="quiet"
+        size="small"
+        startIcon="randomize-reset-reload"
+        onClick={() => retry?.(index, isChoice)}
+        tooltip={
+          !hasContent
+            ? 'We currently do not support retrying functions'
+            : 'Retry'
+        }
+        disabled={!hasContent}>
+        Retry
+      </Button>
+      <Button
+        variant="quiet"
+        size="small"
+        startIcon="pencil-edit"
+        onClick={() => {
+          setEditorHeight(
+            contentRef?.current?.clientHeight
+              ? // Accounts for padding and save buttons
+                contentRef.current.clientHeight - 56
+              : null
+          );
+        }}
+        tooltip={
+          !hasContent ? 'We currently do not support editing functions' : 'Edit'
+        }
+        disabled={!hasContent}>
+        Edit
+      </Button>
+      <Button
+        variant="quiet"
+        size="small"
+        startIcon="delete"
+        onClick={() => {
+          if (isChoice) {
+            deleteChoice?.(index);
+          } else {
+            deleteMessage?.(index, responseIndexes);
+          }
+        }}
+        tooltip={isTool ? 'Tool responses cannot be deleted' : 'Delete message'}
+        disabled={isTool}>
+        Delete
+      </Button>
+    </div>
+  );
+};