From 75b7000b1b3f5329841d6293ffdcc1bf1e993be9 Mon Sep 17 00:00:00 2001
From: Xander Song <axiomofjoy@gmail.com>
Date: Wed, 16 Oct 2024 11:41:48 -0700
Subject: [PATCH] feat(playground): wire up tool calling ui (#5029)

---
 .../generative/ToolChoiceSelector.tsx         |  4 +-
 .../playground/PlaygroundChatTemplate.tsx     |  4 +-
 app/src/pages/playground/PlaygroundOutput.tsx | 77 +++++++++++++++++--
 app/src/pages/playground/PlaygroundTool.tsx   |  1 +
 .../PlaygroundOutputSubscription.graphql.ts   | 29 +++++--
 .../__tests__/playgroundUtils.test.ts         |  2 +-
 app/src/store/playground/playgroundStore.tsx  |  9 ++-
 app/src/store/playground/types.ts             | 43 ++++++++++-
 src/phoenix/server/api/subscriptions.py       |  4 +-
 9 files changed, 146 insertions(+), 27 deletions(-)

diff --git a/app/src/components/generative/ToolChoiceSelector.tsx b/app/src/components/generative/ToolChoiceSelector.tsx
index b0a89d095c..32a0404fc3 100644
--- a/app/src/components/generative/ToolChoiceSelector.tsx
+++ b/app/src/components/generative/ToolChoiceSelector.tsx
@@ -35,7 +35,7 @@ type ToolChoicePickerProps = {
   /**
    * The current choice including the default {@link ToolChoice} and any user defined tools
    */
-  choice: ToolChoice;
+  choice: ToolChoice | undefined;
   /**
    * Callback for when the tool choice changes
    */
@@ -52,7 +52,7 @@ export function ToolChoicePicker({
   toolNames,
 }: ToolChoicePickerProps) {
   const currentKey =
-    typeof choice === "string"
+    choice == null || typeof choice === "string"
       ? choice
       : addToolNamePrefix(choice.function.name);
   return (
diff --git a/app/src/pages/playground/PlaygroundChatTemplate.tsx b/app/src/pages/playground/PlaygroundChatTemplate.tsx
index 531064f5f5..bea1b63f1d 100644
--- a/app/src/pages/playground/PlaygroundChatTemplate.tsx
+++ b/app/src/pages/playground/PlaygroundChatTemplate.tsx
@@ -243,7 +243,9 @@ function SortableMessageItem({
         }
         extra={
           <Flex direction="row" gap="size-100">
-            <CopyToClipboardButton text={message.content} />
+            {message.content != null && (
+              <CopyToClipboardButton text={message.content} />
+            )}
             <Button
               aria-label="Delete message"
               icon={<Icon svg={<Icons.TrashOutline />} />}
diff --git a/app/src/pages/playground/PlaygroundOutput.tsx b/app/src/pages/playground/PlaygroundOutput.tsx
index e541c83cff..f0a7f65b4d 100644
--- a/app/src/pages/playground/PlaygroundOutput.tsx
+++ b/app/src/pages/playground/PlaygroundOutput.tsx
@@ -1,12 +1,14 @@
 import React, { useMemo, useState } from "react";
 import { useSubscription } from "react-relay";
 import { graphql, GraphQLSubscriptionConfig } from "relay-runtime";
+import { css } from "@emotion/react";
 
 import { Card, Flex, Icon, Icons } from "@arizeai/components";
 
 import { useCredentialsContext } from "@phoenix/contexts/CredentialsContext";
 import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
 import { useChatMessageStyles } from "@phoenix/hooks/useChatMessageStyles";
+import type { ToolCall } from "@phoenix/store";
 import { ChatMessage, generateMessageId } from "@phoenix/store";
 import { assertUnreachable } from "@phoenix/typeUtils";
 
@@ -24,11 +26,37 @@ import { PlaygroundInstanceProps } from "./types";
 interface PlaygroundOutputProps extends PlaygroundInstanceProps {}
 
 function PlaygroundOutputMessage({ message }: { message: ChatMessage }) {
-  const styles = useChatMessageStyles(message.role);
+  const { role, content, toolCalls } = message;
+  const styles = useChatMessageStyles(role);
 
   return (
-    <Card title={message.role} {...styles} variant="compact">
-      {message.content}
+    <Card title={role} {...styles} variant="compact">
+      {content != null && (
+        <Flex direction="column" alignItems="start">
+          {content}
+        </Flex>
+      )}
+      {toolCalls && toolCalls.length > 0
+        ? toolCalls.map((toolCall) => {
+            return (
+              <pre
+                key={toolCall.id}
+                css={css`
+                  text-wrap: wrap;
+                  margin: var(--ac-global-dimension-static-size-100) 0;
+                `}
+              >
+                {toolCall.function.name}(
+                {JSON.stringify(
+                  JSON.parse(toolCall.function.arguments),
+                  null,
+                  2
+                )}
+                )
+              </pre>
+            );
+          })
+        : null}
     </Card>
   );
 }
@@ -106,6 +134,7 @@ function useChatCompletionSubscription({
           $messages: [ChatCompletionMessageInput!]!
           $model: GenerativeModelInput!
           $invocationParameters: InvocationParameters!
+          $tools: [JSON!]
           $apiKey: String
         ) {
           chatCompletion(
@@ -113,6 +142,7 @@ function useChatCompletionSubscription({
               messages: $messages
               model: $model
               invocationParameters: $invocationParameters
+              tools: $tools
               apiKey: $apiKey
             }
           ) {
@@ -196,7 +226,8 @@ function PlaygroundOutputText(props: PlaygroundInstanceProps) {
     throw new Error("We only support chat templates for now");
   }
 
-  const [output, setOutput] = useState<string>("");
+  const [output, setOutput] = useState<string | undefined>(undefined);
+  const [toolCalls, setToolCalls] = useState<ToolCall[]>([]);
 
   useChatCompletionSubscription({
     params: {
@@ -206,15 +237,46 @@ function PlaygroundOutputText(props: PlaygroundInstanceProps) {
         name: instance.model.modelName || "",
       },
       invocationParameters: {
-        temperature: 0.1, // TODO: add invocation parameters
+        toolChoice: instance.toolChoice,
       },
+      tools: instance.tools.map((tool) => tool.definition),
       apiKey: credentials[instance.model.provider],
     },
     runId: instance.activeRunId,
     onNext: (response) => {
       const chatCompletion = response.chatCompletion;
       if (chatCompletion.__typename === "TextChunk") {
-        setOutput((acc) => acc + chatCompletion.content);
+        setOutput((acc) => (acc || "") + chatCompletion.content);
+      } else if (chatCompletion.__typename === "ToolCallChunk") {
+        setToolCalls((toolCalls) => {
+          let toolCallExists = false;
+          const updated = toolCalls.map((toolCall) => {
+            if (toolCall.id === chatCompletion.id) {
+              toolCallExists = true;
+              return {
+                ...toolCall,
+                function: {
+                  ...toolCall.function,
+                  arguments:
+                    toolCall.function.arguments +
+                    chatCompletion.function.arguments,
+                },
+              };
+            } else {
+              return toolCall;
+            }
+          });
+          if (!toolCallExists) {
+            updated.push({
+              id: chatCompletion.id,
+              function: {
+                name: chatCompletion.function.name,
+                arguments: chatCompletion.function.arguments,
+              },
+            });
+          }
+          return updated;
+        });
       }
     },
     onCompleted: () => {
@@ -222,7 +284,7 @@ function PlaygroundOutputText(props: PlaygroundInstanceProps) {
     },
   });
 
-  if (!output) {
+  if (!output && (toolCalls.length === 0 || instance.isRunning)) {
     return (
       <Flex direction="row" gap="size-100" alignItems="center">
         <Icon svg={<Icons.LoadingOutline />} />
@@ -236,6 +298,7 @@ function PlaygroundOutputText(props: PlaygroundInstanceProps) {
         id: generateMessageId(),
         content: output,
         role: "ai",
+        toolCalls: toolCalls,
       }}
     />
   );
diff --git a/app/src/pages/playground/PlaygroundTool.tsx b/app/src/pages/playground/PlaygroundTool.tsx
index de0705b772..d3e30bca2e 100644
--- a/app/src/pages/playground/PlaygroundTool.tsx
+++ b/app/src/pages/playground/PlaygroundTool.tsx
@@ -81,6 +81,7 @@ export function PlaygroundTool({
                 instanceId: playgroundInstanceId,
                 patch: {
                   tools: instanceTools.filter((t) => t.id !== tool.id),
+                  toolChoice: undefined,
                 },
               });
             }}
diff --git a/app/src/pages/playground/__generated__/PlaygroundOutputSubscription.graphql.ts b/app/src/pages/playground/__generated__/PlaygroundOutputSubscription.graphql.ts
index 5203f36a5e..d560c6a348 100644
--- a/app/src/pages/playground/__generated__/PlaygroundOutputSubscription.graphql.ts
+++ b/app/src/pages/playground/__generated__/PlaygroundOutputSubscription.graphql.ts
@@ -1,5 +1,5 @@
 /**
- * @generated SignedSource<<fed32a7465da8f81b8fda04674f80001>>
+ * @generated SignedSource<<d1f4c21ebf3b85b639c7ff17beee84ca>>
  * @lightSyntaxTransform
  * @nogrep
  */
@@ -33,6 +33,7 @@ export type PlaygroundOutputSubscription$variables = {
   invocationParameters: InvocationParameters;
   messages: ReadonlyArray<ChatCompletionMessageInput>;
   model: GenerativeModelInput;
+  tools?: ReadonlyArray<any> | null;
 };
 export type PlaygroundOutputSubscription$data = {
   readonly chatCompletion: {
@@ -77,7 +78,12 @@ v3 = {
   "kind": "LocalArgument",
   "name": "model"
 },
-v4 = [
+v4 = {
+  "defaultValue": null,
+  "kind": "LocalArgument",
+  "name": "tools"
+},
+v5 = [
   {
     "alias": null,
     "args": [
@@ -102,6 +108,11 @@ v4 = [
             "kind": "Variable",
             "name": "model",
             "variableName": "model"
+          },
+          {
+            "kind": "Variable",
+            "name": "tools",
+            "variableName": "tools"
           }
         ],
         "kind": "ObjectValue",
@@ -183,12 +194,13 @@ return {
       (v0/*: any*/),
       (v1/*: any*/),
       (v2/*: any*/),
-      (v3/*: any*/)
+      (v3/*: any*/),
+      (v4/*: any*/)
     ],
     "kind": "Fragment",
     "metadata": null,
     "name": "PlaygroundOutputSubscription",
-    "selections": (v4/*: any*/),
+    "selections": (v5/*: any*/),
     "type": "Subscription",
     "abstractKey": null
   },
@@ -198,23 +210,24 @@ return {
       (v2/*: any*/),
       (v3/*: any*/),
       (v1/*: any*/),
+      (v4/*: any*/),
       (v0/*: any*/)
     ],
     "kind": "Operation",
     "name": "PlaygroundOutputSubscription",
-    "selections": (v4/*: any*/)
+    "selections": (v5/*: any*/)
   },
   "params": {
-    "cacheID": "ce54188370b1d1657129ab3f4a4baa28",
+    "cacheID": "924d84f911c5156af0abb2a371d098f2",
     "id": null,
     "metadata": {},
     "name": "PlaygroundOutputSubscription",
     "operationKind": "subscription",
-    "text": "subscription PlaygroundOutputSubscription(\n  $messages: [ChatCompletionMessageInput!]!\n  $model: GenerativeModelInput!\n  $invocationParameters: InvocationParameters!\n  $apiKey: String\n) {\n  chatCompletion(input: {messages: $messages, model: $model, invocationParameters: $invocationParameters, apiKey: $apiKey}) {\n    __typename\n    ... on TextChunk {\n      content\n    }\n    ... on ToolCallChunk {\n      id\n      function {\n        name\n        arguments\n      }\n    }\n  }\n}\n"
+    "text": "subscription PlaygroundOutputSubscription(\n  $messages: [ChatCompletionMessageInput!]!\n  $model: GenerativeModelInput!\n  $invocationParameters: InvocationParameters!\n  $tools: [JSON!]\n  $apiKey: String\n) {\n  chatCompletion(input: {messages: $messages, model: $model, invocationParameters: $invocationParameters, tools: $tools, apiKey: $apiKey}) {\n    __typename\n    ... on TextChunk {\n      content\n    }\n    ... on ToolCallChunk {\n      id\n      function {\n        name\n        arguments\n      }\n    }\n  }\n}\n"
   }
 };
 })();
 
-(node as any).hash = "5c2ef45511941f179e8b09d4ff53503b";
+(node as any).hash = "30b9973d6ea69054a907549af97c0e5f";
 
 export default node;
diff --git a/app/src/pages/playground/__tests__/playgroundUtils.test.ts b/app/src/pages/playground/__tests__/playgroundUtils.test.ts
index 1d71fa48d8..2f505c0a19 100644
--- a/app/src/pages/playground/__tests__/playgroundUtils.test.ts
+++ b/app/src/pages/playground/__tests__/playgroundUtils.test.ts
@@ -33,7 +33,7 @@ const expectedPlaygroundInstanceWithIO: PlaygroundInstance = {
   },
   input: { variableKeys: [], variablesValueCache: {} },
   tools: [],
-  toolChoice: "auto",
+  toolChoice: undefined,
   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
diff --git a/app/src/store/playground/playgroundStore.tsx b/app/src/store/playground/playgroundStore.tsx
index fb3e347b2d..88a437a2c6 100644
--- a/app/src/store/playground/playgroundStore.tsx
+++ b/app/src/store/playground/playgroundStore.tsx
@@ -97,7 +97,7 @@ export function createPlaygroundInstance(): PlaygroundInstance {
     template: generateChatCompletionTemplate(),
     model: { provider: DEFAULT_MODEL_PROVIDER, modelName: "gpt-4o" },
     tools: [],
-    toolChoice: "auto",
+    toolChoice: undefined,
     // TODO(apowell) - use datasetId if in dataset mode
     input: { variablesValueCache: {}, variableKeys: [] },
     output: undefined,
@@ -310,9 +310,10 @@ export const createPlaygroundStore = (
             // for each chat message in the instance
             instance.template.messages.forEach((message) => {
               // extract variables from the message content
-              const extractedVariables = utils.extractVariables(
-                message.content
-              );
+              const extractedVariables =
+                message.content == null
+                  ? []
+                  : utils.extractVariables(message.content);
               extractedVariables.forEach((variable) => {
                 variables.add(variable);
               });
diff --git a/app/src/store/playground/types.ts b/app/src/store/playground/types.ts
index c7aa70caf1..01d6357f82 100644
--- a/app/src/store/playground/types.ts
+++ b/app/src/store/playground/types.ts
@@ -8,14 +8,51 @@ export type GenAIOperationType = "chat" | "text_completion";
  */
 export type PlaygroundInputMode = "manual" | "dataset";
 
+/**
+ * A tool call that invokes a function with JSON arguments
+ * @example
+ * ```typescript
+ *  {
+ *   id: "1",
+ *   function: {
+ *     name: "getCurrentWeather",
+ *     arguments: "{ \"city\": \"San Francisco\" }"
+ *   }
+ * }
+ * ```
+ */
+export type ToolCall = {
+  id: string;
+  function: {
+    name: string;
+    arguments: string;
+  };
+};
+
 /**
  * A chat message with a role and content
- * @example { role: "user", content: "What is the meaning of life?" }
+ * @example { role: "user", content: "What is the weather in San Francisco?" }
+ * @example
+ * ```typescript
+ * {
+ *   "role": "assistant",
+ *   "toolCalls": [
+ *     {
+ *       "id": "1",
+ *       "function": {
+ *         "name": "getCurrentWeather",
+ *         "arguments": "{ \"city\": \"San Francisco\" }"
+ *       }
+ *     }
+ *   ]
+ * }
+ * ```
  */
 export type ChatMessage = {
   id: number;
   role: ChatMessageRole;
-  content: string;
+  content?: string;
+  toolCalls?: ToolCall[];
 };
 
 /**
@@ -82,7 +119,7 @@ export interface PlaygroundInstance {
   id: number;
   template: PlaygroundTemplate;
   tools: Tool[];
-  toolChoice: ToolChoice;
+  toolChoice: ToolChoice | undefined;
   input: PlaygroundInput;
   model: ModelConfig;
   output: ChatMessage[] | undefined | string;
diff --git a/src/phoenix/server/api/subscriptions.py b/src/phoenix/server/api/subscriptions.py
index 9272608b5b..638eb1943a 100644
--- a/src/phoenix/server/api/subscriptions.py
+++ b/src/phoenix/server/api/subscriptions.py
@@ -1,5 +1,6 @@
 import json
 from collections import defaultdict
+from dataclasses import fields
 from datetime import datetime
 from itertools import chain
 from typing import (
@@ -280,8 +281,9 @@ def _llm_tools(tools: List[JSONScalarType]) -> Iterator[Tuple[str, Any]]:
 
 
 def _input_value_and_mime_type(input: ChatCompletionInput) -> Iterator[Tuple[str, Any]]:
+    assert any(field.name == (api_key := "api_key") for field in fields(ChatCompletionInput))
     yield INPUT_MIME_TYPE, JSON
-    yield INPUT_VALUE, safe_json_dumps(jsonify(input))
+    yield INPUT_VALUE, safe_json_dumps({k: v for k, v in jsonify(input).items() if k != api_key})
 
 
 def _output_value_and_mime_type(output: Any) -> Iterator[Tuple[str, Any]]: