feat(playground): add tool role messages to ui (#5103)

* get tool cals working for non tool roles

* renaming

* plumb through tool calls to server

* add tool role message ui

* mikeldking: small style changes

* fix: Align vertical edge padding across playground message editors (#5134)

* update tool and chat copy to clipboard

* update message types for subscriptions on server, change message tool calls and mode when switching off of ai

* buld gql

* add descriptions to gql fields

* update to_openai_tool_call to take in a single tool and return a single tool call param

* build graphql

---------

Co-authored-by: Mikyo King <mikyo@arize.com>
Co-authored-by: Anthony Powell <apowell@arize.com>

app/schema.graphql

@@ -77,8 +77,16 @@ input ChatCompletionInput {
 input ChatCompletionMessageInput {
   role: ChatCompletionMessageRole!
 
-  """The content of the message as JSON to support text and tools"""
+  """The content of the message as JSON to support various kinds of text"""
   content: JSON!
+
+  """The tool calls that were made in the message"""
+  toolCalls: [JSON!]
+
+  """
+  The ID that corresponds to a prior tool call. Used to link a tool message to a pre-existing tool call.
+  """
+  toolCallId: String
 }
 
 enum ChatCompletionMessageRole {

app/src/components/trace/SpanKindIcon.tsx

@@ -36,7 +36,7 @@ const ToolSVG = () => (
   </svg>
 );
 
-const ToolFilledSVG = () => (
+export const ToolFilledSVG = () => (
   <svg
     width="20"
     height="20"

app/src/pages/playground/ChatMessageToolCallsEditor.tsx

@@ -0,0 +1,75 @@
+import React, { useCallback, useState } from "react";
+import { JSONSchema7 } from "json-schema";
+
+import { JSONEditor } from "@phoenix/components/code";
+import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
+import {
+  openAIToolCallsJSONSchema,
+  openAIToolCallsSchema,
+} from "@phoenix/schemas";
+import { ChatMessage } from "@phoenix/store";
+import { safelyParseJSON } from "@phoenix/utils/jsonUtils";
+
+import { PlaygroundInstanceProps } from "./types";
+
+/**
+ * Editor for message tool calls
+ */
+export function ChatMessageToolCallsEditor({
+  playgroundInstanceId,
+  toolCalls,
+  templateMessages,
+  messageId,
+}: PlaygroundInstanceProps & {
+  toolCalls: ChatMessage["toolCalls"];
+  templateMessages: ChatMessage[];
+  messageId: number;
+}) {
+  const updateInstance = usePlaygroundContext((state) => state.updateInstance);
+  const [toolCallsValue, setToolCallsValue] = useState(() =>
+    JSON.stringify(toolCalls, null, 2)
+  );
+
+  const onChange = useCallback(
+    (value: string) => {
+      setToolCallsValue(value);
+      const { json: definition } = safelyParseJSON(value);
+      if (definition == null) {
+        return;
+      }
+      // Don't use data here returned by safeParse, as we want to allow for extra keys,
+      // there is no "deepPassthrough" to allow for extra keys
+      // at all levels of the schema, so we just use the json parsed value here,
+      // knowing that it is valid with potentially extra keys
+      const { success } = openAIToolCallsSchema.safeParse(definition);
+      if (!success) {
+        return;
+      }
+      updateInstance({
+        instanceId: playgroundInstanceId,
+        patch: {
+          template: {
+            __type: "chat",
+            messages: templateMessages.map((m) =>
+              messageId === m.id
+                ? {
+                    ...m,
+                    toolCalls: definition,
+                  }
+                : m
+            ),
+          },
+        },
+      });
+    },
+    [messageId, playgroundInstanceId, templateMessages, updateInstance]
+  );
+
+  return (
+    <JSONEditor
+      value={toolCallsValue}
+      jsonSchema={openAIToolCallsJSONSchema as JSONSchema7}
+      onChange={onChange}
+    />
+  );
+}

app/src/pages/playground/MessageContentRadioGroup.tsx

@@ -0,0 +1,39 @@
+import React from "react";
+
+import { Icon, Icons, Radio, RadioGroup } from "@arizeai/components";
+
+export type MessageMode = "text" | "toolCalls";
+
+function isMessageMode(value: string): value is MessageMode {
+  return value === "text" || value === "toolCalls";
+}
+
+export function MessageContentRadioGroup({
+  messageMode,
+  onChange,
+}: {
+  messageMode: MessageMode;
+  onChange: (messageMode: MessageMode) => void;
+}) {
+  return (
+    <RadioGroup
+      defaultValue={messageMode}
+      variant="inline-button"
+      size="compact"
+      onChange={(v) => {
+        if (isMessageMode(v)) {
+          onChange(v);
+        } else {
+          throw new Error(`Unknown message mode: ${v}`);
+        }
+      }}
+    >
+      <Radio label="text input" value={"text"}>
+        <Icon svg={<Icons.MessageSquareOutline />} />
+      </Radio>
+      <Radio label="tool calling" value={"toolCalls"}>
+        <Icon svg={<Icons.Code />} />
+      </Radio>
+    </RadioGroup>
+  );
+}

app/src/pages/playground/MessageRolePicker.tsx

@@ -51,6 +51,7 @@ export function MessageRolePicker({
       <Item key="system">System</Item>
       <Item key="user">User</Item>
       <Item key="ai">AI</Item>
+      <Item key="tool">Tool</Item>
     </Picker>
   );
 }