Consuming Web Streams with useState, SWR and React Query

In this blog post I present you three different ways you can consume text streams with React.

Let's imagine we are building a chatbot interface that will send user messages to an API that will stream down the response one token at the time.

If you are interested in a full guide to build a chatbot refer to this older post or look at the GitHub repository for this article.

Streaming

To build the interface we need a function that calls an API and allows the application to read incoming tokens.
We can use a Javascript Generator to iterate through an async stream of tokens using a simple for loop.

const wait = (ms: number) =>
  new Promise((resolve) => setTimeout(resolve, ms));

// This is how the tokens are made available 
async function* getCompletion(prompt: string) {
  await wait(300);
  yield "Ciao, ";   // token: Ciao, 
  await wait(300);  
  yield "Come ";    // token: Come 
  await wait(300);
  yield "Stai?";    // token: Stai?
}

// This is how the consumer can read the tokens
for await (const token of getCompletion("Ciao")) {
  console.log("token:", token);
}

Here is the actual implementation:

// src/getCompletion.ts
async function* getCompletion(
  prompt: string,
  signal?: AbortSignal,
) {
  const url = new URL(
    "/completion",
    "http://localhost:4000",
  );
  url.searchParams.append("prompt", prompt);

  const res = await fetch(url, {
    method: "GET",
    signal,
  });

  const reader = res.body?.getReader();
  if (!reader) throw new Error("No reader");
  const decoder = new TextDecoder();

  let i = 0;
  while (i < 1000) {
    i++;
    const { done, value } = await reader.read();
    if (done) return;
    const token = decoder.decode(value);
    yield token;

    if (signal?.aborted) {
      await reader.cancel();
      return;
    }
  }
}

export default getCompletion;

https://github.com/fibonacid/streaming-use-effect-react-query-swr/blob/main/src/lib/shared.ts

useEffect + useState

To render the tokens in a React application we need to wrap the getCompletion function in a hook.
Let's create a custom hook that uses only react primitives (useState, useCallback);

// src/useCompletion.ts
import { useCallback, useState } from "react";
import getCompletion from "getCompletion";

export default function useCompletion() {
  const [tokens, setTokens] = useState<string[]>([]);
  const [isLoading, setIsLoading] = useState(false);
  const [error, setError] = useState<unknown>();
  const [abortController, setAbortController] =
    useState<AbortController | null>(null);

  const mutate = useCallback(
    async (prompt: string) => {
      setIsLoading(true);
      setTokens([]);

      if (abortController) {
        abortController.abort();
      }
      const controller = new AbortController();
      const signal = controller.signal;
      setAbortController(controller);

      try {
        for await (const token of getCompletion(
          prompt,
          signal,
        )) {
          console.log("token", token);
          setTokens((prev) => [...prev, token]);
        }
      } catch (err) {
        if (
          err instanceof Error &&
          err.name === "AbortError"
        ) {
          return; // abort errors are expected
        }
        setError(err);
      }
      setIsLoading(false);
      setAbortController(null);
    },
    [abortController],
  );

  const data = tokens.join("");
  return [mutate, { data, isLoading, error }] as const;
}

https://github.com/fibonacid/streaming-use-effect-react-query-swr/blob/main/src/lib/custom.tsx

Now, let's create a simple component that allow users to send a prompt and read the incoming response.

// src/App.tsx
import useCompletion from "./useCompletion"; 

export default function App() {
  const [prompt, setPrompt] = useState("");
  const [mutate, { data }] = useCompletion();
  return (
    <main>
      <h1>Chat</h1>
      <form
        onClick={() => {
          setPrompt("");
          void mutate(prompt);
        }}
      >
        <label htmlFor="prompt">Prompt</label>
        <input
          id="prompt"
          onChange={(e) => setPrompt(e.currentTarget.value)}
          value={prompt}
        />
        <button>Send</button>
      </form>
      {data && <p>{data}</p>}
    </main>
  );
}

SWR

This next implementation uses the swr package to store the streaming data. This implementation has been copied and simplified from the ai package from Vercel.

// src/useCompletionSWR.ts

import { useId } from "react";
import useSWR from "swr";
import useSWRMutation from "swr/mutation";
import getCompletion from "./getCompletion";

export default function useCompletionSWR() {
  const id = useId();
  const { data, mutate } = useSWR<string>(
    ["completion", id],
    null,
  );
  const [abortController, setAbortController] =
    useState<AbortController | null>();

  const { trigger, isMutating, error } = useSWRMutation<
    void,
    unknown,
    [string, string],
    string
  >(["completion", id], async (_, { arg: prompt }) => {
    void mutate("", false);
    if (abortController) {
      abortController.abort();
    }
    const controller = new AbortController();
    const signal = controller.signal;
    setAbortController(controller);

    for await (const token of getCompletion(
      prompt,
      signal,
    )) {
      void mutate(
        (prev) => (prev ? prev + token : token),
        false,
      );
    }
    setAbortController(null);
  });

  return [
    trigger,
    { data, error, isLoading: isMutating },
  ] as const;
}

https://github.com/fibonacid/streaming-use-effect-react-query-swr/blob/main/src/lib/custom.tsx

React Query

This last example is very similar to the swr one, but uses React Query.

// src/useCompletionRQ
import { useId } from "react";
import { useQueryClient, useQuery, useMutation } from "@tanstack/react-query";

export default function useCompletionRQ() {
  const id = useId();

  const queryClient = useQueryClient();
  const { data } = useQuery<string>({
    queryKey: ["completion", id],
  });
  const [abortController, setAbortController] =
    useState<AbortController | null>();

  const { mutate, error, isLoading } = useMutation({
    mutationKey: ["completion", id],
    mutationFn: async (prompt: string) => {
      if (abortController) {
        abortController.abort();
      }
      const controller = new AbortController();
      const signal = controller.signal;
      setAbortController(controller);

      for await (const token of getCompletion(
        prompt,
        signal,
      )) {
        queryClient.setQueryData<string>(
          ["completion", id],
          (prev) => (prev ? prev + token : token),
        );
      }
      setAbortController(null);
    },
  });

  return [mutate, { data, error, isLoading }] as const;
}

https://github.com/fibonacid/streaming-use-effect-react-query-swr/blob/main/src/lib/react-query.tsx

We can render the same <App /> element but React Query requires a global store:

// src/App.tsx
import {
  QueryClient,
  QueryClientProvider,
} from "@tanstack/react-query";
import useCompletion from "./useCompletionRQ"; 

const queryClient = new QueryClient();

function Chat() {
  const [mutate, { data }] = useCompletionRQ();
  return <main>{/* ... */}</main>
}

function App() {
  return (
    <QueryClientProvider client={queryClient}>
      <Chat />
    </QueryClientProvider />
  )
}

export default App;

Conclusion

The option that uses useState is very simple and provides the same functionality as the other ones.
The React Query and SWR versions might be useful if your application heavily relies on one of those libraries and you would like to keep things consistent.

Blog

Consuming Web Streams with useState, SWR and React Query

Lorenzo Rivosecchi

Streaming

useEffect + useState

SWR

React Query

Conclusion

Join Our Newsletter. No Spam, Only the good stuff.

Related