feat: dynamic model discovery from vercel ai gateway (#1353)

lib/ai/entitlements.ts

@@ -1,9 +1,7 @@
 import type { UserType } from "@/app/(auth)/auth";
-import type { ChatModel } from "./models";
 
 type Entitlements = {
   maxMessagesPerDay: number;
-  availableChatModelIds: ChatModel["id"][];
 };
 
 export const entitlementsByUserType: Record<UserType, Entitlements> = {
@@ -11,16 +9,14 @@ export const entitlementsByUserType: Record<UserType, Entitlements> = {
    * For users without an account
    */
   guest: {
-    maxMessagesPerDay: 20,
-    availableChatModelIds: ["chat-model", "chat-model-reasoning"],
+    maxMessagesPerDay: 10,
   },
 
   /*
    * For users with an account
    */
   regular: {
-    maxMessagesPerDay: 100,
-    availableChatModelIds: ["chat-model", "chat-model-reasoning"],
+    maxMessagesPerDay: 50,
   },
 
   /*

lib/ai/models.mock.ts

@@ -1,5 +1,28 @@
 import type { LanguageModel } from "ai";
 
+const mockResponses: Record<string, string> = {
+  default: "This is a mock response for testing.",
+  weather: "The weather in San Francisco is sunny and 72°F.",
+  greeting: "Hello! How can I help you today?",
+};
+
+function getResponseForPrompt(prompt: unknown): string {
+  const promptStr = JSON.stringify(prompt).toLowerCase();
+
+  if (promptStr.includes("weather") || promptStr.includes("temperature")) {
+    return mockResponses.weather;
+  }
+  if (
+    promptStr.includes("hello") ||
+    promptStr.includes("hi") ||
+    promptStr.includes("hey")
+  ) {
+    return mockResponses.greeting;
+  }
+
+  return mockResponses.default;
+}
+
 const createMockModel = (): LanguageModel => {
   return {
     specificationVersion: "v2",
@@ -9,20 +32,116 @@ const createMockModel = (): LanguageModel => {
     supportedUrls: [],
     supportsImageUrls: false,
     supportsStructuredOutputs: false,
+    doGenerate: async ({ prompt }: { prompt: unknown }) => ({
+      rawCall: { rawPrompt: null, rawSettings: {} },
+      finishReason: "stop",
+      usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 },
+      content: [{ type: "text", text: getResponseForPrompt(prompt) }],
+      warnings: [],
+    }),
+    doStream: ({ prompt }: { prompt: unknown }) => {
+      const response = getResponseForPrompt(prompt);
+      const words = response.split(" ");
+
+      return {
+        stream: new ReadableStream({
+          async start(controller) {
+            for (const word of words) {
+              controller.enqueue({
+                type: "text-delta",
+                textDelta: `${word} `,
+              });
+              await new Promise((resolve) => {
+                setTimeout(resolve, 10);
+              });
+            }
+            controller.enqueue({
+              type: "finish",
+              finishReason: "stop",
+              usage: { inputTokens: 10, outputTokens: 20 },
+            });
+            controller.close();
+          },
+        }),
+        rawCall: { rawPrompt: null, rawSettings: {} },
+      };
+    },
+  } as unknown as LanguageModel;
+};
+
+const createMockReasoningModel = (): LanguageModel => {
+  return {
+    specificationVersion: "v2",
+    provider: "mock",
+    modelId: "mock-reasoning-model",
+    defaultObjectGenerationMode: "tool",
+    supportedUrls: [],
+    supportsImageUrls: false,
+    supportsStructuredOutputs: false,
     doGenerate: async () => ({
       rawCall: { rawPrompt: null, rawSettings: {} },
       finishReason: "stop",
       usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 },
-      content: [{ type: "text", text: "Hello, world!" }],
+      content: [{ type: "text", text: "This is a reasoned response." }],
+      reasoning: [
+        { type: "text", text: "Let me think through this step by step..." },
+      ],
       warnings: [],
     }),
-    doStream: async () => ({
+    doStream: () => ({
+      stream: new ReadableStream({
+        async start(controller) {
+          controller.enqueue({
+            type: "reasoning",
+            textDelta: "Let me think through this step by step... ",
+          });
+          await new Promise((resolve) => {
+            setTimeout(resolve, 10);
+          });
+          controller.enqueue({
+            type: "text-delta",
+            textDelta: "This is a reasoned response.",
+          });
+          controller.enqueue({
+            type: "finish",
+            finishReason: "stop",
+            usage: { inputTokens: 10, outputTokens: 20 },
+          });
+          controller.close();
+        },
+      }),
+      rawCall: { rawPrompt: null, rawSettings: {} },
+    }),
+  } as unknown as LanguageModel;
+};
+
+const createMockTitleModel = (): LanguageModel => {
+  return {
+    specificationVersion: "v2",
+    provider: "mock",
+    modelId: "mock-title-model",
+    defaultObjectGenerationMode: "tool",
+    supportedUrls: [],
+    supportsImageUrls: false,
+    supportsStructuredOutputs: false,
+    doGenerate: async () => ({
+      rawCall: { rawPrompt: null, rawSettings: {} },
+      finishReason: "stop",
+      usage: { inputTokens: 5, outputTokens: 5, totalTokens: 10 },
+      content: [{ type: "text", text: "Test Conversation" }],
+      warnings: [],
+    }),
+    doStream: () => ({
       stream: new ReadableStream({
         start(controller) {
           controller.enqueue({
             type: "text-delta",
-            id: "mock-id",
-            delta: "Mock response",
+            textDelta: "Test Conversation",
+          });
+          controller.enqueue({
+            type: "finish",
+            finishReason: "stop",
+            usage: { inputTokens: 5, outputTokens: 5 },
           });
           controller.close();
         },
@@ -33,6 +152,6 @@ const createMockModel = (): LanguageModel => {
 };
 
 export const chatModel = createMockModel();
-export const reasoningModel = createMockModel();
-export const titleModel = createMockModel();
+export const reasoningModel = createMockReasoningModel();
+export const titleModel = createMockTitleModel();
 export const artifactModel = createMockModel();

lib/ai/models.ts

@@ -1,21 +1,89 @@
-export const DEFAULT_CHAT_MODEL: string = "chat-model";
+// Curated list of top models from Vercel AI Gateway
+export const DEFAULT_CHAT_MODEL = "google/gemini-2.5-flash-lite";
 
 export type ChatModel = {
   id: string;
   name: string;
+  provider: string;
   description: string;
 };
 
 export const chatModels: ChatModel[] = [
+  // Anthropic
   {
-    id: "chat-model",
-    name: "Grok Vision",
-    description: "Advanced multimodal model with vision and text capabilities",
+    id: "anthropic/claude-haiku-4.5",
+    name: "Claude Haiku 4.5",
+    provider: "anthropic",
+    description: "Fast and affordable, great for everyday tasks",
   },
   {
-    id: "chat-model-reasoning",
-    name: "Grok Reasoning",
-    description:
-      "Uses advanced chain-of-thought reasoning for complex problems",
+    id: "anthropic/claude-sonnet-4.5",
+    name: "Claude Sonnet 4.5",
+    provider: "anthropic",
+    description: "Best balance of speed, intelligence, and cost",
+  },
+  {
+    id: "anthropic/claude-opus-4.5",
+    name: "Claude Opus 4.5",
+    provider: "anthropic",
+    description: "Most capable Anthropic model",
+  },
+  // OpenAI
+  {
+    id: "openai/gpt-4.1-mini",
+    name: "GPT-4.1 Mini",
+    provider: "openai",
+    description: "Fast and cost-effective for simple tasks",
+  },
+  {
+    id: "openai/gpt-5.2",
+    name: "GPT-5.2",
+    provider: "openai",
+    description: "Most capable OpenAI model",
+  },
+  // Google
+  {
+    id: "google/gemini-2.5-flash-lite",
+    name: "Gemini 2.5 Flash Lite",
+    provider: "google",
+    description: "Ultra fast and affordable",
+  },
+  {
+    id: "google/gemini-3-pro-preview",
+    name: "Gemini 3 Pro",
+    provider: "google",
+    description: "Most capable Google model",
+  },
+  // xAI
+  {
+    id: "xai/grok-4.1-fast-non-reasoning",
+    name: "Grok 4.1 Fast",
+    provider: "xai",
+    description: "Fast with 30K context",
+  },
+  // Reasoning models (extended thinking)
+  {
+    id: "anthropic/claude-3.7-sonnet-thinking",
+    name: "Claude 3.7 Sonnet",
+    provider: "reasoning",
+    description: "Extended thinking for complex problems",
+  },
+  {
+    id: "xai/grok-code-fast-1-thinking",
+    name: "Grok Code Fast",
+    provider: "reasoning",
+    description: "Reasoning optimized for code",
   },
 ];
+
+// Group models by provider for UI
+export const modelsByProvider = chatModels.reduce(
+  (acc, model) => {
+    if (!acc[model.provider]) {
+      acc[model.provider] = [];
+    }
+    acc[model.provider].push(model);
+    return acc;
+  },
+  {} as Record<string, ChatModel[]>
+);

lib/ai/prompts.ts

@@ -59,7 +59,11 @@ export const systemPrompt = ({
 }) => {
   const requestPrompt = getRequestPromptFromHints(requestHints);
 
-  if (selectedChatModel === "chat-model-reasoning") {
+  // reasoning models don't need artifacts prompt (they can't use tools)
+  if (
+    selectedChatModel.includes("reasoning") ||
+    selectedChatModel.includes("thinking")
+  ) {
     return `${regularPrompt}\n\n${requestPrompt}`;
   }
 

lib/ai/providers.ts

@@ -6,6 +6,8 @@ import {
 } from "ai";
 import { isTestEnvironment } from "../constants";
 
+const THINKING_SUFFIX_REGEX = /-thinking$/;
+
 export const myProvider = isTestEnvironment
   ? (() => {
       const {
@@ -23,14 +25,38 @@ export const myProvider = isTestEnvironment
         },
       });
     })()
-  : customProvider({
-      languageModels: {
-        "chat-model": gateway.languageModel("xai/grok-2-vision-1212"),
-        "chat-model-reasoning": wrapLanguageModel({
-          model: gateway.languageModel("xai/grok-3-mini"),
-          middleware: extractReasoningMiddleware({ tagName: "think" }),
-        }),
-        "title-model": gateway.languageModel("xai/grok-2-1212"),
-        "artifact-model": gateway.languageModel("xai/grok-2-1212"),
-      },
+  : null;
+
+export function getLanguageModel(modelId: string) {
+  if (isTestEnvironment && myProvider) {
+    return myProvider.languageModel(modelId);
+  }
+
+  const isReasoningModel =
+    modelId.includes("reasoning") || modelId.endsWith("-thinking");
+
+  if (isReasoningModel) {
+    const gatewayModelId = modelId.replace(THINKING_SUFFIX_REGEX, "");
+
+    return wrapLanguageModel({
+      model: gateway.languageModel(gatewayModelId),
+      middleware: extractReasoningMiddleware({ tagName: "thinking" }),
     });
+  }
+
+  return gateway.languageModel(modelId);
+}
+
+export function getTitleModel() {
+  if (isTestEnvironment && myProvider) {
+    return myProvider.languageModel("title-model");
+  }
+  return gateway.languageModel("anthropic/claude-haiku-4.5");
+}
+
+export function getArtifactModel() {
+  if (isTestEnvironment && myProvider) {
+    return myProvider.languageModel("artifact-model");
+  }
+  return gateway.languageModel("anthropic/claude-haiku-4.5");
+}

lib/ai/tools/request-suggestions.ts

@@ -5,7 +5,7 @@ import { getDocumentById, saveSuggestions } from "@/lib/db/queries";
 import type { Suggestion } from "@/lib/db/schema";
 import type { ChatMessage } from "@/lib/types";
 import { generateUUID } from "@/lib/utils";
-import { myProvider } from "../providers";
+import { getArtifactModel } from "../providers";
 
 type RequestSuggestionsProps = {
   session: Session;
@@ -38,7 +38,7 @@ export const requestSuggestions = ({
       >[] = [];
 
       const { elementStream } = streamObject({
-        model: myProvider.languageModel("artifact-model"),
+        model: getArtifactModel(),
         system:
           "You are a help writing assistant. Given a piece of writing, please offer suggestions to improve the piece of writing and describe the change. It is very important for the edits to contain full sentences instead of just words. Max 5 suggestions.",
         prompt: document.content,