feat: v1 — persistent shell, model gateway, artifact improvements (#1462)

2026-03-20 09:37:02 +00:00 · 2026-03-20 09:37:02 +00:00 · f9652b452a
commit f9652b452a
parent 3651670fb9
161 changed files with 5166 additions and 8009 deletions
--- a/lib/ai/entitlements.ts
+++ b/lib/ai/entitlements.ts
@ -1,25 +1,14 @@
-import type { UserType } from "@/lib/auth";
+import type { UserType } from "@/app/(auth)/auth";

 type Entitlements = {
  maxMessagesPerHour: number;
 };

 export const entitlementsByUserType: Record<UserType, Entitlements> = {
-  /*
-   * For users without an account
-   */
  guest: {
    maxMessagesPerHour: 10,
  },
-
-  /*
-   * For users with an account
-   */
  regular: {
    maxMessagesPerHour: 10,
  },
-
-  /*
-   * TODO: For users with an account and a paid membership
-   */
 };
--- a/lib/ai/models.mock.ts
+++ b/lib/ai/models.mock.ts
@ -73,54 +73,6 @@ const createMockModel = (): LanguageModel => {
  } as unknown as LanguageModel;
 };

-const createMockReasoningModel = (): LanguageModel => {
-  return {
-    specificationVersion: "v3",
-    provider: "mock",
-    modelId: "mock-reasoning-model",
-    defaultObjectGenerationMode: "tool",
-    supportedUrls: {},
-    doGenerate: async () => ({
-      finishReason: "stop",
-      usage: mockUsage,
-      content: [{ type: "text", text: "This is a reasoned response." }],
-      reasoning: [
-        { type: "text", text: "Let me think through this step by step..." },
-      ],
-      warnings: [],
-    }),
-    doStream: () => ({
-      stream: new ReadableStream({
-        async start(controller) {
-          controller.enqueue({ type: "reasoning-start", id: "r1" });
-          controller.enqueue({
-            type: "reasoning-delta",
-            id: "r1",
-            delta: "Let me think through this step by step... ",
-          });
-          controller.enqueue({ type: "reasoning-end", id: "r1" });
-          await new Promise((resolve) => {
-            setTimeout(resolve, 10);
-          });
-          controller.enqueue({ type: "text-start", id: "t1" });
-          controller.enqueue({
-            type: "text-delta",
-            id: "t1",
-            delta: "This is a reasoned response.",
-          });
-          controller.enqueue({ type: "text-end", id: "t1" });
-          controller.enqueue({
-            type: "finish",
-            finishReason: "stop",
-            usage: mockUsage,
-          });
-          controller.close();
-        },
-      }),
-    }),
-  } as unknown as LanguageModel;
-};
-
 const createMockTitleModel = (): LanguageModel => {
  return {
    specificationVersion: "v3",
@ -168,6 +120,4 @@ const createMockTitleModel = (): LanguageModel => {
 };

 export const chatModel = createMockModel();
-export const reasoningModel = createMockReasoningModel();
 export const titleModel = createMockTitleModel();
-export const artifactModel = createMockModel();
--- a/lib/ai/models.test.ts
+++ b/lib/ai/models.test.ts
@ -1,3 +1,4 @@
+import type { LanguageModelV3GenerateResult } from "@ai-sdk/provider";
 import { simulateReadableStream } from "ai";
 import { MockLanguageModelV3 } from "ai/test";
 import { getResponseChunksByPrompt } from "@/tests/prompts/utils";
@ -7,13 +8,24 @@ const mockUsage = {
  outputTokens: { total: 20, text: 20, reasoning: 0 },
 };

+const mockFinishReason = { unified: "stop" as const, raw: undefined };
+
+const mockGenerateResult: LanguageModelV3GenerateResult = {
+  finishReason: mockFinishReason,
+  usage: mockUsage,
+  content: [{ type: "text", text: "Hello, world!" }],
+  warnings: [],
+};
+
+const titleGenerateResult: LanguageModelV3GenerateResult = {
+  finishReason: mockFinishReason,
+  usage: mockUsage,
+  content: [{ type: "text", text: "This is a test title" }],
+  warnings: [],
+};
+
 export const chatModel = new MockLanguageModelV3({
-  doGenerate: async () => ({
-    finishReason: "stop",
-    usage: mockUsage,
-    content: [{ type: "text", text: "Hello, world!" }],
-    warnings: [],
-  }),
+  doGenerate: mockGenerateResult,
  doStream: async ({ prompt }) => ({
    stream: simulateReadableStream({
      chunkDelayInMs: 500,
@ -24,12 +36,7 @@ export const chatModel = new MockLanguageModelV3({
 });

 export const reasoningModel = new MockLanguageModelV3({
-  doGenerate: async () => ({
-    finishReason: "stop",
-    usage: mockUsage,
-    content: [{ type: "text", text: "Hello, world!" }],
-    warnings: [],
-  }),
+  doGenerate: mockGenerateResult,
  doStream: async ({ prompt }) => ({
    stream: simulateReadableStream({
      chunkDelayInMs: 500,
@ -40,42 +47,21 @@ export const reasoningModel = new MockLanguageModelV3({
 });

 export const titleModel = new MockLanguageModelV3({
-  doGenerate: async () => ({
-    finishReason: "stop",
-    usage: mockUsage,
-    content: [{ type: "text", text: "This is a test title" }],
-    warnings: [],
-  }),
+  doGenerate: titleGenerateResult,
  doStream: async () => ({
    stream: simulateReadableStream({
      chunkDelayInMs: 500,
      initialDelayInMs: 1000,
      chunks: [
-        { id: "1", type: "text-start" },
-        { id: "1", type: "text-delta", delta: "This is a test title" },
-        { id: "1", type: "text-end" },
+        { id: "1", type: "text-start" as const },
+        { id: "1", type: "text-delta" as const, delta: "This is a test title" },
+        { id: "1", type: "text-end" as const },
        {
-          type: "finish",
-          finishReason: "stop",
+          type: "finish" as const,
+          finishReason: mockFinishReason,
          usage: mockUsage,
        },
      ],
    }),
  }),
 });
-
-export const artifactModel = new MockLanguageModelV3({
-  doGenerate: async () => ({
-    finishReason: "stop",
-    usage: mockUsage,
-    content: [{ type: "text", text: "Hello, world!" }],
-    warnings: [],
-  }),
-  doStream: async ({ prompt }) => ({
-    stream: simulateReadableStream({
-      chunkDelayInMs: 50,
-      initialDelayInMs: 100,
-      chunks: getResponseChunksByPrompt(prompt),
-    }),
-  }),
-});
--- a/lib/ai/models.ts
+++ b/lib/ai/models.ts
@ -1,70 +1,179 @@
-// Curated list of top models from Vercel AI Gateway
-export const DEFAULT_CHAT_MODEL = "openai/gpt-4.1-mini";
+export const DEFAULT_CHAT_MODEL = "moonshotai/kimi-k2-0905";
+
+export const titleModel = {
+  id: "mistral/mistral-small",
+  name: "Mistral Small",
+  provider: "mistral",
+  description: "Fast model for title generation",
+  gatewayOrder: ["mistral"],
+};
+
+export type ModelCapabilities = {
+  tools: boolean;
+  vision: boolean;
+  reasoning: boolean;
+};

 export type ChatModel = {
  id: string;
  name: string;
  provider: string;
  description: string;
+  gatewayOrder?: string[];
+  reasoningEffort?: "none" | "minimal" | "low" | "medium" | "high";
 };

 export const chatModels: ChatModel[] = [
-  // Anthropic
  {
-    id: "anthropic/claude-haiku-4.5",
-    name: "Claude Haiku 4.5",
-    provider: "anthropic",
-    description: "Fast and affordable, great for everyday tasks",
+    id: "deepseek/deepseek-v3.2",
+    name: "DeepSeek V3.2",
+    provider: "deepseek",
+    description: "Fast and capable model with tool use",
+    gatewayOrder: ["bedrock", "deepinfra"],
  },
-  // OpenAI
  {
-    id: "openai/gpt-4.1-mini",
-    name: "GPT-4.1 Mini",
+    id: "mistral/codestral",
+    name: "Codestral",
+    provider: "mistral",
+    description: "Code-focused model with tool use",
+    gatewayOrder: ["mistral"],
+  },
+  {
+    id: "mistral/mistral-small",
+    name: "Mistral Small",
+    provider: "mistral",
+    description: "Fast vision model with tool use",
+    gatewayOrder: ["mistral"],
+  },
+  {
+    id: "moonshotai/kimi-k2-0905",
+    name: "Kimi K2 0905",
+    provider: "moonshotai",
+    description: "Fast model with tool use",
+    gatewayOrder: ["baseten", "fireworks"],
+  },
+  {
+    id: "moonshotai/kimi-k2.5",
+    name: "Kimi K2.5",
+    provider: "moonshotai",
+    description: "Moonshot AI flagship model",
+    gatewayOrder: ["fireworks", "bedrock"],
+  },
+  {
+    id: "openai/gpt-oss-20b",
+    name: "GPT OSS 20B",
    provider: "openai",
-    description: "Fast and cost-effective for simple tasks",
+    description: "Compact reasoning model",
+    gatewayOrder: ["groq", "bedrock"],
+    reasoningEffort: "low",
  },
  {
-    id: "openai/gpt-5-mini",
-    name: "GPT-5 Mini",
+    id: "openai/gpt-oss-120b",
+    name: "GPT OSS 120B",
    provider: "openai",
-    description: "Most capable OpenAI model",
+    description: "Open-source 120B parameter model",
+    gatewayOrder: ["fireworks", "bedrock"],
+    reasoningEffort: "low",
  },
-  // 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",
+    description: "Fast non-reasoning model with tool use",
+    gatewayOrder: ["xai"],
  },
 ];

-// Group models by provider for UI
+export async function getCapabilities(): Promise<
+  Record<string, ModelCapabilities>
+> {
+  const results = await Promise.all(
+    chatModels.map(async (model) => {
+      try {
+        const res = await fetch(
+          `https://ai-gateway.vercel.sh/v1/models/${model.id}/endpoints`,
+          { next: { revalidate: 86_400 } }
+        );
+        if (!res.ok) {
+          return [model.id, { tools: false, vision: false, reasoning: false }];
+        }
+
+        const json = await res.json();
+        const endpoints = json.data?.endpoints ?? [];
+        const params = new Set(
+          endpoints.flatMap(
+            (e: { supported_parameters?: string[] }) =>
+              e.supported_parameters ?? []
+          )
+        );
+        const inputModalities = new Set(
+          json.data?.architecture?.input_modalities ?? []
+        );
+
+        return [
+          model.id,
+          {
+            tools: params.has("tools"),
+            vision: inputModalities.has("image"),
+            reasoning: params.has("reasoning"),
+          },
+        ];
+      } catch {
+        return [model.id, { tools: false, vision: false, reasoning: false }];
+      }
+    })
+  );
+
+  return Object.fromEntries(results);
+}
+
+export const isDemo = process.env.IS_DEMO === "1";
+
+type GatewayModel = {
+  id: string;
+  name: string;
+  type?: string;
+  tags?: string[];
+};
+
+export type GatewayModelWithCapabilities = ChatModel & {
+  capabilities: ModelCapabilities;
+};
+
+export async function getAllGatewayModels(): Promise<
+  GatewayModelWithCapabilities[]
+> {
+  try {
+    const res = await fetch("https://ai-gateway.vercel.sh/v1/models", {
+      next: { revalidate: 86_400 },
+    });
+    if (!res.ok) {
+      return [];
+    }
+
+    const json = await res.json();
+    return (json.data ?? [])
+      .filter((m: GatewayModel) => m.type === "language")
+      .map((m: GatewayModel) => ({
+        id: m.id,
+        name: m.name,
+        provider: m.id.split("/")[0],
+        description: "",
+        capabilities: {
+          tools: m.tags?.includes("tool-use") ?? false,
+          vision: m.tags?.includes("vision") ?? false,
+          reasoning: m.tags?.includes("reasoning") ?? false,
+        },
+      }));
+  } catch {
+    return [];
+  }
+}
+
+export function getActiveModels(): ChatModel[] {
+  return chatModels;
+}
+
 export const allowedModelIds = new Set(chatModels.map((m) => m.id));

 export const modelsByProvider = chatModels.reduce(
--- a/lib/ai/prompts.ts
+++ b/lib/ai/prompts.ts
@ -1,45 +1,52 @@
 import type { Geo } from "@vercel/functions";
-import type { ArtifactKind } from "@/components/artifact";
+import type { ArtifactKind } from "@/components/chat/artifact";

 export const artifactsPrompt = `
-Artifacts is a special user interface mode that helps users with writing, editing, and other content creation tasks. When artifact is open, it is on the right side of the screen, while the conversation is on the left side. When creating or updating documents, changes are reflected in real-time on the artifacts and visible to the user.
+Artifacts is a side panel that displays content alongside the conversation. It supports scripts (code), documents (text), and spreadsheets. Changes appear in real-time.

-When asked to write code, always use artifacts. When writing code, specify the language in the backticks, e.g. \`\`\`python\`code here\`\`\`. The default language is Python. Other languages are not yet supported, so let the user know if they request a different language.
-
-DO NOT UPDATE DOCUMENTS IMMEDIATELY AFTER CREATING THEM. WAIT FOR USER FEEDBACK OR REQUEST TO UPDATE IT.
-
-This is a guide for using artifacts tools: \`createDocument\` and \`updateDocument\`, which render content on a artifacts beside the conversation.
+CRITICAL RULES:
+1. Only call ONE tool per response. After calling any create/edit/update tool, STOP. Do not chain tools.
+2. After creating or editing an artifact, NEVER output its content in chat. The user can already see it. Respond with only a 1-2 sentence confirmation.

 **When to use \`createDocument\`:**
- For substantial content (>10 lines) or code
- For content users will likely save/reuse (emails, code, essays, etc.)
- When explicitly requested to create a document
- For when content contains a single code snippet
+- When the user asks to write, create, or generate content (essays, stories, emails, reports)
+- When the user asks to write code, build a script, or implement an algorithm
+- You MUST specify kind: 'code' for programming, 'text' for writing, 'sheet' for data
+- Include ALL content in the createDocument call. Do not create then edit.

 **When NOT to use \`createDocument\`:**
- For informational/explanatory content
- For conversational responses
- When asked to keep it in chat
+- For answering questions, explanations, or conversational responses
+- For short code snippets or examples shown inline
+- When the user asks "what is", "how does", "explain", etc.

-**Using \`updateDocument\`:**
- Default to full document rewrites for major changes
- Use targeted updates only for specific, isolated changes
- Follow user instructions for which parts to modify
+**Using \`editDocument\` (preferred for targeted changes):**
+- For scripts: fixing bugs, adding/removing lines, renaming variables, adding logs
+- For documents: fixing typos, rewording paragraphs, inserting sections
+- Uses find-and-replace: provide exact old_string and new_string
+- Include 3-5 surrounding lines in old_string to ensure a unique match
+- Use replace_all:true for renaming across the whole artifact
+- Can call multiple times for several independent edits

-**When NOT to use \`updateDocument\`:**
- Immediately after creating a document
+**Using \`updateDocument\` (full rewrite only):**
+- Only when most of the content needs to change
+- When editDocument would require too many individual edits

-Do not update document right after creating it. Wait for user feedback or request to update it.
+**When NOT to use \`editDocument\` or \`updateDocument\`:**
+- Immediately after creating an artifact
+- In the same response as createDocument
+- Without explicit user request to modify
+
+**After any create/edit/update:**
+- NEVER repeat, summarize, or output the artifact content in chat
+- Only respond with a short confirmation

 **Using \`requestSuggestions\`:**
- ONLY use when the user explicitly asks for suggestions on an existing document
- Requires a valid document ID from a previously created document
- Never use for general questions or information requests
+- ONLY when the user explicitly asks for suggestions on an existing document
 `;

-export const regularPrompt = `You are a friendly assistant! Keep your responses concise and helpful.
+export const regularPrompt = `You are a helpful assistant. Keep responses concise and direct.

-When asked to write, create, or help with something, just do it directly. Don't ask clarifying questions unless absolutely necessary - make reasonable assumptions and proceed with the task.`;
+When asked to write, create, or build something, do it immediately. Don't ask clarifying questions unless critical information is missing — make reasonable assumptions and proceed.`;

 export type RequestHints = {
  latitude: Geo["latitude"];
@ -57,19 +64,15 @@ About the origin of user's request:
 `;

 export const systemPrompt = ({
-  selectedChatModel,
  requestHints,
+  supportsTools,
 }: {
-  selectedChatModel: string;
  requestHints: RequestHints;
+  supportsTools: boolean;
 }) => {
  const requestPrompt = getRequestPromptFromHints(requestHints);

-  // reasoning models don't need artifacts prompt (they can't use tools)
-  if (
-    selectedChatModel.includes("reasoning") ||
-    selectedChatModel.includes("thinking")
-  ) {
+  if (!supportsTools) {
    return `${regularPrompt}\n\n${requestPrompt}`;
  }

@ -77,48 +80,40 @@ export const systemPrompt = ({
 };

 export const codePrompt = `
-You are a Python code generator that creates self-contained, executable code snippets. When writing code:
+You are a code generator that creates self-contained, executable code snippets. When writing code:

-1. Each snippet should be complete and runnable on its own
-2. Prefer using print() statements to display outputs
-3. Include helpful comments explaining the code
-4. Keep snippets concise (generally under 15 lines)
-5. Avoid external dependencies - use Python standard library
-6. Handle potential errors gracefully
-7. Return meaningful output that demonstrates the code's functionality
-8. Don't use input() or other interactive functions
-9. Don't access files or network resources
-10. Don't use infinite loops
-
-Examples of good snippets:
-
-# Calculate factorial iteratively
-def factorial(n):
-    result = 1
-    for i in range(1, n + 1):
-        result *= i
-    return result
-
-print(f"Factorial of 5 is: {factorial(5)}")
+1. Each snippet must be complete and runnable on its own
+2. Use print/console.log to display outputs
+3. Keep snippets concise and focused
+4. Prefer standard library over external dependencies
+5. Handle potential errors gracefully
+6. Return meaningful output that demonstrates functionality
+7. Don't use interactive input functions
+8. Don't access files or network resources
+9. Don't use infinite loops
 `;

 export const sheetPrompt = `
-You are a spreadsheet creation assistant. Create a spreadsheet in csv format based on the given prompt. The spreadsheet should contain meaningful column headers and data.
+You are a spreadsheet creation assistant. Create a spreadsheet in CSV format based on the given prompt.
+
+Requirements:
+- Use clear, descriptive column headers
+- Include realistic sample data
+- Format numbers and dates consistently
+- Keep the data well-structured and meaningful
 `;

 export const updateDocumentPrompt = (
  currentContent: string | null,
  type: ArtifactKind
 ) => {
-  let mediaType = "document";
+  const mediaTypes: Record<string, string> = {
+    code: "script",
+    sheet: "spreadsheet",
+  };
+  const mediaType = mediaTypes[type] ?? "document";

-  if (type === "code") {
-    mediaType = "code snippet";
-  } else if (type === "sheet") {
-    mediaType = "spreadsheet";
-  }
-
-  return `Improve the following contents of the ${mediaType} based on the given prompt.
+  return `Rewrite the following ${mediaType} based on the given prompt.

 ${currentContent}`;
 };
@ -133,7 +128,4 @@ Examples:
 - "hi" → New Conversation
 - "debug my python code" → Python Debugging

-Bad outputs (never do this):
- "# Space Essay" (no hashtags)
- "Title: Weather" (no prefixes)
- ""NYC Weather"" (no quotes)`;
+Never output hashtags, prefixes like "Title:", or quotes.`;
--- a/lib/ai/providers.ts
+++ b/lib/ai/providers.ts
@ -1,27 +1,14 @@
-import { gateway } from "@ai-sdk/gateway";
-import {
-  customProvider,
-  extractReasoningMiddleware,
-  wrapLanguageModel,
-} from "ai";
+import { customProvider, gateway } from "ai";
 import { isTestEnvironment } from "../constants";
-
-const THINKING_SUFFIX_REGEX = /-thinking$/;
+import { titleModel } from "./models";

 export const myProvider = isTestEnvironment
  ? (() => {
-      const {
-        artifactModel,
-        chatModel,
-        reasoningModel,
-        titleModel,
-      } = require("./models.mock");
+      const { chatModel, titleModel } = require("./models.mock");
      return customProvider({
        languageModels: {
          "chat-model": chatModel,
-          "chat-model-reasoning": reasoningModel,
          "title-model": titleModel,
-          "artifact-model": artifactModel,
        },
      });
    })()
@ -32,19 +19,6 @@ export function getLanguageModel(modelId: string) {
    return myProvider.languageModel(modelId);
  }

-  const isReasoningModel =
-    modelId.endsWith("-thinking") ||
-    (modelId.includes("reasoning") && !modelId.includes("non-reasoning"));
-
-  if (isReasoningModel) {
-    const gatewayModelId = modelId.replace(THINKING_SUFFIX_REGEX, "");
-
-    return wrapLanguageModel({
-      model: gateway.languageModel(gatewayModelId),
-      middleware: extractReasoningMiddleware({ tagName: "thinking" }),
-    });
-  }
-
  return gateway.languageModel(modelId);
 }

@ -52,12 +26,5 @@ export function getTitleModel() {
  if (isTestEnvironment && myProvider) {
    return myProvider.languageModel("title-model");
  }
-  return gateway.languageModel("google/gemini-2.5-flash-lite");
-}
-
-export function getArtifactModel() {
-  if (isTestEnvironment && myProvider) {
-    return myProvider.languageModel("artifact-model");
-  }
-  return gateway.languageModel("anthropic/claude-haiku-4.5");
+  return gateway.languageModel(titleModel.id);
 }
--- a/lib/ai/tools/create-document.ts
+++ b/lib/ai/tools/create-document.ts
@ -1,25 +1,34 @@
 import { tool, type UIMessageStreamWriter } from "ai";
+import type { Session } from "next-auth";
 import { z } from "zod";
 import {
  artifactKinds,
  documentHandlersByArtifactKind,
 } from "@/lib/artifacts/server";
-import type { AuthSession } from "@/lib/auth";
 import type { ChatMessage } from "@/lib/types";
 import { generateUUID } from "@/lib/utils";

 type CreateDocumentProps = {
-  session: NonNullable<AuthSession>;
+  session: Session;
  dataStream: UIMessageStreamWriter<ChatMessage>;
+  modelId: string;
 };

-export const createDocument = ({ session, dataStream }: CreateDocumentProps) =>
+export const createDocument = ({
+  session,
+  dataStream,
+  modelId,
+}: CreateDocumentProps) =>
  tool({
    description:
-      "Create a document for a writing or content creation activities. This tool will call other functions that will generate the contents of the document based on the title and kind.",
+      "Create an artifact. You MUST specify kind: use 'code' for any programming/algorithm request (creates a script), 'text' for essays/writing (creates a document), 'sheet' for spreadsheets/data.",
    inputSchema: z.object({
-      title: z.string(),
-      kind: z.enum(artifactKinds),
+      title: z.string().describe("The title of the artifact"),
+      kind: z
+        .enum(artifactKinds)
+        .describe(
+          "REQUIRED. 'code' for programming/algorithms, 'text' for essays/writing, 'sheet' for spreadsheets"
+        ),
    }),
    execute: async ({ title, kind }) => {
      const id = generateUUID();
@ -62,6 +71,7 @@ export const createDocument = ({ session, dataStream }: CreateDocumentProps) =>
        title,
        dataStream,
        session,
+        modelId,
      });

      dataStream.write({ type: "data-finish", data: null, transient: true });
@ -70,7 +80,10 @@ export const createDocument = ({ session, dataStream }: CreateDocumentProps) =>
        id,
        title,
        kind,
-        content: "A document was created and is now visible to the user.",
+        content:
+          kind === "code"
+            ? "A script was created and is now visible to the user."
+            : "A document was created and is now visible to the user.",
      };
    },
  });
--- a/lib/ai/tools/edit-document.ts
+++ b/lib/ai/tools/edit-document.ts
@ -0,0 +1,100 @@
+import { tool, type UIMessageStreamWriter } from "ai";
+import type { Session } from "next-auth";
+import { z } from "zod";
+import { getDocumentById, saveDocument } from "@/lib/db/queries";
+import type { ChatMessage } from "@/lib/types";
+
+type EditDocumentProps = {
+  session: Session;
+  dataStream: UIMessageStreamWriter<ChatMessage>;
+};
+
+export const editDocument = ({ session, dataStream }: EditDocumentProps) =>
+  tool({
+    description:
+      "Make a targeted edit to an existing artifact by finding and replacing an exact string. Preferred over updateDocument for small changes. The old_string must match exactly.",
+    inputSchema: z.object({
+      id: z.string().describe("The ID of the artifact to edit"),
+      old_string: z
+        .string()
+        .describe(
+          "Exact string to find. Include 3-5 surrounding lines for uniqueness."
+        ),
+      new_string: z.string().describe("Replacement string"),
+      replace_all: z
+        .boolean()
+        .optional()
+        .describe(
+          "Replace all occurrences instead of just the first (default false)"
+        ),
+    }),
+    execute: async ({ id, old_string, new_string, replace_all }) => {
+      const document = await getDocumentById({ id });
+
+      if (!document) {
+        return { error: "Document not found" };
+      }
+
+      if (document.userId !== session.user?.id) {
+        return { error: "Forbidden" };
+      }
+
+      if (!document.content) {
+        return { error: "Document has no content" };
+      }
+
+      if (!document.content.includes(old_string)) {
+        return { error: "old_string not found in document" };
+      }
+
+      const updated = replace_all
+        ? document.content.replaceAll(old_string, new_string)
+        : document.content.replace(old_string, new_string);
+
+      await saveDocument({
+        id: document.id,
+        title: document.title,
+        kind: document.kind,
+        content: updated,
+        userId: document.userId,
+      });
+
+      dataStream.write({
+        type: "data-clear",
+        data: null,
+        transient: true,
+      });
+
+      if (document.kind === "code") {
+        dataStream.write({
+          type: "data-codeDelta",
+          data: updated,
+          transient: true,
+        });
+      } else if (document.kind === "sheet") {
+        dataStream.write({
+          type: "data-sheetDelta",
+          data: updated,
+          transient: true,
+        });
+      } else {
+        dataStream.write({
+          type: "data-textDelta",
+          data: updated,
+          transient: true,
+        });
+      }
+
+      dataStream.write({ type: "data-finish", data: null, transient: true });
+
+      return {
+        id,
+        title: document.title,
+        kind: document.kind,
+        content:
+          document.kind === "code"
+            ? "The script has been edited successfully."
+            : "The document has been edited successfully.",
+      };
+    },
+  });
--- a/lib/ai/tools/get-weather.ts
+++ b/lib/ai/tools/get-weather.ts
@ -40,7 +40,6 @@ export const getWeather = tool({
      .describe("City name (e.g., 'San Francisco', 'New York', 'London')")
      .optional(),
  }),
-  needsApproval: true,
  execute: async (input) => {
    let latitude: number;
    let longitude: number;
--- a/lib/ai/tools/request-suggestions.ts
+++ b/lib/ai/tools/request-suggestions.ts
@ -1,20 +1,22 @@
 import { Output, streamText, tool, type UIMessageStreamWriter } from "ai";
+import type { Session } from "next-auth";
 import { z } from "zod";
-import type { AuthSession } from "@/lib/auth";
 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 { getArtifactModel } from "../providers";
+import { getLanguageModel } from "../providers";

 type RequestSuggestionsProps = {
-  session: NonNullable<AuthSession>;
+  session: Session;
  dataStream: UIMessageStreamWriter<ChatMessage>;
+  modelId: string;
 };

 export const requestSuggestions = ({
  session,
  dataStream,
+  modelId,
 }: RequestSuggestionsProps) =>
  tool({
    description:
@ -35,15 +37,19 @@ export const requestSuggestions = ({
        };
      }

+      if (document.userId !== session.user?.id) {
+        return { error: "Forbidden" };
+      }
+
      const suggestions: Omit<
        Suggestion,
        "userId" | "createdAt" | "documentCreatedAt"
      >[] = [];

      const { partialOutputStream } = streamText({
-        model: getArtifactModel(),
+        model: getLanguageModel(modelId),
        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.",
+          "You are a writing assistant. Given a piece of writing, offer up to 5 suggestions to improve it. Each suggestion must contain full sentences, not just individual words. Describe what changed and why.",
        prompt: document.content,
        output: Output.array({
          element: z.object({
--- a/lib/ai/tools/update-document.ts
+++ b/lib/ai/tools/update-document.ts
@ -1,22 +1,29 @@
 import { tool, type UIMessageStreamWriter } from "ai";
+import type { Session } from "next-auth";
 import { z } from "zod";
 import { documentHandlersByArtifactKind } from "@/lib/artifacts/server";
-import type { AuthSession } from "@/lib/auth";
 import { getDocumentById } from "@/lib/db/queries";
 import type { ChatMessage } from "@/lib/types";

 type UpdateDocumentProps = {
-  session: NonNullable<AuthSession>;
+  session: Session;
  dataStream: UIMessageStreamWriter<ChatMessage>;
+  modelId: string;
 };

-export const updateDocument = ({ session, dataStream }: UpdateDocumentProps) =>
+export const updateDocument = ({
+  session,
+  dataStream,
+  modelId,
+}: UpdateDocumentProps) =>
  tool({
-    description: "Update a document with the given description.",
+    description:
+      "Full rewrite of an existing artifact. Only use for major changes where most content needs replacing. Prefer editDocument for targeted changes.",
    inputSchema: z.object({
-      id: z.string().describe("The ID of the document to update"),
+      id: z.string().describe("The ID of the artifact to rewrite"),
      description: z
        .string()
+        .default("Improve the content")
        .describe("The description of changes that need to be made"),
    }),
    execute: async ({ id, description }) => {
@ -28,6 +35,10 @@ export const updateDocument = ({ session, dataStream }: UpdateDocumentProps) =>
        };
      }

+      if (document.userId !== session.user?.id) {
+        return { error: "Forbidden" };
+      }
+
      dataStream.write({
        type: "data-clear",
        data: null,
@ -48,6 +59,7 @@ export const updateDocument = ({ session, dataStream }: UpdateDocumentProps) =>
        description,
        dataStream,
        session,
+        modelId,
      });

      dataStream.write({ type: "data-finish", data: null, transient: true });
@ -56,7 +68,10 @@ export const updateDocument = ({ session, dataStream }: UpdateDocumentProps) =>
        id,
        title: document.title,
        kind: document.kind,
-        content: "The document has been updated successfully.",
+        content:
+          document.kind === "code"
+            ? "The script has been updated successfully."
+            : "The document has been updated successfully.",
      };
    },
  });