Improve chat navigation, model selection, and streaming

components/tools/registry.ts

@@ -0,0 +1,12 @@
+'use client';
+
+import type { ComponentType } from 'react';
+import { FirecrawlWebFetchTool } from './firecrawl-web-fetch-tool';
+import { FirecrawlWebSearchTool } from './firecrawl-web-search-tool';
+import type { ToolComponentProps } from './types';
+
+// Add a renderer here for each tool that needs a custom presentation.
+export const toolComponents: Record<string, ComponentType<ToolComponentProps> | undefined> = {
+  firecrawlWebSearch: FirecrawlWebSearchTool,
+  firecrawlWebFetch: FirecrawlWebFetchTool,
+};

components/tools/tool-part-renderer.tsx

@@ -0,0 +1,45 @@
+'use client';
+
+import { getToolName } from 'ai';
+import {
+  Tool,
+  ToolContent,
+  ToolHeader,
+  ToolInput,
+  ToolOutput,
+} from '@/components/ai-elements/tool';
+import { Loader } from '@/components/ai-elements/loader';
+import { toolComponents } from './registry';
+import type { ToolComponentProps } from './types';
+
+export function ToolPartRenderer({ part }: Pick<ToolComponentProps, 'part'>) {
+  const toolName = getToolName(part);
+  const CustomRenderer = toolComponents[toolName];
+
+  if (CustomRenderer) {
+    return <CustomRenderer part={part} toolName={toolName} />;
+  }
+
+  if (part.state === 'input-streaming' && !part.input) {
+    return (
+      <div className="mb-4 flex items-center gap-2 text-muted-foreground text-sm">
+        <Loader size={14} />
+        <span>Running {toolName}...</span>
+      </div>
+    );
+  }
+
+  return (
+    <Tool>
+      <ToolHeader
+        title={toolName}
+        type={part.type === 'dynamic-tool' ? `tool-${toolName}` : part.type}
+        state={part.state}
+      />
+      <ToolContent>
+        {part.input != null && <ToolInput input={part.input} />}
+        <ToolOutput output={part.output} errorText={part.errorText} />
+      </ToolContent>
+    </Tool>
+  );
+}

components/tools/types.ts

@@ -1,8 +1,6 @@
-import type { UIMessage } from "ai";
-
-type MessagePart = UIMessage["parts"][number];
+import type { DynamicToolUIPart, ToolUIPart } from "ai";
 
 export interface ToolComponentProps {
-  part: MessagePart;
+  part: ToolUIPart | DynamicToolUIPart;
   toolName: string;
 }