style(weave): chatview styling

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

@@ -1,8 +1,7 @@
-import React, {useEffect, useRef} from 'react';
+import React, {useEffect, useMemo, useRef} from 'react';
 
 import {useDeepMemo} from '../../../../../../hookUtils';
 import {ChoicesView} from './ChoicesView';
-import {HorizontalRuleWithLabel} from './HorizontalRuleWithLabel';
 import {MessageList} from './MessageList';
 import {Chat} from './types';
 
@@ -15,24 +14,38 @@ export const ChatView = ({chat}: ChatViewProps) => {
 
   const chatResult = useDeepMemo(chat.result);
 
+  const scrollLastMessage = useMemo(
+    () => !(outputRef.current && chatResult && chatResult.choices),
+    [chatResult]
+  );
+
   useEffect(() => {
     if (outputRef.current && chatResult && chatResult.choices) {
       outputRef.current.scrollIntoView();
     }
   }, [chatResult]);
 
   return (
-    <div>
-      <HorizontalRuleWithLabel label="Input" />
-      <MessageList messages={chat.request.messages} />
+    <div className="flex flex-col pb-32">
+      <span className="mb-[-16px] text-sm font-semibold text-moon-800">
+        Messages
+      </span>
+      <MessageList
+        messages={chat.request?.messages || []}
+        scrollLastMessage={scrollLastMessage}
+      />
       {chatResult && chatResult.choices && (
-        <div className="mt-12" ref={outputRef}>
-          <HorizontalRuleWithLabel label="Output" />
-          <ChoicesView
-            isStructuredOutput={chat.isStructuredOutput}
-            choices={chatResult.choices}
-          />
-        </div>
+        <>
+          <span className="mt-16 text-sm font-semibold text-moon-800">
+            Response
+          </span>
+          <div ref={outputRef}>
+            <ChoicesView
+              isStructuredOutput={chat.isStructuredOutput}
+              choices={chatResult.choices}
+            />
+          </div>
+        </>
       )}
     </div>
   );

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

@@ -11,6 +11,10 @@ type ChoiceViewProps = {
 export const ChoiceView = ({choice, isStructuredOutput}: ChoiceViewProps) => {
   const {message} = choice;
   return (
-    <MessagePanel message={message} isStructuredOutput={isStructuredOutput} />
+    <MessagePanel
+      index={choice.index}
+      message={message}
+      isStructuredOutput={isStructuredOutput}
+    />
   );
 };

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

@@ -1,18 +1,81 @@
-import React from 'react';
+import React, {useEffect, useRef} from 'react';
 
 import {MessagePanel} from './MessagePanel';
-import {Messages} from './types';
+import {Message, Messages, ToolCall} from './types';
 
 type MessageListProps = {
   messages: Messages;
+  scrollLastMessage?: boolean;
 };
 
-export const MessageList = ({messages}: MessageListProps) => {
+export const MessageList = ({
+  messages,
+  scrollLastMessage = false,
+}: MessageListProps) => {
+  const lastMessageRef = useRef<HTMLDivElement>(null);
+  const processedMessages = processToolCallMessages(messages);
+
+  useEffect(() => {
+    if (lastMessageRef.current && scrollLastMessage) {
+      lastMessageRef.current.scrollIntoView();
+    }
+  }, [messages.length, scrollLastMessage]);
+
   return (
-    <div className="flex flex-col gap-36">
-      {messages.map((m, i) => (
-        <MessagePanel key={i} message={m} />
+    <div className="flex flex-col">
+      {processedMessages.map((m, i) => (
+        <div
+          ref={i === processedMessages.length - 1 ? lastMessageRef : null}
+          key={i}>
+          <MessagePanel index={m.original_index ?? i} message={m} />
+        </div>
       ))}
     </div>
   );
 };
+
+// Associates tool calls with their responses
+const processToolCallMessages = (messages: Messages): Messages => {
+  const processedMessages: Message[] = [];
+  for (let i = 0; i < messages.length; i++) {
+    const message = messages[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,
+      });
+      continue;
+    }
+
+    // Otherwise, we need to associate the tool calls with their responses.
+    // Get all the next messages where role = tool, these are all the responses
+    const toolMessages: Message[] = [];
+    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,
+      });
+      i++;
+    }
+
+    const toolCallsWithResponses: ToolCall[] = message.tool_calls.map(
+      toolCall => ({
+        ...toolCall,
+        response: toolMessages.find(
+          toolMessage => toolMessage.tool_call_id === toolCall.id
+        ),
+      })
+    );
+
+    processedMessages.push({
+      ...message,
+      tool_calls: toolCallsWithResponses,
+    });
+  }
+  return processedMessages;
+};

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

@@ -1,3 +1,4 @@
+import {Callout} from '@wandb/weave/components/Callout';
 import classNames from 'classnames';
 import _ from 'lodash';
 import React from 'react';
@@ -7,34 +8,96 @@ import {ToolCalls} from './ToolCalls';
 import {Message} from './types';
 
 type MessagePanelProps = {
+  index: number;
   message: Message;
   isStructuredOutput?: boolean;
+  isNested?: boolean;
 };
 
 export const MessagePanel = ({
   message,
   isStructuredOutput,
+  isNested,
 }: MessagePanelProps) => {
   const isUser = message.role === 'user';
-  const bg = isUser ? 'bg-cactus-300/[0.48]' : 'bg-moon-100';
+  const isSystemPrompt = message.role === 'system';
+  const isTool = message.role === 'tool';
+  const hasToolCalls =
+    message.tool_calls != null && message.tool_calls.length > 0;
+  const hasContent = message.content != null && message.content.length > 0;
+
   return (
-    <div className={classNames('rounded-[8px] px-16 py-8', bg)}>
-      <div style={{fontVariantCaps: 'all-small-caps'}}>{message.role}</div>
-      {message.content && (
-        <div>
-          {_.isString(message.content) ? (
-            <MessagePanelPart
-              value={message.content}
-              isStructuredOutput={isStructuredOutput}
+    <div className={classNames('flex gap-8', {'mt-24': !isTool})}>
+      {!isNested && !isSystemPrompt && (
+        <div className="w-32">
+          {!isUser && !isTool && (
+            <Callout
+              size="small"
+              icon="robot-service-member"
+              color="moon"
+              className="h-32 w-32"
             />
-          ) : (
-            message.content.map((p, i) => (
-              <MessagePanelPart key={i} value={p} />
-            ))
           )}
         </div>
       )}
-      {message.tool_calls && <ToolCalls toolCalls={message.tool_calls} />}
+
+      <div
+        className={classNames('relative overflow-visible rounded-lg', {
+          'border-t border-moon-250': isTool,
+          'bg-moon-100': isSystemPrompt,
+          'bg-cactus-300/[0.24]': isUser,
+          'w-full': !isUser,
+          'max-w-3xl': isUser,
+          'ml-auto': isUser,
+          'mr-auto': !isUser,
+          'py-8': hasContent,
+        })}>
+        <div>
+          {isSystemPrompt && (
+            <div className="flex justify-between px-16">
+              <div className="text-sm text-moon-500">
+                {message.role.charAt(0).toUpperCase() + message.role.slice(1)}
+              </div>
+            </div>
+          )}
+
+          {isTool && (
+            <div className={classNames({'px-16': isNested}, 'pb-8')}>
+              <div className="text-sm font-semibold text-moon-500">
+                Response
+              </div>
+            </div>
+          )}
+
+          <div className={classNames('w-full overflow-y-hidden')}>
+            {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>
+            )}
+            {hasToolCalls && (
+              <div
+                className={classNames({
+                  'border-t border-moon-250 pt-8': hasContent,
+                })}>
+                <ToolCalls toolCalls={message.tool_calls!} />
+              </div>
+            )}
+          </div>
+        </div>
+      </div>
     </div>
   );
 };