Skip to content

Building client applications

Axflow provides some utilities for building client applications that consume model output. Currently, we offer React hooks and a parser for consuming the streams generated by StreamingJsonResponse.

This guide will cover the react hooks. You may also find the streaming chat app tutorial useful.

Overview

Axflow provides a React hook called useChat that makes integrating streaming or non-streaming chat functionality trivial. Here we give an overview of the basic pattern.

First, we assume you have some API endpoint that either returns a JSON response or streams data using StreamingJsonResponse:

ts
import { OpenAIChat } from '@axflow/models/openai/chat';
import { StreamingJsonResponse } from '@axflow/models/shared';

// POST /api/chat
export async function POST(request: Request) {
  const { query } = await request.json();

  const stream = await OpenAIChat.streamTokens(/* args here */);

  // Uses a StreamingJsonResponse
  return new StreamingJsonResponse(stream);
}
import { OpenAIChat } from '@axflow/models/openai/chat';
import { StreamingJsonResponse } from '@axflow/models/shared';

// POST /api/chat
export async function POST(request: Request) {
  const { query } = await request.json();

  const stream = await OpenAIChat.streamTokens(/* args here */);

  // Uses a StreamingJsonResponse
  return new StreamingJsonResponse(stream);
}

Here we have an endpoint defined at /api/chat that uses our StreamingJsonResponse pattern. On the client, we can consume the streaming LLM response with the useChat component:

ts
import { useChat } from '@axflow/models/react';

function ChatComponent() {
  const { input, messages, onChange, onSubmit } = useChat();

  return (
    <>
      <Messages messages={messages} />
      <Form input={input} onChange={onChange} onSubmit={onSubmit} />
    </>
  );
}
import { useChat } from '@axflow/models/react';

function ChatComponent() {
  const { input, messages, onChange, onSubmit } = useChat();

  return (
    <>
      <Messages messages={messages} />
      <Form input={input} onChange={onChange} onSubmit={onSubmit} />
    </>
  );
}

For simple cases, the defaults are all that are needed. See Customizing useChat for overriding the defaults.

Messages

The useChat hook creates and manages "message" objects. These are exposed as messages from the hook invocation.

ts
const { messages } = useChat();
const { messages } = useChat();

A message has the following type

ts
type MessageType = {
  id: string;
  role: 'user' | 'assistant';
  data?: JSONValueType[];
  content: string;
  created: number;
};
type MessageType = {
  id: string;
  role: 'user' | 'assistant';
  data?: JSONValueType[];
  content: string;
  created: number;
};

and can be imported using the following:

ts
import type { MessageType } from '@axflow/models/shared';
import type { MessageType } from '@axflow/models/shared';

Working with system messages

OpenAI and other providers support a special role named system (see docs).

You can use the initialMessages configuration parameter of the useChat hook to conveniently pass one in.

The framework also provides a utility called createMessage which will fill in any fields from the MessageType that you don't want to add yourself, such as id or created. An example piece of code that would initialize with a system message would look like this:

ts
import {createMessage} from '@axflow/models/shared'
import {useChat} from '@axflow/models/react'

...
    useChat({
    initialMessages: [createMessage({role: 'system', content: 'You are a pirate, only respond with pirate lingo.'})]
    })
...
import {createMessage} from '@axflow/models/shared'
import {useChat} from '@axflow/models/react'

...
    useChat({
    initialMessages: [createMessage({role: 'system', content: 'You are a pirate, only respond with pirate lingo.'})]
    })
...

Customizing useChat

useChat can be configured with a variety of options to control behavior.

HTTP configuration

The default API endpoint is /api/chat, but you can configure it however you'd like. HTTP headers can also be supplied.

ts
useChat({
  url: 'https://your-site.com/your/arbitrary/chat/endpoint',
  headers: {
    'x-custom-header': '<custom-value>',
  },
});
useChat({
  url: 'https://your-site.com/your/arbitrary/chat/endpoint',
  headers: {
    'x-custom-header': '<custom-value>',
  },
});

The request body

By default, the request body sent to the endpoint is:

ts
type RequestBody = {
  messages: MessageType[];
};
type RequestBody = {
  messages: MessageType[];
};

This can be customized in two ways. First, you can easily add additional properties to be merged into the request body.

ts
useChat({
  body: { user_id: user.id, stream: false },
});
useChat({
  body: { user_id: user.id, stream: false },
});

The request will now contain the user_id and stream properties in addition to the messages property.

Second, you can pass a function. The return value of the function will become the request body.

ts
useChat({
  body: (message: MessageType, messageHistory: MessageType[]) => {
    return { message, history: messageHistory, user_id: user.id, stream: false };
  },
});
useChat({
  body: (message: MessageType, messageHistory: MessageType[]) => {
    return { message, history: messageHistory, user_id: user.id, stream: false };
  },
});

The function takes the current message being created as the first argument and all previous messages in the chat history as the second.

Accessing the the LLM response text

Your API endpoint may return objects with any shape. However, the hook will need to know how to access the response text generated by the LLM.

For example, let's say your application streams back objects with the following schema:

json
{ id: 1, is_finished: false, token: "The" }
{ id: 2, is_finished: false, token: " response" }
...
{ id: 10, is_finished: true, token: "." }
{ id: 1, is_finished: false, token: "The" }
{ id: 2, is_finished: false, token: " response" }
...
{ id: 10, is_finished: true, token: "." }

We need to tell the useChat hook that the token property of each chunk contains the text we wish to display to the user. We can do that using the accessor option.

