Merge pull request #1 from chm3778/add-websearch

prisma/components/settings/ThinkingSection.tsx

@@ -1,6 +1,6 @@
 
 import React from 'react';
-import { RefreshCw } from 'lucide-react';
+import { RefreshCw, Search } from 'lucide-react';
 import { AppConfig, ModelOption } from '../../types';
 import { getValidThinkingLevels } from '../../config';
 import LevelSelect from './LevelSelect';
@@ -39,6 +39,25 @@ const ThinkingSection = ({ config, setConfig, model }: ThinkingSectionProps) =>
           </label>
        </div>
 
+       <div className="p-3 bg-indigo-50 border border-indigo-100 rounded-lg flex items-center justify-between">
+          <div className="flex items-center gap-2">
+             <Search size={16} className="text-indigo-600" />
+             <div>
+                <p className="text-sm font-medium text-indigo-900">Web Search</p>
+                <p className="text-[10px] text-indigo-600/80">Allow the model to use grounded web search when supported.</p>
+             </div>
+          </div>
+          <label className="relative inline-flex items-center cursor-pointer">
+            <input
+              type="checkbox"
+              checked={config.enableWebSearch ?? false}
+              onChange={(e) => setConfig({ ...config, enableWebSearch: e.target.checked })}
+              className="sr-only peer"
+            />
+            <div className="w-9 h-5 bg-slate-300 peer-focus:outline-none rounded-full peer peer-checked:after:translate-x-full peer-checked:after:border-white after:content-[''] after:absolute after:top-[2px] after:start-[2px] after:bg-white after:border-gray-300 after:border after:rounded-full after:h-4 after:w-4 after:transition-all peer-checked:bg-indigo-600"></div>
+          </label>
+       </div>
+
       <LevelSelect 
         label="Manager: Planning Strategy" 
         value={config.planningLevel} 

prisma/config.ts

@@ -36,6 +36,7 @@ export const DEFAULT_CONFIG: AppConfig = {
   customBaseUrl: '',
   enableCustomApi: false,
   enableRecursiveLoop: false,
+  enableWebSearch: false,
   apiProvider: 'google',
   customModels: []
 };

prisma/hooks/useDeepThink.ts

@@ -27,6 +27,8 @@ export const useDeepThink = () => {
     appendExperts
   } = useDeepThinkState();
 
+  let enableWebSearch = false;
+
   /**
    * Orchestrates a single expert's lifecycle (Start -> Stream -> End)
    */
