Skip to content

Commit

Permalink
feat(playground): make chat messages match style in traces (#4931)
Browse files Browse the repository at this point in the history
* feat(playground): make chat messages match style in traces

* small fixes

* WIP
  • Loading branch information
mikeldking committed Oct 11, 2024
1 parent 784606a commit ed6e05e
Show file tree
Hide file tree
Showing 6 changed files with 112 additions and 66 deletions.
35 changes: 35 additions & 0 deletions app/src/hooks/useChatMessageStyles.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
import { useMemo } from "react";

import { ViewStyleProps } from "@arizeai/components";

export function useChatMessageStyles(
role: string
): Pick<ViewStyleProps, "backgroundColor" | "borderColor"> {
return useMemo<ViewStyleProps>(() => {
if (role === "user" || role === "human") {
return {
backgroundColor: "grey-100",
borderColor: "grey-500",
};
} else if (role === "assistant" || role === "ai") {
return {
backgroundColor: "blue-100",
borderColor: "blue-700",
};
} else if (role === "system") {
return {
backgroundColor: "indigo-100",
borderColor: "indigo-700",
};
} else if (["function", "tool"].includes(role)) {
return {
backgroundColor: "yellow-100",
borderColor: "yellow-700",
};
}
return {
backgroundColor: "grey-100",
borderColor: "grey-700",
};
}, [role]);
}
6 changes: 3 additions & 3 deletions app/src/pages/playground/Playground.tsx
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import React from "react";
import React, { Fragment } from "react";
import { Panel, PanelGroup, PanelResizeHandle } from "react-resizable-panels";

import { Button, Flex, Heading, View } from "@arizeai/components";
Expand Down Expand Up @@ -48,12 +48,12 @@ function PlaygroundInstances() {
<Flex direction="row" alignItems="stretch" height="100%">
<PanelGroup direction="horizontal">
{instances.map((instance, i) => (
<>
<Fragment key={i}>
{i !== 0 && <PanelResizeHandle css={resizeHandleCSS} />}
<Panel defaultSize={50}>
<PlaygroundInstance key={i} playgroundInstanceId={instance.id} />
</Panel>
</>
</Fragment>
))}
</PanelGroup>
</Flex>
Expand Down
75 changes: 52 additions & 23 deletions app/src/pages/playground/PlaygroundChatTemplate.tsx
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
import React from "react";
import React, { PropsWithChildren } from "react";
import { css } from "@emotion/react";

import { Card, TextArea } from "@arizeai/components";
import { Card, CardProps, TextArea } from "@arizeai/components";

import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
import { useChatMessageStyles } from "@phoenix/hooks/useChatMessageStyles";

import { MessageRolePicker } from "./MessageRolePicker";
import { PlaygroundInstanceProps } from "./types";
Expand Down Expand Up @@ -35,7 +36,8 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
{template.messages.map((message, index) => {
return (
<li key={index}>
<Card
<ChatMessageCard
role={message.role}
title={
<MessageRolePicker
includeLabel={false}
Expand All @@ -55,31 +57,58 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
}}
/>
}
variant="compact"
backgroundColor="light"
borderColor="light"
>
<TextArea
height={100}
value={message.content}
onChange={(val) => {
updateInstance({
instanceId: id,
patch: {
template: {
__type: "chat",
messages: template.messages.map((message, i) =>
i === index ? { ...message, content: val } : message
),
<div
css={css`
// TODO: remove these styles once the codemiror editor is added
.ac-textfield {
border: none !important;
border-radius: 0;
textarea {
padding: var(--ac-global-dimension-size-200);
}
}
`}
>
<TextArea
value={message.content}
height={200}
variant="quiet"
onChange={(val) => {
updateInstance({
instanceId: id,
patch: {
template: {
__type: "chat",
messages: template.messages.map((message, i) =>
i === index ? { ...message, content: val } : message
),
},
},
},
});
}}
/>
</Card>
});
}}
/>
</div>
</ChatMessageCard>
</li>
);
})}
</ul>
);
}

