feat(playground): parse span attribute llm.input_messages for playg…

…round span replay (#4906)

* begin building zod schemas for llm attributes

* feat(playground): parse span input messages

* move to utils, update role

* update initial instance id

* update id generation, add todo, add test files

* add tests

* memoize

* move jest-canvas-mock to dev dependencies

* update span not found error message

* parse roles as strings and corece to ChatMessageRole

* update ChatRoleMap comment

* fix typos

* update prop type to be InitialPlaygroundState

* update naming

* fix naming conflict

app/jest.config.js

@@ -5,6 +5,8 @@ module.exports = {
   preset: "ts-jest",
   testEnvironment: "jsdom",
   prettierPath: null,
+  setupFiles: ["<rootDir>/jest.setup.ts"],
+  testMatch: ["**/__tests__/*.test.ts?(x)"],
   transform: {
     "^.+\\.[jt]sx?$": ["esbuild-jest"],
   },

app/jest.setup.ts

@@ -0,0 +1,10 @@
+import "jest-canvas-mock";
+jest.mock("@phoenix/config");
+
+Object.defineProperty(window, "Config", {
+  value: {
+    authenticationEnabled: true,
+    basename: "/",
+    platformVersion: "1.0.0",
+  },
+});

app/package.json

@@ -47,6 +47,7 @@
     "three-stdlib": "^2.30.4",
     "use-deep-compare-effect": "^1.8.1",
     "use-zustand": "^0.0.4",
+    "zod": "^3.23.8",
     "zustand": "^4.5.4"
   },
   "devDependencies": {
@@ -76,6 +77,7 @@
     "eslint-plugin-simple-import-sort": "^10.0.0",
     "graphql": "^16.9.0",
     "jest": "^29.7.0",
+    "jest-canvas-mock": "^2.5.2",
     "jest-environment-jsdom": "^29.7.0",
     "only-allow": "^1.2.1",
     "prettier": "^3.3.3",

app/src/pages/playground/Playground.tsx

@@ -8,13 +8,14 @@ import {
   PlaygroundProvider,
   usePlaygroundContext,
 } from "@phoenix/contexts/PlaygroundContext";
+import { InitialPlaygroundState } from "@phoenix/store";
 
 import { PlaygroundInstance } from "./PlaygroundInstance";
 import { PlaygroundOperationTypeRadioGroup } from "./PlaygroundOperationTypeRadioGroup";
 
-export function Playground() {
+export function Playground(props: InitialPlaygroundState) {
   return (
-    <PlaygroundProvider>
+    <PlaygroundProvider {...props}>
       <View
         borderBottomColor="dark"
         borderBottomWidth="thin"

app/src/pages/playground/SpanPlaygroundPage.tsx

@@ -1,10 +1,31 @@
-import React from "react";
+import React, { useMemo } from "react";
+import { useLoaderData } from "react-router";
 
-// import { useLoaderData } from "react-router";
+import { spanPlaygroundPageLoaderQuery$data } from "./__generated__/spanPlaygroundPageLoaderQuery.graphql";
 import { Playground } from "./Playground";
+import { transformSpanAttributesToPlaygroundInstance } from "./playgroundUtils";
 
 export function SpanPlaygroundPage() {
-  // const data = useLoaderData();
+  const data = useLoaderData() as spanPlaygroundPageLoaderQuery$data;
+  const span = useMemo(() => {
+    if (data.span.__typename === "Span") {
+      return data.span;
+    }
+    return null;
+  }, [data.span]);
 
-  return <Playground />;
+  if (!span) {
+    throw new Error("Span not found");
+  }
+
+  const playgroundInstance = useMemo(
+    () => transformSpanAttributesToPlaygroundInstance(span),
+    [span]
+  );
+
+  return (
+    <Playground
+      instances={playgroundInstance != null ? [playgroundInstance] : undefined}
+    />
+  );
 }

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

@@ -0,0 +1,40 @@
+import { PlaygroundSpan } from "../spanPlaygroundPageLoader";
+
+export const basePlaygroundSpan: PlaygroundSpan = {
+  __typename: "Span",
+  context: {
+    spanId: "test",
+  },
+  attributes: "",
+};
+export const spanAttributesWithInputMessages = {
+  llm: {
+    output_messages: [
+      {
+        message: {
+          content: "This is an AI Answer",
+          role: "assistant",
+        },
+      },
+    ],
+    model_name: "gpt-3.5-turbo",
+    token_count: { completion: 9.0, prompt: 1881.0, total: 1890.0 },
+    input_messages: [
+      {
+        message: {
+          content: "You are a chatbot",
+          role: "system",
+        },
+      },
+      {
+        message: {
+          content: "Anser me the following question. Are you sentient?",
+          role: "user",
+        },
+      },
+    ],
+    invocation_parameters:
+      '{"context_window": 16384, "num_output": -1, "is_chat_model": true, "is_function_calling_model": true, "model_name": "gpt-3.5-turbo"}',
+  },
+  openinference: { span: { kind: "LLM" } },
+} as const;

app/src/pages/playground/__tests__/playgroundUtils.test.ts

@@ -0,0 +1,96 @@
+import { _resetInstanceId } from "@phoenix/store";
+
+import {
+  getChatRole,
+  transformSpanAttributesToPlaygroundInstance,
+} from "../playgroundUtils";
+
+import {
+  basePlaygroundSpan,
+  spanAttributesWithInputMessages,
+} from "./fixtures";
+
+const expectedPlaygroundInstance = {
+  id: 0,
+  activeRunId: null,
+  isRunning: false,
+  input: {
+    variables: {},
+  },
+  template: {
+    __type: "chat",
+    messages: spanAttributesWithInputMessages.llm.input_messages.map(
+      ({ message }) => message
+    ),
+  },
+  output: spanAttributesWithInputMessages.llm.output_messages,
+  tools: undefined,
+};
+
+describe("transformSpanAttributesToPlaygroundInstance", () => {
+  beforeEach(() => {
+    _resetInstanceId();
+  });
+  it("should throw if the attributes are not parsable", () => {
+    const span = {
+      ...basePlaygroundSpan,
+      attributes: "invalid json",
+    };
+    expect(() => transformSpanAttributesToPlaygroundInstance(span)).toThrow(
+      "Invalid span attributes, attributes must be valid JSON"
+    );
+  });
+
+  it("should return null if the attributes do not match the schema", () => {
+    const span = {
+      ...basePlaygroundSpan,
+      attributes: JSON.stringify({}),
+    };
+    expect(transformSpanAttributesToPlaygroundInstance(span)).toBeNull();
+  });
+
+  it("should return a PlaygroundInstance if the attributes contain llm.input_messages", () => {
+    const span = {
+      ...basePlaygroundSpan,
+      attributes: JSON.stringify(spanAttributesWithInputMessages),
+    };
+
+    expect(transformSpanAttributesToPlaygroundInstance(span)).toEqual(
+      expectedPlaygroundInstance
+    );
+  });
+
+  it("should return a PlaygroundInstance if the attributes contain llm.input_messages, even if output_messages are not present", () => {
+    const span = {
+      ...basePlaygroundSpan,
+      attributes: JSON.stringify({
+        ...spanAttributesWithInputMessages,
+        llm: {
+          ...spanAttributesWithInputMessages.llm,
+          output_messages: undefined,
+        },
+      }),
+    };
+    expect(transformSpanAttributesToPlaygroundInstance(span)).toEqual({
+      ...expectedPlaygroundInstance,
+      output: undefined,
+    });
+  });
+});
+
+describe("getChatRole", () => {
+  it("should return the role if it is a valid ChatMessageRole", () => {
+    expect(getChatRole("user")).toEqual("user");
+  });
+
+  it("should return the ChatMessageRole if the role is included in ChatRoleMap", () => {
+    expect(getChatRole("assistant")).toEqual("ai");
+    // expect(getChatRole("bot")).toEqual("ai");
+    // expect(getChatRole("system")).toEqual("system");
+    // expect(getChatRole("human:")).toEqual("user");
+  });
+
+  it("should return DEFAULT_CHAT_ROLE if the role is not found", () => {
+    expect(getChatRole("invalid")).toEqual("user");
+  });
+});

app/src/pages/playground/constants.tsx

@@ -1 +1,16 @@
+import { ChatMessageRole } from "@phoenix/store";
+
 export const NUM_MAX_PLAYGROUND_INSTANCES = 2;
+
+export const DEFAULT_CHAT_ROLE = "user";
+
+/**
+ * Map of {@link ChatMessageRole} to potential role values.
+ * Used to map roles to a canonical role.
+ */
+export const ChatRoleMap: Record<ChatMessageRole, string[]> = {
+  user: ["user", "human"],
+  ai: ["assistant", "bot", "ai"],
+  system: ["system"],
+  tool: ["tool"],
+};