Merge branch 'master' into tim/basic_batch_actions_UI

weave-js/src/components/PagePanelComponents/Home/Browse2/Browse2RootObjectVersionOutputOf.tsx

@@ -1,11 +1,11 @@
 import Typography from '@mui/material/Typography';
-import {useNodeValue} from '@wandb/weave/react';
+import {parseRefMaybe, useNodeValue} from '@wandb/weave/react';
 import React, {FC, useMemo} from 'react';
 import {useParams} from 'react-router-dom';
 
 import {callsTableFilter, callsTableNode, callsTableOpCounts} from './callTree';
 import {Browse2RootObjectVersionItemParams} from './CommonLib';
-import {parseRefMaybe, SmallRef} from './SmallRef';
+import {SmallRef} from './SmallRef';
 
 export const Browse2RootObjectVersionOutputOf: FC<{uri: string}> = ({uri}) => {
   const params = useParams<Browse2RootObjectVersionItemParams>();

weave-js/src/components/PagePanelComponents/Home/Browse2/SpanDetails.tsx

@@ -4,6 +4,7 @@ import * as globals from '@wandb/weave/common/css/globals.styles';
 import * as _ from 'lodash';
 import React, {FC} from 'react';
 
+import {parseRefMaybe} from '../../../../react';
 import {StatusChip} from '../Browse3/pages/common/StatusChip';
 import {Call} from './callTree';
 import {DisplayControlChars} from './CommonLib';
@@ -13,7 +14,7 @@ import {
   OpenAIChatInputView,
   OpenAIChatOutputView,
 } from './openai';
-import {parseRefMaybe, SmallRef} from './SmallRef';
+import {SmallRef} from './SmallRef';
 
 const ObjectView: FC<{obj: any}> = ({obj}) => {
   if (_.isPlainObject(obj)) {

weave-js/src/components/PagePanelComponents/Home/Browse3/filters/common.ts

@@ -8,10 +8,9 @@ import {
   GridFilterItem,
   GridFilterModel,
 } from '@mui/x-data-grid-pro';
-import {isWeaveObjectRef} from '@wandb/weave/react';
+import {isWeaveObjectRef, parseRefMaybe} from '@wandb/weave/react';
 import _ from 'lodash';
 
-import {parseRefMaybe} from '../../Browse2/SmallRef';
 import {WEAVE_REF_PREFIX} from '../pages/wfReactInterface/constants';
 import {TraceCallSchema} from '../pages/wfReactInterface/traceServerClientTypes';
 

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

@@ -252,9 +252,9 @@ const CallPageInnerVertical: FC<{
           </Box>
         </Box>
       }
-      isSidebarOpen={showTraceTree}
       headerContent={<CallOverview call={currentCall} />}
-      leftSidebar={
+      isLeftSidebarOpen={showTraceTree}
+      leftSidebarContent={
         <Tailwind style={{display: 'contents'}}>
           <div className="h-full bg-moon-50">
             {loading ? (

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

@@ -1,9 +1,10 @@
 import _ from 'lodash';
 import React from 'react';
 
+import {parseRefMaybe} from '../../../../../../react';
 import {Timestamp} from '../../../../../Timestamp';
 import {UserLink} from '../../../../../UserLink';
-import {parseRefMaybe, SmallRef} from '../../../Browse2/SmallRef';
+import {SmallRef} from '../../../Browse2/SmallRef';
 import {SimpleKeyValueTable} from '../common/SimplePageLayout';
 import {CallSchema} from '../wfReactInterface/wfDataModelHooksInterface';
 import {CostTable} from './cost';

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

@@ -28,10 +28,13 @@ import React, {
 } from 'react';
 import {useHistory} from 'react-router-dom';
 
-import {isWeaveObjectRef, parseRef} from '../../../../../../react';
+import {
+  isWeaveObjectRef,
+  parseRef,
+  parseRefMaybe,
+} from '../../../../../../react';
 import {flattenObjectPreservingWeaveTypes} from '../../../Browse2/browse2Util';
 import {CellValue} from '../../../Browse2/CellValue';
-import {parseRefMaybe} from '../../../Browse2/SmallRef';
 import {
   useWeaveflowCurrentRouteContext,
   WeaveflowPeekContext,

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

@@ -18,9 +18,9 @@ import React, {
   useState,
 } from 'react';
 
+import {parseRefMaybe} from '../../../../../../react';
 import {LoadingDots} from '../../../../../LoadingDots';
 import {Browse2OpDefCode} from '../../../Browse2/Browse2OpDefCode';
-import {parseRefMaybe} from '../../../Browse2/SmallRef';
 import {isWeaveRef} from '../../filters/common';
 import {StyledDataGrid} from '../../StyledDataGrid';
 import {isCustomWeaveTypePayload} from '../../typeViews/customWeaveType.types';

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

@@ -1,7 +1,11 @@
 import React, {useMemo} from 'react';
 
-import {isWeaveObjectRef, parseRef} from '../../../../../../react';
-import {parseRefMaybe, SmallRef} from '../../../Browse2/SmallRef';
+import {
+  isWeaveObjectRef,
+  parseRef,
+  parseRefMaybe,
+} from '../../../../../../react';
+import {SmallRef} from '../../../Browse2/SmallRef';
 import {isWeaveRef} from '../../filters/common';
 import {isCustomWeaveTypePayload} from '../../typeViews/customWeaveType.types';
 import {CustomWeaveTypeDispatcher} from '../../typeViews/CustomWeaveTypeDispatcher';

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;
+};