diff --git a/src/lib/agents/search/researcher/actions/registry.ts b/src/lib/agents/search/researcher/actions/registry.ts index 8e0530c..a8de513 100644 --- a/src/lib/agents/search/researcher/actions/registry.ts +++ b/src/lib/agents/search/researcher/actions/registry.ts @@ -5,6 +5,7 @@ import { ClassifierOutput, ResearchAction, SearchAgentConfig, + SearchSources, } from '../../types'; class ActionRegistry { @@ -22,6 +23,7 @@ class ActionRegistry { classification: ClassifierOutput; fileIds: string[]; mode: SearchAgentConfig['mode']; + sources: SearchSources[]; }): ResearchAction[] { return Array.from( this.actions.values().filter((action) => action.enabled(config)), @@ -32,6 +34,7 @@ class ActionRegistry { classification: ClassifierOutput; fileIds: string[]; mode: SearchAgentConfig['mode']; + sources: SearchSources[]; }): Tool[] { const availableActions = this.getAvailableActions(config); @@ -46,6 +49,7 @@ class ActionRegistry { classification: ClassifierOutput; fileIds: string[]; mode: SearchAgentConfig['mode']; + sources: SearchSources[]; }): string { const availableActions = this.getAvailableActions(config); diff --git a/src/lib/agents/search/researcher/actions/webSearch.ts b/src/lib/agents/search/researcher/actions/webSearch.ts index 5eb3af4..3c712a2 100644 --- a/src/lib/agents/search/researcher/actions/webSearch.ts +++ b/src/lib/agents/search/researcher/actions/webSearch.ts @@ -82,6 +82,7 @@ const webSearchAction: ResearchAction = { return prompt; }, enabled: (config) => + config.sources.includes('web') && config.classification.classification.skipSearch === false, execute: async (input, additionalConfig) => { input.queries = input.queries.slice(0, 3); diff --git a/src/lib/agents/search/researcher/index.ts b/src/lib/agents/search/researcher/index.ts index 334ba74..814b0e3 100644 --- a/src/lib/agents/search/researcher/index.ts +++ b/src/lib/agents/search/researcher/index.ts @@ -23,6 +23,7 @@ class Researcher { classification: input.classification, fileIds: input.config.fileIds, mode: input.config.mode, + sources: input.config.sources, }); const availableActionsDescription = @@ -30,6 +31,7 @@ class Researcher { classification: input.classification, fileIds: input.config.fileIds, mode: input.config.mode, + sources: input.config.sources, }); const researchBlockId = crypto.randomUUID(); diff --git a/src/lib/agents/search/types.ts b/src/lib/agents/search/types.ts index 0733de3..04cbe19 100644 --- a/src/lib/agents/search/types.ts +++ b/src/lib/agents/search/types.ts @@ -107,6 +107,7 @@ export interface ResearchAction< classification: ClassifierOutput; fileIds: string[]; mode: SearchAgentConfig['mode']; + sources: SearchSources[]; }) => boolean; execute: ( params: z.infer,