HOWTO: use a custom Streaming REST API with assistant-ui from scratch

README.md

HOWTO: use a custom Streaming REST api with assistant-ui from scratch

Instructions

1. Wrap your Thread component within the ChatRuntimeProvider.
2. Update the endpoint and configure appropriate auth headers

For more info on the SSE protocol, checkout the event stream protocol docs.

*Why so much code?

This example aims to show how to integrate a custom streaming REST api endpoint with assistant-ui without using additional packages. If you like, you can use existing packages for dealing with SSE parsing and communication.

ChatRuntimeProvider.tsx

import { ReactNode } from "react";
import {
  AssistantRuntimeProvider,
  useLocalRuntime,
  type ChatModelAdapter,
} from "@assistant-ui/react";

import { parseEventChunk } from "./sse-util";
import { generateStream } from "./stream-utils";

const TRUSTED_SSEVENTS = ["message"];

const CustomChatModelAdapter: ChatModelAdapter = {
  async *run({ messages, abortSignal, runConfig }) {
    const lastUserMessage = messages.filter((m) => m.role === "user").pop();
    /**
     * Call your REST API streaming endpoint here
     * Replace YOUR_ENDPOINT with the actual endpoint
     * Replace AUTH_HEADERS with the actual headers
     */
    const stream = await generateStream(YOUR_ENDPOINT, lastUserMessage, AUTH_HEADERS, abortSignal);

    let text = "";
    for await (const chunk of stream) {
      const event = parseEventChunk(chunk);
      // Ignore empty events or events that are not of type "message"
      if (!event || !TRUSTED_SSEVENTS.includes(event.type)) {
        continue;
      }

      // If your data is just a string, you can yield it directly
      const data = event.data;
      text += data;
      yield {
        content: [{ type: "text", text }],
      };
    }
  },
};

export function ChatRuntimeProvider({
  children,
}: Readonly<{
  children: ReactNode;
}>) {
  const runtime = useLocalRuntime(CustomChatModelAdapter);

  return <AssistantRuntimeProvider runtime={runtime}>{children}</AssistantRuntimeProvider>;
}

sse-util.ts

export type CustomEventDataType = CustomEvent & {
  data: string[];
};
const FIELD_SEPARATOR = ":";
/**
 * Parse a received SSE event chunk into a constructed event object.
 * See https://github.com/smartmuki/sse-ts/blob/master/lib/sse.ts#L222
 */
export const parseEventChunk = (chunk: string) => {
  if (!chunk || !chunk.length) {
    return null;
  }

  const e: any = { data: [], event: "message" };

  chunk.split(/\n|\r\n|\r/).forEach((line: string) => {
    line = line.trimEnd();
    const index = line.indexOf(FIELD_SEPARATOR);
    if (index <= 0) {
      // Line was either empty, or started with a separator and is a comment.
      // Either way, ignore.
      return;
    }

    const field = line.substring(0, index);
    if (!(field in e)) {
      return;
    }

    const value = line.substring(index + 1).trimStart();
    if (field === "data") {
      e[field].push(value);
    } else {
      e[field] = value;
    }
  });

  const event = new CustomEvent(e.event) as CustomEventDataType;
  event.data = e.data;
  return event;
};

stream-utils.ts

import { AxiosHeaders } from "axios";

export async function generateStream(
  url: string,
  data: any,
  headers: AxiosHeaders,
  abortSignal?: AbortSignal,
) {
  const response = await fetch(url, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      ...headers,
    },
    body: JSON.stringify(data),
    signal: abortSignal,
  });
  if (response.status !== 200) throw new Error(response.status.toString());
  if (!response.body) {
    throw new Error("No response body");
  }
  return getIterableStream(response.body);
}

export async function* getIterableStream(body: ReadableStream<Uint8Array>): AsyncIterable<string> {
  const reader = body.getReader();
  const decoder = new TextDecoder();

  while (true) {
    const { value, done } = await reader.read();
    if (done) {
      break;
    }
    const decodedChunk = decoder.decode(value, { stream: true });
    yield decodedChunk;
  }
}