Consuming Web Streams with useState, SWR and React Query
Tags: javascript, frontend, llm, react
Published: 2023-07-10 12:46:34

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](https://dev.to/fibonacid/quick-guide-for-langchain-streaming-via-expressjs-iod) or look at the GitHub [repository](https://github.com/fibonacid/streaming-use-effect-react-query-swr) 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](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Guide/Iterators_and_Generators) to iterate through an async stream of tokens using a simple for loop.

```typescript
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:

```typescript
// 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](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`);


```typescript
// 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](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.

```tsx
// 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`](https://swr.vercel.app/) package to store the streaming data. This implementation has been copied and simplified from the [`ai`](https://github.com/vercel-labs/ai) package from Vercel.

```typescript
// 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](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](https://tanstack.com/query/v3/).

```tsx
// 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](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:

```tsx
// 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.