Custom Tool UI

当默认的 ToolFallback 不能满足需求时，可以通过 Thread 的 components prop 自定义工具 UI。

场景

假设后端有一个 getWeather 工具调用，我们希望在聊天界面中以卡片形式展示天气信息，而不是默认的折叠面板。

后端定义

app/api/chat/route.ts

import { openai } from "@ai-sdk/openai";
import { streamText, tool } from "ai";
import { z } from "zod";

export async function POST(req: Request) {
  const { messages } = await req.json();

  const result = streamText({
    model: openai("gpt-4o"),
    messages,
    tools: {
      getWeather: tool({
        description: "获取指定城市的天气信息",
        parameters: z.object({
          city: z.string().describe("城市名称"),
        }),
        execute: async ({ city }) => {
          // 模拟天气数据
          return {
            city,
            temperature: 24,
            condition: "晴",
            humidity: 65,
            windSpeed: "3级",
          };
        },
      }),
    },
  });

  return result.toUIMessageStreamResponse();
}

自定义 ToolUI 组件

components/weather-tool-ui.tsx

"use client";

import type { ToolInvocation } from "ai";

// 定义天气数据的类型
interface IWeatherResult {
  city: string;
  temperature: number;
  condition: string;
  humidity: number;
  windSpeed: string;
}

export function WeatherToolUI({ part }: { part: ToolInvocation }) {
  // 工具调用中的 "result" 是执行结果
  const result = part.result as IWeatherResult;

  // 运行中状态
  if (part.state === "call") {
    return (
      <div className="rounded-lg border bg-card p-4">
        <div className="flex items-center gap-2 text-muted-foreground">
          <span className="animate-spin">&#9696;</span>
          <span>正在查询天气...</span>
        </div>
      </div>
    );
  }

  // 错误状态
  if (part.state === "result" && "error" in result) {
    return (
      <div className="rounded-lg border border-destructive/50 bg-destructive/10 p-4">
        <p className="text-destructive text-sm">天气查询失败</p>
      </div>
    );
  }

  // 成功状态 — 展示天气卡片
  return (
    <div className="rounded-lg border bg-gradient-to-br from-sky-50 to-blue-50 p-4 dark:from-sky-950 dark:to-blue-950">
      <div className="flex items-center justify-between">
        <div>
          <h3 className="text-lg font-semibold">{result.city}</h3>
          <p className="text-sm text-muted-foreground">{result.condition}</p>
        </div>
        <div className="text-right">
          <p className="text-3xl font-bold">{result.temperature}°C</p>
        </div>
      </div>
      {/* 详细信息 */}
      <div className="mt-3 flex gap-4 border-t pt-3">
        <div>
          <p className="text-xs text-muted-foreground">湿度</p>
          <p className="text-sm font-medium">{result.humidity}%</p>
        </div>
        <div>
          <p className="text-xs text-muted-foreground">风速</p>
          <p className="text-sm font-medium">{result.windSpeed}</p>
        </div>
      </div>
    </div>
  );
}

注册工具 UI

通过 Thread 的 components prop，将 WeatherToolUI 注册为 getWeather 工具的自定义 UI：

components/chat.tsx

"use client";

import { AssistantRuntimeProvider, useChatRuntime } from "@assistant-ui/react";
import { Thread } from "aui-components";
import { WeatherToolUI } from "./weather-tool-ui";

export function Chat() {
  const runtime = useChatRuntime({
    api: "/api/chat",
  });

  return (
    <AssistantRuntimeProvider aui={runtime.aui} runtime={runtime}>
      <div className="h-screen">
        <Thread
          components={{
            // 注意：key 必须和 tool name 完全一致
            getWeather: WeatherToolUI,
          }}
        />
      </div>
    </AssistantRuntimeProvider>
  );
}

效果

当模型调用 getWeather 工具时，聊天界面不再显示默认的 ToolFallback 折叠面板，而是直接渲染我们自定义的天气卡片。

多个工具 UI

如果后端有多个工具，可以为每个工具注册独立的 UI 组件：

<Thread
  components={{
    getWeather: WeatherToolUI,
    searchWeb: SearchToolUI,
    createReminder: ReminderToolUI,
  }}
/>

对于没有注册自定义 UI 的工具，会回退到默认的 ToolFallback 组件。

Custom Tool UI

场景

后端定义

自定义 ToolUI 组件

注册工具 UI

效果

多个工具 UI

相关链接

On this page