system role (#3098)

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

@@ -6,7 +6,7 @@ import React, {useState} from 'react';
 import {CallChat} from '../../CallPage/CallChat';
 import {TraceCallSchema} from '../../wfReactInterface/traceServerClientTypes';
 import {PlaygroundContext} from '../PlaygroundContext';
-import {PlaygroundState} from '../types';
+import {PlaygroundMessageRole, PlaygroundState} from '../types';
 import {PlaygroundCallStats} from './PlaygroundCallStats';
 import {PlaygroundChatInput} from './PlaygroundChatInput';
 import {PlaygroundChatTopBar} from './PlaygroundChatTopBar';
@@ -51,7 +51,7 @@ export const PlaygroundChat = ({
   const {deleteMessage, editMessage, deleteChoice, editChoice, addMessage} =
     useChatFunctions(setPlaygroundStateField);
 
-  const handleAddMessage = (role: 'assistant' | 'user', text: string) => {
+  const handleAddMessage = (role: PlaygroundMessageRole, text: string) => {
     for (let i = 0; i < playgroundStates.length; i++) {
       addMessage(i, {role, content: text});
     }
@@ -160,7 +160,7 @@ export const PlaygroundChat = ({
                           retry: (messageIndex: number, isChoice?: boolean) =>
                             handleRetry(idx, messageIndex, isChoice),
                           sendMessage: (
-                            role: 'assistant' | 'user' | 'tool',
+                            role: PlaygroundMessageRole,
                             content: string,
                             toolCallId?: string
                           ) => {

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/PlaygroundChatInput.tsx

@@ -4,13 +4,14 @@ import {Button} from '@wandb/weave/components/Button';
 import React, {useState} from 'react';
 
 import {StyledTextArea} from '../StyledTextarea';
+import {PlaygroundMessageRole} from '../types';
 
 type PlaygroundChatInputProps = {
   chatText: string;
   setChatText: (text: string) => void;
   isLoading: boolean;
-  onSend: (role: 'assistant' | 'user') => void;
-  onAdd: (role: 'assistant' | 'user', text: string) => void;
+  onSend: (role: PlaygroundMessageRole) => void;
+  onAdd: (role: PlaygroundMessageRole, text: string) => void;
   settingsTab: number | null;
 };
 
@@ -32,22 +33,21 @@ export const PlaygroundChatInput: React.FC<PlaygroundChatInputProps> = ({
   onAdd,
   settingsTab,
 }) => {
-  const [addMessageRole, setAddMessageRole] = useState<'assistant' | 'user'>(
-    'user'
-  );
+  const [addMessageRole, setAddMessageRole] =
+    useState<PlaygroundMessageRole>('user');
   const [shouldReset, setShouldReset] = useState(false);
 
   const handleReset = () => {
     setShouldReset(true);
     setTimeout(() => setShouldReset(false), 0);
   };
 
-  const handleSend = (role: 'assistant' | 'user') => {
+  const handleSend = (role: PlaygroundMessageRole) => {
     onSend(role);
     handleReset();
   };
 
-  const handleAdd = (role: 'assistant' | 'user', text: string) => {
+  const handleAdd = (role: PlaygroundMessageRole, text: string) => {
     onAdd(role, text);
     handleReset();
   };
@@ -115,6 +115,14 @@ export const PlaygroundChatInput: React.FC<PlaygroundChatInputProps> = ({
                 className="ml-4 rounded-r-none"
                 variant="secondary"
                 size="medium"
+                active={addMessageRole === 'system'}
+                onClick={() => setAddMessageRole('system')}>
+                System
+              </Button>
+              <Button
+                className="rounded-none"
+                variant="secondary"
+                size="medium"
                 active={addMessageRole === 'assistant'}
                 onClick={() => setAddMessageRole('assistant')}>
                 Assistant

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/useChatCompletionFunctions.tsx

@@ -6,6 +6,7 @@ import {Message} from '../../ChatView/types';
 import {useGetTraceServerClientContext} from '../../wfReactInterface/traceServerClientContext';
 import {CompletionsCreateRes} from '../../wfReactInterface/traceServerClientTypes';
 import {PlaygroundState} from '../types';
+import {PlaygroundMessageRole} from '../types';
 import {getInputFromPlaygroundState} from '../usePlaygroundState';
 import {clearTraceCall} from './useChatFunctions';
 
@@ -57,7 +58,7 @@ export const useChatCompletionFunctions = (
   };
 
   const handleSend = async (
-    role: 'assistant' | 'user' | 'tool',
+    role: PlaygroundMessageRole,
     callIndex?: number,
     content?: string,
     toolCallId?: string
@@ -135,7 +136,7 @@ export const useChatCompletionFunctions = (
 
 // Helper functions
 const createMessage = (
-  role: 'assistant' | 'user' | 'tool',
+  role: PlaygroundMessageRole,
   content: string,
   toolCallId?: string
 ): Message | undefined => {

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

@@ -1,6 +1,7 @@
 import {createContext, useContext} from 'react';
 
 import {Message} from '../ChatView/types';
+import {PlaygroundMessageRole} from './types';
 
 export type PlaygroundContextType = {
   isPlayground: boolean;
@@ -13,7 +14,7 @@ export type PlaygroundContextType = {
 
   retry: (messageIndex: number, isChoice?: boolean) => void;
   sendMessage: (
-    role: 'assistant' | 'user' | 'tool',
+    role: PlaygroundMessageRole,
     content: string,
     toolCallId?: string
   ) => void;

weave-js/src/components/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/types.ts

@@ -28,3 +28,5 @@ export type PlaygroundState = {
 export type PlaygroundStateKey = keyof PlaygroundState;
 
 export type OptionalTraceCallSchema = Partial<TraceCallSchema>;
+
+export type PlaygroundMessageRole = 'assistant' | 'user' | 'system' | 'tool';