feat(playground): accumulate errors while parsing span attributes (#4941

)

* feat(playground): accumulate errors while parsing span attributes

* update tests to include message id's / tools

* update output for runs

* update comments / test names

* move parsing errors into span playground banners

* cleanup spanplaygroundpage

* add function to verify zod schema matches type

app/src/pages/playground/PlaygroundChatTemplate.tsx

@@ -30,6 +30,7 @@ import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
 import { useChatMessageStyles } from "@phoenix/hooks/useChatMessageStyles";
 import {
   ChatMessage,
+  ChatMessageRole,
   generateMessageId,
   PlaygroundChatTemplate as PlaygroundChatTemplateType,
 } from "@phoenix/store";
@@ -143,7 +144,7 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
                       ...template.messages,
                       {
                         id: generateMessageId(),
-                        role: "user",
+                        role: ChatMessageRole.user,
                         content: "",
                       },
                     ],

app/src/pages/playground/PlaygroundOutput.tsx

@@ -5,7 +5,12 @@ import { graphql, GraphQLSubscriptionConfig } from "relay-runtime";
 import { Card, Flex, Icon, Icons } from "@arizeai/components";
 
 import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
-import { ChatMessage, ChatMessageRole } from "@phoenix/store";
+import { useChatMessageStyles } from "@phoenix/hooks/useChatMessageStyles";
+import {
+  ChatMessage,
+  ChatMessageRole,
+  generateMessageId,
+} from "@phoenix/store";
 import { assertUnreachable } from "@phoenix/typeUtils";
 
 import {
@@ -15,11 +20,22 @@ import {
   PlaygroundOutputSubscription$data,
   PlaygroundOutputSubscription$variables,
 } from "./__generated__/PlaygroundOutputSubscription.graphql";
+import { isChatMessages } from "./playgroundUtils";
 import { TitleWithAlphabeticIndex } from "./TitleWithAlphabeticIndex";
 import { PlaygroundInstanceProps } from "./types";
 
 interface PlaygroundOutputProps extends PlaygroundInstanceProps {}
 
+function PlaygroundOutputMessage({ message }: { message: ChatMessage }) {
+  const styles = useChatMessageStyles(message.role);
+
+  return (
+    <Card title={message.role} {...styles} variant="compact">
+      {message.content}
+    </Card>
+  );
+}
+
 export function PlaygroundOutput(props: PlaygroundOutputProps) {
   const instanceId = props.playgroundInstanceId;
   const instance = usePlaygroundContext((state) =>
@@ -29,22 +45,46 @@ export function PlaygroundOutput(props: PlaygroundOutputProps) {
     state.instances.findIndex((instance) => instance.id === instanceId)
   );
   if (!instance) {
-    return null;
+    throw new Error("Playground instance not found");
   }
 
   const runId = instance.activeRunId;
   const hasRunId = runId !== null;
+
+  const OutputEl = useMemo(() => {
+    if (hasRunId) {
+      return (
+        <PlaygroundOutputText key={runId} playgroundInstanceId={instanceId} />
+      );
+    }
+    if (isChatMessages(instance.output)) {
+      const messages = instance.output;
+
+      return messages.map((message, index) => {
+        return <PlaygroundOutputMessage key={index} message={message} />;
+      });
+    }
+    if (typeof instance.output === "string") {
+      return (
+        <PlaygroundOutputMessage
+          message={{
+            id: generateMessageId(),
+            content: instance.output,
+            role: ChatMessageRole.ai,
+          }}
+        />
+      );
+    }
+    return "click run to see output";
+  }, [hasRunId, instance.output, instanceId, runId]);
+
   return (
     <Card
       title={<TitleWithAlphabeticIndex index={index} title="Output" />}
       collapsible
       variant="compact"
     >
-      {hasRunId ? (
-        <PlaygroundOutputText key={runId} playgroundInstanceId={instanceId} />
-      ) : (
-        "click run to see output"
-      )}
+      {OutputEl}
     </Card>
   );
 }
@@ -104,27 +144,27 @@ function toGqlChatCompletionRole(
   role: ChatMessageRole
 ): ChatCompletionMessageRole {
   switch (role) {
-    case "system":
+    case ChatMessageRole.system:
       return "SYSTEM";
-    case "user":
+    case ChatMessageRole.user:
       return "USER";
-    case "tool":
+    case ChatMessageRole.tool:
       return "TOOL";
-    case "ai":
+    case ChatMessageRole.ai:
       return "AI";
     default:
       assertUnreachable(role);
   }
 }
 
 function PlaygroundOutputText(props: PlaygroundInstanceProps) {
-  const instance = usePlaygroundContext(
-    (state) => state.instances[props.playgroundInstanceId]
+  const instances = usePlaygroundContext((state) => state.instances);
+  const instance = instances.find(
+    (instance) => instance.id === props.playgroundInstanceId
   );
   const markPlaygroundInstanceComplete = usePlaygroundContext(
     (state) => state.markPlaygroundInstanceComplete
   );
-  const [output, setOutput] = useState<string>("");
   if (!instance) {
     throw new Error("No instance found");
   }
@@ -136,6 +176,8 @@ function PlaygroundOutputText(props: PlaygroundInstanceProps) {
     throw new Error("We only support chat templates for now");
   }
 
+  const [output, setOutput] = useState<string>("");
+
   useChatCompletionSubscription({
     params: {
       messages: instance.template.messages.map(toGqlChatCompletionMessage),
@@ -157,5 +199,13 @@ function PlaygroundOutputText(props: PlaygroundInstanceProps) {
       </Flex>
     );
   }
-  return <span>{output}</span>;
+  return (
+    <PlaygroundOutputMessage
+      message={{
+        id: generateMessageId(),
+        content: output,
+        role: ChatMessageRole.ai,
+      }}
+    />
+  );
 }

app/src/pages/playground/SpanPlaygroundPage.tsx

@@ -3,8 +3,6 @@ import { useLoaderData, useNavigate } from "react-router";
 
 import { Alert, Button, Flex, Icon, Icons } from "@arizeai/components";
 
-import { createPlaygroundInstance } from "@phoenix/store";
-
 import { spanPlaygroundPageLoaderQuery$data } from "./__generated__/spanPlaygroundPageLoaderQuery.graphql";
 import { Playground } from "./Playground";
 import { transformSpanAttributesToPlaygroundInstance } from "./playgroundUtils";
@@ -22,37 +20,37 @@ export function SpanPlaygroundPage() {
     throw new Error("Span not found");
   }
 
-  const playgroundInstance = useMemo(
+  const { playgroundInstance, parsingErrors } = useMemo(
     () => transformSpanAttributesToPlaygroundInstance(span),
     [span]
   );
 
   return (
     <Flex direction="column" height="100%">
-      <SpanPlaygroundBanners span={span} />
-      <Playground
-        instances={
-          playgroundInstance != null
-            ? [playgroundInstance]
-            : [createPlaygroundInstance()]
-        }
-      />
+      <SpanPlaygroundBanners span={span} parsingErrors={parsingErrors} />
+      <Playground instances={[playgroundInstance]} />
     </Flex>
   );
 }
 
 function SpanPlaygroundBanners({
   span,
+  parsingErrors,
 }: {
   span: Extract<
     NonNullable<spanPlaygroundPageLoaderQuery$data["span"]>,
     { __typename: "Span" }
   >;
+
+  parsingErrors?: string[];
 }) {
   const navigate = useNavigate();
+  const hasParsingErrors = parsingErrors && parsingErrors.length > 0;
   const [showBackBanner, setShowBackBanner] = useState(true);
+  const [showParsingErrorsBanner, setShowParsingErrorsBanner] =
+    useState(hasParsingErrors);
   return (
-    <div>
+    <Flex direction={"column"} gap={"size-50"}>
       {showBackBanner && (
         <Alert
           variant="info"
@@ -77,6 +75,23 @@ function SpanPlaygroundBanners({
           }
         >{`Replay and iterate on your LLM call from your ${span.project.name} project`}</Alert>
       )}
-    </div>
+      {showParsingErrorsBanner && hasParsingErrors && (
+        <Alert
+          variant="warning"
+          banner
+          dismissable
+          onDismissClick={() => {
+            setShowParsingErrorsBanner(false);
+          }}
+          title="The following errors occurred when parsing span attributes:"
+        >
+          <ul>
+            {parsingErrors.map((error) => (
+              <li key={error}>{error}</li>
+            ))}
+          </ul>
+        </Alert>
+      )}
+    </Flex>
   );
 }

app/src/pages/playground/__tests__/fixtures.ts

@@ -34,7 +34,7 @@ export const spanAttributesWithInputMessages = {
       },
       {
         message: {
-          content: "Anser me the following question. Are you sentient?",
+          content: "hello?",
           role: "user",
         },
       },