feat(tracing): pretty print JSON strings for structured data outputs (#…

…5811)
Arize-ai · Dec 21, 2024 · b371f7c · b371f7c
1 parent aabe535
commit b371f7c
Show file tree

Hide file tree

Showing 7 changed files with 54 additions and 29 deletions.
diff --git a/app/src/components/code/styles.tsx b/app/src/components/code/styles.tsx
@@ -1,9 +1,6 @@
 import { css } from "@emotion/react";
 
 export const readOnlyCodeMirrorCSS = css`
-  .cm-content {
-    padding: var(--ac-global-dimension-static-size-200) 0;
-  }
   .cm-editor,
   .cm-gutters {
     background-color: transparent;

diff --git a/app/src/components/markdown/MarkdownBlock.tsx b/app/src/components/markdown/MarkdownBlock.tsx
@@ -3,6 +3,8 @@ import Markdown from "react-markdown";
 import remarkGfm from "remark-gfm";
 import { css } from "@emotion/react";
 
+import { PrettyText } from "../utility";
+
 import { useMarkdownMode } from "./MarkdownDisplayContext";
 import { MarkdownDisplayMode } from "./types";
 
@@ -26,15 +28,7 @@ export function MarkdownBlock({
       <Markdown remarkPlugins={[remarkGfm]}>{children}</Markdown>
     </div>
   ) : (
-    <pre
-      css={css`
-        white-space: pre-wrap;
-        text-wrap: wrap;
-        margin: 0;
-      `}
-    >
-      {children}
-    </pre>
+    <PrettyText>{children}</PrettyText>
   );
 }
 

diff --git a/app/src/components/utility/PrettyText.tsx b/app/src/components/utility/PrettyText.tsx
@@ -0,0 +1,27 @@
+import React from "react";
+import { css } from "@emotion/react";
+
+import { JSONBlock } from "@phoenix/components/code";
+import { usePrettyText } from "@phoenix/hooks/usePrettyText";
+import { assertUnreachable } from "@phoenix/typeUtils";
+
+export function PrettyText({ children }: { children: string }) {
+  const { text, textType } = usePrettyText(children);
+  if (textType === "string") {
+    return (
+      <pre
+        css={css`
+          white-space: pre-wrap;
+          text-wrap: wrap;
+          margin: 0;
+        `}
+      >
+        {text}
+      </pre>
+    );
+  }
+  if (textType === "json") {
+    return <JSONBlock value={text} />;
+  }
+  assertUnreachable(textType);
+}
diff --git a/app/src/components/utility/index.tsx b/app/src/components/utility/index.tsx
@@ -0,0 +1 @@
+export * from "./PrettyText";
diff --git a/app/src/hooks/usePrettyText.ts b/app/src/hooks/usePrettyText.ts
@@ -0,0 +1,21 @@
+import { useMemo } from "react";
+
+type TextType = "string" | "json";
+/**
+ * A hook that takes in text, detects if it is a JSON string, and formats it
+ * @param text
+ * @returns formatted text
+ */
+export function usePrettyText(text: string): {
+  text: string;
+  textType: TextType;
+} {
+  return useMemo(() => {
+    try {
+      const parsed = JSON.parse(text);
+      return { text: JSON.stringify(parsed, null, 2), textType: "json" };
+    } catch (e) {
+      return { text, textType: "string" };
+    }
+  }, [text]);
+}
diff --git a/app/src/pages/playground/PlaygroundOutput.tsx b/app/src/pages/playground/PlaygroundOutput.tsx
@@ -7,14 +7,13 @@ import {
   requestSubscription,
 } from "relay-runtime";
 
-import { Card, Flex, View } from "@arizeai/components";
+import { Card, View } from "@arizeai/components";
 
 import { Loading } from "@phoenix/components";
 import {
   ConnectedMarkdownBlock,
   ConnectedMarkdownModeRadioGroup,
   MarkdownDisplayProvider,
-  useMarkdownMode,
 } from "@phoenix/components/markdown";
 import { useNotifyError } from "@phoenix/contexts";
 import { useCredentialsContext } from "@phoenix/contexts/CredentialsContext";
@@ -84,7 +83,6 @@ function PlaygroundOutputMessage({
 }) {
   const { role, content, toolCalls } = message;
   const styles = useChatMessageStyles(role);
-  const { mode: markdownMode } = useMarkdownMode();
 
   return (
     <Card
@@ -94,15 +92,7 @@ function PlaygroundOutputMessage({
       extra={<ConnectedMarkdownModeRadioGroup />}
     >
       {content != null && !Array.isArray(content) && (
-        <Flex direction="column" alignItems="start">
-          {markdownMode === "text" ? (
-            content
-          ) : (
-            <View overflow="auto" maxWidth="100%">
-              <ConnectedMarkdownBlock>{content}</ConnectedMarkdownBlock>
-            </View>
-          )}
-        </Flex>
+        <ConnectedMarkdownBlock>{content}</ConnectedMarkdownBlock>
       )}
       {toolCalls && toolCalls.length > 0
         ? toolCalls.map((toolCall) => {

diff --git a/app/src/pages/trace/SpanDetails.tsx b/app/src/pages/trace/SpanDetails.tsx
@@ -310,7 +310,6 @@ export function SpanDetails({
               {...defaultCardProps}
               titleExtra={attributesContextualHelp}
               extra={<CopyToClipboardButton text={span.attributes} />}
-              bodyStyle={{ padding: 0 }}
             >
               <JSONBlock>{span.attributes}</JSONBlock>
             </Card>
@@ -1587,7 +1586,6 @@ function SpanIO({ span }: { span: Span }) {
           title="All Attributes"
           titleExtra={attributesContextualHelp}
           {...defaultCardProps}
-          bodyStyle={{ padding: 0 }}
           extra={<CopyToClipboardButton text={span.attributes} />}
         >
           <JSONBlock>{span.attributes}</JSONBlock>
@@ -1598,9 +1596,6 @@ function SpanIO({ span }: { span: Span }) {
 }
 
 const codeMirrorCSS = css`
-  .cm-content {
-    padding: var(--ac-global-dimension-static-size-200) 0;
-  }
   .cm-editor,
   .cm-gutters {
     background-color: transparent;