@@ -57,6 +59,7 @@ export const useDeepThink = () => {
         context,
         attachments,
         budget,
+        enableWebSearch,
         signal,
         (textChunk, thoughtChunk) => {
           fullContent += textChunk;
@@ -98,6 +101,7 @@ export const useDeepThink = () => {
     if (abortControllerRef.current) abortControllerRef.current.abort();
     abortControllerRef.current = new AbortController();
     const signal = abortControllerRef.current.signal;
+    enableWebSearch = config.enableWebSearch ?? false;
 
     logger.info('System', 'Starting DeepThink Process', { model, provider: getAIProvider(model) });
 
@@ -137,7 +141,7 @@ export const useDeepThink = () => {
         query, 
         recentHistory,
         currentAttachments,
-        getThinkingBudget(config.planningLevel, model)
+        getThinkingBudget(config.planningLevel, model), enableWebSearch
       );
 
       const primaryExpert: ExpertResult = {
@@ -197,7 +201,7 @@ export const useDeepThink = () => {
 
           const reviewResult = await executeManagerReview(
             ai, model, query, expertsDataRef.current,
-            getThinkingBudget(config.planningLevel, model)
+            getThinkingBudget(config.planningLevel, model), enableWebSearch
           );
 
           if (signal.aborted) return;
@@ -243,7 +247,7 @@ export const useDeepThink = () => {
       await streamSynthesisResponse(
         ai, model, query, recentHistory, expertsDataRef.current,
         currentAttachments,
-        getThinkingBudget(config.synthesisLevel, model), signal,
+        getThinkingBudget(config.synthesisLevel, model), enableWebSearch, signal,
         (textChunk, thoughtChunk) => {
             fullFinalText += textChunk;
             fullFinalThoughts += thoughtChunk;

prisma/services/deepThink/expert.ts

@@ -16,12 +16,28 @@ export const streamExpertResponse = async (
   context: string,
   attachments: MessageAttachment[],
   budget: number,
+  enableWebSearch: boolean,
   signal: AbortSignal,
   onChunk: (text: string, thought: string) => void
 ): Promise<void> => {
   const isGoogle = isGoogleProvider(ai);
 
   if (isGoogle) {
+    const tools = enableWebSearch ? [{ googleSearch: {} }] : undefined;
+    const sourceMap = new Map<string, string | undefined>();
+
+    const addSources = (resp: any) => {
+      const chunks = resp?.candidates?.[0]?.groundingMetadata?.groundingChunks;
+      if (!Array.isArray(chunks)) return;
+      for (const chunk of chunks) {
+        const uri = chunk?.web?.uri;
+        const title = chunk?.web?.title;
+        if (typeof uri === 'string' && uri.trim() && !sourceMap.has(uri)) {
+          sourceMap.set(uri, typeof title === 'string' && title.trim() ? title : undefined);
+        }
+      }
+    };
+
     const contents: any = {
       role: 'user',
       parts: [{ text: expert.prompt }]
@@ -38,7 +54,20 @@ export const streamExpertResponse = async (
       });
     }
 
-    const streamResult = await withRetry(() => ai.models.generateContentStream({
+    let streamResult: any;
+    if (tools) {
+      try {
+        streamResult = await withRetry(() => ai.models.generateContentStream({
+          model: model,
+          contents: contents,
+          config: { systemInstruction: getExpertSystemInstruction(expert.role, expert.description, context), temperature: expert.temperature, tools, thinkingConfig: { thinkingBudget: budget, includeThoughts: true } }
+        }));
+      } catch (e) {
+        logger.warn("Expert", `Web search tool failed for expert ${expert.role}; retrying without it`, e);
+      }
+    }
+
+    streamResult = streamResult || await withRetry(() => ai.models.generateContentStream({
       model: model,
       contents: contents,
       config: {
@@ -58,6 +87,8 @@ export const streamExpertResponse = async (
         let chunkText = "";
         let chunkThought = "";
 
+        if (enableWebSearch) addSources(chunk);
+
         if (chunk.candidates?.[0]?.content?.parts) {
           for (const part of chunk.candidates[0].content.parts) {
             if (part.thought) {
@@ -69,6 +100,11 @@ export const streamExpertResponse = async (
           onChunk(chunkText, chunkThought);
         }
       }
+
+      if (!signal.aborted && enableWebSearch && sourceMap.size > 0) {
+        const sourcesMd = Array.from(sourceMap.entries()).map(([uri, title]) => (title ? `- [${title}](${uri})` : `- ${uri}`)).join('\n');
+        onChunk(`\n\n---\n\n**Sources**\n${sourcesMd}\n`, '');
+      }
     } catch (streamError) {
       logger.error("Expert", `Stream interrupted for expert ${expert.role}`, streamError);
       throw streamError;

prisma/services/deepThink/manager.ts

@@ -17,12 +17,14 @@ export const executeManagerAnalysis = async (
   query: string,
   context: string,
   attachments: MessageAttachment[],
-  budget: number
+  budget: number,
+  enableWebSearch: boolean
 ): Promise<AnalysisResult> => {
   const isGoogle = isGoogleProvider(ai);
   const textPrompt = `Context:\n${context}\n\nCurrent Query: "${query}"`;
 
   if (isGoogle) {
+    const tools = enableWebSearch ? [{ googleSearch: {} }] : undefined;
     const managerSchema = {
       type: Type.OBJECT,
       properties: {
@@ -61,7 +63,20 @@ export const executeManagerAnalysis = async (
     }
 
     try {
-      const analysisResp = await withRetry(() => ai.models.generateContent({
+      let analysisResp: any;
+      if (tools) {
+        try {
+          analysisResp = await withRetry(() => ai.models.generateContent({
+            model: model,
+            contents: contents,
+            config: { systemInstruction: MANAGER_SYSTEM_PROMPT, responseMimeType: "application/json", responseSchema: managerSchema, tools, thinkingConfig: { includeThoughts: true, thinkingBudget: budget } }
+          }));
+        } catch (e) {
+          logger.warn("Manager", "Web search tool failed; retrying without it", e);
+        }
+      }
+
+      analysisResp = analysisResp || await withRetry(() => ai.models.generateContent({
         model: model,
         contents: contents,
         config: {
@@ -151,7 +166,8 @@ export const executeManagerReview = async (
   model: ModelOption,
   query: string,
   currentExperts: ExpertResult[],
-  budget: number
+  budget: number,
+  enableWebSearch: boolean
 ): Promise<ReviewResult> => {
   const isGoogle = isGoogleProvider(ai);
   const expertOutputs = currentExperts.map(e =>
@@ -161,6 +177,7 @@ export const executeManagerReview = async (
   const content = `User Query: "${query}"\n\nCurrent Expert Outputs:\n${expertOutputs}`;
 
   if (isGoogle) {
+    const tools = enableWebSearch ? [{ googleSearch: {} }] : undefined;
     const reviewSchema = {
       type: Type.OBJECT,
       properties: {
@@ -186,7 +203,20 @@ export const executeManagerReview = async (
     };
 
     try {
-      const resp = await withRetry(() => ai.models.generateContent({
+      let resp: any;
+      if (tools) {
+        try {
+          resp = await withRetry(() => ai.models.generateContent({
+            model: model,
+            contents: content,
+            config: { systemInstruction: MANAGER_REVIEW_SYSTEM_PROMPT, responseMimeType: "application/json", responseSchema: reviewSchema, tools, thinkingConfig: { includeThoughts: true, thinkingBudget: budget } }
+          }));
+        } catch (e) {
+          logger.warn("Manager", "Web search tool failed during review; retrying without it", e);
+        }
+      }
+
+      resp = resp || await withRetry(() => ai.models.generateContent({
         model: model,
         contents: content,
         config: {

prisma/services/deepThink/synthesis.ts

@@ -17,13 +17,15 @@ export const streamSynthesisResponse = async (
   expertResults: ExpertResult[],
   attachments: MessageAttachment[],
   budget: number,
+  enableWebSearch: boolean,
   signal: AbortSignal,
   onChunk: (text: string, thought: string) => void
 ): Promise<void> => {
   const prompt = getSynthesisPrompt(historyContext, query, expertResults);
   const isGoogle = isGoogleProvider(ai);
 
   if (isGoogle) {
+    const tools = enableWebSearch ? [{ googleSearch: {} }] : undefined;
     const contents: any = {
       role: 'user',
       parts: [{ text: prompt }]
@@ -40,7 +42,20 @@ export const streamSynthesisResponse = async (
       });
     }
 
-    const synthesisStream = await withRetry(() => ai.models.generateContentStream({
+    let synthesisStream: any;
+    if (tools) {
+      try {
+        synthesisStream = await withRetry(() => ai.models.generateContentStream({
+          model: model,
+          contents: contents,
+          config: { tools, thinkingConfig: { thinkingBudget: budget, includeThoughts: true } }
+        }));
+      } catch (e) {
+        logger.warn("Synthesis", "Web search tool failed; retrying without it", e);
+      }
+    }
+
+    synthesisStream = synthesisStream || await withRetry(() => ai.models.generateContentStream({
       model: model,
       contents: contents,
       config: {

prisma/types.ts

@@ -51,6 +51,7 @@ export type AppConfig = {
   customBaseUrl?: string;
   enableCustomApi?: boolean;
   enableRecursiveLoop?: boolean;
+  enableWebSearch?: boolean;
   apiProvider?: ApiProvider;
   customModels?: CustomModel[];
 };