Remove unused ai-reply file and tools; refactor chat API to integrate new prompt management and enhance email thread handling. Update AI chat component for improved thread rendering and error handling.

apps/mail/actions/ai-reply.ts

@@ -1,91 +0,0 @@
-'use server';
-
-import { generateCompletions } from '@/lib/groq';
-import { headers } from 'next/headers';
-import { auth } from '@/lib/auth';
-
-// Generates an AI response for an email reply based on the thread content
-export async function generateAIResponse(
-  threadContent: string,
-  originalSender: string,
-): Promise<string> {
-  const headersList = await headers();
-  const session = await auth.api.getSession({ headers: headersList });
-
-  if (!session?.user) return '';
-
-  if (!process.env.GROQ_API_KEY) {
-    throw new Error('Groq API key is not configured');
-  }
-
-  // Use a more aggressive content reduction approach
-  const processedContent = threadContent;
-
-  // Create the system message
-  const systemPrompt = `You are an email assistant helping ${session.user.name} write professional and concise email replies.
-  
-  CRITICAL INSTRUCTIONS:
-  - Return ONLY the email content itself
-  - DO NOT include ANY explanatory text or meta-text like "Here's a draft" or "Here's a reply"
-  - DO NOT include ANY text before or after the email content
-  - Start directly with the greeting (e.g. "Hi John,")
-  - End with just the name
-  - Generate a real, ready-to-send email reply, not a template
-  - Do not include placeholders like [Recipient], [discount percentage], etc.
-  - Do not include formatting instructions or explanations
-  - Do not include "Subject:" lines
-  - Write as if this email is ready to be sent immediately
-  - Use real, specific content instead of placeholders
-  - Address the recipient directly without using [brackets]
-  - Be concise but thorough (2-3 paragraphs maximum)
-  - Write in the first person as if you are ${session.user.name}
-  - Double space paragraphs (2 newlines)
-  - Add two spaces below the sign-off
-  - End with the name: ${session.user.name}`;
-
-  // Create the user message - keep it shorter
-  const prompt = `
-  Here's the context of the email thread (some parts may be summarized or truncated due to length):
-  ${processedContent}
-
-  Write a professional, helpful, and concise email reply to ${originalSender}.
-  Keep your response under 200 words.
-  
-  CRITICAL: Return ONLY the email content itself. DO NOT include ANY explanatory text or meta-text.
-  Start directly with the greeting and end with just the name.
-
-  Important instructions:
-  - Return ONLY the email content itself
-  - DO NOT include ANY explanatory text or meta-text
-  - Start directly with the greeting (e.g. "Hi John,")
-  - End with just the name
-  - Generate a real, ready-to-send email reply, not a template
-  - Do not include placeholders like [Recipient], [discount percentage], etc.
-  - Do not include formatting instructions or explanations
-  - Do not include "Subject:" lines
-  - Write as if this email is ready to be sent immediately
-  - Use real, specific content instead of placeholders
-  - Address the recipient directly without using [brackets]
-  - Be concise but thorough (2-3 paragraphs maximum)
-  - Write in the first person as if you are ${session.user.name}
-  - Double space paragraphs (2 newlines)
-  - Add two spaces below the sign-off
-  - End with the name: ${session.user.name}`;
-
-  try {
-    console.log('Generating AI response with prompt:', prompt);
-    // Use direct fetch to the Groq API
-    const { completion } = await generateCompletions({
-      model: 'gpt-4o-mini',
-      systemPrompt: process.env.AI_SYSTEM_PROMPT || systemPrompt,
-      prompt,
-      temperature: 0.6,
-      userName: session.user.name,
-    });
-
-    return completion;
-  } catch (error: any) {
-    console.error('Error generating AI response:', error);
-    throw error;
-  }
-}

apps/mail/app/api/chat/prompt.ts

