Skip to content

Commit

Permalink
feat(playground): add tools ui (#5002)
Browse files Browse the repository at this point in the history
* feat(playground): add tools ui

* add tools back after rebase

* update tool type to be partial in store for editing

* make tool editor uncontrolled, add tool choice selector

* make fields pass through

* allow for extra keys in the json schema

* support json schema in jsonEditor, remove jsonToolEditor

* update descriptions

* fix types

* fix key collision on choice picker

* update comment

* more comments

* WIP

* styling

* cleanup

---------

Co-authored-by: Mikyo King <mikyo@arize.com>
  • Loading branch information
Parker-Stafford and mikeldking authored Oct 15, 2024
1 parent 415ec6f commit 767bd37
Show file tree
Hide file tree
Showing 17 changed files with 981 additions and 42 deletions.
5 changes: 4 additions & 1 deletion app/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
"license": "None",
"private": true,
"dependencies": {
"@arizeai/components": "^1.8.6",
"@arizeai/components": "^1.8.7",
"@arizeai/openinference-semantic-conventions": "^0.10.0",
"@arizeai/point-cloud": "^3.0.6",
"@codemirror/autocomplete": "6.12.0",
Expand All @@ -27,6 +27,7 @@
"@uiw/codemirror-theme-github": "^4.23.5",
"@uiw/codemirror-theme-nord": "^4.23.0",
"@uiw/react-codemirror": "^4.23.0",
"codemirror-json-schema": "^0.7.8",
"copy-to-clipboard": "^3.3.3",
"d3-format": "^3.1.0",
"d3-scale-chromatic": "^3.1.0",
Expand Down Expand Up @@ -55,6 +56,7 @@
"use-deep-compare-effect": "^1.8.1",
"use-zustand": "^0.0.4",
"zod": "^3.23.8",
"zod-to-json-schema": "^3.23.3",
"zustand": "^4.5.4"
},
"devDependencies": {
Expand All @@ -64,6 +66,7 @@
"@types/d3-format": "^3.0.4",
"@types/d3-scale-chromatic": "^3.0.3",
"@types/d3-time-format": "^4.0.3",
"@types/json-schema": "^7.0.15",
"@types/lodash": "^4.17.7",
"@types/node": "^22.5.4",
"@types/react": "18.3.10",
Expand Down
462 changes: 457 additions & 5 deletions app/pnpm-lock.yaml

Large diffs are not rendered by default.

11 changes: 11 additions & 0 deletions app/src/@types/generative.d.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,3 +4,14 @@ declare type ModelProvider = "OPENAI" | "AZURE_OPENAI" | "ANTHROPIC";
* The role of a chat message
*/
declare type ChatMessageRole = "user" | "system" | "ai" | "tool";

/**
* The tool picking mechanism for an LLM
* Either "auto", "required", "none", or a specific tool
* @see https://platform.openai.com/docs/api-reference/chat/create#chat-create-tool_choice
*/
declare type ToolChoice =
| "auto"
| "required"
| "none"
| { type: "function"; function: { name: string } };
1 change: 1 addition & 0 deletions app/src/components/code/CodeEditorFieldWrapper.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ const codeEditorFormWrapperCSS = css`
.cm-editor {
border-radius: var(--ac-global-rounding-small);
}
box-sizing: border-box;
.cm-focused {
outline: none;
}
Expand Down
44 changes: 38 additions & 6 deletions app/src/components/code/JSONEditor.tsx
Original file line number Diff line number Diff line change
@@ -1,24 +1,56 @@
import React from "react";
import { json, jsonParseLinter } from "@codemirror/lang-json";
import React, { useMemo } from "react";
import { json, jsonLanguage, jsonParseLinter } from "@codemirror/lang-json";
import { linter } from "@codemirror/lint";
import { EditorView } from "@codemirror/view";
import { EditorView, hoverTooltip } from "@codemirror/view";
import { githubLight } from "@uiw/codemirror-theme-github";
import { nord } from "@uiw/codemirror-theme-nord";
import CodeMirror, { ReactCodeMirrorProps } from "@uiw/react-codemirror";
import {
handleRefresh,
jsonCompletion,
jsonSchemaHover,
jsonSchemaLinter,
stateExtensions,
} from "codemirror-json-schema";
import { JSONSchema7 } from "json-schema";

import { useTheme } from "@phoenix/contexts";

export type JSONEditorProps = Omit<
ReactCodeMirrorProps,
"theme" | "extensions" | "editable"
>;
> & {
/**
* JSON Schema to use for validation, if provided will enable JSON Schema validation with tooltips in the editor
*/
jsonSchema?: JSONSchema7;
};

export function JSONEditor(props: JSONEditorProps) {
const { theme } = useTheme();
const codeMirrorTheme = theme === "light" ? undefined : nord;
const codeMirrorTheme = theme === "light" ? githubLight : nord;
const extensions = useMemo(() => {
const baseExtensions = [
json(),
EditorView.lineWrapping,
linter(jsonParseLinter()),
];
if (props.jsonSchema) {
baseExtensions.push(
linter(jsonSchemaLinter(), { needsRefresh: handleRefresh }),
jsonLanguage.data.of({
autocomplete: jsonCompletion(),
}),
hoverTooltip(jsonSchemaHover()),
stateExtensions(props.jsonSchema)
);
}
return baseExtensions;
}, [props.jsonSchema]);
return (
<CodeMirror
value={props.value}
extensions={[json(), EditorView.lineWrapping, linter(jsonParseLinter())]}
extensions={extensions}
editable
theme={codeMirrorTheme}
{...props}
Expand Down
102 changes: 102 additions & 0 deletions app/src/components/generative/ToolChoiceSelector.tsx
Original file line number Diff line number Diff line change
@@ -0,0 +1,102 @@
import React from "react";

import { Flex, Item, Label, Picker } from "@arizeai/components";

type DefaultToolChoice = Extract<ToolChoice, "auto" | "required" | "none">;

const isDefaultToolChoice = (choice: string): choice is DefaultToolChoice => {
return choice === "auto" || choice === "required" || choice === "none";
};

/**
* A prefix to add to user defined tools in the picker to avoid picker key collisions with default {@link ToolChoice} keys
*/
const TOOL_NAME_PREFIX = "tool_";

/**
* Adds a prefix to user defined tool names to avoid conflicts picker key collisions with default {@link ToolChoice} keys
* @param toolName The name of a tool
* @returns The tool name with the "TOOL_NAME_PREFIX" prefix added
*/
const addToolNamePrefix = (toolName: string) =>
`${TOOL_NAME_PREFIX}${toolName}`;

/**
* Removes the "TOOL_NAME_PREFIX" prefix from a tool name so that it can be used as a choice that corresponds to an actual tool
* @param toolName The name of a tool with the "TOOL_NAME_PREFIX" prefix
* @returns The tool name with the "TOOL_NAME_PREFIX" prefix removed
*/
const removeToolNamePrefix = (toolName: string) =>
toolName.startsWith(TOOL_NAME_PREFIX)
? toolName.slice(TOOL_NAME_PREFIX.length)
: toolName;

type ToolChoicePickerProps = {
/**
* The current choice including the default {@link ToolChoice} and any user defined tools
*/
choice: ToolChoice;
/**
* Callback for when the tool choice changes
*/
onChange: (choice: ToolChoice) => void;
/**
* A list of user defined tool names
*/
toolNames: string[];
};

export function ToolChoicePicker({
choice,
onChange,
toolNames,
}: ToolChoicePickerProps) {
const currentKey =
typeof choice === "string"
? choice
: addToolNamePrefix(choice.function.name);
return (
<Picker
selectedKey={currentKey}
label="Tool Choice"
aria-label="Tool Choice for an LLM"
onSelectionChange={(choice) => {
if (typeof choice !== "string") {
return;
}
if (choice.startsWith(TOOL_NAME_PREFIX)) {
onChange({
type: "function",
function: {
name: removeToolNamePrefix(choice),
},
});
} else if (isDefaultToolChoice(choice)) {
onChange(choice);
}
}}
>
{[
<Item key="auto">
<Flex gap={"size-100"}>
Tools auto-selected by LLM <Label>auto</Label>
</Flex>
</Item>,
<Item key="required">
<Flex gap={"size-100"}>
Use at least one tool <Label>required</Label>
</Flex>
</Item>,
<Item key="none">
<Flex gap={"size-100"}>
Don&apos;t use any tools <Label>none</Label>
</Flex>
</Item>,
// Add "TOOL_NAME_PREFIX" prefix to user defined tool names to avoid conflicts with default keys
...toolNames.map((toolName) => (
<Item key={addToolNamePrefix(toolName)}>{toolName}</Item>
)),
]}
</Picker>
);
}
1 change: 1 addition & 0 deletions app/src/components/generative/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
export * from "./ToolChoiceSelector";
23 changes: 8 additions & 15 deletions app/src/pages/playground/Playground.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -87,27 +87,25 @@ const playgroundPromptPanelContentCSS = css`
flex-direction: column;
height: 100%;
overflow: hidden;
.ac-accordion {
& > .ac-accordion {
display: flex;
flex-direction: column;
height: 100%;
overflow: hidden;
flex: 1 1 auto;
.ac-accordion-item {
& > .ac-accordion-item {
height: 100%;
overflow: hidden;
flex: 1 1 auto;
.ac-accordion-itemContent {
height: 100%;
overflow: hidden;
flex: 1 1 auto;
& > * {
& > .ac-view {
height: 100%;
flex: 1 1 auto;
overflow: auto;
box-sizing: border-box;
// Fix padding issue with flexbox
padding-bottom: 57px !important;
}
}
}
Expand Down Expand Up @@ -147,20 +145,15 @@ function PlaygroundContent() {
</Flex>
}
>
<View height="100%" padding="size-200">
<View height="100%" padding="size-200" paddingBottom="size-900">
<Flex direction="row" gap="size-200">
{instances.map((instance, i) => (
<div
key={i}
css={css`
flex: 1 1 0px;
`}
>
{instances.map((instance) => (
<View key={instance.id} flex="1 1 0px">
<PlaygroundTemplate
key={i}
key={instance.id}
playgroundInstanceId={instance.id}
/>
</div>
</View>
))}
</Flex>
</View>
Expand Down
28 changes: 26 additions & 2 deletions app/src/pages/playground/PlaygroundChatTemplate.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,13 @@ import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
import { useChatMessageStyles } from "@phoenix/hooks/useChatMessageStyles";
import {
ChatMessage,
createTool,
generateMessageId,
PlaygroundChatTemplate as PlaygroundChatTemplateType,
} from "@phoenix/store";

import { MessageRolePicker } from "./MessageRolePicker";
import { PlaygroundTools } from "./PlaygroundTools";
import { PlaygroundInstanceProps } from "./types";

const MESSAGE_Z_INDEX = 1;
Expand All @@ -54,6 +56,7 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
if (!playgroundInstance) {
throw new Error(`Playground instance ${id} not found`);
}
const hasTools = playgroundInstance.tools.length > 0;
const { template } = playgroundInstance;
if (template.__type !== "chat") {
throw new Error(`Invalid template type ${template.__type}`);
Expand Down Expand Up @@ -123,10 +126,30 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
paddingEnd="size-200"
paddingTop="size-100"
paddingBottom="size-100"
borderTopColor="dark"
borderColor="dark"
borderTopWidth="thin"
borderBottomWidth={hasTools ? "thin" : undefined}
>
<Flex direction="row" justifyContent="end">
<Flex direction="row" justifyContent="end" gap="size-100">
<Button
variant="default"
aria-label="add tool"
size="compact"
icon={<Icon svg={<Icons.PlusOutline />} />}
onClick={() => {
updateInstance({
instanceId: id,
patch: {
tools: [
...playgroundInstance.tools,
createTool(playgroundInstance.tools.length + 1),
],
},
});
}}
>
Tool
</Button>
<Button
variant="default"
aria-label="add message"
Expand Down Expand Up @@ -155,6 +178,7 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
</Button>
</Flex>
</View>
{hasTools ? <PlaygroundTools {...props} /> : null}
</DndContext>
);
}
Expand Down
2 changes: 1 addition & 1 deletion app/src/pages/playground/PlaygroundOutput.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ export function PlaygroundOutput(props: PlaygroundOutputProps) {
state.instances.findIndex((instance) => instance.id === instanceId)
);
if (!instance) {
throw new Error("Playground instance not found");
throw new Error(`Playground instance ${instanceId} not found`);
}

const runId = instance.activeRunId;
Expand Down
Loading

0 comments on commit 767bd37

Please sign in to comment.