feat(playground): add tools ui

app/package.json

@@ -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": {

app/src/components/code/JSONToolEditor.tsx

@@ -0,0 +1,45 @@
+import React from "react";
+import { json, jsonLanguage, jsonParseLinter } from "@codemirror/lang-json";
+import { linter } from "@codemirror/lint";
+import { EditorView, hoverTooltip } from "@codemirror/view";
+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 { useTheme } from "@phoenix/contexts";
+import { toolJSONSchema } from "@phoenix/schemas/toolSchema";
+
+export type JSONToolEditorProps = Omit<
+  ReactCodeMirrorProps,
+  "theme" | "extensions" | "editable"
+>;
+
+export function JSONToolEditor(props: JSONToolEditorProps) {
+  const { theme } = useTheme();
+  const codeMirrorTheme = theme === "light" ? undefined : nord;
+  return (
+    <CodeMirror
+      value={props.value}
+      extensions={[
+        json(),
+        EditorView.lineWrapping,
+        linter(jsonParseLinter()),
+        linter(jsonSchemaLinter(), { needsRefresh: handleRefresh }),
+        jsonLanguage.data.of({
+          autocomplete: jsonCompletion(),
+        }),
+        hoverTooltip(jsonSchemaHover()),
+        stateExtensions(toolJSONSchema as JSONSchema7),
+      ]}
+      editable
+      theme={codeMirrorTheme}
+      {...props}
+    />
+  );
+}

app/src/components/code/index.tsx

@@ -6,3 +6,4 @@ export * from "./CodeWrap";
 export * from "./TypeScriptBlock";
 export * from "./CodeEditorFieldWrapper";
 export * from "./CodeLanguageRadioGroup";
+export * from "./JSONToolEditor";

app/src/pages/playground/PlaygroundChatTemplate.tsx

@@ -27,6 +27,7 @@ import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
 import { useChatMessageStyles } from "@phoenix/hooks/useChatMessageStyles";
 import {
   ChatMessage,
+  createTool,
   generateMessageId,
   PlaygroundChatTemplate as PlaygroundChatTemplateType,
 } from "@phoenix/store";
@@ -124,7 +125,26 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
         borderTopColor="dark"
         borderTopWidth="thin"
       >
-        <Flex direction="row" justifyContent="end">
+        <Flex direction="row" justifyContent="end" gap="size-50">
+          <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"

app/src/pages/playground/PlaygroundOutput.tsx

@@ -41,7 +41,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;

app/src/pages/playground/PlaygroundTools.tsx

@@ -0,0 +1,85 @@
+import React from "react";
+
+import { Button, Card, Flex, Icon, Icons } from "@arizeai/components";
+
+import { JSONToolEditor } from "@phoenix/components/code";
+import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
+
+import { TitleWithAlphabeticIndex } from "./TitleWithAlphabeticIndex";
+import { PlaygroundInstanceProps } from "./types";
+
+interface PlaygroundToolsProps extends PlaygroundInstanceProps {}
+
+export function PlaygroundTools(props: PlaygroundToolsProps) {
+  const instanceId = props.playgroundInstanceId;
+  const instance = usePlaygroundContext((state) =>
+    state.instances.find(
+      (instance) => instance.id === props.playgroundInstanceId
+    )
+  );
+  if (instance == null) {
+    throw new Error(`Playground instance ${instanceId} not found`);
+  }
+  const updateInstance = usePlaygroundContext((state) => state.updateInstance);
+  if (instance.tools == null) {
+    throw new Error(`Playground instance ${instanceId} does not have tools`);
+  }
+  const index = usePlaygroundContext((state) =>
+    state.instances.findIndex((instance) => instance.id === instanceId)
+  );
+
+  return (
+    <Card
+      title={<TitleWithAlphabeticIndex index={index} title="Tools" />}
+      collapsible
+      variant="compact"
+    >
+      <Flex direction="column" gap="size-50">
+        {instance.tools.map((tool) => (
+          <Card
+            collapsible
+            variant="compact"
+            key={tool.id}
+            title={tool.definition.function.name}
+            bodyStyle={{ padding: 0 }}
+            extra={
+              <Button
+                aria-label="Delete tool"
+                icon={<Icon svg={<Icons.TrashOutline />} />}
+                variant="default"
+                size="compact"
+                onClick={() => {
+                  updateInstance({
+                    instanceId,
+                    patch: {
+                      tools: instance.tools.filter((t) => t.id !== tool.id),
+                    },
+                  });
+                }}
+              />
+            }
+          >
+            <JSONToolEditor
+              value={JSON.stringify(tool.definition, null, 2)}
+              onChange={(value) => {
+                updateInstance({
+                  instanceId,
+                  patch: {
+                    tools: instance.tools.map((t) =>
+                      t.id === tool.id
+                        ? {
+                            ...t,
+                            definition: JSON.parse(value),
+                          }
+                        : t
+                    ),
+                  },
+                });
+              }}
+            />
+          </Card>
+        ))}
+      </Flex>
+    </Card>
+  );
+}

app/src/pages/playground/__tests__/playgroundUtils.test.ts

@@ -26,7 +26,7 @@ const expectedPlaygroundInstanceWithIO: PlaygroundInstance = {
     provider: "OPENAI",
     modelName: "gpt-4o",
   },
