PlaygroundOutput.tsx•18.8 kB
import { Key, Suspense, useCallback, useEffect } from "react";
import { useMutation, useRelayEnvironment } from "react-relay";
import {
graphql,
GraphQLSubscriptionConfig,
PayloadError,
requestSubscription,
} from "relay-runtime";
import {
Card,
Flex,
Icon,
Icons,
Loading,
Text,
Tooltip,
TooltipTrigger,
TriggerWrap,
View,
} from "@phoenix/components";
import {
ConnectedMarkdownBlock,
ConnectedMarkdownModeSelect,
MarkdownDisplayProvider,
} from "@phoenix/components/markdown";
import { useNotifyError } from "@phoenix/contexts";
import { useCredentialsContext } from "@phoenix/contexts/CredentialsContext";
import {
usePlaygroundContext,
usePlaygroundStore,
} from "@phoenix/contexts/PlaygroundContext";
import { useChatMessageStyles } from "@phoenix/hooks/useChatMessageStyles";
import {
ChatMessage,
generateMessageId,
PlaygroundRepetition,
} from "@phoenix/store";
import { isStringKeyedObject } from "@phoenix/typeUtils";
import {
getErrorMessagesFromRelayMutationError,
getErrorMessagesFromRelaySubscriptionError,
} from "@phoenix/utils/errorUtils";
import { ExperimentRepetitionSelector } from "../experiment/ExperimentRepetitionSelector";
import PlaygroundOutputMutation, {
PlaygroundOutputMutation as PlaygroundOutputMutationType,
PlaygroundOutputMutation$data,
} from "./__generated__/PlaygroundOutputMutation.graphql";
import PlaygroundOutputSubscription, {
PlaygroundOutputSubscription as PlaygroundOutputSubscriptionType,
PlaygroundOutputSubscription$data,
} from "./__generated__/PlaygroundOutputSubscription.graphql";
import { PlaygroundErrorWrap } from "./PlaygroundErrorWrap";
import { PlaygroundOutputMoveButton } from "./PlaygroundOutputMoveButton";
import {
PartialOutputToolCall,
PlaygroundToolCall,
} from "./PlaygroundToolCall";
import { getChatCompletionInput, isChatMessages } from "./playgroundUtils";
import { RunMetadataFooter } from "./RunMetadataFooter";
import { TitleWithAlphabeticIndex } from "./TitleWithAlphabeticIndex";
import { PlaygroundInstanceProps } from "./types";
interface PlaygroundOutputProps extends PlaygroundInstanceProps {}
/**
* A chat message with potentially partial tool calls, for when tool calls are being streamed back to the client
*/
type PlaygroundOutputMessageType = Omit<ChatMessage, "toolCalls"> & {
toolCalls?: ChatMessage["toolCalls"] | readonly PartialOutputToolCall[];
};
const getToolCallKey = (
toolCall:
| NonNullable<ChatMessage["toolCalls"]>[number]
| PartialOutputToolCall[]
): Key => {
if (
isStringKeyedObject(toolCall) &&
"id" in toolCall &&
(typeof toolCall.id === "string" || typeof toolCall.id === "number")
) {
return toolCall.id;
} else if (
isStringKeyedObject(toolCall) &&
"toolUse" in toolCall &&
isStringKeyedObject(toolCall.toolUse) &&
"toolUseId" in toolCall.toolUse &&
(typeof toolCall.toolUse.toolUseId === "string" ||
typeof toolCall.toolUse.toolUseId === "number")
) {
return toolCall.toolUse.toolUseId;
}
return JSON.stringify(toolCall);
};
function PlaygroundOutputMessage({
message,
}: {
message: PlaygroundOutputMessageType;
}) {
const { role, content, toolCalls } = message;
const styles = useChatMessageStyles(role);
return (
<Card title={role} {...styles} extra={<ConnectedMarkdownModeSelect />}>
{content != null && !Array.isArray(content) && (
<ConnectedMarkdownBlock>{content}</ConnectedMarkdownBlock>
)}
{toolCalls && toolCalls.length > 0
? toolCalls.map((toolCall) => {
return (
<View
key={`tool-call-${getToolCallKey(toolCall)}`}
paddingX="size-200"
paddingY="size-200"
borderTopWidth="thin"
borderTopColor="blue-500"
>
<PlaygroundToolCall
key={getToolCallKey(toolCall)}
toolCall={toolCall}
/>
</View>
);
})
: null}
</Card>
);
}
function PlaygroundOutputContent({
output,
partialToolCalls,
}: {
output: PlaygroundRepetition["output"];
partialToolCalls: readonly PartialOutputToolCall[];
}) {
if (isChatMessages(output)) {
return output.map((message, index) => {
return <PlaygroundOutputMessage key={index} message={message} />;
});
}
if (typeof output === "string" || partialToolCalls.length > 0) {
return (
<PlaygroundOutputMessage
message={{
id: generateMessageId(),
content: output ?? undefined,
role: "ai",
toolCalls: partialToolCalls,
}}
/>
);
}
return "click run to see output";
}
export function PlaygroundOutput(props: PlaygroundOutputProps) {
const instanceId = props.playgroundInstanceId;
const instances = usePlaygroundContext((state) => state.instances);
const instance = instances.find((instance) => instance.id === instanceId);
if (!instance) {
throw new Error(`No instance found for id ${instanceId}`);
}
if (instance.template.__type !== "chat") {
throw new Error("We only support chat templates for now");
}
const streaming = usePlaygroundContext((state) => state.streaming);
const credentials = useCredentialsContext((state) => state);
const index = usePlaygroundContext((state) =>
state.instances.findIndex((instance) => instance.id === instanceId)
);
const {
appendRepetitionOutput,
setSelectedRepetitionNumber,
setRepetitionSpanId,
setRepetitionError,
setRepetitionStatus,
setRepetitionToolCalls,
addRepetitionPartialToolCall,
clearRepetitions,
markPlaygroundInstanceComplete,
} = usePlaygroundContext((state) => ({
appendRepetitionOutput: state.appendRepetitionOutput,
setSelectedRepetitionNumber: state.setSelectedRepetitionNumber,
setRepetitionSpanId: state.setRepetitionSpanId,
setRepetitionError: state.setRepetitionError,
setRepetitionStatus: state.setRepetitionStatus,
setRepetitionToolCalls: state.setRepetitionToolCalls,
addRepetitionPartialToolCall: state.addRepetitionPartialToolCall,
clearRepetitions: state.clearRepetitions,
markPlaygroundInstanceComplete: state.markPlaygroundInstanceComplete,
}));
const environment = useRelayEnvironment();
const playgroundStore = usePlaygroundStore();
const numInstanceRepetitions = Object.keys(instance.repetitions).length;
const numRepetitionErrors = Object.values(instance.repetitions).filter(
(output) => output?.error != null
).length;
const selectedRepetitionNumber = instance.selectedRepetitionNumber;
const selectedRepetition = instance.repetitions[selectedRepetitionNumber];
const selectedRepetitionError = selectedRepetition?.error ?? null;
const selectedRepetitionToolCalls = Object.values(
selectedRepetition?.toolCalls ?? {}
);
const selectedRepetitionSpanId = selectedRepetition?.spanId;
const selectedRepetitionSuccessfullyCompleted =
selectedRepetition?.status === "finished" &&
selectedRepetition?.error == null;
const [generateChatCompletion] = useMutation<PlaygroundOutputMutationType>(
PlaygroundOutputMutation
);
const runInProgress = instances.some(
(instance) => instance.activeRunId != null
);
const notifyErrorToast = useNotifyError();
const notifyError = useCallback(
({ title, message, ...rest }: Parameters<typeof notifyErrorToast>[0]) => {
notifyErrorToast({
title,
message,
...rest,
});
},
[notifyErrorToast]
);
const handleChatCompletionSubscriptionPayload = useCallback(
({ chatCompletion }: PlaygroundOutputSubscription$data) => {
if (chatCompletion.__typename === "TextChunk") {
const content = chatCompletion.content;
if (content == null || chatCompletion.repetitionNumber == null) {
return;
}
setRepetitionStatus(
instanceId,
chatCompletion.repetitionNumber,
"streamInProgress"
);
appendRepetitionOutput(
instanceId,
chatCompletion.repetitionNumber,
content
);
return;
} else if (chatCompletion.__typename === "ToolCallChunk") {
const chatCompletionId = chatCompletion.id;
const chatCompletionFunction = chatCompletion.function;
if (
chatCompletionFunction == null ||
chatCompletionId == null ||
chatCompletion.repetitionNumber == null
) {
return;
}
setRepetitionStatus(
instanceId,
chatCompletion.repetitionNumber,
"streamInProgress"
);
addRepetitionPartialToolCall(
instanceId,
chatCompletion.repetitionNumber,
{
id: chatCompletionId,
function: {
name: chatCompletionFunction.name,
arguments: chatCompletionFunction.arguments,
},
}
);
}
if (chatCompletion.__typename === "ChatCompletionSubscriptionResult") {
if (chatCompletion.repetitionNumber == null) {
return;
}
setRepetitionStatus(
instanceId,
chatCompletion.repetitionNumber,
"finished"
);
if (chatCompletion.span != null) {
setRepetitionSpanId(
instanceId,
chatCompletion.repetitionNumber,
chatCompletion.span.id
);
}
return;
}
if (chatCompletion.__typename === "ChatCompletionSubscriptionError") {
if (chatCompletion.repetitionNumber == null) {
return;
}
setRepetitionStatus(
instanceId,
chatCompletion.repetitionNumber,
"finished"
);
setRepetitionError(instanceId, chatCompletion.repetitionNumber, {
title: "Chat completion failed",
message: chatCompletion.message,
});
}
},
[
addRepetitionPartialToolCall,
instanceId,
appendRepetitionOutput,
setRepetitionSpanId,
setRepetitionStatus,
setRepetitionError,
]
);
const handleChatCompletionMutationPayload = useCallback(
(
response: PlaygroundOutputMutation$data,
errors: PayloadError[] | null
) => {
markPlaygroundInstanceComplete(props.playgroundInstanceId);
if (errors != null && errors.length > 0) {
notifyError({
title: "Chat completion failed",
message: errors[0].message,
});
return;
}
const instance = playgroundStore
.getState()
.instances.find((inst) => inst.id === instanceId);
if (instance == null) {
return;
}
response.chatCompletion.repetitions.forEach((repetition) => {
const repetitionNumber = repetition.repetitionNumber;
setRepetitionStatus(instanceId, repetitionNumber, "finished");
if (repetition.content != null) {
appendRepetitionOutput(
instanceId,
repetitionNumber,
repetition.content
);
}
if (repetition.toolCalls.length > 0) {
setRepetitionToolCalls(instanceId, repetitionNumber, [
...repetition.toolCalls,
]);
}
if (repetition.span != null) {
setRepetitionSpanId(instanceId, repetitionNumber, repetition.span.id);
}
if (repetition.errorMessage != null) {
setRepetitionError(instanceId, repetitionNumber, {
title: "Chat completion failed",
message: repetition.errorMessage,
});
}
});
},
[
instanceId,
markPlaygroundInstanceComplete,
notifyError,
setRepetitionToolCalls,
playgroundStore,
props.playgroundInstanceId,
appendRepetitionOutput,
setRepetitionSpanId,
setRepetitionStatus,
setRepetitionError,
]
);
useEffect(() => {
if (!runInProgress) {
return;
}
const input = getChatCompletionInput({
playgroundStore,
instanceId,
credentials,
});
if (streaming) {
const config: GraphQLSubscriptionConfig<PlaygroundOutputSubscriptionType> =
{
subscription: PlaygroundOutputSubscription,
variables: {
input,
},
onNext: (response) => {
if (response) {
handleChatCompletionSubscriptionPayload(response);
}
},
onCompleted: () => {
markPlaygroundInstanceComplete(props.playgroundInstanceId);
},
onError: (error) => {
markPlaygroundInstanceComplete(props.playgroundInstanceId);
const instance = playgroundStore
.getState()
.instances.find((inst) => inst.id === instanceId);
if (instance != null) {
Object.keys(instance.repetitions).forEach((repetitionNumber) => {
setRepetitionStatus(
instanceId,
parseInt(repetitionNumber),
"finished"
);
});
}
const errorMessages =
getErrorMessagesFromRelaySubscriptionError(error);
if (errorMessages != null && errorMessages.length > 0) {
notifyError({
title: "Failed to get output",
message: errorMessages.join("\n"),
});
} else {
notifyError({
title: "Failed to get output",
message: error.message,
});
}
},
};
const subscription = requestSubscription(environment, config);
return subscription.dispose;
}
const disposable = generateChatCompletion({
variables: {
input,
},
onCompleted: handleChatCompletionMutationPayload,
onError(error) {
markPlaygroundInstanceComplete(props.playgroundInstanceId);
clearRepetitions(instanceId);
const errorMessages = getErrorMessagesFromRelayMutationError(error);
if (errorMessages != null && errorMessages.length > 0) {
notifyError({
title: "Failed to get output",
message: errorMessages.join("\n"),
});
} else {
notifyError({
title: "Failed to get output",
message: error.message,
});
}
},
});
return disposable.dispose;
}, [
credentials,
environment,
generateChatCompletion,
instanceId,
clearRepetitions,
markPlaygroundInstanceComplete,
notifyError,
handleChatCompletionMutationPayload,
handleChatCompletionSubscriptionPayload,
runInProgress,
playgroundStore,
props.playgroundInstanceId,
setRepetitionStatus,
streaming,
]);
return (
<Card
title={<TitleWithAlphabeticIndex index={index} title="Output" />}
extra={
<Flex direction="row" gap="size-150" alignItems="center">
{numInstanceRepetitions > 1 && numRepetitionErrors > 0 && (
<TooltipTrigger>
<TriggerWrap>
<Icon svg={<Icons.AlertTriangleOutline />} color="danger" />
</TriggerWrap>
<Tooltip>
<Text>{`${numRepetitionErrors} repetition ${numRepetitionErrors > 1 ? "s" : ""} failed`}</Text>
</Tooltip>
</TooltipTrigger>
)}
{numInstanceRepetitions > 1 && (
<ExperimentRepetitionSelector
repetitionNumber={selectedRepetitionNumber}
totalRepetitions={numInstanceRepetitions}
setRepetitionNumber={(n) => {
let repetitionNumber: number;
if (typeof n === "function") {
repetitionNumber = n(selectedRepetitionNumber);
} else {
repetitionNumber = n;
}
setSelectedRepetitionNumber(instanceId, repetitionNumber);
}}
/>
)}
<PlaygroundOutputMoveButton
isDisabled={!selectedRepetitionSuccessfullyCompleted}
output={selectedRepetition?.output}
toolCalls={selectedRepetitionToolCalls}
instance={instance}
cleanupOutput={() => {
clearRepetitions(instanceId);
}}
/>
</Flex>
}
collapsible
>
{(() => {
switch (true) {
case selectedRepetition?.status === "pending":
return (
<View padding="size-200">
<Loading message="Running..." />
</View>
);
case selectedRepetitionError != null:
return (
<View padding="size-200">
<PlaygroundErrorWrap>
{selectedRepetitionError.message}
</PlaygroundErrorWrap>
</View>
);
default:
return (
<>
<View padding="size-200">
<MarkdownDisplayProvider>
<PlaygroundOutputContent
output={selectedRepetition?.output ?? null}
partialToolCalls={selectedRepetitionToolCalls}
/>
</MarkdownDisplayProvider>
</View>
<Suspense>
{selectedRepetitionSpanId ? (
<RunMetadataFooter spanId={selectedRepetitionSpanId} />
) : null}
</Suspense>
</>
);
}
})()}
</Card>
);
}
// eslint-disable-next-line @typescript-eslint/no-unused-expressions
graphql`
subscription PlaygroundOutputSubscription($input: ChatCompletionInput!) {
chatCompletion(input: $input) {
__typename
repetitionNumber
... on TextChunk {
content
}
... on ToolCallChunk {
id
function {
name
arguments
}
}
... on ChatCompletionSubscriptionResult {
span {
id
}
}
... on ChatCompletionSubscriptionError {
message
}
}
}
`;
// eslint-disable-next-line @typescript-eslint/no-unused-expressions
graphql`
mutation PlaygroundOutputMutation($input: ChatCompletionInput!) {
chatCompletion(input: $input) {
__typename
repetitions {
repetitionNumber
content
errorMessage
span {
id
}
toolCalls {
id
function {
name
arguments
}
}
}
}
}
`;