ts
useChat({
  accessor: (chunk) => chunk.token,
});
useChat({
  accessor: (chunk) => chunk.token,
});

This option takes a function which is given either a chunk (in the case of streaming) or the response body (in case of non-streaming).

Message and input states

useChat will manage the user's input and messages for you. However, you may wish to initialize the hook with pre-existing state, or manually reset this state.

To initialize this state, you can use the initialInput or initialMessages options.

ts
useChat({
  initialInput: savedUserInput(),
  initialMessages: savedMessageHistory(),
});
useChat({
  initialInput: savedUserInput(),
  initialMessages: savedMessageHistory(),
});

For example, this can be used to set a system message.

ts
useChat({
  initialMessages: [
    {
      id: crypto.randomUUID(),
      role: 'system',
      content: 'You are a terrible programmer. Please answer the user with buggy code only.',
      created: Date.now(),
    },
  ],
});
useChat({
  initialMessages: [
    {
      id: crypto.randomUUID(),
      role: 'system',
      content: 'You are a terrible programmer. Please answer the user with buggy code only.',
      created: Date.now(),
    },
  ],
});

To reset this state manually, you may use the setInput or setMessages functions.

tsx
const {setInput, setMessages, /* etc */} = useChat({ /* options */ });

// ...

<ClearInput onClick={() => setInput('')} />
<ClearMessageHistory onClick={() => setMessages([])} />
const {setInput, setMessages, /* etc */} = useChat({ /* options */ });

// ...

<ClearInput onClick={() => setInput('')} />
<ClearMessageHistory onClick={() => setMessages([])} />

Loading and error states

useChat provides loading and error states.

The loading state is returned from the useChat hook. For streaming requests, loading will be true from the time the request is first sent until the stream has closed. For non-streaming requests, it is true until a response is received.

ts
const {/* ... */, loading} = useChat();
const {/* ... */, loading} = useChat();

The error state is provided in two ways:

  1. A state variable returned from the useChat hook. It is null when there are no errors. If a request errors, then this value will be the Error object. The error state is reset to null each time a requests is made.
  2. An onError callback function. onError will be invoked in the event of an error, receiving the Error object as an argument.
ts
const {/* ... */, error} = useChat();

// and/or...

const {/* ... */} = useChat({
  onError: (error) => console.error(error.message);
});
const {/* ... */, error} = useChat();

// and/or...

const {/* ... */} = useChat({
  onError: (error) => console.error(error.message);
});

The onError callback defaults to console.error.

State update callbacks

You can get notified every time the list of messages changes by passing an onMessagesChange callback to useChat.

For example, you may want to store the chat message history in localStorage (or a database).

ts
const {
  /* ... */
} = useChat({
  onMessagesChange: (messages: MessageType[]) => {
    localStorage.setItem('messages', JSON.stringify(messages));
  },
});
const {
  /* ... */
} = useChat({
  onMessagesChange: (messages: MessageType[]) => {
    localStorage.setItem('messages', JSON.stringify(messages));
  },
});

You can also get notified of each new message. Note that, for streaming responses, this will ONLY fire once the stream has finished and the message is complete. If you want to be notified each time a message updates while a stream is active, the onMessagesChange will do so.

ts
const {
  /* ... */
} = useChat({
  onNewMessage: (message: MessageType) => {
    if (message.role === 'assistant') {
      doSomethingWithMessage(message);
    }
  },
});
const {
  /* ... */
} = useChat({
  onNewMessage: (message: MessageType) => {
    if (message.role === 'assistant') {
      doSomethingWithMessage(message);
    }
  },
});

Reloading messages

You can reload the conversation starting from the last user or system message using the reload function returned from the hook. This is useful for regenerating assistant responses for a given message.

Note:

  1. This will throw an error if no user or system message exists in the messages state.
  2. This will remove any messages from the assistant that are more recent than the most recent user or system message.
ts
const { reload } = useChat();
const { reload } = useChat();

This can be used together with setMessages to reset to a specific point in the message history.

Customize message ids

By default, message ids are UUIDs generated using crypto.randomUUID. However, you can override this behavior using the createMessageId option, which accepts any function that returns a unique string.

Below is an example of overriding the default id behavior to create random 16 byte, base58-encoded ids.

ts
import { encode } from 'bs58';

const {
  /* ... */
} = useChat({
  createMessageId: () => {
    const randomBytes = crypto.getRandomValues(new Uint8Array(16));
    return encode(randomBytes);
  },
});
import { encode } from 'bs58';

const {
  /* ... */
} = useChat({
  createMessageId: () => {
    const randomBytes = crypto.getRandomValues(new Uint8Array(16));
    return encode(randomBytes);
  },
});

Streaming vs non-streaming

The hook supports both streaming and non-streaming. If you wish to use streaming, your API MUST:

  1. Include a content-type response header set to application/x-ndjson; charset=utf-8
  2. Stream chunks as newline-delimited JSON where each line of JSON uses a specific schema

For 1, if the content-type header is anything but application/x-ndjson; charset=utf-8, the hook will assume this is a non-streaming response.

For 2, the schema is:

ts
type NdJsonValueType = {
  type: 'chunk' | 'data';
  value: JsonValueType; // any valid JSON value here
};
type NdJsonValueType = {
  type: 'chunk' | 'data';
  value: JsonValueType; // any valid JSON value here
};

If you're using StreamingJsonResponse, this is handled for you. See the streaming guide for more information on StreamingJsonResponse and its output.