@@ -0,0 +1,59 @@
+const prompt = `You are an intelligent email management assistant with access to powerful Gmail operations. You can help users organize their inbox by searching, analyzing, and performing actions on their emails.
+
+Core Capabilities:
+1. Search & Analysis
+   - Search through email threads using complex queries
+   - Analyze email content, subjects, and patterns
+   - Identify email categories and suggested organizations
+
+2. Label Management
+   - Create new labels with custom colors
+   - View existing labels
+   - Apply labels to emails based on content analysis
+   - Suggest label hierarchies for better organization
+
+3. Email Organization
+   - Archive emails that don't need immediate attention
+   - Mark emails as read/unread strategically 
+   - Apply bulk actions to similar emails
+   - Help maintain inbox zero principles
+
+Available Tools:
+- listThreads: Search and retrieve email threads
+- archiveThreads: Move emails out of inbox
+- markThreadsRead/Unread: Manage read status
+- createLabel: Create new organizational labels
+- addLabelsToThreads: Apply labels to emails
+- getUserLabels: View existing label structure
+
+Best Practices:
+1. Always confirm actions before processing large numbers of emails
+2. Suggest organizational strategies based on user's email patterns
+3. Explain your reasoning when recommending actions
+4. Be cautious with permanent actions like deletion
+5. Consider email importance and urgency when organizing
+
+Examples of how you can help:
+- "Find all my unread newsletter emails and help me organize them"
+- "Create a systematic way to handle my recruitment emails"
+- "Help me clean up my inbox by identifying and archiving non-critical emails"
+- "Set up a label system for my project-related emails"
+
+When suggesting actions, consider:
+- Email importance and time sensitivity
+- Natural groupings and categories
+- Workflow optimization
+- Future searchability
+- Maintenance requirements
+
+Response Format Rules:
+1. NEVER include tool call results in your text response
+2. NEVER start responses with phrases like "Here is", "I found", etc.
+3. ONLY respond with exactly one of these two options:
+   - "Done." (when the action is completed successfully)
+   - "Could not complete action." (when the action fails or cannot be completed)
+
+Remember: Your goal is to help users maintain an organized, efficient, and stress-free email system while preserving important information and accessibility.
+`;
+
+export default prompt;

apps/mail/app/api/chat/route.ts

@@ -1,10 +1,10 @@
 import { getActiveConnection } from '@/actions/utils';
 import { ToolInvocation, streamText } from 'ai';
-import { type IConfig } from '../driver/types';
 import { NextResponse } from 'next/server';
 import { createDriver } from '../driver';
 import { openai } from '@ai-sdk/openai';
-import { listThreads } from './tools';
+import prompt from './prompt';
+import { z } from 'zod';
 
 interface Message {
   role: 'user' | 'assistant';
@@ -20,28 +20,105 @@ export async function POST(req: Request) {
     console.error('Unauthorized: No valid connection found');
     return NextResponse.json({}, { status: 401 });
   }
+
   const driver = await createDriver(connection.providerId, {
     auth: {
       access_token: connection.accessToken,
-      refresh_token: connection.refreshToken,
+      refresh_token: connection.refreshToken!,
       email: connection.email,
     },
   });
 
   const result = streamText({
     model: openai('gpt-4o'),
-    system: `
-        You are a helpful assistant.
-        You are able to list email threads. Use the listThreads tool to do so.
-        Parameters:
-        - folder: the folder to list threads from
-        - query: the search query
-        - maxResults: the maximum number of results
-        - labelIds: the label IDs to filter by
-    `,
+    system: prompt,
     messages,
     tools: {
-      listThreads: listThreads(driver),
+      listThreads: {
+        description: 'List email threads',
+        parameters: z.object({
+          folder: z
+            .string()
+            .optional()
+            .default('inbox')
+            .describe('The folder to list threads from'),
+          query: z.string().optional().describe('The search query'),
+          maxResults: z.number().optional().default(20).describe('The maximum number of results'),
+          labelIds: z.array(z.string()).optional().describe('The label IDs to filter by'),
+        }),
+        execute: async ({ folder, query, maxResults, labelIds }) => {
+          return driver.list(folder, query, maxResults, labelIds, undefined);
+        },
+      },
+      archiveThreads: {
+        description: 'Archive email threads by removing them from inbox',
+        parameters: z.object({
+          threadIds: z.array(z.string()).describe('Array of thread IDs to archive'),
+        }),
+        execute: async ({ threadIds }) => {
+          await driver.modifyLabels(threadIds, {
+            removeLabels: ['INBOX'],
+            addLabels: [],
+          });
+          return { archived: threadIds.length };
+        },
+      },
+      markThreadsRead: {
+        description: 'Mark email threads as read',
+        parameters: z.object({
+          threadIds: z.array(z.string()).describe('Array of thread IDs to mark as read'),
+        }),
+        execute: async ({ threadIds }) => {
+          await driver.markAsRead(threadIds);
+          return { marked: threadIds.length };
+        },
+      },
+      markThreadsUnread: {
+        description: 'Mark email threads as unread',
+        parameters: z.object({
+          threadIds: z.array(z.string()).describe('Array of thread IDs to mark as unread'),
+        }),
+        execute: async ({ threadIds }) => {
+          await driver.markAsUnread(threadIds);
+          return { marked: threadIds.length };
+        },
+      },
+      createLabel: {
+        description: 'Create a new label',
+        parameters: z.object({
+          name: z.string().describe('Name of the label to create'),
+          backgroundColor: z.string().optional().describe('Background color for the label'),
+          textColor: z.string().optional().describe('Text color for the label'),
+        }),
+        execute: async ({ name, backgroundColor = '#e3e3e3', textColor = '#666666' }) => {
+          const label = await driver.createLabel({
+            name,
+            color: { backgroundColor: '#FFFFFF', textColor: '#000000' },
+          });
+          return { created: label.id };
+        },
+      },
+      addLabelsToThreads: {
+        description: 'Add labels to email threads',
+        parameters: z.object({
+          threadIds: z.array(z.string()).describe('Array of thread IDs to label'),
+          labelIds: z.array(z.string()).describe('Array of label IDs to add'),
+        }),
+        execute: async ({ threadIds, labelIds }) => {
+          await driver.modifyLabels(threadIds, {
+            addLabels: labelIds,
+            removeLabels: [],
+          });
+          return { labeled: threadIds.length };
+        },
+      },
+      getUserLabels: {
+        description: 'Get all user labels',
+        parameters: z.object({}),
+        execute: async () => {
+          return driver.getUserLabels();
+        },
+      },
     },
   });
 

apps/mail/app/api/chat/tools.ts

@@ -1,35 +0,0 @@
-import type { MailManager } from '../driver/types';
-
-export const listThreads = (driver: MailManager) => ({
-  description: 'List email threads',
-  parameters: {
-    type: 'object',
-    properties: {
-      folder: { type: 'string', description: 'Folder to list threads from' },
-      query: { type: 'string', description: 'Search query' },
-      maxResults: { type: 'number', description: 'Maximum number of results' },
-      labelIds: {
-        type: 'array',
-        items: { type: 'string' },
-        description: 'Label IDs to filter by',
-      },
-      pageToken: { type: 'string', description: 'Page token for pagination' },
-    },
-    required: ['folder'],
-  },
-  execute: async ({
-    folder,
-    query,
-    maxResults,
-    labelIds,
-  }: {
-    folder: string;
-    query?: string;
-    maxResults?: number;
-    labelIds?: string[];
-  }) => {
-    console.log('trying to list threads');
-
-    return driver.list(folder, query, maxResults, labelIds, undefined);
-  },
-});

apps/mail/components/create/ai-chat.tsx

@@ -1,15 +1,17 @@
 'use client';
 
-import { ArrowUpIcon, Mic, CheckIcon, XIcon, Plus, Command } from 'lucide-react';
+import { ArrowUpIcon, Mic, CheckIcon, XIcon, Plus, Command, ArrowDownCircle } from 'lucide-react';
 import { usePathname, useRouter, useSearchParams } from 'next/navigation';
 import { useSearchValue } from '@/hooks/use-search-value';
 import { useConnections } from '@/hooks/use-connections';
 import { useRef, useCallback, useEffect } from 'react';
 import { Button } from '@/components/ui/button';
+import { useThread } from '@/hooks/use-threads';
 import { useSession } from '@/lib/auth-client';
 import { CurvedArrow } from '../icons/icons';
 import { AITextarea } from './ai-textarea';
 import { useChat } from '@ai-sdk/react';
+import { useQueryState } from 'nuqs';
 import { cn } from '@/lib/utils';
 import { useState } from 'react';
 import VoiceChat from './voice';
@@ -45,6 +47,29 @@ interface AIChatProps {
   onReset?: () => void;
 }
 
+const renderThread = (thread: { id: string; title: string; snippet: string }) => {
+  const [, setThreadId] = useQueryState('threadId');
+  const { data: getThread } = useThread(thread.id);
+  return getThread ? (
+    <div
+      onClick={() => setThreadId(thread.id)}
+      key={thread.id}
+      className="bg-subtleBlack cursor-pointer rounded-md border p-2 hover:bg-black"
+    >
+      <p>{getThread.latest?.subject}</p>
+    </div>
+  ) : null;
+};
+
+const RenderThreads = ({
+  threads,
+}: {
+  threads: { id: string; title: string; snippet: string }[];
+}) => {
+  const [, setThreadId] = useQueryState('threadId');
+  return threads.map(renderThread);
+};
+
 export function AIChat({ editor, onMessagesChange, onReset }: AIChatProps) {
   const [value, setValue] = useState('');
   const [showVoiceChat, setShowVoiceChat] = useState(false);
@@ -59,7 +84,7 @@ export function AIChat({ editor, onMessagesChange, onReset }: AIChatProps) {
   const { data: session } = useSession();
   const { data: connections } = useConnections();
 
-  const { messages, input, setInput, append, handleSubmit, status } = useChat({
+  const { messages, input, setInput, error, handleSubmit, status } = useChat({
     api: '/api/chat',
     maxSteps: 5,
   });
@@ -274,9 +299,35 @@ export function AIChat({ editor, onMessagesChange, onReset }: AIChatProps) {
                     : 'overflow-wrap-anywhere mr-auto break-words bg-[#f0f0f0] p-3 dark:bg-[#313131]', // Assistant messages aligned to left
                 )}
               >
-                <div className="prose dark:prose-invert overflow-wrap-anywhere break-words text-sm font-medium">
+                {/* <div className="prose dark:prose-invert overflow-wrap-anywhere break-words text-sm font-medium">
                   {message.content}
-                </div>
+                </div> */}
+
+                {message.parts.map((part) => {
+                  if (part.type === 'text') {
+                    return <p>{part.text}</p>;
+                  }
+                  if (part.type === 'reasoning') {
+                    return <p>Reasoning: {part.reasoning}</p>;
+                  }
+                  if (part.type === 'tool-invocation') {
+                    return (
+                      'result' in part.toolInvocation &&
+                      ('threads' in part.toolInvocation.result ? (
+                        <RenderThreads threads={part.toolInvocation.result.threads} />
+                      ) : (
+                        <p>No threads found</p>
+                      ))
+                    );
+                  }
+                  if (part.type === 'source') {
+                    return <p>Source: {part.source.title}</p>;
+                  }
+                  //   if (part.type === 'step-start') {
+                  //     return <ArrowDownCircle className="mx-auto h-4 w-4" />;
+                  //   }
+                  //   return <p>{part.type}</p>;
+                })}
 
                 {/* {message.type === 'search' &&
                 message.searchContent &&
@@ -362,6 +413,7 @@ export function AIChat({ editor, onMessagesChange, onReset }: AIChatProps) {
           {/* Invisible element to scroll to */}
           <div ref={messagesEndRef} />
 
+          {JSON.stringify(error)}
           {/* Loading indicator */}
           {status === 'submitted' && (
             <div className="flex flex-col gap-2 rounded-lg">

apps/mail/lib/providers.tsx

@@ -15,16 +15,16 @@ export function Providers({ children, ...props }: React.ComponentProps<typeof Ne
   const theme = settings?.colorTheme || 'system';
 
   return (
-    <AISidebarProvider>
-      <JotaiProvider>
-        <NuqsAdapter>
+    <NuqsAdapter>
+      <AISidebarProvider>
+        <JotaiProvider>
           <NextThemesProvider {...props} defaultTheme={theme}>
             <SidebarProvider>
               <PostHogProvider>{children}</PostHogProvider>
             </SidebarProvider>
           </NextThemesProvider>
-        </NuqsAdapter>
-      </JotaiProvider>
-    </AISidebarProvider>
+        </JotaiProvider>
+      </AISidebarProvider>
+    </NuqsAdapter>
   );
 }