chore(weave): Llm dropdown (#2939)

* add settings drawer with simple toggle * lint * add top bar * add llm dropdown * prttier
wandb · Nov 8, 2024 · 7c9a7b1 · 7c9a7b1
1 parent e9c0ee4
commit 7c9a7b1
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 0 deletions.
diff --git a/...ents/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/LLMDropdown.tsx b/...ents/PagePanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/LLMDropdown.tsx
@@ -0,0 +1,47 @@
+import {Box} from '@mui/material';
+import {Select} from '@wandb/weave/components/Form/Select';
+import React from 'react';
+
+import {LLM_MAX_TOKENS} from '../llmMaxTokens';
+
+interface LLMDropdownProps {
+  value: string;
+  onChange: (value: string, maxTokens: number) => void;
+}
+
+export const LLMDropdown: React.FC<LLMDropdownProps> = ({value, onChange}) => {
+  const options = Object.keys(LLM_MAX_TOKENS).map(llm => ({
+    value: llm,
+    label: llm,
+  }));
+
+  return (
+    <Box
+      sx={{
+        width: '150px',
+        '& #react-select-2-listbox': {
+          width: '300px',
+          maxHeight: '500px',
+        },
+        '& #react-select-2-listbox > div': {
+          maxHeight: '500px',
+        },
+      }}>
+      <Select
+        value={options.find(opt => opt.value === value)}
+        onChange={option => {
+          if (option) {
+            const maxTokens =
+              LLM_MAX_TOKENS[
+                (option as {value: string}).value as keyof typeof LLM_MAX_TOKENS
+              ]?.max_tokens || 0;
+            onChange((option as {value: string}).value, maxTokens);
+          }
+        }}
+        options={options}
+        size="medium"
+        isSearchable
+      />
+    </Box>
+  );
+};
diff --git a/...PanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/PlaygroundChatTopBar.tsx b/...PanelComponents/Home/Browse3/pages/PlaygroundPage/PlaygroundChat/PlaygroundChatTopBar.tsx
@@ -7,6 +7,7 @@ import {useHistory} from 'react-router-dom';
 import {CopyableId} from '../../common/Id';
 import {OptionalTraceCallSchema, PlaygroundState} from '../types';
 import {DEFAULT_SYSTEM_MESSAGE} from '../usePlaygroundState';
+import {LLMDropdown} from './LLMDropdown';
 
 type PlaygroundChatTopBarProps = {
   idx: number;
@@ -57,6 +58,16 @@ export const PlaygroundChatTopBar: React.FC<PlaygroundChatTopBarProps> = ({
     }
   };
 
+  const handleModelChange = (
+    index: number,
+    model: string,
+    maxTokens: number
+  ) => {
+    setPlaygroundStateField(index, 'model', model);
+    setPlaygroundStateField(index, 'maxTokensLimit', maxTokens);
+    setPlaygroundStateField(index, 'maxTokens', maxTokens / 2);
+  };
+
   return (
     <Box
       sx={{
@@ -73,6 +84,12 @@ export const PlaygroundChatTopBar: React.FC<PlaygroundChatTopBarProps> = ({
           backgroundColor: 'white',
         }}>
         {!onlyOneChat && <Tag label={`${idx + 1}`} />}
+        <LLMDropdown
+          value={playgroundStates[idx].model}
+          onChange={(model, maxTokens) =>
+            handleModelChange(idx, model, maxTokens)
+          }
+        />
         {playgroundStates[idx].traceCall?.id && (
           <CopyableId id={playgroundStates[idx]!.traceCall!.id!} type="Call" />
         )}