function ChatMessageCard(
props: PropsWithChildren<{ title: CardProps["title"]; role: string }>
) {
const styles = useChatMessageStyles(props.role);
return (
<Card
title={props.title}
variant="compact"
{...styles}
bodyStyle={{ padding: 0 }}
>
{props.children}
</Card>
);
}
8 changes: 7 additions & 1 deletion app/src/pages/playground/SpanPlaygroundPage.tsx
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import React, { useMemo } from "react";
import { useLoaderData } from "react-router";

import { createPlaygroundInstance } from "@phoenix/store";

import { spanPlaygroundPageLoaderQuery$data } from "./__generated__/spanPlaygroundPageLoaderQuery.graphql";
import { Playground } from "./Playground";
import { transformSpanAttributesToPlaygroundInstance } from "./playgroundUtils";
Expand All @@ -25,7 +27,11 @@ export function SpanPlaygroundPage() {

return (
<Playground
instances={playgroundInstance != null ? [playgroundInstance] : undefined}
instances={
playgroundInstance != null
? [playgroundInstance]
: [createPlaygroundInstance()]
}
/>
);
}
30 changes: 2 additions & 28 deletions app/src/pages/trace/SpanDetails.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,6 @@ import {
TooltipTrigger,
View,
ViewProps,
ViewStyleProps,
} from "@arizeai/components";
import {
DocumentAttributePostfixes,
Expand All @@ -65,6 +64,7 @@ import { SpanKindLabel } from "@phoenix/components/trace/SpanKindLabel";
import { useNotifySuccess, useTheme } from "@phoenix/contexts";
import { useFeatureFlag } from "@phoenix/contexts/FeatureFlagsContext";
import { usePreferencesContext } from "@phoenix/contexts/PreferencesContext";
import { useChatMessageStyles } from "@phoenix/hooks/useChatMessageStyles";
import {
AttributeDocument,
AttributeEmbedding,
Expand Down Expand Up @@ -1292,33 +1292,7 @@ function LLMMessage({ message }: { message: AttributeMessage }) {
message[MessageAttributePostfixes.function_call_arguments_json] &&
message[MessageAttributePostfixes.function_call_name];
const role = message[MessageAttributePostfixes.role] || "unknown";
const messageStyles = useMemo<ViewStyleProps>(() => {
if (role === "user") {
return {
backgroundColor: "grey-100",
borderColor: "grey-500",
};
} else if (role === "assistant") {
return {
backgroundColor: "blue-100",
borderColor: "blue-700",
};
} else if (role === "system") {
return {
backgroundColor: "indigo-100",
borderColor: "indigo-700",
};
} else if (["function", "tool"].includes(role)) {
return {
backgroundColor: "yellow-100",
borderColor: "yellow-700",
};
}
return {
backgroundColor: "grey-100",
borderColor: "grey-700",
};
}, [role]);
const messageStyles = useChatMessageStyles(role);

return (
<MarkdownDisplayProvider>
Expand Down
24 changes: 13 additions & 11 deletions app/src/store/playgroundStore.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -193,24 +193,26 @@ const DEFAULT_TEXT_COMPLETION_TEMPLATE: PlaygroundTextCompletionTemplate = {
prompt: "{{question}}",
};

export function createPlaygroundInstance(): PlaygroundInstance {
return {
id: generateInstanceId(),
template: DEFAULT_CHAT_COMPLETION_TEMPLATE,
tools: {},
input: { variables: {} },
output: {},
activeRunId: null,
isRunning: false,
};
}

export const createPlaygroundStore = (
initialProps?: InitialPlaygroundState
) => {
const playgroundStore: StateCreator<PlaygroundState> = (set, get) => ({
operationType: "chat",
inputMode: "manual",
setInputMode: (inputMode: PlaygroundInputMode) => set({ inputMode }),
instances: [
{
id: generateInstanceId(),
template: DEFAULT_CHAT_COMPLETION_TEMPLATE,
tools: {},
input: { variables: {} },
output: {},
activeRunId: null,
isRunning: false,
},
],
instances: [createPlaygroundInstance()],
setOperationType: (operationType: GenAIOperationType) => {
if (operationType === "chat") {
// TODO: this is incorrect, it should only change the template
Expand Down

0 comments on commit ed6e05e

Please sign in to comment.