feat: Implement a new tool integration system with dedicated UI components and server-side utilities, replacing the previous assistant tool and model selection mechanisms.

Author: arnaudpfu

apps/mobile/components/ai-chat/thread.tsx

@@ -1,5 +1,5 @@
 import { FC, useCallback, type ReactNode } from "react";
-import { View } from "react-native";
+import { Pressable, View } from "react-native";
 
 import { Icon, IconName } from "../ui/icon";
 import { Text } from "../ui/text";
@@ -299,6 +299,37 @@ const ModelSelector: React.FC = () => {
   );
 };
 
+const ToolsSelector: React.FC = () => {
+	return (
+    <>
+      <ComposerPrimitive.SelectTools 
+        header={
+          <>
+            <Divider />
+            <Text className="font-bold">Tools</Text>
+          </>
+        }
+        toolUI={({ toolName, toggle, isEnabled }) => (
+          <Pressable className="flex-1 active:opacity-50" onPress={toggle}>
+            <View className="flex-1 p-2">
+              <View className="mb-1 flex flex-row gap-2 text-lg">
+                <Text className={cn("font-bold", isEnabled ? "text-cyan-500" : "")}>{toolName.replace(/-/g, ' ').replace(/\b\w/g, c => c.toUpperCase())}</Text>
+                <Icon name="Brain" size={20} />
+              </View>
+              <Text className={cn("text-muted-foreground text-sm", isEnabled ? "text-cyan-500" : "")}>
+                Groundbreaking advancements in reasoning
+              </Text>
+            </View>
+            {isEnabled && (
+              <Icon name="Check" size={20} color={isEnabled ? "#00b8db" : undefined} />
+            )}
+            
+          </Pressable>
+      )} />
+    </>
+	);
+};
+
 const ComposerAction: FC = () => {
   const textMutedForeground = useCSSVariable("--color-muted-foreground");
   const borderColor = useCSSVariable("--color-border");
@@ -322,11 +353,8 @@ const ComposerAction: FC = () => {
           <ComposerAddAttachmentSheet />
           <Divider />
           <ModelSelector />
+          <ToolsSelector />
         </View>
-        {/* <View className="h-40 px-4">
-          <ComposerAddAttachmentFile />
-          <ComposerAddAttachmentImage />
-        </View> */}
       </ActionSheetContent>
     </ActionSheet>
   );

apps/mobile/components/tools/weather-tool-ui.tsx

@@ -1,9 +1,10 @@
 import { FC, useMemo } from "react";
 import { View } from "react-native";
-import { useAssistantToolUI } from "@creatorem/ai-chat";
+import { useToolUI } from "@creatorem/ai-chat";
 import type { ToolCallMessagePartProps } from "@creatorem/ai-chat/types/message-part-component-types";
 import { Text } from "../ui/text";
 import { cn } from "~/lib/cn";
+import { IconName } from "../ui/icon";
 
 const HOUR_LABELS = ["1PM", "2PM", "3PM", "4PM", "5PM"] as const;
 
@@ -121,8 +122,13 @@ const WeatherToolCard: FC<ToolCallMessagePartProps> = ({
 };
 
 export const WeatherToolRegistration: FC = () => {
-  useAssistantToolUI({
+  useToolUI({
     toolName: "weather",
+    display: {
+      icon: 'CloudSun' as IconName,
+      title: 'Weather',
+      description: 'Get the weather forecast for a location',
+    },
     render: WeatherToolCard,
   });
   return null;

apps/nextjs-example/app/api/chat/route.ts

@@ -1,7 +1,7 @@
 import { createOpenAI } from "@ai-sdk/openai";
 import { streamText, convertToModelMessages, type UIMessage, smoothStream } from "ai";
 import { weatherTool } from "../../../lib/tools/weather-tool";
-import { getModel, getDisabledToolsFiter } from "@creatorem/ai-chat/server";
+import { getServerUtils } from "@creatorem/ai-chat/server";
 import { DEFAULT_AI_MODEL } from "@/lib/ai-constants";
 
 const groq = createOpenAI({
@@ -11,14 +11,15 @@ const groq = createOpenAI({
 
 export async function POST(req: Request) {
   const { messages, ...body }: { messages: UIMessage[] } = await req.json();
+  const { getModel, toolsFilter, frontendTools } = getServerUtils(body);
 
-  const model = groq(getModel(body) ?? DEFAULT_AI_MODEL)
-  const toolsFilter = getDisabledToolsFiter(body);
+  const model = groq(getModel() ?? DEFAULT_AI_MODEL)
 
   const result = streamText({
     model,
     messages: await convertToModelMessages(messages),
     tools: toolsFilter({
+      ...frontendTools,
       weather: weatherTool,
     }),
     experimental_transform: smoothStream({

apps/nextjs-example/app/test2/page.tsx

@@ -16,6 +16,7 @@ import {
 	MoreHorizontalIcon,
 	PencilIcon,
 	RefreshCwIcon,
+	Settings2Icon,
 	SquareIcon,
 } from "lucide-react";
 import { TooltipIconButton } from "@/components/ai-chat/tooltip-icon-button";
@@ -62,8 +63,7 @@ import * as BranchPickerPrimitive from "@creatorem/ai-react/primitives/branch-pi
 import * as ThreadPrimitive from "@creatorem/ai-react/primitives/thread";
 import * as MessagePrimitive from "@creatorem/ai-react/primitives/message";
 import { AI_MODELS, DEFAULT_AI_MODEL } from "@/lib/ai-constants";
-import { useCallback } from "react";
-import { InputGroup } from "@/components/ui/input-group";
+import { Switch } from "@/components/ui/switch";
 
 export default function Chat() {
 	return (
@@ -321,11 +321,50 @@ const ModelSelector: React.FC = () => {
 	);
 };
 
+const ToolsSelector: React.FC = () => {
+	return (
+		<DropdownMenu>
+			<DropdownMenuTrigger asChild>
+				<Button variant="ghost">
+					<Settings2Icon className="size-4" />
+					Tools
+				</Button>
+			</DropdownMenuTrigger>
+			<DropdownMenuContent side="top" align="start" className="z-90 w-64 [--radius:0.95rem]">
+				<DropdownMenuLabel>AI Tools</DropdownMenuLabel>
+				<DropdownMenuSeparator />
+				<ComposerPrimitive.SelectTools empty={(
+					<DropdownMenuItem disabled>No tools available</DropdownMenuItem>
+				)} toolUI={({ toolName, toggle, isEnabled }) => (
+					<DropdownMenuItem
+						key={toolName}
+						onSelect={toggle}
+						className="flex cursor-pointer items-center justify-between gap-2"
+					>
+						<div className="flex flex-col items-start">
+							{toolName.replace(/-/g, ' ').replace(/\b\w/g, c => c.toUpperCase())}
+						</div>
+						{/* {isEnabled && (
+							<CheckIcon className="text-primary size-4" />
+						)} */}
+						<Switch
+							checked={isEnabled}
+							onCheckedChange={() => toggle()}
+							onClick={(e) => e.stopPropagation()}
+						/>
+					</DropdownMenuItem>
+				)} />
+			</DropdownMenuContent>
+		</DropdownMenu>
+	);
+};
+
 const ComposerAction: FC = () => {
 	return (
 		<div className="aui-composer-action-wrapper relative mx-2 mb-2 flex items-center justify-between">
 			<div className="flex">
 				<ComposerAddAttachment />
+				<ToolsSelector />
 			</div>
 
 			<div className="flex gap-2">

apps/nextjs-example/app/test2/weather-tool-ui.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import { FC } from "react";
-import { useAssistantToolUI } from "@creatorem/ai-react";
+import { useToolUI } from "@creatorem/ai-react";
 import type { ToolCallMessagePartProps } from "@creatorem/ai-chat/types/message-part-component-types";
 
 const HOUR_LABELS = ["1PM", "2PM", "3PM", "4PM", "5PM"] as const;
@@ -177,7 +177,7 @@ const WeatherToolCard: FC<ToolCallMessagePartProps> = ({
 };
 
 export const WeatherToolRegistration: FC = () => {
-  useAssistantToolUI({
+  useToolUI({
     toolName: "weather",
     render: WeatherToolCard,
   });

apps/nextjs-example/components/ui/switch.tsx

@@ -0,0 +1,26 @@
+"use client"
+
+import * as React from "react"
+import * as SwitchPrimitives from '@radix-ui/react-switch';
+
+import { cn } from "@/lib/utils"
+
+const Switch: React.FC<React.ComponentPropsWithoutRef<typeof SwitchPrimitives.Root>> = ({ className, ...props }) => (
+    <SwitchPrimitives.Root
+        className={cn(
+            'focus-visible:ring-ring focus-visible:ring-offset-background data-[state=checked]:bg-primary data-[state=unchecked]:bg-input peer inline-flex h-5 w-9 shrink-0 cursor-pointer items-center rounded-full border-2 border-transparent shadow-xs transition-colors focus-visible:ring-2 focus-visible:ring-offset-2 focus-visible:outline-hidden disabled:cursor-not-allowed disabled:opacity-50',
+            className,
+        )}
+        {...props}
+    >
+        <SwitchPrimitives.Thumb
+            className={cn(
+                'bg-background pointer-events-none block h-4 w-4 rounded-full shadow-lg ring-0 transition-transform data-[state=checked]:translate-x-4 data-[state=unchecked]:translate-x-0',
+            )}
+        />
+    </SwitchPrimitives.Root>
+);
+Switch.displayName = SwitchPrimitives.Root.displayName;
+
+export { Switch };
+

apps/nextjs-example/package.json

@@ -29,6 +29,7 @@
     "@radix-ui/react-dropdown-menu": "^2.1.16",
     "@radix-ui/react-separator": "^1.1.8",
     "@radix-ui/react-slot": "^1.2.4",
+    "@radix-ui/react-switch": "^1.2.6",
     "@radix-ui/react-tooltip": "^1.2.8",
     "ai": "^6.0.50",
     "class-variance-authority": "^0.7.1",

packages-test-3/ai-chat/src/ai-provider.tsx

@@ -6,9 +6,9 @@ import { createStore, useStore, type StoreApi } from 'zustand';
 import { ThreadAdapter, AttachmentAdapter, DictationAdapter } from "./types/adapters";
 import { emptyStorageThreadAdapter } from "./adapters/empty-storage-adapter";
 import { Thread, Threads } from "./types/entities";
-import { useChat, UseChatOptions } from "@ai-sdk/react";
+import { UseChatOptions } from "@ai-sdk/react";
 import { AiChatEventHandler, AiChatEvents } from "./primitives/events";
-import type { Toolkit } from "./types/toolbox";
+import { UITool } from "./types/tools-types";
 
 type LanguageModelV1CallSettings = {
     maxTokens?: number;
@@ -26,6 +26,7 @@ type LanguageModelConfig = {
     modelName?: string;
 };
 
+
 export type AiContextType = {
     adapters?: {
         attachment?: AttachmentAdapter;
@@ -34,18 +35,19 @@ export type AiContextType = {
     }
     priority?: number | undefined;
     system?: string | undefined;
-    tools?: Record<string, Tool<any, any>> | undefined;
-    toolkit?: Toolkit | undefined;
     callSettings?: LanguageModelV1CallSettings | undefined;
     config?: LanguageModelConfig | undefined;
     selectedModel: string | null;
     setSelectedModel: (selectedModel: string | null) => void;
-    disabledTools: string[];
-    setDisabledTools: (disabledTools: string[]) => void;
     chatOptions?: Omit<UseChatOptions<Thread['messages'][0]> & ChatInit<Thread['messages'][0]>, 'id' | 'transport'> & {
         transportOptions?: HttpChatTransportInitOptions<Thread['messages'][0]>
     }
     eventHandler: AiChatEventHandler
+    
+    tools: Record<string,UITool<any, any>>
+    upsertTool: (
+        tool: UITool<any, any>
+    ) => void;
 };
 
 const AiContextStoreCtx = React.createContext<StoreApi<AiContextType> | null>(null);
@@ -72,16 +74,24 @@ export const useAiEvent = <TEvent extends keyof AiChatEvents>(name: TEvent, p: (
     }, [eventHandler, name, p])
 };
 
-export function AiProvider({ children, ...value }: { children: React.ReactNode } & Omit<AiContextType, 'eventHandler' | 'selectedModel' | 'disabledTools' | 'setSelectedModel' | 'setDisabledTools'> & Partial<Pick<AiContextType, 'selectedModel' | 'disabledTools'>>) {
+export function AiProvider({ children, ...value }: { children: React.ReactNode } & Omit<AiContextType, 'eventHandler' | 'selectedModel' | 'tools' | 'upsertTool'>) {
     // Create store once
     const storeRef = useRef<StoreApi<AiContextType> | null>(null);
     if (storeRef.current === null) {
         storeRef.current = createStore<AiContextType>((set) => ({
             ...value,
-            selectedModel: value.selectedModel ?? null,
+            selectedModel: null,
             setSelectedModel: (selectedModel) => set({ selectedModel }),
-            disabledTools: value.disabledTools ?? [],
-            setDisabledTools: (disabledTools) => set({ disabledTools }),
+            tools: {},
+            upsertTool: (tool) => {
+                const toolName = tool.toolName;
+                set((prev) => ({
+                    tools: {
+                        ...prev.tools,
+                        [toolName]: tool,
+                    },
+                }));
+            },
             eventHandler: new AiChatEventHandler(),
         }));
     }

packages-test-3/ai-chat/src/model-context/index.ts

@@ -1,2 +1 @@
-export { useAssistantTool } from "./use-assistant-tool";
-export { useAssistantToolUI } from "./use-assistant-tool-ui";
+export { useToolUI } from "./use-tool-ui";