fix(weave): fix anthropic chats in playground (#3213)

* fix anthropic chats in playground * deep clone * lint again * idk why this isnt auto sorting
wandb · Dec 12, 2024 · cdf6574 · cdf6574
1 parent 32fc0da
commit cdf6574
Show file tree

Hide file tree

Showing 5 changed files with 53 additions and 8 deletions.
diff --git a/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/ChoicesView.tsx b/weave-js/src/components/PagePanelComponents/Home/Browse3/pages/ChatView/ChoicesView.tsx
@@ -31,7 +31,11 @@ export const ChoicesView = ({
   }
   if (choices.length === 1) {
     return (
-      <ChoiceView choice={choices[0]} isStructuredOutput={isStructuredOutput} />
+      <ChoiceView
+        choice={choices[0]}
+        isStructuredOutput={isStructuredOutput}
+        choiceIndex={0}
+      />
     );
   }
   return (

diff --git a/...s/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/PlaygroundChat.tsx b/...s/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/PlaygroundChat.tsx
@@ -35,7 +35,6 @@ export const PlaygroundChat = ({
   setSettingsTab,
   settingsTab,
 }: PlaygroundChatProps) => {
-  console.log('playgroundStates', playgroundStates);
   const [chatText, setChatText] = useState('');
   const [isLoading, setIsLoading] = useState(false);
 

diff --git a/...PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/useChatFunctions.tsx b/...PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/useChatFunctions.tsx
@@ -43,8 +43,6 @@ export const useChatFunctions = (
     messageIndex: number,
     newMessage: Message
   ) => {
-    console.log('editMessage', callIndex, messageIndex, newMessage);
-
     setPlaygroundStateField(callIndex, 'traceCall', prevTraceCall => {
       const newTraceCall = clearTraceCall(
         cloneDeep(prevTraceCall as OptionalTraceCallSchema)
@@ -108,7 +106,6 @@ export const useChatFunctions = (
     choiceIndex: number,
     newChoice: Message
   ) => {
-    console.log('editChoice', callIndex, choiceIndex, newChoice);
     setPlaygroundStateField(callIndex, 'traceCall', prevTraceCall => {
       const newTraceCall = clearTraceCall(
         cloneDeep(prevTraceCall as OptionalTraceCallSchema)

diff --git a/...s/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundPage.tsx b/...s/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundPage.tsx
@@ -7,7 +7,11 @@ import {SimplePageLayoutWithHeader} from '../common/SimplePageLayout';
 import {useWFHooks} from '../wfReactInterface/context';
 import {PlaygroundChat} from './PlaygroundChat/PlaygroundChat';
 import {PlaygroundSettings} from './PlaygroundSettings/PlaygroundSettings';
-import {DEFAULT_SYSTEM_MESSAGE, usePlaygroundState} from './usePlaygroundState';
+import {
+  DEFAULT_SYSTEM_MESSAGE,
+  parseTraceCall,
+  usePlaygroundState,
+} from './usePlaygroundState';
 
 export type PlaygroundPageProps = {
   entity: string;
@@ -89,7 +93,10 @@ export const PlaygroundPageInner = (props: PlaygroundPageProps) => {
       for (const [idx, state] of newStates.entries()) {
         for (const c of calls || []) {
           if (state.traceCall.id === c.callId) {
-            newStates[idx] = {...state, traceCall: c.traceCall || {}};
+            newStates[idx] = {
+              ...state,
+              traceCall: parseTraceCall(c.traceCall || {}),
+            };
             break;
           }
         }

diff --git a/...rc/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/usePlaygroundState.ts b/...rc/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/usePlaygroundState.ts
@@ -1,5 +1,11 @@
+import {cloneDeep} from 'lodash';
 import {SetStateAction, useCallback, useState} from 'react';
 
+import {
+  anthropicContentBlocksToChoices,
+  hasStringProp,
+  isAnthropicCompletionFormat,
+} from '../ChatView/hooks';
 import {LLM_MAX_TOKENS_KEYS, LLMMaxTokensKey} from './llmMaxTokens';
 import {
   OptionalTraceCallSchema,
@@ -77,7 +83,7 @@ export const usePlaygroundState = () => {
       setPlaygroundStates(prevState => {
         const newState = {...prevState[0]};
 
-        newState.traceCall = traceCall;
+        newState.traceCall = parseTraceCall(traceCall);
 
         if (!inputs) {
           return [newState];
@@ -155,3 +161,35 @@ export const getInputFromPlaygroundState = (state: PlaygroundState) => {
     tools: tools.length > 0 ? tools : undefined,
   };
 };
+
+// This is a helper function to parse the trace call output for anthropic
+// so that the playground can display the choices
+export const parseTraceCall = (traceCall: OptionalTraceCallSchema) => {
+  const parsedTraceCall = cloneDeep(traceCall);
+
+  // Handles anthropic outputs
+  // Anthropic has content and stop_reason as top-level fields
+  if (isAnthropicCompletionFormat(parsedTraceCall.output)) {
+    const {content, stop_reason, ...outputs} = parsedTraceCall.output as any;
+    parsedTraceCall.output = {
+      ...outputs,
+      choices: anthropicContentBlocksToChoices(content, stop_reason),
+    };
+  }
+  // Handles anthropic inputs
+  // Anthropic has system message as a top-level request field
+  if (hasStringProp(parsedTraceCall.inputs, 'system')) {
+    const {messages, system, ...inputs} = parsedTraceCall.inputs as any;
+    parsedTraceCall.inputs = {
+      ...inputs,
+      messages: [
+        {
+          role: 'system',
+          content: system,
+        },
+        ...messages,
+      ],
+    };
+  }
+  return parsedTraceCall;
+};