Building Custom UI

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

// 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>
  );
}

"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>
  );
}

"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>
  );
}

"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>
  );
}

"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>
  );
}

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>;
}