From e5919487465a8975534d65748c1bf2a801d83e67 Mon Sep 17 00:00:00 2001 From: chm3778 <206088075+chm3778@users.noreply.github.com> Date: Wed, 21 Jan 2026 15:06:50 +0800 Subject: [PATCH] Merge pull request #1 from chm3778/add-websearch Add web search --- .../components/settings/ThinkingSection.tsx | 21 +++++++++- prisma/config.ts | 1 + prisma/hooks/useDeepThink.ts | 10 +++-- prisma/services/deepThink/expert.ts | 38 ++++++++++++++++++- prisma/services/deepThink/manager.ts | 38 +++++++++++++++++-- prisma/services/deepThink/synthesis.ts | 17 ++++++++- prisma/types.ts | 1 + 7 files changed, 116 insertions(+), 10 deletions(-) diff --git a/prisma/components/settings/ThinkingSection.tsx b/prisma/components/settings/ThinkingSection.tsx index b2513de..ca74398 100644 --- a/prisma/components/settings/ThinkingSection.tsx +++ b/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) => +
+
+ +
+

Web Search

+

Allow the model to use grounded web search when supported.

+
+
+ +
+ { 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; diff --git a/prisma/services/deepThink/expert.ts b/prisma/services/deepThink/expert.ts index 562a578..0a5dba9 100644 --- a/prisma/services/deepThink/expert.ts +++ b/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 => { const isGoogle = isGoogleProvider(ai); if (isGoogle) { + const tools = enableWebSearch ? [{ googleSearch: {} }] : undefined; + const sourceMap = new Map(); + + 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; diff --git a/prisma/services/deepThink/manager.ts b/prisma/services/deepThink/manager.ts index e23e65e..901ceab 100644 --- a/prisma/services/deepThink/manager.ts +++ b/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 => { 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 => { 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: { diff --git a/prisma/services/deepThink/synthesis.ts b/prisma/services/deepThink/synthesis.ts index 2499fd6..02a6d0f 100644 --- a/prisma/services/deepThink/synthesis.ts +++ b/prisma/services/deepThink/synthesis.ts @@ -17,6 +17,7 @@ export const streamSynthesisResponse = async ( expertResults: ExpertResult[], attachments: MessageAttachment[], budget: number, + enableWebSearch: boolean, signal: AbortSignal, onChunk: (text: string, thought: string) => void ): Promise => { @@ -24,6 +25,7 @@ export const streamSynthesisResponse = async ( 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: { diff --git a/prisma/types.ts b/prisma/types.ts index 39fbb03..5b7a8e5 100644 --- a/prisma/types.ts +++ b/prisma/types.ts @@ -51,6 +51,7 @@ export type AppConfig = { customBaseUrl?: string; enableCustomApi?: boolean; enableRecursiveLoop?: boolean; + enableWebSearch?: boolean; apiProvider?: ApiProvider; customModels?: CustomModel[]; };