feat(playground): add tools ui

app/package.json

@@ -6,7 +6,7 @@
   "license": "None",
   "private": true,
   "dependencies": {
-    "@arizeai/components": "^1.8.6",
+    "@arizeai/components": "^1.8.7",
     "@arizeai/openinference-semantic-conventions": "^0.10.0",
     "@arizeai/point-cloud": "^3.0.6",
     "@codemirror/autocomplete": "6.12.0",
@@ -27,6 +27,7 @@
     "@uiw/codemirror-theme-github": "^4.23.5",
     "@uiw/codemirror-theme-nord": "^4.23.0",
     "@uiw/react-codemirror": "^4.23.0",
+    "codemirror-json-schema": "^0.7.8",
     "copy-to-clipboard": "^3.3.3",
     "d3-format": "^3.1.0",
     "d3-scale-chromatic": "^3.1.0",
@@ -55,6 +56,7 @@
     "use-deep-compare-effect": "^1.8.1",
     "use-zustand": "^0.0.4",
     "zod": "^3.23.8",
+    "zod-to-json-schema": "^3.23.3",
     "zustand": "^4.5.4"
   },
   "devDependencies": {
@@ -64,6 +66,7 @@
     "@types/d3-format": "^3.0.4",
     "@types/d3-scale-chromatic": "^3.0.3",
     "@types/d3-time-format": "^4.0.3",
+    "@types/json-schema": "^7.0.15",
     "@types/lodash": "^4.17.7",
     "@types/node": "^22.5.4",
     "@types/react": "18.3.10",

app/src/@types/generative.d.ts

@@ -4,3 +4,14 @@ declare type ModelProvider = "OPENAI" | "AZURE_OPENAI" | "ANTHROPIC";
  * The role of a chat message
  */
 declare type ChatMessageRole = "user" | "system" | "ai" | "tool";
+
+/**
+ * The tool picking mechanism for an LLM
+ * Either "auto", "required", "none", or a specific tool
+ * @see https://platform.openai.com/docs/api-reference/chat/create#chat-create-tool_choice
+ */
+declare type ToolChoice =
+  | "auto"
+  | "required"
+  | "none"
+  | { type: "function"; function: { name: string } };

app/src/components/code/CodeEditorFieldWrapper.tsx

@@ -27,6 +27,7 @@ const codeEditorFormWrapperCSS = css`
   .cm-editor {
     border-radius: var(--ac-global-rounding-small);
   }
+  box-sizing: border-box;
   .cm-focused {
     outline: none;
   }

app/src/components/code/JSONEditor.tsx

@@ -1,24 +1,56 @@
-import React from "react";
-import { json, jsonParseLinter } from "@codemirror/lang-json";
+import React, { useMemo } from "react";
+import { json, jsonLanguage, jsonParseLinter } from "@codemirror/lang-json";
 import { linter } from "@codemirror/lint";
-import { EditorView } from "@codemirror/view";
+import { EditorView, hoverTooltip } from "@codemirror/view";
+import { githubLight } from "@uiw/codemirror-theme-github";
 import { nord } from "@uiw/codemirror-theme-nord";
 import CodeMirror, { ReactCodeMirrorProps } from "@uiw/react-codemirror";
+import {
+  handleRefresh,
+  jsonCompletion,
+  jsonSchemaHover,
+  jsonSchemaLinter,
+  stateExtensions,
+} from "codemirror-json-schema";
+import { JSONSchema7 } from "json-schema";
 
 import { useTheme } from "@phoenix/contexts";
 
 export type JSONEditorProps = Omit<
   ReactCodeMirrorProps,
   "theme" | "extensions" | "editable"
->;
+> & {
+  /**
+   * JSON Schema to use for validation, if provided will enable JSON Schema validation with tooltips in the editor
+   */
+  jsonSchema?: JSONSchema7;
+};
 
 export function JSONEditor(props: JSONEditorProps) {
   const { theme } = useTheme();
-  const codeMirrorTheme = theme === "light" ? undefined : nord;
+  const codeMirrorTheme = theme === "light" ? githubLight : nord;
+  const extensions = useMemo(() => {
+    const baseExtensions = [
+      json(),
+      EditorView.lineWrapping,
+      linter(jsonParseLinter()),
+    ];
+    if (props.jsonSchema) {
+      baseExtensions.push(
+        linter(jsonSchemaLinter(), { needsRefresh: handleRefresh }),
+        jsonLanguage.data.of({
+          autocomplete: jsonCompletion(),
+        }),
+        hoverTooltip(jsonSchemaHover()),
+        stateExtensions(props.jsonSchema)
+      );
+    }
+    return baseExtensions;
+  }, [props.jsonSchema]);
   return (
     <CodeMirror
       value={props.value}
-      extensions={[json(), EditorView.lineWrapping, linter(jsonParseLinter())]}
+      extensions={extensions}
       editable
       theme={codeMirrorTheme}
       {...props}

app/src/components/generative/ToolChoiceSelector.tsx

@@ -0,0 +1,102 @@
+import React from "react";
+
+import { Flex, Item, Label, Picker } from "@arizeai/components";
+
+type DefaultToolChoice = Extract<ToolChoice, "auto" | "required" | "none">;
+
+const isDefaultToolChoice = (choice: string): choice is DefaultToolChoice => {
+  return choice === "auto" || choice === "required" || choice === "none";
+};
+
+/**
+ * A prefix to add to user defined tools in the picker to avoid picker key collisions with default {@link ToolChoice} keys
+ */
+const TOOL_NAME_PREFIX = "tool_";
+
+/**
+ * Adds a prefix to user defined tool names to avoid conflicts picker key collisions with default {@link ToolChoice} keys
+ * @param toolName The name of a tool
+ * @returns  The tool name with the "TOOL_NAME_PREFIX" prefix added
+ */
+const addToolNamePrefix = (toolName: string) =>
+  `${TOOL_NAME_PREFIX}${toolName}`;
+
+/**
+ * Removes the "TOOL_NAME_PREFIX" prefix from a tool name so that it can be used as a choice that corresponds to an actual tool
+ * @param toolName The name of a tool with the "TOOL_NAME_PREFIX" prefix
+ * @returns The tool name with the "TOOL_NAME_PREFIX" prefix removed
+ */
+const removeToolNamePrefix = (toolName: string) =>
+  toolName.startsWith(TOOL_NAME_PREFIX)
+    ? toolName.slice(TOOL_NAME_PREFIX.length)
+    : toolName;
+
+type ToolChoicePickerProps = {
+  /**
+   * The current choice including the default {@link ToolChoice} and any user defined tools
+   */
+  choice: ToolChoice;
+  /**
+   * Callback for when the tool choice changes
+   */
+  onChange: (choice: ToolChoice) => void;
+  /**
+   * A list of user defined tool names
+   */
+  toolNames: string[];
+};
+
+export function ToolChoicePicker({
+  choice,
+  onChange,
+  toolNames,
+}: ToolChoicePickerProps) {
+  const currentKey =
+    typeof choice === "string"
+      ? choice
+      : addToolNamePrefix(choice.function.name);
+  return (
+    <Picker
+      selectedKey={currentKey}
+      label="Tool Choice"
+      aria-label="Tool Choice for an LLM"
+      onSelectionChange={(choice) => {
+        if (typeof choice !== "string") {
+          return;
+        }
+        if (choice.startsWith(TOOL_NAME_PREFIX)) {
+          onChange({
+            type: "function",
+            function: {
+              name: removeToolNamePrefix(choice),
+            },
+          });
+        } else if (isDefaultToolChoice(choice)) {
+          onChange(choice);
+        }
+      }}
+    >
+      {[
+        <Item key="auto">
+          <Flex gap={"size-100"}>
+            Tools auto-selected by LLM <Label>auto</Label>
+          </Flex>
+        </Item>,
+        <Item key="required">
+          <Flex gap={"size-100"}>
+            Use at least one tool <Label>required</Label>
+          </Flex>
+        </Item>,
+        <Item key="none">
+          <Flex gap={"size-100"}>
+            Don&apos;t use any tools <Label>none</Label>
+          </Flex>
+        </Item>,
+        // Add "TOOL_NAME_PREFIX" prefix to user defined tool names to avoid conflicts with default keys
+        ...toolNames.map((toolName) => (
+          <Item key={addToolNamePrefix(toolName)}>{toolName}</Item>
+        )),
+      ]}
+    </Picker>
+  );
+}

app/src/components/generative/index.ts

@@ -0,0 +1 @@
+export * from "./ToolChoiceSelector";

app/src/pages/playground/Playground.tsx

@@ -87,27 +87,25 @@ const playgroundPromptPanelContentCSS = css`
   flex-direction: column;
   height: 100%;
   overflow: hidden;
-  .ac-accordion {
+  & > .ac-accordion {
     display: flex;
     flex-direction: column;
     height: 100%;
     overflow: hidden;
     flex: 1 1 auto;
-    .ac-accordion-item {
+    & > .ac-accordion-item {
       height: 100%;
       overflow: hidden;
       flex: 1 1 auto;
       .ac-accordion-itemContent {
         height: 100%;
         overflow: hidden;
         flex: 1 1 auto;
-        & > * {
+        & > .ac-view {
           height: 100%;
           flex: 1 1 auto;
           overflow: auto;
           box-sizing: border-box;
-          // Fix padding issue with flexbox
-          padding-bottom: 57px !important;
         }
       }
     }
@@ -147,20 +145,15 @@ function PlaygroundContent() {
                 </Flex>
               }
             >
-              <View height="100%" padding="size-200">
+              <View height="100%" padding="size-200" paddingBottom="size-900">
                 <Flex direction="row" gap="size-200">
-                  {instances.map((instance, i) => (
-                    <div
-                      key={i}
-                      css={css`
-                        flex: 1 1 0px;
-                      `}
-                    >
+                  {instances.map((instance) => (
+                    <View key={instance.id} flex="1 1 0px">
                       <PlaygroundTemplate
-                        key={i}
+                        key={instance.id}
                         playgroundInstanceId={instance.id}
                       />
-                    </div>
+                    </View>
                   ))}
                 </Flex>
               </View>

app/src/pages/playground/PlaygroundChatTemplate.tsx

@@ -25,11 +25,13 @@ import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
 import { useChatMessageStyles } from "@phoenix/hooks/useChatMessageStyles";
 import {
   ChatMessage,
+  createTool,
   generateMessageId,
   PlaygroundChatTemplate as PlaygroundChatTemplateType,
 } from "@phoenix/store";
 
 import { MessageRolePicker } from "./MessageRolePicker";
+import { PlaygroundTools } from "./PlaygroundTools";
 import { PlaygroundInstanceProps } from "./types";
 
 const MESSAGE_Z_INDEX = 1;
@@ -54,6 +56,7 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
   if (!playgroundInstance) {
     throw new Error(`Playground instance ${id} not found`);
   }
+  const hasTools = playgroundInstance.tools.length > 0;
   const { template } = playgroundInstance;
   if (template.__type !== "chat") {
     throw new Error(`Invalid template type ${template.__type}`);
@@ -123,10 +126,30 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
         paddingEnd="size-200"
         paddingTop="size-100"
         paddingBottom="size-100"
-        borderTopColor="dark"
+        borderColor="dark"
         borderTopWidth="thin"
+        borderBottomWidth={hasTools ? "thin" : undefined}
       >
-        <Flex direction="row" justifyContent="end">
+        <Flex direction="row" justifyContent="end" gap="size-100">
+          <Button
+            variant="default"
+            aria-label="add tool"
+            size="compact"
+            icon={<Icon svg={<Icons.PlusOutline />} />}
+            onClick={() => {
+              updateInstance({
+                instanceId: id,
+                patch: {
+                  tools: [
+                    ...playgroundInstance.tools,
+                    createTool(playgroundInstance.tools.length + 1),
+                  ],
+                },
+              });
+            }}
+          >
+            Tool
+          </Button>
           <Button
             variant="default"
             aria-label="add message"
@@ -155,6 +178,7 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
           </Button>
         </Flex>
       </View>
+      {hasTools ? <PlaygroundTools {...props} /> : null}
     </DndContext>
   );
 }

app/src/pages/playground/PlaygroundOutput.tsx

@@ -42,7 +42,7 @@ export function PlaygroundOutput(props: PlaygroundOutputProps) {
     state.instances.findIndex((instance) => instance.id === instanceId)
   );
   if (!instance) {
-    throw new Error("Playground instance not found");
+    throw new Error(`Playground instance ${instanceId} not found`);
   }
 
   const runId = instance.activeRunId;