GraphQL LLM Streaming Examples

index.ts

import { createServer } from 'node:http'
import { createSchema, createYoga, Repeater } from 'graphql-yoga'
import { useDeferStream } from '@graphql-yoga/plugin-defer-stream'
import OpenAI from 'openai'

const openai = new OpenAI({
  // apiKey: 'my api key', // defaults to process.env["OPENAI_API_KEY"]
})

const generateCompletion = (prompt: string): Repeater<string> => {
  return new Repeater<string>(async (push, stop) => {
    const stream = await openai.chat.completions.create({
      model: 'gpt-3.5-turbo',
      messages: [{ role: 'user', content: prompt }],
      stream: true,
      max_tokens: 100,
    })
    for await (const part of stream) {
      const delta = part.choices[0]?.delta?.content ?? ''
      push(delta)
    }
    stop()

    await stop.then(() => {
      console.log('cancel')
      stream.controller.abort()
    })
  })
}

const yoga = createYoga({
  schema: createSchema({
    typeDefs: /* GraphQL */ `
      input GenerateCompletionInput {
        prompt: String!
      }

      type GenerateCompletionPayload {
        tokens: [String!]!
      }

      type GenerateCompletionSubPayload {
        token: String!
        finished: Boolean!
      }

      type Query {
        isHealthy: Boolean!
      }

      type Mutation {
        generateCompletion(input: GenerateCompletionInput!): GenerateCompletionPayload!
      }

      type Subscription {
        generateCompletionSub(prompt: String!): GenerateCompletionSubPayload!
      }
    `,
    resolvers: {
      Query: {
        isHealthy: () => true,
      },
      Mutation: {
        generateCompletion: (_, { input }) => ({ tokens: [], prompt: input.prompt }),
      },
      GenerateCompletionPayload: {
        tokens: (parent) => generateCompletion(parent.prompt),
      },
      Subscription: {
        generateCompletionSub: {
          subscribe: (_, { prompt }) => generateCompletion(prompt),
          resolve: (token) => ({ token, finished: false }),
        },
      },
    },
  }),
  plugins: [useDeferStream()],
})

const server = createServer(yoga)

server.listen(4000, () => {
  console.info('Server is running on http://localhost:4000/graphql')
})

package.json

{
  "name": "graphql-llm",
  "version": "1.0.0",
  "description": "",
  "main": "index.js",
  "type": "module",
  "scripts": {
    "test": "echo \"Error: no test specified\" && exit 1",
    "start": "ts-node-esm src/index.ts"
  },
  "author": "",
  "license": "ISC",
  "dependencies": {
    "@graphql-yoga/plugin-defer-stream": "^2.0.4",
    "graphql": "^16.8.1",
    "graphql-yoga": "^4.0.4",
    "openai": "^4.8.0"
  },
  "devDependencies": {
    "@types/node": "^20.6.3",
    "ts-node": "^10.9.1",
    "typescript": "^5.2.2"
  }
}

tsconfig.json

{
  "compilerOptions": {
    "target": "ESNext",
    "module": "NodeNext",
    "esModuleInterop": true,
    "forceConsistentCasingInFileNames": true,
    "skipLibCheck": true
  }
}