Skip to content

Commit

Permalink
feat(playground): parse span attribute llm.input_messages for playg…
Browse files Browse the repository at this point in the history
…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
  • Loading branch information
Parker-Stafford authored and mikeldking committed Oct 9, 2024
1 parent 1fd02cd commit 6091c19
Show file tree
Hide file tree
Showing 15 changed files with 399 additions and 17 deletions.
2 changes: 2 additions & 0 deletions app/jest.config.js
Original file line number Diff line number Diff line change
Expand Up @@ -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"],
},
Expand Down
10 changes: 10 additions & 0 deletions app/jest.setup.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
import "jest-canvas-mock";
jest.mock("@phoenix/config");

Object.defineProperty(window, "Config", {
value: {
authenticationEnabled: true,
basename: "/",
platformVersion: "1.0.0",
},
});
2 changes: 2 additions & 0 deletions app/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
Expand Down Expand Up @@ -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",
Expand Down
26 changes: 26 additions & 0 deletions app/pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

5 changes: 3 additions & 2 deletions app/src/pages/playground/Playground.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
29 changes: 25 additions & 4 deletions app/src/pages/playground/SpanPlaygroundPage.tsx
Original file line number Diff line number Diff line change
@@ -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}
/>
);
}
40 changes: 40 additions & 0 deletions app/src/pages/playground/__tests__/fixtures.ts
Original file line number Diff line number Diff line change
@@ -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;
96 changes: 96 additions & 0 deletions app/src/pages/playground/__tests__/playgroundUtils.test.ts
Original file line number Diff line number Diff line change
@@ -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");
});
});
15 changes: 15 additions & 0 deletions app/src/pages/playground/constants.tsx
Original file line number Diff line number Diff line change
@@ -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"],
};
Loading

0 comments on commit 6091c19

Please sign in to comment.