From a4975408797f501144734213bcf7fee9f5d41d03 Mon Sep 17 00:00:00 2001
From: Xander Song <axiomofjoy@gmail.com>
Date: Tue, 1 Oct 2024 16:42:05 -0700
Subject: [PATCH] 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                              |   5 +-
 app/pnpm-lock.yaml                            |  83 +++++++-------
 app/schema.graphql                            |   8 ++
 app/src/RelayEnvironment.ts                   |  28 ++++-
 .../components/auth/OneTimeAPIKeyDialog.tsx   |   2 +-
 app/src/pages/dataset/DatasetCodeDropdown.tsx |   6 +-
 .../pages/playground/MessageRolePicker.tsx    |  46 ++++++++
 .../playground/PlaygroundChatTemplate.tsx     |  20 +++-
 .../PlaygroundInputModeRadioGroup.tsx         |   2 +-
 .../pages/playground/PlaygroundInstance.tsx   |   2 +-
 app/src/pages/playground/PlaygroundOutput.tsx | 101 +++++++++++++++++-
 .../pages/playground/PlaygroundTemplate.tsx   |  71 +++++++++---
 .../PlaygroundOutputSubscription.graphql.ts   |  82 ++++++++++++++
 app/src/pages/settings/SettingsPage.tsx       |   2 +-
 app/src/pages/trace/SpanCodeDropdown.tsx      |   4 +-
 app/src/pages/trace/SpanDetails.tsx           |   3 -
 app/src/store/playgroundStore.tsx             |  71 +++++++++++-
 pyproject.toml                                |   1 +
 src/phoenix/server/api/schema.py              |   2 +
 src/phoenix/server/api/subscriptions.py       |  76 +++++++++++++
 src/phoenix/server/app.py                     |   2 +
 21 files changed, 544 insertions(+), 73 deletions(-)
 create mode 100644 app/src/pages/playground/MessageRolePicker.tsx
 create mode 100644 app/src/pages/playground/__generated__/PlaygroundOutputSubscription.graphql.ts
 create mode 100644 src/phoenix/server/api/subscriptions.py

diff --git a/app/package.json b/app/package.json
index cd16447d6a9..03095faa6c6 100644
--- 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
index f3a0ca56610..bab64f61191 100644
--- a/app/pnpm-lock.yaml
+++ b/app/pnpm-lock.yaml
@@ -12,14 +12,14 @@ importers:
   .:
     dependencies:
       '@arizeai/components':
-        specifier: ^1.8.1
-        version: 1.8.1(@types/react@18.2.48)(eslint@8.57.0)(react-dom@18.2.0(react@18.2.0))(react@18.2.0)
+        specifier: ^1.8.3
+        version: 1.8.3(@types/react@18.3.10)(eslint@8.57.0)(react-dom@18.2.0(react@18.2.0))(react@18.2.0)
       '@arizeai/openinference-semantic-conventions':
         specifier: ^0.10.0
         version: 0.10.0
       '@arizeai/point-cloud':
         specifier: ^3.0.6
