feat(playground): parse span input messages

Arize-ai · Oct 7, 2024 · 182c053 · 182c053
1 parent a9bd2d2
commit 182c053
Show file tree

Hide file tree

Showing 6 changed files with 62 additions and 88 deletions.
diff --git a/app/src/pages/playground/Playground.tsx b/app/src/pages/playground/Playground.tsx
@@ -8,13 +8,14 @@ import {
   PlaygroundProvider,
   usePlaygroundContext,
 } from "@phoenix/contexts/PlaygroundContext";
+import { PlaygroundProps } from "@phoenix/store";
 
 import { PlaygroundInstance } from "./PlaygroundInstance";
 import { PlaygroundOperationTypeRadioGroup } from "./PlaygroundOperationTypeRadioGroup";
 
-export function Playground() {
+export function Playground(props: Partial<PlaygroundProps>) {
   return (
-    <PlaygroundProvider>
+    <PlaygroundProvider {...props}>
       <View
         borderBottomColor="dark"
         borderBottomWidth="thin"

diff --git a/app/src/pages/playground/SpanPlaygroundPage.tsx b/app/src/pages/playground/SpanPlaygroundPage.tsx
@@ -1,50 +1,18 @@
 import React, { useMemo } from "react";
 import { useLoaderData } from "react-router";
-import { z } from "zod";
 
 import { PlaygroundInstance } from "@phoenix/store";
+import { safelyParseJSON } from "@phoenix/utils/jsonUtils";
 
 import { spanPlaygroundPageLoaderQuery$data } from "./__generated__/spanPlaygroundPageLoaderQuery.graphql";
 import { Playground } from "./Playground";
+import { llmAttributesSchema } from "./schemas";
 
 type PlaygroundSpan = Extract<
   spanPlaygroundPageLoaderQuery$data["span"],
   { __typename: "Span" }
 >;
 
-const toolCallSchema = z
-  .object({
-    function: z
-      .object({
-        name: z.string(),
-        arguments: z.string(),
-      })
-      .partial(),
-  })
-  .partial();
-
-const messageSchema = z
-  .object({
-    role: z.string(),
-    content: z.string(),
-    name: z.string(),
-    tool_calls: z.array(toolCallSchema),
-  })
-  .partial();
-
-const llmAttributesSchema = z
-  .object({
-    model_name: z.string(),
-    prompts: z.array(z.string()),
-    prompt_template: z
-      .object({
-        template: z.string(),
-        variables: z.record(z.string()),
-      })
-      .partial(),
-  })
-  .partial();
-
 export function SpanPlaygroundPage() {
   const data = useLoaderData() as spanPlaygroundPageLoaderQuery$data;
   const span = useMemo(() => {
@@ -58,37 +26,34 @@ export function SpanPlaygroundPage() {
     throw new Error("Invalid span");
   }
 
-  const PlaygroundInstance = z.object({
-    template: z.string(),
-  });
+  const playgroundInstance = transformSpanAttributesToPlaygroundInstance(span);
 
-  return <Playground />;
+  return <Playground instances={[playgroundInstance]} />;
 }
 
-type MessageTemplate = {
-  messages: { role: string; message: string }[];
-};
-
-const a = z.object({
-  messages: z.array(
-    z.object({
-      role: z.string(),
-      message: z.string(),
-    })
-  ),
-});
 function transformSpanAttributesToPlaygroundInstance(
   span: PlaygroundSpan
 ): PlaygroundInstance {
-  const result = ZMessageTemplate.safeParse(span);
-  if (!result.success) {
+  const { json: parsedAttributes, parseError } = safelyParseJSON(
+    span.attributes
+  );
+  if (parseError) {
+    throw new Error("Invalid attributes");
+  }
+  const { data, success } = llmAttributesSchema.safeParse(parsedAttributes);
+  if (!success) {
     throw new Error("Invalid data");
   }
-  return result.data;
-  // return {
-  //   messages: data.messages.map(message => ({
-  //     role: message.role,
-  //     message: message.message
-  //   }))
-  // };
+  return {
+    id: 1,
+    input: {
+      variables: {},
+    },
+    template: {
+      __type: "chat",
+      messages: data["llm.input_messages"],
+    },
+    output: data["llm.output_messages"],
+    tools: undefined,
+  };
 }
diff --git a/app/src/pages/playground/constants.tsx b/app/src/pages/playground/constants.tsx
@@ -1 +1,3 @@
 export const NUM_MAX_PLAYGROUND_INSTANCES = 2;
+
+export const DEFAULT_CHAT_ROLE = "system";
diff --git a/app/src/pages/playground/schemas.ts b/app/src/pages/playground/schemas.ts
@@ -2,10 +2,10 @@ import { z } from "zod";
 
 import {
   ImageAttributesPostfixes,
-  LLMAttributePostfixes,
   MessageAttributePostfixes,
   MessageContentsAttributePostfixes,
   SemanticAttributePrefixes,
+  SemanticConventions,
 } from "@arizeai/openinference-semantic-conventions";
 
 /**
@@ -22,6 +22,7 @@ const toolCallSchema = z
       .partial(),
   })
   .partial();
+
 /**
  * The zod schema for llm message contents
  * @see {@link https://github.com/Arize-ai/openinference/blob/main/spec/semantic_conventions.md|Semantic Conventions}
@@ -43,25 +44,25 @@ const messageContentSchema = z.object({
     })
     .partial(),
 });
-const messageSchema = z
-  .object({
-    [MessageAttributePostfixes.role]: z.string(),
-    [MessageAttributePostfixes.content]: z.string(),
-    [MessageAttributePostfixes.name]: z.string(),
-    [MessageAttributePostfixes.tool_calls]: z.array(toolCallSchema),
-    [MessageAttributePostfixes.contents]: z.array(messageContentSchema),
-  })
-  .partial();
+/**
+ * The zod schema for llm messages
+ * @see {@link https://github.com/Arize-ai/openinference/blob/main/spec/semantic_conventions.md|Semantic Conventions}
+ */
+const messageSchema = z.object({
+  [MessageAttributePostfixes.role]: z.string(),
+  [MessageAttributePostfixes.content]: z.string(),
+  [MessageAttributePostfixes.name]: z.string().optional(),
+  [MessageAttributePostfixes.tool_calls]: z.array(toolCallSchema).optional(),
+  [MessageAttributePostfixes.contents]: z
+    .array(messageContentSchema)
+    .optional(),
+});
 
-const llmAttributesSchema = z
-  .object({
-    [LLMAttributePostfixes.model_name]: z.string(),
-    [LLMAttributePostfixes.prompts]: z.array(z.string()),
-    prompt_template: z
-      .object({
-        template: z.string(),
-        variables: z.record(z.string()),
-      })
-      .partial(),
-  })
-  .partial();
+/**
+ * The zod schema for llm attributes
+ * @see {@link https://github.com/Arize-ai/openinference/blob/main/spec/semantic_conventions.md|Semantic Conventions}
+ */
+export const llmAttributesSchema = z.object({
+  [SemanticConventions.LLM_INPUT_MESSAGES]: z.array(messageSchema),
+  [SemanticConventions.LLM_OUTPUT_MESSAGES]: z.optional(z.array(messageSchema)),
+});
diff --git a/app/src/pages/trace/SpanDetails.tsx b/app/src/pages/trace/SpanDetails.tsx
@@ -79,6 +79,7 @@ import {
   AttributeTool,
 } from "@phoenix/openInference/tracing/types";
 import { assertUnreachable, isStringArray } from "@phoenix/typeUtils";
+import { safelyParseJSON } from "@phoenix/utils/jsonUtils";
 import { formatFloat, numberFormatter } from "@phoenix/utils/numberFormatUtils";
 
 import { RetrievalEvaluationLabel } from "../project/RetrievalEvaluationLabel";
@@ -117,11 +118,7 @@ const useSafelyParsedJSON = (
   jsonStr: string
 ): { json: { [key: string]: unknown } | null; parseError?: unknown } => {
   return useMemo(() => {
-    try {
-      return { json: JSON.parse(jsonStr) };
-    } catch (e) {
-      return { json: null, parseError: e };
-    }
+    return safelyParseJSON(jsonStr);
   }, [jsonStr]);
 };
 

diff --git a/app/src/utils/jsonUtils.ts b/app/src/utils/jsonUtils.ts
@@ -32,3 +32,11 @@ export function isJSONString({
 export function isJSONObjectString(str: string) {
   return isJSONString({ str, excludeArray: true, excludePrimitives: true });
 }
+
+export function safelyParseJSON(str: string) {
+  try {
+    return { json: JSON.parse(str) };
+  } catch (e) {
+    return { json: null, parseError: e };
+  }
+}
Original file line number	Diff line number	Diff line change
		@@ -1 +1,3 @@
		export const NUM_MAX_PLAYGROUND_INSTANCES = 2;

		export const DEFAULT_CHAT_ROLE = "system";