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

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.",
       };
     },
   });

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.",
+      };
+    },
+  });

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;

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({

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.",
       };
     },
   });