From 711da0b94b3ae9f102eb54c4053700cd7cc4e0d4 Mon Sep 17 00:00:00 2001
From: Jeremy <git@jrmyphlmn.com>
Date: Mon, 3 Feb 2025 20:41:32 +0530
Subject: [PATCH] feat: switch prompts based on selected chat model (#754)

---
 app/(chat)/api/chat/route.ts | 28 ++++++++++------------------
 lib/ai/prompts.ts            | 12 +++++++++++-
 2 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/app/(chat)/api/chat/route.ts b/app/(chat)/api/chat/route.ts
index 9736e3c..ec404ec 100644
--- a/app/(chat)/api/chat/route.ts
+++ b/app/(chat)/api/chat/route.ts
@@ -3,7 +3,6 @@ import {
   createDataStreamResponse,
   smoothStream,
   streamText,
-  wrapLanguageModel,
 } from 'ai';
 
 import { auth } from '@/app/(auth)/auth';
@@ -29,21 +28,6 @@ import { getWeather } from '@/lib/ai/tools/get-weather';
 
 export const maxDuration = 60;
 
-type AllowedTools =
-  | 'createDocument'
-  | 'updateDocument'
-  | 'requestSuggestions'
-  | 'getWeather';
-
-const blocksTools: AllowedTools[] = [
-  'createDocument',
-  'updateDocument',
-  'requestSuggestions',
-];
-
-const weatherTools: AllowedTools[] = ['getWeather'];
-const allTools: AllowedTools[] = [...blocksTools, ...weatherTools];
-
 export async function POST(request: Request) {
   const {
     id,
@@ -79,10 +63,18 @@ export async function POST(request: Request) {
     execute: (dataStream) => {
       const result = streamText({
         model: myProvider.languageModel(selectedChatModel),
-        system: systemPrompt,
+        system: systemPrompt({ selectedChatModel }),
         messages,
         maxSteps: 5,
-        experimental_activeTools: allTools,
+        experimental_activeTools:
+          selectedChatModel === 'chat-model-reasoning'
+            ? []
+            : [
+                'getWeather',
+                'createDocument',
+                'updateDocument',
+                'requestSuggestions',
+              ],
         experimental_transform: smoothStream({ chunking: 'word' }),
         experimental_generateMessageId: generateUUID,
         tools: {
diff --git a/lib/ai/prompts.ts b/lib/ai/prompts.ts
index 07ae18f..fa3e546 100644
--- a/lib/ai/prompts.ts
+++ b/lib/ai/prompts.ts
@@ -34,7 +34,17 @@ Do not update document right after creating it. Wait for user feedback or reques
 export const regularPrompt =
   'You are a friendly assistant! Keep your responses concise and helpful.';
 
-export const systemPrompt = `${regularPrompt}\n\n${blocksPrompt}`;
+export const systemPrompt = ({
+  selectedChatModel,
+}: {
+  selectedChatModel: string;
+}) => {
+  if (selectedChatModel === 'chat-model-reasoning') {
+    return regularPrompt;
+  } else {
+    return `${regularPrompt}\n\n${blocksPrompt}`;
+  }
+};
 
 export const codePrompt = `
 You are a Python code generator that creates self-contained, executable code snippets. When writing code: