feat: improve chat behavior (#308)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: ben <ben@prologe.io>

packages/plugin/settings.ts

@@ -32,9 +32,10 @@ export class FileOrganizerSettings {
   enableFabric = false;
   useFolderEmbeddings = false;
   useVaultTitles = true;
+  enableSearchGrounding = false;
   showLocalLLMInChat = false;
   customFolderInstructions = "";
-  selectedModel: "gpt-4o" | "llama3.2" | "gemini-2.0-flash-exp" = "gpt-4o";
+  selectedModel: "gpt-4o" | "llama3.2" = "gpt-4o";
   customModelName = "llama3.2";
   tagScoreThreshold = 70;
   formatBehavior: "override" | "newFile" = "override";

packages/plugin/views/assistant/ai-chat/chat.tsx

@@ -18,6 +18,7 @@ import { ModelSelector } from "./model-selector";
 import { ModelType } from "./types";
 import { AudioRecorder } from "./audio-recorder";
 import { logger } from "../../../services/logger";
+import { SearchToggle } from "./components/search-toggle";
 import { SubmitButton } from "./submit-button";
 import { getUniqueReferences, useContextItems } from "./use-context-items";
 import { ContextItems } from "./components/context-items";
@@ -127,6 +128,7 @@ export const ChatComponent: React.FC<ChatComponentProps> = ({
     enableScreenpipe: plugin.settings.enableScreenpipe,
     newUnifiedContext: contextString,
     model: plugin.settings.selectedModel, // Pass selected model to server
+    enableSearchGrounding: plugin.settings.enableSearchGrounding,
   };
 
   const [groundingMetadata, setGroundingMetadata] =
@@ -161,8 +163,7 @@ export const ChatComponent: React.FC<ChatComponentProps> = ({
       // Handle different model types
       if (
         !plugin.settings.showLocalLLMInChat ||
-        selectedModel === "gpt-4o" ||
-        selectedModel === "gemini-2.0-flash-exp"
+        selectedModel === "gpt-4o"
       ) {
         // Use server fetch for non-local models
         return fetch(url, options);
@@ -500,10 +501,13 @@ export const ChatComponent: React.FC<ChatComponentProps> = ({
             unifiedContext={contextString}
             maxContextSize={maxContextSize}
           />
-          <ModelSelector
-            selectedModel={selectedModel}
-            onModelSelect={setSelectedModel}
-          />
+          <div className="flex items-center space-x-2">
+            <SearchToggle selectedModel={selectedModel} />
+            <ModelSelector
+              selectedModel={selectedModel}
+              onModelSelect={setSelectedModel}
+            />
+          </div>
         </div>
       </div>
     </div>

packages/plugin/views/assistant/ai-chat/components/SourcesSection.tsx

@@ -47,7 +47,7 @@ export function SourcesSection({ groundingMetadata }: SourcesSectionProps) {
   const displayedSources = showAll ? sources : sources.slice(0, 3);
 
   return (
-    <div className="mt-6 space-y-4">
+    <div className="mt-6 space-y-4 m-2 z-50">
       <div className="flex justify-between items-center">
         <h3 className="text-base font-medium text-gray-900 dark:text-gray-100">Sources</h3>
         {sources.length > 3 && (
@@ -80,7 +80,14 @@ export function SourcesSection({ groundingMetadata }: SourcesSectionProps) {
             </div>
             <div className="flex-grow min-w-0">
               <h4 className="font-medium text-sm text-gray-900 dark:text-gray-100 truncate">
-                {source.title || source.domain}
+                <a 
+                  href={source.url} 
+                  target="_blank" 
+                  rel="noopener noreferrer" 
+                  className="text-blue-600 hover:underline"
+                >
+                  {source.title || source.domain}
+                </a>
               </h4>
               <p className="text-xs text-gray-500 dark:text-gray-400 truncate">
                 {source.domain}

packages/plugin/views/assistant/ai-chat/components/search-toggle.tsx

@@ -0,0 +1,49 @@
+import React, { useState } from 'react';
+import { usePlugin } from '../../provider';
+import { ModelType } from '../types';
+
+interface SearchToggleProps {
+  selectedModel: ModelType;
+}
+
+export function SearchToggle({ selectedModel }: SearchToggleProps) {
+  const plugin = usePlugin();
+  const [isEnabled, setIsEnabled] = useState(plugin.settings.enableSearchGrounding);
+
+  const handleToggle = async () => {
+    plugin.settings.enableSearchGrounding = !plugin.settings.enableSearchGrounding;
+    await plugin.saveSettings();
+    setIsEnabled(!isEnabled);
+  };
+
+  if (selectedModel !== 'gpt-4o') {
+    return null;
+  }
+
+  return (
+    <button
+      onClick={handleToggle}
+      className={`flex items-center space-x-1 text-sm px-2 py-1 rounded transition-colors ${
+        isEnabled 
+          ? "bg-blue-600 text-white hover:bg-blue-700" 
+          : "bg-[--background-primary-alt] text-[--text-muted] hover:text-[--text-normal] hover:bg-[--background-modifier-hover]"
+      }`}
+      title={isEnabled ? "Disable internet search" : "Enable internet search"}
+    >
+      <svg
+        className="w-4 h-4"
+        viewBox="0 0 24 24"
+        fill="none"
+        stroke="currentColor"
+        strokeWidth="2"
+        strokeLinecap="round"
+        strokeLinejoin="round"
+      >
+        <circle cx="12" cy="12" r="10" />
+        <line x1="2" y1="12" x2="22" y2="12" />
+        <path d="M12 2a15.3 15.3 0 0 1 4 10 15.3 15.3 0 0 1-4 10 15.3 15.3 0 0 1-4-10 15.3 15.3 0 0 1 4-10z" />
+      </svg>
+      {isEnabled && <span>Search</span>}
+    </button>
+  );
+}

packages/plugin/views/assistant/ai-chat/model-selector.tsx

@@ -6,7 +6,6 @@ import { usePlugin } from '../provider';
 const MODEL_DISPLAY_NAMES: Record<ModelType, string> = {
   'gpt-4o': 'Classic',
   'llama3.2': 'Local LLM',
-  'gemini-2.0-flash-exp': 'Internet Search',
   'custom': 'Ollama Model'
 } as const;
 
@@ -35,7 +34,7 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
       return;
     }
     onModelSelect(model);
-    if (model === "gpt-4o" || model === "llama3.2" || model === "gemini-2.0-flash-exp") {
+    if (model === "gpt-4o" || model === "llama3.2") {
       plugin.settings.selectedModel = model;
     }
     await plugin.saveSettings();
@@ -85,12 +84,7 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
               >
                 {getDisplayName("gpt-4o")}
               </div>
-              <div
-                onClick={() => handleModelSelect("gemini-2.0-flash-exp")}
-                className="cursor-pointer block w-full text-left px-4 py-2 text-sm text-[--text-normal] hover:bg-[--background-modifier-hover]"
-              >
-                {getDisplayName("gemini-2.0-flash-exp")}
-              </div>
+
               {isCustomizing ? (
                 <div className="px-4 py-2">
                   <input

packages/plugin/views/assistant/ai-chat/tool-handlers/search-annotation-handler.tsx

@@ -38,12 +38,13 @@ export const SearchAnnotationHandler: React.FC<SearchAnnotationProps> = ({
   if (!groundingMetadata?.groundingSupports?.length) return null;
 
   return (
-    <div className="flex flex-col gap-2 p-2 rounded-md bg-[--background-primary-alt]">
+    <div className="flex flex-col gap-2 p-4 rounded-md bg-[--background-primary-alt] m-2">
       <div className="text-[--text-muted] text-sm">Search Results:</div>
       {groundingMetadata.groundingSupports.map((result, index) => {
-        const sources = result.groundingChunkIndices.map(
-          idx => groundingMetadata.groundingChunks[idx]?.web.title
-        ).filter(Boolean);
+        const sources = result.groundingChunkIndices.map(idx => {
+          const chunk = groundingMetadata.groundingChunks[idx]?.web;
+          return chunk ? { title: chunk.title, uri: chunk.uri } : null;
+        }).filter(Boolean);
 
         const maxScore = Math.max(...result.confidenceScores);
 
@@ -54,7 +55,19 @@ export const SearchAnnotationHandler: React.FC<SearchAnnotationProps> = ({
           >
             <div className="flex justify-between items-center">
               <span className="text-[--text-accent] text-sm">
-                {sources.join(', ')}
+                {sources.map((source, i) => (
+                  <React.Fragment key={i}>
+                    {i > 0 && ', '}
+                    <a 
+                      href={source.uri} 
+                      target="_blank" 
+                      rel="noopener noreferrer"
+                      className="hover:underline"
+                    >
+                      {source.title}
+                    </a>
+                  </React.Fragment>
+                ))}
               </span>
               <span className="text-[--text-muted] text-xs">
                 Score: {(maxScore * 100).toFixed(1)}%

packages/plugin/views/assistant/ai-chat/types.ts

@@ -1 +1 @@
-export type ModelType = "gpt-4o" | "llama3.2" | "gemini-2.0-flash-exp" | string;
+export type ModelType = "gpt-4o" | "llama3.2" | string;

packages/web/app/api/(newai)/chat/route.ts

@@ -29,9 +29,24 @@ export async function POST(req: NextRequest) {
           currentDatetime,
           unifiedContext: oldUnifiedContext,
           model: bodyModel,
+          enableSearchGrounding = false,
         } = await req.json();
 
-        const chosenModelName = bodyModel;
+        let chosenModelName;
+        if (enableSearchGrounding) {
+          console.log("Enabling search grounding");
+          console.log("GOOGLE_API_KEY:", process.env.GOOGLE_API_KEY);
+          console.log("GOOGLE_GENERATIVE_AI_API_KEY:", process.env.GOOGLE_GENERATIVE_AI_API_KEY);
+          if (process.env.GOOGLE_API_KEY || process.env.GOOGLE_GENERATIVE_AI_API_KEY) {
+            console.log("Using Gemini model");
+            chosenModelName = "gemini-2.0-flash-exp";
+          } else {
+            console.log("Using GPT-4o model");
+            chosenModelName = "gpt-4o";
+          }
+        } else {
+          chosenModelName = "gpt-4o";
+        }
         console.log("Chat using model:", chosenModelName);
         const model = getModel(chosenModelName);