Remove console logs and simplify AI chat prompt (#1769)

# Cleanup: Remove console logs and refactor AI chat prompt

## Description

This PR removes unnecessary console logs from the codebase and refactors the AI chat prompt to be more concise and structured. The changes include:

1. Removing debug console.log statements from notification settings and hot key handling
2. Adding error handling to the AI sidebar component
3. Streamlining the AI chat prompt structure for better readability and maintenance
4. Improving the inboxRag function with better error handling and logging

## Type of Change

- 🐛 Bug fix (non-breaking change which fixes an issue)
- ⚡ Performance improvement

## Areas Affected

- [x] User Interface/Experience
- [x] API Endpoints

## Testing Done

- [x] Manual testing performed

## Checklist

- [x] I have performed a self-review of my code
- [x] My changes generate no new warnings
- [x] All tests pass locally

## Additional Notes

The AI chat prompt refactoring significantly improves readability while maintaining the same functionality. The removal of console logs helps keep the browser console clean in production.

---

_By submitting this pull request, I confirm that my contribution is made under the terms of the project's license._

<!-- This is an auto-generated comment: release notes by coderabbit.ai -->

## Summary by CodeRabbit

* **New Features**
  * Enhanced email search with improved AI-generated query refinement for more accurate results.

* **Bug Fixes**
  * Improved error handling and fallback messaging during AI-powered searches to prevent disruptions.

* **Refactor**
  * Streamlined and restructured the AI assistant’s internal prompt for clearer, more direct responses and safer actions.
  * Refined tool orchestration for chat streaming to improve reliability.

* **Chores**
  * Removed unused tools and related features to simplify the system.
  * Eliminated unnecessary console logging for cleaner operation.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: MrgSub

apps/mail/app/(routes)/settings/notifications/page.tsx

@@ -15,8 +15,8 @@ import {
 } from '@/components/ui/select';
 import { SettingsCard } from '@/components/settings/settings-card';
 import { zodResolver } from '@hookform/resolvers/zod';
-import { Button } from '@/components/ui/button';
 import { Switch } from '@/components/ui/switch';
+import { Button } from '@/components/ui/button';
 import { useForm } from 'react-hook-form';
 import { Bell } from 'lucide-react';
 import { useState } from 'react';
@@ -41,7 +41,6 @@ export default function NotificationsPage() {
   function onSubmit(values: z.infer<typeof formSchema>) {
     setIsSaving(true);
     setTimeout(() => {
-      console.log(values);
       setIsSaving(false);
     }, 1000);
   }

apps/mail/components/ui/ai-sidebar.tsx

@@ -349,6 +349,7 @@ function AISidebar({ className }: AISidebarProps) {
     agent: 'ZeroAgent',
     name: activeConnection?.id ? String(activeConnection.id) : 'general',
     host: `${import.meta.env.VITE_PUBLIC_BACKEND_URL}`,
+    onError: (e) => console.log(e),
   });
 
   const chatState = useAgentChat({

apps/mail/hooks/use-hot-key.ts

@@ -9,49 +9,44 @@ function initKeyListeners() {
     window.addEventListener('keydown', (e) => {
       // Store the key state
       keyStates.set(e.key, true);
-      
+
       // Also store specific states for modifier keys
       if (e.altKey) {
         keyStates.set('Alt', true);
         keyStates.set('AltLeft', true);
         keyStates.set('AltRight', true);
       }
-      
+
       if (e.shiftKey) {
         keyStates.set('Shift', true);
         keyStates.set('ShiftLeft', true);
         keyStates.set('ShiftRight', true);
       }
-      
-      console.log('Key down:', e.key, 'Alt:', e.altKey, 'Shift:', e.shiftKey);
     });
 
     window.addEventListener('keyup', (e) => {
       // Clear the key state
       keyStates.set(e.key, false);
-      
+
       // Also clear specific states for modifier keys
       if (e.key === 'Alt' || e.key === 'AltLeft' || e.key === 'AltRight') {
         keyStates.set('Alt', false);
         keyStates.set('AltLeft', false);
         keyStates.set('AltRight', false);
       }
-      
+
       if (e.key === 'Shift' || e.key === 'ShiftLeft' || e.key === 'ShiftRight') {
         keyStates.set('Shift', false);
         keyStates.set('ShiftLeft', false);
         keyStates.set('ShiftRight', false);
       }
-      
-      console.log('Key up:', e.key);
     });
 
     window.addEventListener('blur', () => {
       // Clear all key states when window loses focus
       keyStates.forEach((_, key) => {
         keyStates.set(key, false);
       });
-      console.log('Window blur - all keys reset');
     });
 
     listenersInit = true;

apps/server/src/lib/prompts.ts

@@ -19,6 +19,7 @@ import {
   createDataStreamResponse,
   streamText,
   appendResponseMessages,
+  generateText,
 } from 'ai';
 import {
   IncomingMessageType,
@@ -27,6 +28,7 @@ import {
   type OutgoingMessage,
 } from './types';
 import { DurableObjectOAuthClientProvider } from 'agents/mcp/do-oauth-client-provider';
+import { AiChatPrompt, GmailSearchAssistantSystemPrompt } from '../../lib/prompts';
 import { EPrompts, type IOutgoingMessage, type ParsedMessage } from '../../types';
 import type { MailManager, IGetThreadResponse } from '../../lib/driver/types';
 import { connectionToDriver } from '../../lib/server-utils';
@@ -35,7 +37,6 @@ import { withRetry } from '../../lib/gmail-rate-limit';
 import { getPrompt } from '../../pipelines.effect';
 import { AIChatAgent } from 'agents/ai-chat-agent';
 import { ToolOrchestrator } from './orchestrator';
-import { AiChatPrompt } from '../../lib/prompts';
 import { getPromptName } from '../../pipelines';
 import { anthropic } from '@ai-sdk/anthropic';
 import { connection } from '../../db/schema';
@@ -44,6 +45,7 @@ import { tools as authTools } from './tools';
 import { processToolCalls } from './utils';
 import { env } from 'cloudflare:workers';
 import type { Connection } from 'agents';
+import { openai } from '@ai-sdk/openai';
 import { createDb } from '../../db';
 import { DriverRpcDO } from './rpc';
 import { eq } from 'drizzle-orm';
@@ -493,22 +495,48 @@ export class ZeroDriver extends AIChatAgent<typeof env> {
   }
 
   async inboxRag(query: string) {
-    if (!env.AUTORAG_ID) return { result: 'Not enabled', data: [] };
-    const answer = await env.AI.autorag(env.AUTORAG_ID).aiSearch({
-      query: query,
-      //   rewrite_query: true,
-      max_num_results: 3,
-      ranking_options: {
+    if (!env.AUTORAG_ID) {
+      console.warn('[inboxRag] AUTORAG_ID not configured - RAG search disabled');
+      return { result: 'Not enabled', data: [] };
+    }
+
+    try {
+      const startTime = Date.now();
+      console.log(`[inboxRag] Executing AI search with parameters:`, {
+        query,
+        max_num_results: 3,
         score_threshold: 0.3,
-      },
-      //   stream: true,
-      filters: {
-        type: 'eq',
-        key: 'folder',
-        value: `${this.name}/`,
-      },
-    });
-    return { result: answer.response, data: answer.data };
+        folder_filter: `${this.name}/`,
+      });
+
+      const answer = await env.AI.autorag(env.AUTORAG_ID).aiSearch({
+        query: query,
+        //   rewrite_query: true,
+        max_num_results: 3,
+        ranking_options: {
+          score_threshold: 0.3,
+        },
+        //   stream: true,
+        filters: {
+          type: 'eq',
+          key: 'folder',
+          value: `${this.name}/`,
+        },
+      });
+
+      const duration = Date.now() - startTime;
+
+      return { result: answer.response, data: answer.data };
+    } catch (error) {
+      console.error(`[inboxRag] Search failed for query: "${query}"`, {
+        error: error instanceof Error ? error.message : 'Unknown error',
+        stack: error instanceof Error ? error.stack : undefined,
+        user: this.name,
+      });
+
+      // Return empty result on error to prevent breaking the flow
+      return { result: 'Search failed', data: [] };
+    }
   }
 
   async searchThreads(params: {
@@ -532,10 +560,20 @@ export class ZeroDriver extends AIChatAgent<typeof env> {
       }),
     ).pipe(Effect.catchAll(() => Effect.succeed([])));
 
+    const genQueryEffect = Effect.tryPromise(() =>
+      generateText({
+        model: openai(env.OPENAI_MODEL || 'gpt-4o'),
+        system: GmailSearchAssistantSystemPrompt(),
+        prompt: params.query,
+      }).then((response) => response.text),
+    ).pipe(Effect.catchAll(() => Effect.succeed(query)));
+
+    const genQueryResult = await Effect.runPromise(genQueryEffect);
+
     const rawEffect = Effect.tryPromise(() =>
       this.driver!.list({
         folder,
-        query,
+        query: genQueryResult,
         labelIds,
         maxResults,
         pageToken,
@@ -847,7 +885,7 @@ export class ZeroAgent extends AIChatAgent<typeof env> {
         const rawTools = {
           ...(await authTools(connectionId)),
         };
-        const tools = orchestrator.processTools({});
+        const tools = orchestrator.processTools(rawTools);
         const processedMessages = await processToolCalls(
           {
             messages: this.messages,
@@ -861,7 +899,7 @@ export class ZeroAgent extends AIChatAgent<typeof env> {
           model: anthropic(env.OPENAI_MODEL || 'claude-3-5-haiku-latest'),
           maxSteps: 10,
           messages: processedMessages,
-          tools: rawTools,
+          tools,
           onFinish,
           onError: (error) => {
             console.error('Error in streamText', error);

apps/server/src/routes/agent/index.ts

@@ -14,10 +14,9 @@
  * Reuse or distribution of this file requires a license from Zero Email Inc.
  */
 
-import { GmailSearchAssistantSystemPrompt, getCurrentDateContext } from '../../lib/prompts';
 import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
+import { getCurrentDateContext } from '../../lib/prompts';
 import { getZeroAgent } from '../../lib/server-utils';
-import { anthropic } from '@ai-sdk/anthropic';
 import { connection } from '../../db/schema';
 import { FOLDERS } from '../../lib/utils';
 import { env } from 'cloudflare:workers';
@@ -118,31 +117,6 @@ export class ZeroMCP extends McpAgent<typeof env, Record<string, unknown>, { use
       },
     );
 
-    this.server.registerTool(
-      'buildGmailSearchQuery',
-      {
-        description: 'Build Gmail search query using AI assistance',
-        inputSchema: {
-          query: z.string(),
-        },
-      },
-      async (s) => {
-        const result = await generateText({
-          model: anthropic(env.OPENAI_MODEL || 'claude-3-5-haiku-latest'),
-          system: GmailSearchAssistantSystemPrompt(),
-          prompt: s.query,
-        });
-        return {
-          content: [
-            {
-              type: 'text',
-              text: result.text,
-            },
-          ],
-        };
-      },
-    );
-
     const agent = await getZeroAgent(_connection.id);
 
     this.server.registerTool(

apps/server/src/routes/agent/mcp.ts

@@ -2,9 +2,8 @@ import { composeEmail } from '../../trpc/routes/ai/compose';
 import { perplexity } from '@ai-sdk/perplexity';
 import { generateText, tool } from 'ai';
 
-import { colors, GmailSearchAssistantSystemPrompt } from '../../lib/prompts';
 import { getZeroAgent } from '../../lib/server-utils';
-import { anthropic } from '@ai-sdk/anthropic';
+import { colors } from '../../lib/prompts';
 import { env } from 'cloudflare:workers';
 import { Tools } from '../../types';
 import { z } from 'zod';
@@ -186,18 +185,6 @@ const markAsRead = (connectionId: string) =>
     },
   });
 
-// const inboxRag = (connectionId: string, dataStream?: DataStreamWriter) =>
-//   tool({
-//     description: 'Search the inbox for emails',
-//     parameters: z.object({
-//       query: z.string().describe('The query to search the inbox for'),
-//     }),
-//     execute: async ({ query }) => {
-//       console.log('inboxRag', query);
-//       return await agent.inboxRag(query, dataStream);
-//     },
-//   });
-
 const markAsUnread = (connectionId: string) =>
   tool({
     description: 'Mark emails as unread',
@@ -390,22 +377,6 @@ export const webSearch = () =>
     },
   });
 
-const buildGmailSearchQuery = () =>
-  tool({
-    description: 'Build Gmail search query using AI assistance',
-    parameters: z.object({
-      query: z.string(),
-    }),
-    execute: async ({ query }) => {
-      const result = await generateText({
-        model: anthropic(env.OPENAI_MODEL || 'claude-3-5-haiku-latest'),
-        system: GmailSearchAssistantSystemPrompt(),
-        prompt: query,
-      });
-      return result.text;
-    },
-  });
-
 export const tools = async (connectionId: string) => {
   return {
     [Tools.GetThread]: getEmail(),
@@ -425,7 +396,6 @@ export const tools = async (connectionId: string) => {
         query: z.string().describe('The query to search the web for'),
       }),
     }),
-    [Tools.BuildGmailSearchQuery]: buildGmailSearchQuery(),
     [Tools.InboxRag]: tool({
       description:
         'Search the inbox for emails using natural language. Returns only an array of threadIds.',

apps/server/src/routes/agent/tools.ts

@@ -3,9 +3,9 @@ import {
   OutlookSearchAssistantSystemPrompt,
 } from '../../../lib/prompts';
 import { activeDriverProcedure } from '../../trpc';
+import { env } from 'cloudflare:workers';
 import { openai } from '@ai-sdk/openai';
 import { generateObject } from 'ai';
-import { env } from 'cloudflare:workers';
 import { z } from 'zod';
 
 export const generateSearchQuery = activeDriverProcedure
@@ -28,6 +28,7 @@ export const generateSearchQuery = activeDriverProcedure
       schema: z.object({
         query: z.string(),
       }),
+      output: 'object',
     });
 
     return result.object;