Streaming messages from ChatGPT using Swift AsyncSequence

Zach Waugh Zach Waugh •

I was recently playing with the ChatGPT API (like everyone else in the world) and hit the point of needing to consume their streaming API in Swift. The API uses server-sent events to stream partial responses so you can update the UI in real-time as the response streams in. This gave me a chance to use a new protocol introduced in Swift 5.5 called AsyncSequence for the first time with URLSession. John Sundell has a great article that explains it more in depth, but the gist you can consume a sequence asynchronously as easy as you would iterating through an array:

for try await value in asyncSequence {
  // do something with `value`
}

To use that for HTTP requests to process data as it’s received, you can use a new URLSession method (as of iOS 15.0) that returns a tuple of (URLSession.AsyncBytes, URLResponse). The AsyncBytes type has methods to get an AsyncSequence of individual characters or entire lines. Using lines saves us a lot of work where we can let URLSession buffer until a newline is received and give us the whole message at once. I was surprised how easy it was to get this working in Swift for ChatGPT.

let (bytes, response) = try await URLSession.shared.bytes(for: request)

for try await line in bytes.lines {
  print(line)
}

Each line we get there will be partial response, and all we need to do is parse the message. That is less interesting, but here’s a demo of how it looks and a full code sample that you can run from the command-line to query ChatGPT and stream the messages in a simple Swift script is below.

Demo of streaming from ChatGPT in Swift from the command-line
import Foundation

/// Execute a streaming request to ChatGPT and prints the response as it comes in
/// Set OPENAI_API_KEY environment variable to a valid API key
/// This is simplified version to demonstrate using async sequences from an API in Swift
/// and is not production ready as it has almost zero error handling
///
/// Usage:
/// swift chat.swift <query>
///
guard let apiKey = ProcessInfo.processInfo.environment["OPENAI_API_KEY"] else {
  print("*** Error: No OpenAI API key provided! Set an $OPENAI_API_KEY environment variable and try again ")
  exit(1)
}

guard CommandLine.arguments.count == 2 else { 
  print("Usage: swift chat.swift \"<query>\"")
  exit(1)
}

let content = CommandLine.arguments[1]
let request = try makeRequest(content: content, apiKey: apiKey)
let (stream, _) = try await URLSession.shared.bytes(for: request)

for try await line in stream.lines {
  guard let message = parse(line) else { continue }

  print(message, terminator: "")
  fflush(stdout)
}

/// Make a URL request for the ChatGPT endpoint
/// https://platform.openai.com/docs/api-reference/chat/create
func makeRequest(content: String, apiKey: String) throws -> URLRequest {
  let query = Query(messages: [.init(role: "user", content: content)])
  let url = URL(string: "https://api.openai.com/v1/chat/completions")!

  var request = URLRequest(url: url)
  request.httpMethod = "POST"
  request.httpBody = try JSONEncoder().encode(query)
  request.allHTTPHeaderFields = [
    "Content-Type": "application/json",
    "Authorization": "Bearer \(apiKey)"
  ]

  return request
}

/// Parse a line from the stream and extract the message
func parse(_ line: String) -> String? {
  let components = line.split(separator: ":", maxSplits: 1, omittingEmptySubsequences: true)
  guard components.count == 2, components[0] == "data" else { return nil }

  let message = components[1].trimmingCharacters(in: .whitespacesAndNewlines)

  if message == "[DONE]" {
    return "\n"
  } else {
    let chunk = try? JSONDecoder().decode(Chunk.self, from: message.data(using: .utf8)!)
    return chunk?.choices.first?.delta.content
  }
}

/// Encodes the query
struct Query: Encodable {
  struct Message: Encodable {
    let role: String
    let content: String
  }

  let model = "gpt-3.5-turbo"
  let messages: [Message]
  let stream = true
}

/// Decodes a chunk of the streaming response
struct Chunk: Decodable {
  struct Choice: Decodable {
    struct Delta: Decodable {
      let role: String?
      let content: String?
    }

    let delta: Delta
  }

  let choices: [Choice]
}

Note: you need to set an environment variable for OPENAI_API_KEY to use it, then you can run swift chat.swift <query>