feat(playground): streaming chat completions (#4785)

* add subscriptions

* configure relay

* playground mvp

* align the runs for multiple

* add more functionality

* cleanup the components

---------

Co-authored-by: Mikyo King <mikyo@arize.com>

app/package.json

@@ -6,7 +6,7 @@
   "license": "None",
   "private": true,
   "dependencies": {
-    "@arizeai/components": "^1.8.1",
+    "@arizeai/components": "^1.8.3",
     "@arizeai/openinference-semantic-conventions": "^0.10.0",
     "@arizeai/point-cloud": "^3.0.6",
     "@codemirror/autocomplete": "6.12.0",
@@ -25,6 +25,7 @@
     "d3-scale-chromatic": "^3.1.0",
     "d3-time-format": "^4.1.0",
     "date-fns": "^3.6.0",
+    "graphql-ws": "^5.16.0",
     "lodash": "^4.17.21",
     "normalize.css": "^8.0.1",
     "polished": "^4.3.1",

app/schema.graphql

@@ -65,6 +65,10 @@ enum AuthMethod {
 
 union Bin = NominalBin | IntervalBin | MissingValueBin
 
+input ChatCompletionInput {
+  message: String!
+}
+
 input ClearProjectInput {
   id: GlobalID!
 
@@ -1372,6 +1376,10 @@ enum SpanStatusCode {
   UNSET
 }
 
+type Subscription {
+  chatCompletion(input: ChatCompletionInput!): String!
+}
+
 type SystemApiKey implements ApiKey & Node {
   """Name of the API key."""
   name: String!

app/src/RelayEnvironment.ts

@@ -1,9 +1,15 @@
+import { createClient, Sink } from "graphql-ws";
 import {
   Environment,
   FetchFunction,
+  GraphQLResponse,
   Network,
+  Observable,
   RecordSource,
+  RequestParameters,
   Store,
+  SubscribeFunction,
+  Variables,
 } from "relay-runtime";
 
 import { authFetch } from "@phoenix/authFetch";
@@ -50,9 +56,29 @@ const fetchRelay: FetchFunction = async (params, variables, _cacheConfig) => {
   return json;
 };
 
+const wsClient = createClient({
+  url: "ws://localhost:6006/graphql",
+});
+
+const subscribe: SubscribeFunction = (
+  operation: RequestParameters,
+  variables: Variables
+) => {
+  return Observable.create<GraphQLResponse>((sink) => {
+    return wsClient.subscribe(
+      {
+        operationName: operation.name,
+        query: operation.text as string,
+        variables,
+      },
+      sink as Sink
+    );
+  });
+};
+
 // Export a singleton instance of Relay Environment configured with our network layer:
 export default new Environment({
-  network: Network.create(fetchRelay),
+  network: Network.create(fetchRelay, subscribe),
   store: new Store(new RecordSource(), {
     // This property tells Relay to not immediately clear its cache when the user
     // navigates around the app. Relay will hold onto the specified number of

app/src/components/auth/OneTimeAPIKeyDialog.tsx

@@ -37,7 +37,7 @@ export function OneTimeAPIKeyDialog(props: { jwt: string }) {
         <View padding="size-200">
           <Flex direction="row" gap="size-100" alignItems="end">
             <TextField label="API Key" isReadOnly value={jwt} minWidth="100%" />
-            <CopyToClipboardButton text={jwt} size="normal" />
+            <CopyToClipboardButton text={jwt} size="default" />
           </Flex>
         </View>
         <View padding="size-200" borderTopColor="light" borderTopWidth="thin">

app/src/pages/dataset/DatasetCodeDropdown.tsx

@@ -72,7 +72,7 @@ export function DatasetCodeDropdown() {
                   width="100%"
                 >
                   <TextField label="Dataset ID" isReadOnly value={datasetId} />
-                  <CopyToClipboardButton text={datasetId} size="normal" />
+                  <CopyToClipboardButton text={datasetId} size="default" />
                 </Flex>
                 <Flex direction="row" gap="size-100" alignItems="end">
                   <TextField
@@ -84,7 +84,7 @@ export function DatasetCodeDropdown() {
                   <CopyToClipboardButton
                     text={version?.id || "No Versions"}
                     disabled={!version}
-                    size="normal"
+                    size="default"
                   />
                 </Flex>
               </Form>
@@ -105,7 +105,7 @@ export function DatasetCodeDropdown() {
                   <Form>
                     <Flex direction="row" gap="size-100" alignItems="end">
                       <TextField label="URL" isReadOnly value={datasetURL} />
-                      <CopyToClipboardButton text={datasetURL} size="normal" />
+                      <CopyToClipboardButton text={datasetURL} size="default" />
                     </Flex>
                   </Form>
                 </View>

app/src/pages/playground/MessageRolePicker.tsx

@@ -0,0 +1,46 @@
+import React from "react";
+import { css } from "@emotion/react";
+
+import { Item, Picker } from "@arizeai/components";
+
+import { ChatMessageRole } from "@phoenix/store";
+
+const hiddenLabelCSS = css`
+  .ac-field-label {
+    display: none;
+  }
+`;
+
+type MessageRolePickerProps = {
+  /**
+   * The currently selected message role
+   */
+  role: ChatMessageRole;
+  /**
+   * Whether to display a label for the picker
+   * This may be set to false in cases where the picker is rendered in a table for instance
+   * @default true
+   */
+  includeLabel?: boolean;
+};
+
+export function MessageRolePicker({
+  role,
+  includeLabel = true,
+}: MessageRolePickerProps) {
+  return (
+    <Picker
+      selectedKey={role}
+      css={!includeLabel ? hiddenLabelCSS : undefined}
+      label="Role"
+      data-testid="inferences-time-range"
+      aria-label={`Time range for the primary inferences`}
+      size="compact"
+      onSelectionChange={() => {}}
+    >
+      <Item key="system">System</Item>
+      <Item key="user">User</Item>
+      <Item key="assistant">Assistant</Item>
+    </Picker>
+  );
+}