-  tools: {},
+  tools: [],
   template: {
     __type: "chat",
     // These id's are not 0, 1, 2, because we create a playground instance (including messages) at the top of the transformSpanAttributesToPlaygroundInstance function

app/src/schemas/index.ts

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

app/src/schemas/toolSchema.ts

@@ -0,0 +1,70 @@
+import { z } from "zod";
+import zodToJsonSchema from "zod-to-json-schema";
+
+/**
+ * The schema for a tool definition
+ * @see https://platform.openai.com/docs/guides/structured-outputs/supported-schemas
+ */
+export const toolSchema = z.object({
+  type: z.literal("function"),
+  function: z.object({
+    name: z.string().describe("The name of the function"),
+    description: z
+      .string()
+      .optional()
+      .describe("A description of the function"),
+    parameters: z
+      .object({
+        type: z.literal("object"),
+        properties: z
+          .record(
+            z.object({
+              type: z
+                .enum([
+                  "string",
+                  "number",
+                  "boolean",
+                  "object",
+                  "array",
+                  "null",
+                  "integer",
+                ])
+                .describe("The type of the parameter"),
+              description: z
+                .string()
+                .optional()
+                .describe("A description of the parameter"),
+              enum: z
+                .array(z.string())
+                .optional()
+                .describe("The allowed values"),
+            })
+          )
+          .describe("A map of parameter names to their definitions"),
+        required: z
+          .array(z.string())
+          .optional()
+          .describe("The required parameters"),
+        additionalProperties: z
+          .boolean()
+          .optional()
+          .describe("Whether additional properties are allowed"),
+        strict: z
+          .boolean()
+          .optional()
+          .describe("Whether the object should be strict"),
+      })
+      .describe("The parameters that the function accepts"),
+  }),
+});
+
+/**
+ * The type of a tool definition
+ * @see https://platform.openai.com/docs/guides/structured-outputs/supported-schemas
+ */
+export type ToolDefinition = z.infer<typeof toolSchema>;
+
+/**
+ * The JSON schema for a tool definition
+ */
+export const toolJSONSchema = zodToJsonSchema(toolSchema);

app/src/store/playground/playgroundStore.tsx

@@ -9,30 +9,37 @@ import {
   PlaygroundInstance,
   PlaygroundState,
   PlaygroundTextCompletionTemplate,
+  Tool,
 } from "./types";
 
-let playgroundInstanceIdIndex = 0;
-let playgroundRunIdIndex = 0;
+let playgroundInstanceId = 0;
+let playgroundRunId = 0;
 // This value must be truthy in order for message re-ordering to work
-let playgroundMessageIdIndex = 1;
+let playgroundMessageId = 1;
+let playgroundToolId = 0;
 
 /**
  * Generates a new playground instance ID
  */
-export const generateInstanceId = () => playgroundInstanceIdIndex++;
+export const generateInstanceId = () => playgroundInstanceId++;
 
 /**
  * Generates a new playground message ID
  */
-export const generateMessageId = () => playgroundMessageIdIndex++;
+export const generateMessageId = () => playgroundMessageId++;
+
+/**
+ * Generates a new playground tool ID
+ */
+export const generateToolId = () => playgroundToolId++;
 
 /**
  * Resets the playground instance ID to 0
  *
  * NB: This is only used for testing purposes
  */
 export const _resetInstanceId = () => {
-  playgroundInstanceIdIndex = 0;
+  playgroundInstanceId = 0;
 };
 
 /**
@@ -41,7 +48,16 @@ export const _resetInstanceId = () => {
  * NB: This is only used for testing purposes
  */
 export const _resetMessageId = () => {
-  playgroundMessageIdIndex = 0;
+  playgroundMessageId = 0;
+};
+
+/**
+ * Resets the playground tool ID to 0
+ *
+ * NB: This is only used for testing purposes
+ */
+export const _resetToolId = () => {
+  playgroundToolId = 0;
 };
 
 const generateChatCompletionTemplate = (): PlaygroundChatTemplate => ({
@@ -70,13 +86,37 @@ export function createPlaygroundInstance(): PlaygroundInstance {
     id: generateInstanceId(),
     template: generateChatCompletionTemplate(),
     model: { provider: "OPENAI", modelName: "gpt-4o" },
-    tools: {},
+    tools: [],
+    input: { variables: {} },
     output: undefined,
     activeRunId: null,
     isRunning: false,
   };
 }
 
+export function createTool(toolNumber: number): Tool {
+  return {
+    id: generateToolId(),
+    definition: {
+      type: "function",
+      function: {
+        name: `new_function_${toolNumber}`,
+        parameters: {
+          type: "object",
+          properties: {
+            new_arg: {
+              type: "string",
+            },
+          },
+          required: [],
+          additionalProperties: false,
+          strict: false,
+        },
+      },
+    },
+  };
+}
+
 export const createPlaygroundStore = (
   initialProps?: InitialPlaygroundState
 ) => {
@@ -198,7 +238,7 @@ export const createPlaygroundStore = (
       set({
         instances: instances.map((instance) => ({
           ...instance,
-          activeRunId: playgroundRunIdIndex++,
+          activeRunId: playgroundRunId++,
           isRunning: true,
         })),
       });
@@ -210,7 +250,7 @@ export const createPlaygroundStore = (
           if (instance.id === instanceId) {
             return {
               ...instance,
-              activeRunId: playgroundRunIdIndex++,
+              activeRunId: playgroundRunId++,
               isRunning: true,
             };
           }