Skip to content

Commit

Permalink
feat(playground): pull prompt template variables from span and captur…
Browse files Browse the repository at this point in the history 
…e on playground spans (#5642)

* feat(playground): pull prompt template variables from span and capture on playground spans

* update test

* remove storybook log

* fix test
  • Loading branch information
@Parker-Stafford
Parker-Stafford authored Dec 9, 2024
1 parent 016ca63 commit d6382dc
Show file tree
Hide file tree
Showing 9 changed files with 295 additions and 6 deletions.
6 changes: 5 additions & 1 deletion app/src/pages/playground/SpanPlaygroundPage.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,18 +20,22 @@ export function SpanPlaygroundPage() {
throw new Error("Span not found");
}

const { playgroundInstance, parsingErrors } = useMemo(
const { playgroundInstance, parsingErrors, playgroundInput } = useMemo(
() => transformSpanAttributesToPlaygroundInstance(span),
[span]
);

const additionalProps =
playgroundInput != null ? { input: playgroundInput } : {};

return (
<Flex direction="column" height={"100%"}>
<SpanPlaygroundBanners span={span} parsingErrors={parsingErrors} />
<Playground
// remount the playground when the span changes, resetting all local state, closing dialogs, etc.
key={span.id}
instances={[playgroundInstance]}
{...additionalProps}
/>
</Flex>
);
Expand Down
192 changes: 192 additions & 0 deletions app/src/pages/playground/__tests__/playgroundUtils.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,27 +10,33 @@ import {
PlaygroundInstance,
} from "@phoenix/store";

import { InvocationParameterInput } from "../__generated__/PlaygroundDatasetExamplesTableSubscription.graphql";
import {
INPUT_MESSAGES_PARSING_ERROR,
MODEL_CONFIG_PARSING_ERROR,
MODEL_CONFIG_WITH_INVOCATION_PARAMETERS_PARSING_ERROR,
MODEL_CONFIG_WITH_RESPONSE_FORMAT_PARSING_ERROR,
OUTPUT_MESSAGES_PARSING_ERROR,
OUTPUT_VALUE_PARSING_ERROR,
PROMPT_TEMPLATE_VARIABLES_PARSING_ERROR,
SPAN_ATTRIBUTES_PARSING_ERROR,
TOOLS_PARSING_ERROR,
} from "../constants";
import { InvocationParameter } from "../InvocationParametersFormFields";
import {
areInvocationParamsEqual,
areRequiredInvocationParametersConfigured,
extractVariablesFromInstances,
getBaseModelConfigFromAttributes,
getChatRole,
getModelInvocationParametersFromAttributes,
getModelProviderFromModelName,
getOutputFromAttributes,
getPromptTemplateVariablesFromAttributes,
getTemplateMessagesFromAttributes,
getToolsFromAttributes,
getVariablesMapFromInstances,
mergeInvocationParametersWithDefaults,
processAttributeToolCalls,
transformSpanAttributesToPlaygroundInstance,
} from "../playgroundUtils";
Expand Down Expand Up @@ -150,6 +156,7 @@ describe("transformSpanAttributesToPlaygroundInstance", () => {
template: defaultTemplate,
output: undefined,
},
playgroundInput: undefined,
parsingErrors: [
INPUT_MESSAGES_PARSING_ERROR,
OUTPUT_MESSAGES_PARSING_ERROR,
Expand Down Expand Up @@ -1242,3 +1249,188 @@ describe("areInvocationParamsEqual", () => {
expect(areInvocationParamsEqual(paramA, paramB)).toBe(false);
});
});
describe("getPromptTemplateVariablesFromAttributes", () => {
it("should return parsing errors if prompt template variables are invalid", () => {
const parsedAttributes = { llm: { prompt_template: "invalid" } };
expect(getPromptTemplateVariablesFromAttributes(parsedAttributes)).toEqual({
variables: null,
parsingErrors: [PROMPT_TEMPLATE_VARIABLES_PARSING_ERROR],
});
});

it("should return parsed variables if prompt template variables are valid", () => {
const parsedAttributes = {
llm: {
prompt_template: {
variables: JSON.stringify({
name: "John",
age: 30,
}),
},
},
};
expect(getPromptTemplateVariablesFromAttributes(parsedAttributes)).toEqual({
variables: {
name: "John",
age: "30",
},
parsingErrors: [],
});
});

it("should return null variables and no parsing errors if prompt template is not present", () => {
const parsedAttributes = { llm: {} };
expect(getPromptTemplateVariablesFromAttributes(parsedAttributes)).toEqual({
variables: null,
parsingErrors: [],
});
});
});

describe("areRequiredInvocationParametersConfigured", () => {
it("should return true if all required parameters are configured", () => {
const configuredInvocationParameters: InvocationParameterInput[] = [
{
invocationName: "max_tokens",
canonicalName: "MAX_COMPLETION_TOKENS",
valueInt: 1,
},
{ invocationName: "seed", canonicalName: "RANDOM_SEED", valueInt: 2 },
];
const supportedInvocationParameters: InvocationParameter[] = [
{
invocationName: "max_tokens",
canonicalName: "MAX_COMPLETION_TOKENS",
required: true,
__typename: "IntInvocationParameter",
},
{
invocationName: "random seed",
canonicalName: "RANDOM_SEED",
required: true,
__typename: "IntInvocationParameter",
},
];
expect(
areRequiredInvocationParametersConfigured(
configuredInvocationParameters,
supportedInvocationParameters
)
).toBe(true);
});

it("should return false if not all required parameters are configured", () => {
const configuredInvocationParameters: InvocationParameterInput[] = [
{ invocationName: "seed", canonicalName: "RANDOM_SEED", valueInt: 2 },
];
const supportedInvocationParameters: InvocationParameter[] = [
{
invocationName: "max_tokens",
canonicalName: "MAX_COMPLETION_TOKENS",
required: true,
__typename: "IntInvocationParameter",
},
{
invocationName: "random seed",
canonicalName: "RANDOM_SEED",
required: true,
__typename: "IntInvocationParameter",
},
];
expect(
areRequiredInvocationParametersConfigured(
configuredInvocationParameters,
supportedInvocationParameters
)
).toBe(false);
});
});

describe("mergeInvocationParametersWithDefaults", () => {
it("should merge invocation parameters with default values", () => {
const invocationParameters: InvocationParameterInput[] = [
{
invocationName: "max_tokens",
canonicalName: "MAX_COMPLETION_TOKENS",
valueInt: 1,
},
];
const supportedInvocationParameters: InvocationParameter[] = [
{
invocationName: "max_tokens",
canonicalName: "MAX_COMPLETION_TOKENS",
required: true,
__typename: "IntInvocationParameter",
intDefaultValue: 5,
invocationInputField: "value_int",
},
{
invocationName: "random seed",
canonicalName: "RANDOM_SEED",
required: true,
intDefaultValue: 1000,
__typename: "IntInvocationParameter",
invocationInputField: "value_int",
},
];
expect(
mergeInvocationParametersWithDefaults(
invocationParameters,
supportedInvocationParameters
)
).toEqual([
{
invocationName: "max_tokens",
canonicalName: "MAX_COMPLETION_TOKENS",
valueInt: 1,
},
{
invocationName: "random seed",
canonicalName: "RANDOM_SEED",
valueInt: 1000,
},
]);
});

it("should not overwrite existing values with defaults", () => {
const invocationParameters: InvocationParameterInput[] = [
{
invocationName: "max_tokens",
canonicalName: "MAX_COMPLETION_TOKENS",
valueInt: 1,
},
{ invocationName: "seed", canonicalName: "RANDOM_SEED", valueInt: 2 },
];
const supportedInvocationParameters: InvocationParameter[] = [
{
invocationName: "max_tokens",
canonicalName: "MAX_COMPLETION_TOKENS",
required: true,
__typename: "IntInvocationParameter",
intDefaultValue: 5,
invocationInputField: "value_int",
},
{
invocationName: "random seed",
canonicalName: "RANDOM_SEED",
required: true,
intDefaultValue: 1000,
__typename: "IntInvocationParameter",
invocationInputField: "value_int",
},
];
expect(
mergeInvocationParametersWithDefaults(
invocationParameters,
supportedInvocationParameters
)
).toEqual([
{
invocationName: "max_tokens",
canonicalName: "MAX_COMPLETION_TOKENS",
valueInt: 1,
},
{ invocationName: "seed", canonicalName: "RANDOM_SEED", valueInt: 2 },
]);
});
});
4 changes: 3 additions & 1 deletion app/src/pages/playground/constants.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,13 @@ export const SPAN_ATTRIBUTES_PARSING_ERROR =
export const MODEL_CONFIG_PARSING_ERROR =
"Unable to parse model config, expected llm.model_name to be present.";
export const MODEL_CONFIG_WITH_INVOCATION_PARAMETERS_PARSING_ERROR =
"Unable to parse model config, expected llm.invocation_parameters json string to be present.";
"Unable to parse model config, expected llm.invocation_parameters JSON string to be present.";
export const MODEL_CONFIG_WITH_RESPONSE_FORMAT_PARSING_ERROR =
"Unable to parse invocation parameters response_format, expected llm.invocation_parameters.response_format to be a well formed json object or undefined.";
export const TOOLS_PARSING_ERROR =
"Unable to parse tools, expected tools to be an array of valid tools.";
export const PROMPT_TEMPLATE_VARIABLES_PARSING_ERROR =
"Unable to parse prompt template variables, expected prompt template variables to be a valid JSON object string.";

export const modelProviderToModelPrefixMap: Record<ModelProvider, string[]> = {
AZURE_OPENAI: [],
Expand Down
36 changes: 35 additions & 1 deletion app/src/pages/playground/playgroundUtils.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ import {
modelProviderToModelPrefixMap,
OUTPUT_MESSAGES_PARSING_ERROR,
OUTPUT_VALUE_PARSING_ERROR,
PROMPT_TEMPLATE_VARIABLES_PARSING_ERROR,
RESPONSE_FORMAT_PARAM_CANONICAL_NAME,
RESPONSE_FORMAT_PARAM_NAME,
SPAN_ATTRIBUTES_PARSING_ERROR,
Expand All @@ -75,6 +76,7 @@ import {
modelConfigWithInvocationParametersSchema,
modelConfigWithResponseFormatSchema,
outputSchema,
promptTemplateSchema,
} from "./schemas";
import { PlaygroundSpan } from "./spanPlaygroundPageLoader";

Expand Down Expand Up @@ -469,6 +471,32 @@ export function getToolsFromAttributes(
return { tools: processAttributeTools(data), parsingErrors: [] };
}

export function getPromptTemplateVariablesFromAttributes(
parsedAttributes: unknown
):
| { variables: Record<string, string | undefined>; parsingErrors: never[] }
| { variables: null; parsingErrors: string[] } {
const { success, data } = promptTemplateSchema.safeParse(parsedAttributes);
if (!success) {
return {
variables: null,
parsingErrors: [PROMPT_TEMPLATE_VARIABLES_PARSING_ERROR],
};
}

// If there is no template or llm attributes, we don't want to return parsing errors, it just means the span didn't have a prompt template
if (data?.llm?.prompt_template == null) {
return {
variables: null,
parsingErrors: [],
};
}
return {
variables: data.llm.prompt_template.variables,
parsingErrors: [],
};
}

/**
* Takes a {@link PlaygroundSpan|Span} and attempts to transform it's attributes into various fields on a {@link PlaygroundInstance}.
* @param span the {@link PlaygroundSpan|Span} to transform into a playground instance
Expand All @@ -485,6 +513,7 @@ export function transformSpanAttributesToPlaygroundInstance(
* This field is used to store any issues encountered when parsing to display in the playground.
*/
parsingErrors: string[];
playgroundInput?: PlaygroundInput;
} {
const basePlaygroundInstance = createPlaygroundInstance();
const { json: parsedAttributes, parseError } = safelyParseJSON(
Expand Down Expand Up @@ -523,6 +552,8 @@ export function transformSpanAttributesToPlaygroundInstance(
parsedAttributes,
modelSupportedInvocationParameters
);
const { variables, parsingErrors: promptTemplateVariablesParsingErrors } =
getPromptTemplateVariablesFromAttributes(parsedAttributes);
// parse response format separately so that we can get distinct errors messages from the rest of
// the invocation parameters
const { parsingErrors: responseFormatParsingErrors } =
Expand Down Expand Up @@ -565,13 +596,16 @@ export function transformSpanAttributesToPlaygroundInstance(
spanId: span.id,
tools: tools ?? basePlaygroundInstance.tools,
},
playgroundInput:
variables != null ? { variablesValueCache: variables } : undefined,
parsingErrors: [
...messageParsingErrors,
...outputParsingErrors,
...modelConfigParsingErrors,
...toolsParsingErrors,
...invocationParametersParsingErrors,
...responseFormatParsingErrors,
...promptTemplateVariablesParsingErrors,
],
};
}
Expand Down Expand Up @@ -674,7 +708,7 @@ export const getVariablesMapFromInstances = ({
acc[key] = variableValueCache[key] || "";
return acc;
},
{} as Record<string, string>
{} as NonNullable<PlaygroundInput["variablesValueCache"]>
);
return { variablesMap, variableKeys };
};
Expand Down
Loading

0 comments on commit d6382dc

Please sign in to comment.