From aa8c283ceded7ea4333b74fb02824ba2fade2d55 Mon Sep 17 00:00:00 2001
From: Anthony Powell <apowell@arize.com>
Date: Fri, 22 Nov 2024 19:23:51 -0500
Subject: [PATCH] feat: Implement "Move Output" button on Playground Output
 (#5509)

* feat: Implement "Move Output" button to Playground Output

This allows users to copy a desireable result from the playground back into their conversation,
to continue iteration

* Convert moved tool calls into the correct provider schema
---
 app/src/pages/playground/PlaygroundOutput.tsx | 31 ++++++-
 .../playground/PlaygroundOutputMoveButton.tsx | 93 +++++++++++++++++++
 2 files changed, 120 insertions(+), 4 deletions(-)
 create mode 100644 app/src/pages/playground/PlaygroundOutputMoveButton.tsx

diff --git a/app/src/pages/playground/PlaygroundOutput.tsx b/app/src/pages/playground/PlaygroundOutput.tsx
index e72f0fc22e..9a9bb30406 100644
--- a/app/src/pages/playground/PlaygroundOutput.tsx
+++ b/app/src/pages/playground/PlaygroundOutput.tsx
@@ -28,7 +28,7 @@ import {
   generateMessageId,
   PlaygroundInstance,
 } from "@phoenix/store";
-import { isStringKeyedObject } from "@phoenix/typeUtils";
+import { isStringKeyedObject, Mutable } from "@phoenix/typeUtils";
 import {
   getErrorMessagesFromRelayMutationError,
   getErrorMessagesFromRelaySubscriptionError,
@@ -44,6 +44,7 @@ import {
 } from "./__generated__/PlaygroundOutputSubscription.graphql";
 import PlaygroundOutputSubscription from "./__generated__/PlaygroundOutputSubscription.graphql";
 import { PlaygroundErrorWrap } from "./PlaygroundErrorWrap";
+import { PlaygroundOutputMoveButton } from "./PlaygroundOutputMoveButton";
 import {
   PartialOutputToolCall,
   PlaygroundToolCall,
@@ -316,14 +317,24 @@ export function PlaygroundOutput(props: PlaygroundOutputProps) {
     ]
   );
 
+  const cleanup = useCallback(() => {
+    setOutputContent(undefined);
+    setToolCalls([]);
+    setOutputError(null);
+    updateInstance({
+      instanceId,
+      patch: {
+        spanId: null,
+      },
+    });
+  }, [instanceId, updateInstance]);
+
   useEffect(() => {
     if (!hasRunId) {
       return;
     }
     setLoading(true);
-    setOutputContent(undefined);
-    setToolCalls([]);
-    setOutputError(null);
+    cleanup();
     const input = getChatCompletionInput({
       playgroundStore,
       instanceId,
@@ -390,6 +401,7 @@ export function PlaygroundOutput(props: PlaygroundOutputProps) {
       },
     });
   }, [
+    cleanup,
     credentials,
     environment,
     generateChatCompletion,
@@ -408,6 +420,17 @@ export function PlaygroundOutput(props: PlaygroundOutputProps) {
   return (
     <Card
       title={<TitleWithAlphabeticIndex index={index} title="Output" />}
+      titleExtra={
+        outputContent != null || toolCalls?.length > 0 ? (
+          <PlaygroundOutputMoveButton
+            outputContent={outputContent}
+            toolCalls={toolCalls as Mutable<typeof toolCalls>}
+            instance={instance}
+            cleanupOutput={cleanup}
+          />
+        ) : null
+      }
+      titleSeparator
       collapsible
       variant="compact"
       bodyStyle={{ padding: 0 }}
diff --git a/app/src/pages/playground/PlaygroundOutputMoveButton.tsx b/app/src/pages/playground/PlaygroundOutputMoveButton.tsx
new file mode 100644
index 0000000000..74cfddf4d9
--- /dev/null
+++ b/app/src/pages/playground/PlaygroundOutputMoveButton.tsx
@@ -0,0 +1,93 @@
+import React from "react";
+
+import {
+  Button,
+  Icon,
+  Icons,
+  Tooltip,
+  TooltipTrigger,
+} from "@arizeai/components";
+
+import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
+import {
+  ChatMessage,
+  generateMessageId,
+  PlaygroundInstance,
+} from "@phoenix/store";
+import { safelyParseJSON } from "@phoenix/utils/jsonUtils";
+
+import { PartialOutputToolCall } from "./PlaygroundToolCall";
+import {
+  convertMessageToolCallsToProvider,
+  getChatRole,
+} from "./playgroundUtils";
+
+export const PlaygroundOutputMoveButton = ({
+  instance,
+  outputContent,
+  toolCalls,
+  cleanupOutput,
+}: {
+  instance: PlaygroundInstance;
+  outputContent?: string | ChatMessage[];
+  toolCalls: PartialOutputToolCall[];
+  cleanupOutput: () => void;
+}) => {
+  const instanceId = instance.id;
+  const updateInstance = usePlaygroundContext((state) => state.updateInstance);
+  return (
+    <TooltipTrigger delay={500} offset={10}>
+      <Button
+        variant="default"
+        size="compact"
+        icon={<Icon svg={<Icons.PlusCircleOutline />} />}
+        onClick={(e) => {
+          e.stopPropagation();
+          if (instance.template.__type !== "chat") {
+            return;
+          }
+          const messages = Array.isArray(outputContent)
+            ? outputContent
+            : outputContent
+              ? [
+                  {
+                    id: generateMessageId(),
+                    role: getChatRole("ai"),
+                    content: outputContent,
+                  },
+                ]
+              : [];
+          if (toolCalls.length > 0) {
+            messages.push({
+              id: generateMessageId(),
+              role: getChatRole("ai"),
+              toolCalls: convertMessageToolCallsToProvider({
+                provider: instance.model.provider,
+                toolCalls: toolCalls.map((tc) => ({
+                  ...tc,
+                  function: {
+                    ...tc.function,
+                    arguments:
+                      safelyParseJSON(tc.function.arguments)?.json ??
+                      tc.function.arguments,
+                  },
+                })),
+              }),
+            });
+          }
+          updateInstance({
+            instanceId,
+            patch: {
+              template: {
+                __type: "chat",
+                messages: [...instance.template.messages, ...messages],
+              },
+            },
+          });
+          cleanupOutput();
+        }}
+      />
+      <Tooltip>Move Output message to the end of Prompt</Tooltip>
+    </TooltipTrigger>
+  );
+};