Building Custom UI

Build a fully custom Aomi chat interface from scratch using only @aomi-labs/react hooks. Covers provider setup, message rendering, input, and streaming responses.

This guide walks through building a complete chat interface from scratch using only @aomi-labs/react hooks. You control every pixel.

Setup

npm install @aomi-labs/react @assistant-ui/react

Wrap your app with the runtime provider:

// app/providers.tsx
import { AomiRuntimeProvider } from "@aomi-labs/react";

export function Providers({ children }: { children: React.ReactNode }) {
  return (
    <AomiRuntimeProvider backendUrl={process.env.NEXT_PUBLIC_BACKEND_URL!}>
      {children}
    </AomiRuntimeProvider>
  );
}

Message List Component

Display messages for the current thread:

"use client";

import { useAomiRuntime } from "@aomi-labs/react";

export function MessageList() {
  const { currentThreadId, getMessages, isRunning } = useAomiRuntime();
  const messages = getMessages(currentThreadId);

  return (
    <div className="flex-1 overflow-y-auto p-4 space-y-4">
      {messages.map((msg, i) => (
        <div
          key={i}
          className={`flex ${msg.role === "user" ? "justify-end" : "justify-start"}`}
        >
          <div
            className={`max-w-[80%] rounded-lg px-4 py-2 ${
              msg.role === "user"
                ? "bg-blue-600 text-white"
                : "bg-gray-100 text-gray-900"
            }`}
          >
            {msg.content.map((block, j) => (
              <span key={j}>
                {block.type === "text" ? block.text : null}
              </span>
            ))}
          </div>
        </div>
      ))}

      {isRunning && (
        <div className="flex justify-start">
          <div className="bg-gray-100 rounded-lg px-4 py-2 text-gray-500">
            Thinking...
          </div>
        </div>
      )}
    </div>
  );
}

Chat Input Component

"use client";

import { useState } from "react";
import { useAomiRuntime } from "@aomi-labs/react";

export function ChatInput() {
  const { sendMessage, isRunning, cancelGeneration } = useAomiRuntime();
  const [input, setInput] = useState("");

  const handleSubmit = async (e: React.FormEvent) => {
    e.preventDefault();
    if (!input.trim() || isRunning) return;
    const text = input;
    setInput("");
    await sendMessage(text);
  };

  return (
    <form onSubmit={handleSubmit} className="border-t p-4 flex gap-2">
      <input
        type="text"
        value={input}
        onChange={(e) => setInput(e.target.value)}
        placeholder="Type a message..."
        className="flex-1 rounded-lg border px-4 py-2"
        disabled={isRunning}
      />
      {isRunning ? (
        <button
          type="button"
          onClick={cancelGeneration}
          className="rounded-lg bg-red-500 px-4 py-2 text-white"
        >
          Stop
        </button>
      ) : (
        <button
          type="submit"
          disabled={!input.trim()}
          className="rounded-lg bg-blue-600 px-4 py-2 text-white disabled:opacity-50"
        >
          Send
        </button>
      )}
    </form>
  );
}

Thread Switcher

Manage multiple conversation threads:

"use client";

import { useAomiRuntime } from "@aomi-labs/react";

export function ThreadSwitcher() {
  const { currentThreadId, threadMetadata, selectThread, createThread, deleteThread } =
    useAomiRuntime();

  const threads = Array.from(threadMetadata.entries());

  return (
    <div className="w-64 border-r h-full flex flex-col">
      <div className="p-3 border-b">
        <button
          onClick={() => createThread()}
          className="w-full rounded-lg bg-gray-100 px-3 py-2 text-sm hover:bg-gray-200"
        >
          + New Thread
        </button>
      </div>

      <div className="flex-1 overflow-y-auto">
        {threads.map(([id, meta]) => (
          <div
            key={id}
            onClick={() => selectThread(id)}
            className={`flex items-center justify-between px-3 py-2 cursor-pointer text-sm ${
              id === currentThreadId
                ? "bg-blue-50 text-blue-700"
                : "hover:bg-gray-50"
            }`}
          >
            <span className="truncate">{meta.title || "Untitled"}</span>
            <button
              onClick={(e) => {
                e.stopPropagation();
                deleteThread(id);
              }}
              className="text-gray-400 hover:text-red-500 text-xs"
            >
              x
            </button>
          </div>
        ))}
      </div>
    </div>
  );
}

Compose Everything

"use client";

import { MessageList } from "@/components/message-list";
import { ChatInput } from "@/components/chat-input";
import { ThreadSwitcher } from "@/components/thread-switcher";

export default function ChatPage() {
  return (
    <div className="flex h-screen">
      <ThreadSwitcher />
      <div className="flex flex-1 flex-col">
        <MessageList />
        <ChatInput />
      </div>
    </div>
  );
}

Streaming

The isRunning flag is true while the assistant generates. Messages update in real time as tokens stream in. For granular control, subscribe to SSE events:

import { useEffect } from "react";
import { useEventContext } from "@aomi-labs/react";

function StreamMonitor() {
  const { subscribe, sseStatus } = useEventContext();

  useEffect(() => {
    const unsub = subscribe("streaming_text", (event) => {
      console.log("Streaming chunk:", event.payload);
    });
    return unsub;
  }, [subscribe]);

  return <div>SSE: {sseStatus}</div>;
}

Key Principle

The headless library provides all the state and actions. Your components are purely presentational. There are no layout constraints, no CSS to override, and no component hierarchy to follow. You decide how messages look, where threads are listed, and how controls are arranged.

Next Steps

• Hooks Reference: complete hook API
• Runtime Provider: provider setup
• Custom Tools: extend with your own tools