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>
Arize-ai · Oct 1, 2024 · a497540 · a497540
1 parent 6fc2bcf
commit a497540
Show file tree

Hide file tree

Showing 21 changed files with 544 additions and 73 deletions.
diff --git a/app/package.json b/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",
@@ -57,7 +58,7 @@
     "@types/jest": "^29.5.12",
     "@types/lodash": "^4.17.7",
     "@types/node": "^22.5.4",
-    "@types/react": "18.2.48",
+    "@types/react": "^18.3.10",
     "@types/react-dom": "^18.3.0",
     "@types/react-relay": "^16.0.6",
     "@types/recharts": "^1.8.29",

diff --git a/app/pnpm-lock.yaml b/app/pnpm-lock.yaml
diff --git a/app/schema.graphql b/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!

diff --git a/app/src/RelayEnvironment.ts b/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

diff --git a/app/src/components/auth/OneTimeAPIKeyDialog.tsx b/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">

diff --git a/app/src/pages/dataset/DatasetCodeDropdown.tsx b/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>