-        version: 3.0.6(@react-three/drei@9.108.4(@react-three/fiber@8.0.12(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(@types/react@18.2.48)(@types/three@0.149.0)(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(@react-three/fiber@8.0.12(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(react@18.2.0)(three-stdlib@2.30.4(three@0.139.2))(three@0.139.2)
+        version: 3.0.6(@react-three/drei@9.108.4(@react-three/fiber@8.0.12(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(@types/react@18.3.10)(@types/three@0.149.0)(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(@react-three/fiber@8.0.12(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(react@18.2.0)(three-stdlib@2.30.4(three@0.139.2))(three@0.139.2)
       '@codemirror/autocomplete':
         specifier: 6.12.0
         version: 6.12.0(@codemirror/language@6.10.2)(@codemirror/state@6.4.1)(@codemirror/view@6.28.5)(@lezer/common@1.2.1)
@@ -40,7 +40,7 @@ importers:
         version: 6.28.5
       '@react-three/drei':
         specifier: ^9.108.4
-        version: 9.108.4(@react-three/fiber@8.0.12(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(@types/react@18.2.48)(@types/three@0.149.0)(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2)
+        version: 9.108.4(@react-three/fiber@8.0.12(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(@types/react@18.3.10)(@types/three@0.149.0)(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2)
       '@react-three/fiber':
         specifier: 8.0.12
         version: 8.0.12(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2)
@@ -68,6 +68,9 @@ importers:
       date-fns:
         specifier: ^3.6.0
         version: 3.6.0
+      graphql-ws:
+        specifier: ^5.16.0
+        version: 5.16.0(graphql@16.9.0)
       lodash:
         specifier: ^4.17.21
         version: 4.17.21
@@ -94,7 +97,7 @@ importers:
         version: 4.5.0(react-dom@18.2.0(react@18.2.0))(react@18.2.0)
       react-markdown:
         specifier: ^9.0.1
-        version: 9.0.1(@types/react@18.2.48)(react@18.2.0)
+        version: 9.0.1(@types/react@18.3.10)(react@18.2.0)
       react-relay:
         specifier: ^16.2.0
         version: 16.2.0(react@18.2.0)
@@ -133,11 +136,11 @@ importers:
         version: 0.0.4(react@18.2.0)
       zustand:
         specifier: ^4.5.4
-        version: 4.5.4(@types/react@18.2.48)(react@18.2.0)
+        version: 4.5.4(@types/react@18.3.10)(react@18.2.0)
     devDependencies:
       '@emotion/react':
         specifier: ^11.11.4
-        version: 11.11.4(@types/react@18.2.48)(react@18.2.0)
+        version: 11.11.4(@types/react@18.3.10)(react@18.2.0)
       '@playwright/test':
         specifier: ^1.47.0
         version: 1.47.0
@@ -160,8 +163,8 @@ importers:
         specifier: ^22.5.4
         version: 22.5.4
       '@types/react':
-        specifier: 18.2.48
-        version: 18.2.48
+        specifier: ^18.3.10
+        version: 18.3.10
       '@types/react-dom':
         specifier: ^18.3.0
         version: 18.3.0
@@ -244,8 +247,8 @@ packages:
     resolution: {integrity: sha512-30iZtAPgz+LTIYoeivqYo853f02jBYSd5uGnGpkFV0M3xOt9aN73erkgYAmZU43x4VfqcnLxW9Kpg3R5LC4YYw==}
     engines: {node: '>=6.0.0'}
 
-  '@arizeai/components@1.8.1':
-    resolution: {integrity: sha512-djcZ9noXJRaHsYHYJXAGKeON+vacT3aJvKdFcF0hJeWZ/j7wMGCpvzYwxRgLUooZvj5ifTGBzG9DmF6/sC+v5g==}
+  '@arizeai/components@1.8.3':
+    resolution: {integrity: sha512-U+EV8+GDm0yLNh+xXcEiKY4v7cJgepHMVLRD9pmtYYYdCgiBOLZ7TNQv6gxcRRvlxuqUgLIcP//nwY/roI6MdQ==}
     engines: {node: '>=14'}
     peerDependencies:
       react: '>=18'
@@ -1631,8 +1634,8 @@ packages:
   '@types/react-relay@16.0.6':
     resolution: {integrity: sha512-VTntVQJhlwQYNUlbNgGf8RYy7EtQPRZqsD/w2Si0ygZspJXuNlVdRkklWMFN99EMRhHDpqlNHD8i3wIs7QRz9g==}
 
-  '@types/react@18.2.48':
-    resolution: {integrity: sha512-qboRCl6Ie70DQQG9hhNREz81jqC1cs9EVNcjQ1AU+jH6NFfSAhVVbrrY/+nSF+Bsk4AOwm9Qa61InvMCyV+H3w==}
+  '@types/react@18.3.10':
+    resolution: {integrity: sha512-02sAAlBnP39JgXwkAq3PeU9DVaaGpZyF3MGcC0MKgQVkZor5IiiDAipVaxQHtDJAmO4GIy/rVBy/LzVj76Cyqg==}
 
   '@types/recharts@1.8.29':
     resolution: {integrity: sha512-ulKklaVsnFIIhTQsQw226TnOibrddW1qUQNFVhoQEyY1Z7FRQrNecFCGt7msRuJseudzE9czVawZb17dK/aPXw==}
@@ -1640,9 +1643,6 @@ packages:
   '@types/relay-runtime@14.1.24':
     resolution: {integrity: sha512-ta7vPoFXtEG1wu0Mk7sTngzhmfNGnIe8cDiy3yBEm8pJcGpv55YY/+vWrd9gYd9OQht8rALZpXIYSOLzS/0PVg==}
 
-  '@types/scheduler@0.23.0':
-    resolution: {integrity: sha512-YIoDCTH3Af6XM5VuwGG/QL/CJqga1Zm3NkU3HZ4ZHK2fRMPYP1VczsTUqtsf43PH/iJNVlPHAo2oWX7BSdB2Hw==}
-
   '@types/stack-utils@2.0.3':
     resolution: {integrity: sha512-9aEbYZ3TbYMznPdcdr3SmIrLXwC/AKZXQeCf9Pgao5CKb8CyHuEX5jzWPTkvregvhRJHcpRO6BFoGW9ycaOkYw==}
 
@@ -2789,6 +2789,12 @@ packages:
   graphemer@1.4.0:
     resolution: {integrity: sha512-EtKwoO6kxCL9WO5xipiHTZlSzBm7WLT627TqC/uVRd0HKmq8NXyebnNYxDoBi7wt8eTWrUrKXCOVaFq9x1kgag==}
 
+  graphql-ws@5.16.0:
+    resolution: {integrity: sha512-Ju2RCU2dQMgSKtArPbEtsK5gNLnsQyTNIo/T7cZNp96niC1x0KdJNZV0TIoilceBPQwfb5itrGl8pkFeOUMl4A==}
+    engines: {node: '>=10'}
+    peerDependencies:
+      graphql: '>=0.11 <=16'
+
   graphql@15.3.0:
     resolution: {integrity: sha512-GTCJtzJmkFLWRfFJuoo9RWWa/FfamUHgiFosxi/X1Ani4AVWbeyBenZTNX6dM+7WSbbFfTo/25eh0LLkwHMw2w==}
     engines: {node: '>= 10.x'}
@@ -4879,9 +4885,9 @@ snapshots:
       '@jridgewell/gen-mapping': 0.3.5
       '@jridgewell/trace-mapping': 0.3.25
 
-  '@arizeai/components@1.8.1(@types/react@18.2.48)(eslint@8.57.0)(react-dom@18.2.0(react@18.2.0))(react@18.2.0)':
+  '@arizeai/components@1.8.3(@types/react@18.3.10)(eslint@8.57.0)(react-dom@18.2.0(react@18.2.0))(react@18.2.0)':
     dependencies:
-      '@emotion/react': 11.11.4(@types/react@18.2.48)(react@18.2.0)
+      '@emotion/react': 11.11.4(@types/react@18.3.10)(react@18.2.0)
       '@react-aria/breadcrumbs': 3.5.13(react@18.2.0)
       '@react-aria/button': 3.9.5(react@18.2.0)
       '@react-aria/dialog': 3.5.14(react-dom@18.2.0(react@18.2.0))(react@18.2.0)
@@ -4923,9 +4929,9 @@ snapshots:
 
   '@arizeai/openinference-semantic-conventions@0.10.0': {}
 
-  '@arizeai/point-cloud@3.0.6(@react-three/drei@9.108.4(@react-three/fiber@8.0.12(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(@types/react@18.2.48)(@types/three@0.149.0)(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(@react-three/fiber@8.0.12(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(react@18.2.0)(three-stdlib@2.30.4(three@0.139.2))(three@0.139.2)':
+  '@arizeai/point-cloud@3.0.6(@react-three/drei@9.108.4(@react-three/fiber@8.0.12(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(@types/react@18.3.10)(@types/three@0.149.0)(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(@react-three/fiber@8.0.12(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(react@18.2.0)(three-stdlib@2.30.4(three@0.139.2))(three@0.139.2)':
     dependencies:
-      '@react-three/drei': 9.108.4(@react-three/fiber@8.0.12(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(@types/react@18.2.48)(@types/three@0.149.0)(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2)
+      '@react-three/drei': 9.108.4(@react-three/fiber@8.0.12(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(@types/react@18.3.10)(@types/three@0.149.0)(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2)
       '@react-three/fiber': 8.0.12(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2)
       react: 18.2.0
       three: 0.139.2
@@ -5282,7 +5288,7 @@ snapshots:
 
   '@emotion/memoize@0.8.1': {}
 
-  '@emotion/react@11.11.4(@types/react@18.2.48)(react@18.2.0)':
+  '@emotion/react@11.11.4(@types/react@18.3.10)(react@18.2.0)':
     dependencies:
       '@babel/runtime': 7.24.8
       '@emotion/babel-plugin': 11.11.0
@@ -5294,7 +5300,7 @@ snapshots:
       hoist-non-react-statics: 3.3.2
       react: 18.2.0
     optionalDependencies:
-      '@types/react': 18.2.48
+      '@types/react': 18.3.10
     transitivePeerDependencies:
       - supports-color
 
@@ -6265,7 +6271,7 @@ snapshots:
       '@swc/helpers': 0.5.12
       react: 18.2.0
 
-  '@react-three/drei@9.108.4(@react-three/fiber@8.0.12(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(@types/react@18.2.48)(@types/three@0.149.0)(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2)':
+  '@react-three/drei@9.108.4(@react-three/fiber@8.0.12(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2))(@types/react@18.3.10)(@types/three@0.149.0)(react-dom@18.2.0(react@18.2.0))(react@18.2.0)(three@0.139.2)':
     dependencies:
       '@babel/runtime': 7.24.8
       '@mediapipe/tasks-vision': 0.10.8
@@ -6289,7 +6295,7 @@ snapshots:
       three-mesh-bvh: 0.7.6(three@0.139.2)
       three-stdlib: 2.30.4(three@0.139.2)
       troika-three-text: 0.49.1(three@0.139.2)
-      tunnel-rat: 0.1.2(@types/react@18.2.48)(react@18.2.0)
+      tunnel-rat: 0.1.2(@types/react@18.3.10)(react@18.2.0)
       utility-types: 3.11.0
       uuid: 9.0.1
       zustand: 3.7.2(react@18.2.0)
@@ -6603,32 +6609,29 @@ snapshots:
 
   '@types/react-dom@18.3.0':
     dependencies:
-      '@types/react': 18.2.48
+      '@types/react': 18.3.10
 
   '@types/react-reconciler@0.26.7':
     dependencies:
-      '@types/react': 18.2.48
+      '@types/react': 18.3.10
 
   '@types/react-relay@16.0.6':
     dependencies:
-      '@types/react': 18.2.48
+      '@types/react': 18.3.10
       '@types/relay-runtime': 14.1.24
 
-  '@types/react@18.2.48':
+  '@types/react@18.3.10':
     dependencies:
       '@types/prop-types': 15.7.12
-      '@types/scheduler': 0.23.0
       csstype: 3.1.3
 
   '@types/recharts@1.8.29':
     dependencies:
       '@types/d3-shape': 1.3.12
-      '@types/react': 18.2.48
+      '@types/react': 18.3.10
 
   '@types/relay-runtime@14.1.24': {}
 
-  '@types/scheduler@0.23.0': {}
-
   '@types/stack-utils@2.0.3': {}
 
   '@types/stats.js@0.17.3': {}
@@ -8077,6 +8080,10 @@ snapshots:
 
   graphemer@1.4.0: {}
 
+  graphql-ws@5.16.0(graphql@16.9.0):
+    dependencies:
+      graphql: 16.9.0
+
   graphql@15.3.0: {}
 
   graphql@16.9.0: {}
@@ -9705,10 +9712,10 @@ snapshots:
 
   react-is@18.3.1: {}
 
-  react-markdown@9.0.1(@types/react@18.2.48)(react@18.2.0):
+  react-markdown@9.0.1(@types/react@18.3.10)(react@18.2.0):
     dependencies:
       '@types/hast': 3.0.4
-      '@types/react': 18.2.48
+      '@types/react': 18.3.10
       devlop: 1.1.0
       hast-util-to-jsx-runtime: 2.3.0
       html-url-attributes: 3.0.0
@@ -10320,9 +10327,9 @@ snapshots:
 
   tslib@2.6.3: {}
 
-  tunnel-rat@0.1.2(@types/react@18.2.48)(react@18.2.0):
+  tunnel-rat@0.1.2(@types/react@18.3.10)(react@18.2.0):
     dependencies:
-      zustand: 4.5.4(@types/react@18.2.48)(react@18.2.0)
+      zustand: 4.5.4(@types/react@18.3.10)(react@18.2.0)
     transitivePeerDependencies:
       - '@types/react'
       - immer
@@ -10661,11 +10668,11 @@ snapshots:
     optionalDependencies:
       react: 18.2.0
 
-  zustand@4.5.4(@types/react@18.2.48)(react@18.2.0):
+  zustand@4.5.4(@types/react@18.3.10)(react@18.2.0):
     dependencies:
       use-sync-external-store: 1.2.0(react@18.2.0)
     optionalDependencies:
-      '@types/react': 18.2.48
+      '@types/react': 18.3.10
       react: 18.2.0
 
   zwitch@2.0.4: {}
diff --git a/app/schema.graphql b/app/schema.graphql
index 9d3c39903d6..f7ec737a63d 100644
--- 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
index 6122d551e07..ee4dbe6a51c 100644
--- 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
index 82e1d31c874..a85b3b86d9b 100644
--- 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
index 61480577394..cd338421c6b 100644
--- 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>
diff --git a/app/src/pages/playground/MessageRolePicker.tsx b/app/src/pages/playground/MessageRolePicker.tsx
new file mode 100644
index 00000000000..33307a5539a
--- /dev/null
+++ b/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>
+  );
+}
diff --git a/app/src/pages/playground/PlaygroundChatTemplate.tsx b/app/src/pages/playground/PlaygroundChatTemplate.tsx
index 99623e68c8e..f2103de9873 100644
--- a/app/src/pages/playground/PlaygroundChatTemplate.tsx
+++ b/app/src/pages/playground/PlaygroundChatTemplate.tsx
@@ -1,9 +1,11 @@
 import React from "react";
+import { css } from "@emotion/react";
 
 import { Card, TextArea } from "@arizeai/components";
 
 import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
 
+import { MessageRolePicker } from "./MessageRolePicker";
 import { PlaygroundInstanceProps } from "./types";
 
 interface PlaygroundChatTemplateProps extends PlaygroundInstanceProps {}
@@ -22,11 +24,25 @@ export function PlaygroundChatTemplate(props: PlaygroundChatTemplateProps) {
   }
 
   return (
-    <ul>
+    <ul
+      css={css`
+        display: flex;
+        flex-direction: column;
+        gap: var(--ac-global-dimension-size-200);
+        padding: var(--ac-global-dimension-size-200);
+      `}
+    >
       {template.messages.map((message, index) => {
         return (
           <li key={index}>
-            <Card title={message.role} variant="compact">
+            <Card
+              title={
+                <MessageRolePicker includeLabel={false} role={message.role} />
+              }
+              variant="compact"
+              backgroundColor="light"
+              borderColor="light"
+            >
               <TextArea
                 height={100}
                 value={message.content}
diff --git a/app/src/pages/playground/PlaygroundInputModeRadioGroup.tsx b/app/src/pages/playground/PlaygroundInputModeRadioGroup.tsx
index 9f5940871e8..313b9089c57 100644
--- a/app/src/pages/playground/PlaygroundInputModeRadioGroup.tsx
+++ b/app/src/pages/playground/PlaygroundInputModeRadioGroup.tsx
@@ -25,7 +25,7 @@ export function PlaygroundInputTypeTypeRadioGroup() {
         Manual
       </Radio>
       <Radio label="Dataset" value={"dataset"}>
-        dataset
+        Dataset
       </Radio>
     </RadioGroup>
   );
diff --git a/app/src/pages/playground/PlaygroundInstance.tsx b/app/src/pages/playground/PlaygroundInstance.tsx
index d2d9ccb55d6..c55fd142417 100644
--- a/app/src/pages/playground/PlaygroundInstance.tsx
+++ b/app/src/pages/playground/PlaygroundInstance.tsx
@@ -36,7 +36,7 @@ export function PlaygroundInstance(props: PlaygroundInstanceProps) {
       />
       <Panel defaultSize={50} order={2} css={panelContentCSS}>
         <PlaygroundInput />
-        <PlaygroundOutput />
+        <PlaygroundOutput {...props} />
       </Panel>
     </PanelGroup>
   );
diff --git a/app/src/pages/playground/PlaygroundOutput.tsx b/app/src/pages/playground/PlaygroundOutput.tsx
index 6c521a707df..77401db98aa 100644
--- a/app/src/pages/playground/PlaygroundOutput.tsx
+++ b/app/src/pages/playground/PlaygroundOutput.tsx
@@ -1,11 +1,106 @@
-import React from "react";
+import React, { useMemo, useState } from "react";
+import { useSubscription } from "react-relay";
+import { graphql, GraphQLSubscriptionConfig } from "relay-runtime";
 
 import { Card } from "@arizeai/components";
 
-export function PlaygroundOutput() {
+import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
+
+import {
+  PlaygroundOutputSubscription,
+  PlaygroundOutputSubscription$data,
+} from "./__generated__/PlaygroundOutputSubscription.graphql";
+import { PlaygroundInstanceProps } from "./types";
+
+interface PlaygroundOutputProps extends PlaygroundInstanceProps {}
+
+export function PlaygroundOutput(props: PlaygroundOutputProps) {
+  const instanceId = props.playgroundInstanceId;
+  const instance = usePlaygroundContext((state) => state.instances[instanceId]);
+  if (!instance) {
+    return null;
+  }
+  const runId = instance.activeRunId;
+  const hasRunId = runId !== null;
   return (
     <Card title="Output" collapsible variant="compact">
-      Output goes here
+      {hasRunId ? (
+        <PlaygroundOutputText key={runId} playgroundInstanceId={instanceId} />
+      ) : (
+        "click run to see output"
+      )}
     </Card>
   );
 }
+
+function useChatCompletionSubscription({
+  message,
+  runId,
+  onNext,
+  onCompleted,
+}: {
+  message: string;
+  runId: number;
+  onNext: (response: PlaygroundOutputSubscription$data) => void;
+  onCompleted: () => void;
+}) {
+  const config = useMemo<
+    GraphQLSubscriptionConfig<PlaygroundOutputSubscription>
+  >(
+    () => ({
+      subscription: graphql`
+        subscription PlaygroundOutputSubscription($message: String!) {
+          chatCompletion(input: { message: $message })
+        }
+      `,
+      variables: { message },
+      onNext: (response) => {
+        if (response) {
+          onNext(response);
+        }
+      },
+      onCompleted: () => {
+        onCompleted();
+      },
+    }),
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+    [runId]
+  );
+  return useSubscription(config);
+}
+
+function PlaygroundOutputText(props: PlaygroundInstanceProps) {
+  const instance = usePlaygroundContext(
+    (state) => state.instances[props.playgroundInstanceId]
+  );
+  const markPlaygroundInstanceComplete = usePlaygroundContext(
+    (state) => state.markPlaygroundInstanceComplete
+  );
+  const [output, setOutput] = useState<string>("");
+  if (!instance) {
+    throw new Error("No instance found");
+  }
+  if (typeof instance.activeRunId !== "number") {
+    throw new Error("No message found");
+  }
+
+  if (instance.template.__type !== "chat") {
+    throw new Error("We only support chat templates for now");
+  }
+
+  const message = instance.template.messages.reduce((acc, message) => {
+    return acc + message.content;
+  }, "");
+
+  useChatCompletionSubscription({
+    message: message,
+    runId: instance.activeRunId,
+    onNext: (response) => {
+      setOutput((acc) => acc + response.chatCompletion);
+    },
+    onCompleted: () => {
+      markPlaygroundInstanceComplete(props.playgroundInstanceId);
+    },
+  });
+  return <span>{output}</span>;
+}
diff --git a/app/src/pages/playground/PlaygroundTemplate.tsx b/app/src/pages/playground/PlaygroundTemplate.tsx
index c40ffd693be..4f57c2e6cea 100644
--- a/app/src/pages/playground/PlaygroundTemplate.tsx
+++ b/app/src/pages/playground/PlaygroundTemplate.tsx
@@ -1,6 +1,17 @@
 import React from "react";
 
-import { Button, Card, Icon, Icons } from "@arizeai/components";
+import {
+  Button,
+  Card,
+  Content,
+  Flex,
+  Icon,
+  Icons,
+  Tooltip,
+  TooltipTrigger,
+  TriggerWrap,
+  View,
+} from "@arizeai/components";
 
 import { usePlaygroundContext } from "@phoenix/contexts/PlaygroundContext";
 
@@ -11,19 +22,23 @@ import { PlaygroundInstanceProps } from "./types";
 interface PlaygroundTemplateProps extends PlaygroundInstanceProps {}
 
 export function PlaygroundTemplate(props: PlaygroundTemplateProps) {
-  const id = props.playgroundInstanceId;
+  const instanceId = props.playgroundInstanceId;
   const instances = usePlaygroundContext((state) => state.instances);
-  const playground = instances.find((instance) => instance.id === id);
-  if (!playground) {
-    throw new Error(`Playground instance ${id} not found`);
+  const runPlaygroundInstance = usePlaygroundContext(
+    (state) => state.runPlaygroundInstance
+  );
+  const instance = instances.find((instance) => instance.id === instanceId);
+  if (!instance) {
+    throw new Error(`Playground instance ${instanceId} not found`);
   }
-  const { template } = playground;
+  const { template } = instance;
 
   return (
     <Card
       title="Template"
       collapsible
       variant="compact"
+      bodyStyle={{ padding: 0 }}
       extra={
         instances.length >= NUM_MAX_PLAYGROUND_INSTANCES ? (
           <DeleteButton {...props} />
@@ -37,6 +52,27 @@ export function PlaygroundTemplate(props: PlaygroundTemplateProps) {
       ) : (
         "Completion Template"
       )}
+      <View
+        paddingTop="size-100"
+        paddingBottom="size-100"
+        paddingStart="size-200"
+        paddingEnd="size-200"
+        borderTopColor="dark"
+        borderTopWidth="thin"
+      >
+        <Flex direction="row" gap="size-100" justifyContent="end">
+          <Button
+            variant="primary"
+            icon={<Icon svg={<Icons.PlayCircleOutline />} />}
+            loading={instance.isRunning}
+            onClick={() => {
+              runPlaygroundInstance(instanceId);
+            }}
+          >
+            {instance.isRunning ? "Running..." : "Run"}
+          </Button>
+        </Flex>
+      </View>
     </Card>
   );
 }
@@ -58,13 +94,20 @@ function CompareButton() {
 function DeleteButton(props: PlaygroundInstanceProps) {
   const deleteInstance = usePlaygroundContext((state) => state.deleteInstance);
   return (
-    <Button
-      variant="default"
-      size="compact"
-      icon={<Icon svg={<Icons.TrashOutline />} />}
-      onClick={() => {
-        deleteInstance(props.playgroundInstanceId);
-      }}
-    />
+    <TooltipTrigger>
+      <TriggerWrap>
+        <Button
+          variant="default"
+          size="compact"
+          icon={<Icon svg={<Icons.TrashOutline />} />}
+          onClick={() => {
+            deleteInstance(props.playgroundInstanceId);
+          }}
+        />
+      </TriggerWrap>
+      <Tooltip>
+        <Content>Delete this instance of the playground</Content>
+      </Tooltip>
+    </TooltipTrigger>
   );
 }
diff --git a/app/src/pages/playground/__generated__/PlaygroundOutputSubscription.graphql.ts b/app/src/pages/playground/__generated__/PlaygroundOutputSubscription.graphql.ts
new file mode 100644
index 00000000000..a9a29290f3b
--- /dev/null
+++ b/app/src/pages/playground/__generated__/PlaygroundOutputSubscription.graphql.ts
@@ -0,0 +1,82 @@
+/**
+ * @generated SignedSource<<f13c38a36104dff25a3eb75b2df7893c>>
+ * @lightSyntaxTransform
+ * @nogrep
+ */
+
+/* tslint:disable */
+/* eslint-disable */
+// @ts-nocheck
+
+import { ConcreteRequest, GraphQLSubscription } from 'relay-runtime';
+export type PlaygroundOutputSubscription$variables = {
+  message: string;
+};
+export type PlaygroundOutputSubscription$data = {
+  readonly chatCompletion: string;
+};
+export type PlaygroundOutputSubscription = {
+  response: PlaygroundOutputSubscription$data;
+  variables: PlaygroundOutputSubscription$variables;
+};
+
+const node: ConcreteRequest = (function(){
+var v0 = [
+  {
+    "defaultValue": null,
+    "kind": "LocalArgument",
+    "name": "message"
+  }
+],
+v1 = [
+  {
+    "alias": null,
+    "args": [
+      {
+        "fields": [
+          {
+            "kind": "Variable",
+            "name": "message",
+            "variableName": "message"
+          }
+        ],
+        "kind": "ObjectValue",
+        "name": "input"
+      }
+    ],
+    "kind": "ScalarField",
+    "name": "chatCompletion",
+    "storageKey": null
+  }
+];
+return {
+  "fragment": {
+    "argumentDefinitions": (v0/*: any*/),
+    "kind": "Fragment",
+    "metadata": null,
+    "name": "PlaygroundOutputSubscription",
+    "selections": (v1/*: any*/),
+    "type": "Subscription",
+    "abstractKey": null
+  },
+  "kind": "Request",
+  "operation": {
+    "argumentDefinitions": (v0/*: any*/),
+    "kind": "Operation",
+    "name": "PlaygroundOutputSubscription",
+    "selections": (v1/*: any*/)
+  },
+  "params": {
+    "cacheID": "9f3d7d595d3974bb015d1b73bff91ff5",
+    "id": null,
+    "metadata": {},
+    "name": "PlaygroundOutputSubscription",
+    "operationKind": "subscription",
+    "text": "subscription PlaygroundOutputSubscription(\n  $message: String!\n) {\n  chatCompletion(input: {message: $message})\n}\n"
+  }
+};
+})();
+
+(node as any).hash = "e560f1f0d0df600f9f7eb955c25d4ddc";
+
+export default node;
diff --git a/app/src/pages/settings/SettingsPage.tsx b/app/src/pages/settings/SettingsPage.tsx
index 7cf8da797b1..7abc20eecb0 100644
--- a/app/src/pages/settings/SettingsPage.tsx
+++ b/app/src/pages/settings/SettingsPage.tsx
@@ -82,7 +82,7 @@ export function SettingsPage() {
 function CopyToClipboardButtonWithPadding(props: { text: string }) {
   return (
     <View paddingBottom="19px">
-      <CopyToClipboardButton text={props.text} size="normal" />
+      <CopyToClipboardButton text={props.text} size="default" />
     </View>
   );
 }
diff --git a/app/src/pages/trace/SpanCodeDropdown.tsx b/app/src/pages/trace/SpanCodeDropdown.tsx
index a7de027a05b..e18129acb0b 100644
--- a/app/src/pages/trace/SpanCodeDropdown.tsx
+++ b/app/src/pages/trace/SpanCodeDropdown.tsx
@@ -45,11 +45,11 @@ export function SpanCodeDropdown(props: SpanCodeDropdownProps) {
             <Form>
               <Flex direction="row" gap="size-100" alignItems="end">
                 <TextField label="Span ID" isReadOnly value={spanId} />
-                <CopyToClipboardButton text={spanId} size="normal" />
+                <CopyToClipboardButton text={spanId} size="default" />
               </Flex>
               <Flex direction="row" gap="size-100" alignItems="end">
                 <TextField label="Trace ID" isReadOnly value={traceId} />
-                <CopyToClipboardButton text={traceId} size="normal" />
+                <CopyToClipboardButton text={traceId} size="default" />
               </Flex>
             </Form>
           </View>
diff --git a/app/src/pages/trace/SpanDetails.tsx b/app/src/pages/trace/SpanDetails.tsx
index 628c8c7dd14..4a0b41e40e0 100644
--- a/app/src/pages/trace/SpanDetails.tsx
+++ b/app/src/pages/trace/SpanDetails.tsx
@@ -608,7 +608,6 @@ function LLMSpanInfo(props: { span: Span; spanAttributes: AttributeObject }) {
         }}
         titleSeparator={false}
         variant="compact"
-        // @ts-expect-error force putting the title in as a string
         title={modelNameTitleEl}
       >
         <Tabs>
@@ -1158,7 +1157,6 @@ function DocumentItem({
       bodyStyle={{
         padding: 0,
       }}
-      // @ts-expect-error force putting the title in as a string
       title={
         <Flex direction="row" gap="size-50" alignItems="center">
           <Icon svg={<Icons.FileOutline />} />
@@ -1404,7 +1402,6 @@ function LLMToolSchema({
 
   return (
     <Card
-      // @ts-expect-error force putting the title in as a string
       title={titleEl}
       titleExtra={<Counter variant="light">#{index + 1}</Counter>}
       {...defaultCardProps}
diff --git a/app/src/store/playgroundStore.tsx b/app/src/store/playgroundStore.tsx
index bdccd9e0527..ca925dadb5f 100644
--- a/app/src/store/playgroundStore.tsx
+++ b/app/src/store/playgroundStore.tsx
@@ -4,6 +4,7 @@ import { devtools } from "zustand/middleware";
 export type GenAIOperationType = "chat" | "text_completion";
 
 let playgroundInstanceIdIndex = 0;
+let playgroundRunIdIndex = 0;
 
 /**
  * The input mode for the playground
@@ -18,12 +19,17 @@ export type PlaygroundTemplate =
   | PlaygroundChatTemplate
   | PlaygroundTextCompletionTemplate;
 
+/**
+ * The role of a chat message with a LLM
+ */
+export type ChatMessageRole = "user" | "assistant" | "system" | "tool";
+
 /**
  * A chat message with a role and content
  * @example { role: "user", content: "What is the meaning of life?" }
  */
 export type ChatMessage = {
-  role: string;
+  role: ChatMessageRole;
   content: string;
 };
 
@@ -88,6 +94,11 @@ export interface PlaygroundInstance {
   tools: unknown;
   input: PlaygroundInput;
   output: unknown;
+  activeRunId: number | null;
+  /**
+   * Whether or not the playground instance is actively running or not
+   **/
+  isRunning: boolean;
 }
 
 /**
@@ -128,6 +139,18 @@ export interface PlaygroundState extends PlaygroundProps {
     instanceId: number;
     patch: Partial<PlaygroundInstance>;
   }) => void;
+  /**
+   * Run all the active playground Instances
+   */
+  runPlaygroundInstances: () => void;
+  /**
+   * Run a specific playground instance
+   */
+  runPlaygroundInstance: (instanceId: number) => void;
+  /**
+   * Mark a given playground instance as completed
+   */
+  markPlaygroundInstanceComplete: (instanceId: number) => void;
 }
 
 const DEFAULT_CHAT_COMPLETION_TEMPLATE: PlaygroundChatTemplate = {
@@ -163,6 +186,8 @@ export const createPlaygroundStore = (
         tools: {},
         input: { variables: {} },
         output: {},
+        activeRunId: null,
+        isRunning: false,
       },
     ],
     setOperationType: (operationType: GenAIOperationType) => {
@@ -176,6 +201,8 @@ export const createPlaygroundStore = (
               tools: {},
               input: { variables: {} },
               output: {},
+              activeRunId: null,
+              isRunning: false,
             },
           ],
         });
@@ -188,6 +215,8 @@ export const createPlaygroundStore = (
               tools: {},
               input: { variables: {} },
               output: {},
+              activeRunId: null,
+              isRunning: false,
             },
           ],
         });
@@ -206,6 +235,7 @@ export const createPlaygroundStore = (
           {
             ...instance,
             id: playgroundInstanceIdIndex++,
+            activeRunId: null,
           },
         ],
       });
@@ -253,6 +283,45 @@ export const createPlaygroundStore = (
         }),
       });
     },
+    runPlaygroundInstances: () => {
+      const instances = get().instances;
+      set({
+        instances: instances.map((instance) => ({
+          ...instance,
+          activeRunId: playgroundRunIdIndex++,
+          isRunning: true,
+        })),
+      });
+    },
+    runPlaygroundInstance: (instanceId: number) => {
+      const instances = get().instances;
+      set({
+        instances: instances.map((instance) => {
+          if (instance.id === instanceId) {
+            return {
+              ...instance,
+              activeRunId: playgroundRunIdIndex++,
+              isRunning: true,
+            };
+          }
+          return instance;
+        }),
+      });
+    },
+    markPlaygroundInstanceComplete: (instanceId: number) => {
+      const instances = get().instances;
+      set({
+        instances: instances.map((instance) => {
+          if (instance.id === instanceId) {
+            return {
+              ...instance,
+              isRunning: false,
+            };
+          }
+          return instance;
+        }),
+      });
+    },
     ...initialProps,
   });
   return create(devtools(playgroundStore));
diff --git a/pyproject.toml b/pyproject.toml
index 8bdf7c702bc..e42f105f62b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -61,6 +61,7 @@ dependencies = [
   "fastapi-mail",
   "pydantic>=1.0,!=2.0.*,<3", # exclude 2.0.* since it does not support the `json_encoders` configuration setting
   "authlib",
+  "websockets",
 ]
 dynamic = ["version"]
 
diff --git a/src/phoenix/server/api/schema.py b/src/phoenix/server/api/schema.py
index 0a5ac8b7dec..0f153ba7cdc 100644
--- a/src/phoenix/server/api/schema.py
+++ b/src/phoenix/server/api/schema.py
@@ -3,6 +3,7 @@
 from phoenix.server.api.exceptions import get_mask_errors_extension
 from phoenix.server.api.mutations import Mutation
 from phoenix.server.api.queries import Query
+from phoenix.server.api.subscriptions import Subscription
 
 # This is the schema for generating `schema.graphql`.
 # See https://strawberry.rocks/docs/guides/schema-export
@@ -12,4 +13,5 @@
     query=Query,
     mutation=Mutation,
     extensions=[get_mask_errors_extension()],
+    subscription=Subscription,
 )
diff --git a/src/phoenix/server/api/subscriptions.py b/src/phoenix/server/api/subscriptions.py
new file mode 100644
index 00000000000..e2475a18d50
--- /dev/null
+++ b/src/phoenix/server/api/subscriptions.py
@@ -0,0 +1,76 @@
+from datetime import datetime
+from typing import AsyncIterator
+
+import strawberry
+from sqlalchemy import insert, select
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.server.api.context import Context
+
+
+@strawberry.input
+class ChatCompletionInput:
+    message: str
+
+
+@strawberry.type
+class Subscription:
+    @strawberry.subscription
+    async def chat_completion(
+        self, info: Info[Context, None], input: ChatCompletionInput
+    ) -> AsyncIterator[str]:
+        from openai import AsyncOpenAI
+        from openai.types.chat import ChatCompletionUserMessageParam
+
+        client = AsyncOpenAI()
+        chunk_contents = []
+        start_time = datetime.now()
+        async for chunk in await client.chat.completions.create(
+            messages=[ChatCompletionUserMessageParam(role="user", content=input.message)],
+            model="gpt-4",
+            stream=True,
+        ):
+            choice = chunk.choices[0]
+            if choice.finish_reason is None:
+                assert isinstance(chunk_content := chunk.choices[0].delta.content, str)
+                yield chunk_content
+                chunk_contents.append(chunk_content)
+        end_time = datetime.now()
+        async with info.context.db() as session:
+            # insert dummy data
+            trace_id = str(start_time)
+            span_id = str(end_time)
+            default_project_id = await session.scalar(
+                select(models.Project.id).where(models.Project.name == "default")
+            )
+            trace_rowid = await session.scalar(
+                insert(models.Trace)
+                .returning(models.Trace.id)
+                .values(
+                    project_rowid=default_project_id,
+                    trace_id=trace_id,
+                    start_time=start_time,
+                    end_time=end_time,
+                )
+            )
+            await session.execute(
+                insert(models.Span).values(
+                    trace_rowid=trace_rowid,
+                    span_id=span_id,
+                    parent_id=None,
+                    name="AsyncOpenAI.chat.completion.create",
+                    span_kind="LLM",
+                    start_time=start_time,
+                    end_time=end_time,
+                    attributes={},
+                    events=[],
+                    status_code="OK",
+                    status_message="",
+                    cumulative_error_count=0,
+                    cumulative_llm_token_count_prompt=0,
+                    cumulative_llm_token_count_completion=0,
+                    llm_token_count_prompt=0,
+                    llm_token_count_completion=0,
+                )
+            )
diff --git a/src/phoenix/server/app.py b/src/phoenix/server/app.py
index 816ae611e48..708e6bcde71 100644
--- a/src/phoenix/server/app.py
+++ b/src/phoenix/server/app.py
@@ -45,6 +45,7 @@
 from starlette.types import Scope, StatefulLifespan
 from strawberry.fastapi import GraphQLRouter
 from strawberry.schema import BaseSchema
+from strawberry.subscriptions import GRAPHQL_TRANSPORT_WS_PROTOCOL
 from typing_extensions import TypeAlias
 
 import phoenix
@@ -555,6 +556,7 @@ def get_context() -> Context:
         include_in_schema=False,
         prefix="/graphql",
         dependencies=(Depends(is_authenticated),) if authentication_enabled else (),
+        subscription_protocols=[GRAPHQL_TRANSPORT_WS_PROTOCOL],
     )