feat: Enable dynamic AI model selection and tool filtering in chat interactions.

Author: arnaudpfu

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

@@ -1,22 +1,28 @@
 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";
 
 const groq = createOpenAI({
   apiKey: process.env.GROQ_API_KEY ?? '',
   baseURL: 'https://api.groq.com/openai/v1',
 });
 
 export async function POST(req: Request) {
-  const { messages }: { messages: UIMessage[] } = await req.json();
+  const { messages, ...body }: { messages: UIMessage[] } = await req.json();
+
+  const model = groq(getModel(body) ?? 'llama-3.3-70b-versatile')
+  const toolsFilter = getDisabledToolsFiter(body);
 
   const result = streamText({
     // model: groq('llama-3.3-70b-versatile'),
-    model: groq('meta-llama/llama-4-scout-17b-16e-instruct'),
+    // model: groq('meta-llama/llama-4-scout-17b-16e-instruct'),
+    // model: groq('meta-llama/llama-4-scout-17b-16e-instruct'),
+    model,
     messages: await convertToModelMessages(messages),
-    tools: {
+    tools: toolsFilter({
       weather: weatherTool,
-    },
+    }),
     experimental_transform: smoothStream({
       delayInMs: 30, // optional: defaults to 10ms
       chunking: 'line', // optional: defaults to 'word'

packages-test-3/ai-chat/package.json

@@ -23,6 +23,7 @@
     "./hook-types": "./src/hook-types.ts",
     "./utils": "./src/utils/index.ts",
     "./hooks": "./src/hooks/index.ts",
+    "./server": "./src/server/index.ts",
     "./types/*": "./src/types/*.ts"
   },
   "main": "./dist/index.js",

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

@@ -38,6 +38,8 @@ export type AiContextType = {
     toolkit?: Toolkit | undefined;
     callSettings?: LanguageModelV1CallSettings | undefined;
     config?: LanguageModelConfig | undefined;
+    selectedModel: string | null;
+    disabledTools: string[];
     chatOptions?: Omit<UseChatOptions<Thread['messages'][0]> & ChatInit<Thread['messages'][0]>, 'id' | 'transport'> & {
         transportOptions?: HttpChatTransportInitOptions<Thread['messages'][0]>
     }
@@ -68,11 +70,16 @@ export const useAiEvent = <TEvent extends keyof AiChatEvents>(name: TEvent, p: (
     }, [eventHandler, name, p])
 };
 
-export function AiProvider({ children, ...value }: { children: React.ReactNode } & Omit<AiContextType, 'eventHandler'>) {
+export function AiProvider({ children, ...value }: { children: React.ReactNode } & Omit<AiContextType, 'eventHandler' | 'selectedModel' | 'disabledTools'> & Partial<Pick<AiContextType, 'selectedModel' | 'disabledTools'>>) {
     // Create store once
     const storeRef = useRef<StoreApi<AiContextType> | null>(null);
     if (storeRef.current === null) {
-        storeRef.current = createStore<AiContextType>(() => ({ ...value, eventHandler: new AiChatEventHandler() }));
+        storeRef.current = createStore<AiContextType>(() => ({
+            ...value,
+            selectedModel: value.selectedModel ?? null,
+            disabledTools: value.disabledTools ?? [],
+            eventHandler: new AiChatEventHandler(),
+        }));
     }
 
     return <AiContextStoreCtx.Provider value={storeRef.current}>

packages-test-3/ai-chat/src/primitives/composer/composer-select-model.tsx

@@ -0,0 +1,28 @@
+"use client";
+
+import { useCallback } from "react";
+import { useAiContextStore } from "../../ai-provider";
+import { ActionButtonElement, ActionButtonProps, createActionButton } from "../../utils/create-action-button";
+
+type SelectModelProps = {
+  model: string;
+};
+
+const useComposerSelectModel = ({ model }: SelectModelProps) => {
+  const aiContextStore = useAiContextStore();
+
+  return useCallback(() => {
+    aiContextStore.setState({ selectedModel: model });
+  }, [aiContextStore, model]);
+};
+
+export namespace ComposerPrimitiveSelectModel {
+  export type Element = ActionButtonElement;
+  export type Props = ActionButtonProps<typeof useComposerSelectModel>;
+}
+
+export const ComposerPrimitiveSelectModel = createActionButton(
+  "ComposerPrimitive.SelectModel",
+  useComposerSelectModel,
+  ["model"],
+);

packages-test-3/ai-chat/src/primitives/composer/index.ts

@@ -2,5 +2,6 @@ export { ComposerPrimitiveRoot as Root, useComposer } from "./composer-provider"
 export { ComposerPrimitiveInput as Input } from "./composer-input";
 export { ComposerPrimitiveSend as Send,useComposerSend } from "./composer-send";
 export { ComposerPrimitiveCancel as Cancel } from "./composer-cancel";
+export { ComposerPrimitiveSelectModel as SelectModel } from "./composer-select-model";
 export { ComposerPrimitiveAttachments as Attachments, ComposerPrimitiveAttachmentByIndex as AttachmentByIndex } from "./composer-attachments";
-export * from './composer-provider'
+export * from './composer-provider'

packages-test-3/ai-chat/src/primitives/thread/thread-root.tsx

@@ -21,6 +21,7 @@ import { ToolCallMessagePartComponent } from "../../types/message-part-component
 import { Unsubscribe } from "../../types/unsuscribe";
 import { toToolsJSONSchema } from "../../stream/schema-utils";
 import { Tool } from "../../stream/tool-types";
+import { BODY_KEY } from "../../utils/request-keys";
 
 export type CustomUIDataTypes = {
     textDelta: string;
@@ -190,6 +191,8 @@ export function ThreadPrimitiveRoot({ children, ...value }: { children: React.Re
     const system = useAiContext(s => s.system);
     const callSettings = useAiContext(s => s.callSettings);
     const config = useAiContext(s => s.config);
+    const selectedModel = useAiContext(s => s.selectedModel);
+    const disabledTools = useAiContext(s => s.disabledTools);
     const activeThreadId = useThreads(s => s.activeThreadId);
     const [title, setTitle] = useState('New thread');
     const [status, setStatus] = useState<Thread['status']>('regular');
@@ -228,6 +231,8 @@ export function ThreadPrimitiveRoot({ children, ...value }: { children: React.Re
         system?: string;
         callSettings?: unknown;
         config?: unknown;
+        selectedModel?: string;
+        disabledTools?: string[];
     }>({ tools: {} });
     const addToolOutputRef = useRef<
         null | ((
@@ -259,6 +264,8 @@ export function ThreadPrimitiveRoot({ children, ...value }: { children: React.Re
                         ...(context.system ? { system: context.system } : {}),
                         ...(context.callSettings ? { callSettings: context.callSettings } : {}),
                         ...(context.config ? { config: context.config } : {}),
+                        [BODY_KEY.SELECTED_MODEL]: context.selectedModel ?? "",
+                        [BODY_KEY.DISABLED_TOOLS]: context.disabledTools ?? [],
                         tools: toToolsJSONSchema(Object.values(toolsRef.current) as any) as any,
                     },
                 };
@@ -649,6 +656,8 @@ export function ThreadPrimitiveRoot({ children, ...value }: { children: React.Re
             system,
             callSettings,
             config,
+            selectedModel: selectedModel ?? undefined,
+            disabledTools: disabledTools ?? [],
         };
     }, []);
 
@@ -661,10 +670,12 @@ export function ThreadPrimitiveRoot({ children, ...value }: { children: React.Re
                 system,
                 callSettings,
                 config,
+                selectedModel: selectedModel ?? undefined,
+                disabledTools: disabledTools ?? [],
             };
         });
         return unsubscribe;
-    }, [system, callSettings, config]);
+    }, [system, callSettings, config, selectedModel, disabledTools]);
 
     useEffect(() => {
         if (!storeRef.current) return;

packages-test-3/ai-chat/src/server/index.ts

@@ -0,0 +1,22 @@
+import z from "zod";
+import { BODY_KEY } from "../utils/request-keys";
+import type { streamText, Tool } from "ai";
+
+export type ModelResolver = (model?: string) => string | null;
+
+const defaultModelResolver: ModelResolver = (model) => model ?? null;
+
+export const getModel =
+  (body: Record<string, unknown>, resolveModel: ModelResolver = defaultModelResolver): string | null => {
+    const model = z.string().nullable().optional().safeParse(body?.[BODY_KEY.SELECTED_MODEL]).data || undefined;
+    return resolveModel(model);
+  };
+
+export const getDisabledToolsFiter = (body: Record<string, unknown>): ((tools: Parameters<typeof streamText>[0]['tools']) => Parameters<typeof streamText>[0]['tools']) => {
+  const disabledTools = z.array(z.string()).nullable().optional().safeParse(body?.[BODY_KEY.DISABLED_TOOLS]).data || [];
+  return (serverTools: Parameters<typeof streamText>[0]['tools']) => Object.fromEntries(
+    Object.entries(serverTools as Record<string, Tool<any, any>>).filter(([toolName]) => !disabledTools.includes(toolName)),
+  )
+};
+
+export { BODY_KEY } from "../utils/request-keys";

packages-test-3/ai-chat/src/utils/request-keys.ts

@@ -0,0 +1,4 @@
+export const BODY_KEY = {
+    SELECTED_MODEL: "selected-model",
+    DISABLED_TOOLS: "disabled-tools",
+}