From 62540dc46b50cf999e8e5fe577b548c8edcd611b Mon Sep 17 00:00:00 2001 From: Logan Yang Date: Fri, 6 Mar 2026 19:59:03 -0800 Subject: [PATCH] fix: remove returnAll from agent-facing search tools to prevent token spikes (#2272) (#2273) The agent was triggering returnAll too aggressively, causing massive token usage. This removes returnAll from the tool schema, planning prompt, and all search paths: - Remove returnAll from localSearch Zod schema (no longer agent-settable) - Remove OR stacking that expanded limits when time range or tags present - Maintain similarity score floor at 0.1 (was dropping to 0.0) - Remove returnAll from CopilotPlusChainRunner planning prompt and extraction - Clean up agent prompt in builtinTools.ts - All search paths now use DEFAULT_MAX_SOURCE_CHUNKS (30) consistently Co-authored-by: Claude Opus 4.6 --- .../chainRunner/CopilotPlusChainRunner.ts | 34 +++------- src/tools/SearchTools.schema.test.ts | 68 +------------------ src/tools/SearchTools.ts | 62 +++-------------- src/tools/builtinTools.ts | 6 +- 4 files changed, 25 insertions(+), 145 deletions(-) diff --git a/src/LLMProviders/chainRunner/CopilotPlusChainRunner.ts b/src/LLMProviders/chainRunner/CopilotPlusChainRunner.ts index 29992ccc..dfa912dc 100644 --- a/src/LLMProviders/chainRunner/CopilotPlusChainRunner.ts +++ b/src/LLMProviders/chainRunner/CopilotPlusChainRunner.ts @@ -104,7 +104,7 @@ export class CopilotPlusChainRunner extends BaseChainRunner { private async planToolCalls( userMessage: string, chatModel: BaseChatModel - ): Promise<{ toolCalls: ToolCallWithExecutor[]; salientTerms: string[]; returnAll: boolean }> { + ): Promise<{ toolCalls: ToolCallWithExecutor[]; salientTerms: string[] }> { const availableTools = this.getAvailableToolsForPlanning(); // Check if model supports native tool calling @@ -113,7 +113,6 @@ export class CopilotPlusChainRunner extends BaseChainRunner { return { toolCalls: [], salientTerms: this.extractSalientTermsFromQuery(userMessage), - returnAll: false, }; } @@ -134,10 +133,7 @@ After analyzing, extract key search terms from the user's message that would be - Preserve the EXACT words and language from the user's message (works for any language) - Exclude time expressions (those are handled by tools) -Include your extracted terms as: [SALIENT_TERMS: term1, term2, term3] - -If the user wants ALL matching notes (e.g., "find all my X", "list every Y", "show me all Z", "how many notes about W"), output: [RETURN_ALL: true] -Otherwise omit RETURN_ALL or output: [RETURN_ALL: false]`; +Include your extracted terms as: [SALIENT_TERMS: term1, term2, term3]`; // Create planning request const planningMessages = [ @@ -165,11 +161,8 @@ Otherwise omit RETURN_ALL or output: [RETURN_ALL: false]`; logInfo("[CopilotPlus] Native tool calls:", nativeToolCalls.length); - // Extract salient terms and returnAll intent from response text - const { salientTerms, returnAll } = this.extractPlanningFieldsFromResponse( - responseText, - userMessage - ); + // Extract salient terms from response text + const { salientTerms } = this.extractPlanningFieldsFromResponse(responseText, userMessage); // Convert native tool calls to executor format const toolCalls: ToolCallWithExecutor[] = []; @@ -186,16 +179,16 @@ Otherwise omit RETURN_ALL or output: [RETURN_ALL: false]`; } } - return { toolCalls, salientTerms, returnAll }; + return { toolCalls, salientTerms }; } /** - * Extract salient terms and returnAll intent from model response. + * Extract salient terms from model response. */ private extractPlanningFieldsFromResponse( responseText: string, originalQuery: string - ): { salientTerms: string[]; returnAll: boolean } { + ): { salientTerms: string[] } { // Extract salient terms from [SALIENT_TERMS: ...] format let salientTerms: string[]; const termsMatch = responseText.match(/\[SALIENT_TERMS:\s*([^\]]+?)\s*\]/i); @@ -209,12 +202,7 @@ Otherwise omit RETURN_ALL or output: [RETURN_ALL: false]`; salientTerms = this.extractSalientTermsFromQuery(originalQuery); } - // Extract returnAll from [RETURN_ALL: true/false] format - // Allow optional whitespace/newlines around the value and before closing bracket - const returnAllMatch = responseText.match(/\[RETURN_ALL:\s*(true|false)\s*\]/i); - const returnAll = returnAllMatch ? returnAllMatch[1].toLowerCase() === "true" : false; - - return { salientTerms, returnAll }; + return { salientTerms }; } /** @@ -237,7 +225,7 @@ Otherwise omit RETURN_ALL or output: [RETURN_ALL: false]`; private async processAtCommands( userMessage: string, existingToolCalls: ToolCallWithExecutor[], - context: { salientTerms: string[]; timeRange?: any; returnAll?: boolean } + context: { salientTerms: string[]; timeRange?: any } ): Promise { const message = userMessage.toLowerCase(); const cleanQuery = this.removeAtCommands(userMessage); @@ -254,7 +242,6 @@ Otherwise omit RETURN_ALL or output: [RETURN_ALL: false]`; query: cleanQuery, salientTerms: context.salientTerms, timeRange: context.timeRange, - returnAll: context.returnAll === true ? true : undefined, }, }); } @@ -817,11 +804,10 @@ Otherwise omit RETURN_ALL or output: [RETURN_ALL: false]`; }); // Process @commands - this may add localSearch, webSearch, or updateMemory - // Pass timeRange and returnAll in context so @vault commands can use them + // Pass timeRange in context so @vault commands can use them toolCalls = await this.processAtCommands(messageForAnalysis, filteredToolCalls, { salientTerms: planningResult.salientTerms, timeRange, - returnAll: planningResult.returnAll, }); } catch (error: any) { return this.handleResponse( diff --git a/src/tools/SearchTools.schema.test.ts b/src/tools/SearchTools.schema.test.ts index d19d7a7f..0c634040 100644 --- a/src/tools/SearchTools.schema.test.ts +++ b/src/tools/SearchTools.schema.test.ts @@ -19,14 +19,6 @@ describe("SearchTools Schema Validation", () => { }) .optional() .describe("Time range for search"), - returnAll: z - .preprocess((val) => { - if (typeof val === "string") { - return val.toLowerCase() === "true"; - } - return val; - }, z.boolean().optional()) - .describe("Return all matching notes"), }); test("validates correct input structure", () => { @@ -109,67 +101,11 @@ describe("SearchTools Schema Validation", () => { expect(result.success).toBe(false); }); - test("accepts boolean returnAll: true", () => { + test("ignores unknown fields like returnAll", () => { const input = { query: "find all notes", salientTerms: ["notes"], returnAll: true }; const result = localSearchSchema.safeParse(input); + // Schema strips unknown fields but still parses successfully expect(result.success).toBe(true); - if (result.success) { - expect(result.data.returnAll).toBe(true); - } - }); - - test("accepts boolean returnAll: false", () => { - const input = { query: "find notes", salientTerms: ["notes"], returnAll: false }; - const result = localSearchSchema.safeParse(input); - expect(result.success).toBe(true); - if (result.success) { - expect(result.data.returnAll).toBe(false); - } - }); - - test("coerces string 'True' to boolean true", () => { - const input = { query: "find all notes", salientTerms: ["notes"], returnAll: "True" }; - const result = localSearchSchema.safeParse(input); - expect(result.success).toBe(true); - if (result.success) { - expect(result.data.returnAll).toBe(true); - } - }); - - test("coerces string 'true' to boolean true", () => { - const input = { query: "find all notes", salientTerms: ["notes"], returnAll: "true" }; - const result = localSearchSchema.safeParse(input); - expect(result.success).toBe(true); - if (result.success) { - expect(result.data.returnAll).toBe(true); - } - }); - - test("coerces string 'FALSE' to boolean false", () => { - const input = { query: "find notes", salientTerms: ["notes"], returnAll: "FALSE" }; - const result = localSearchSchema.safeParse(input); - expect(result.success).toBe(true); - if (result.success) { - expect(result.data.returnAll).toBe(false); - } - }); - - test("coerces string 'False' to boolean false", () => { - const input = { query: "find notes", salientTerms: ["notes"], returnAll: "False" }; - const result = localSearchSchema.safeParse(input); - expect(result.success).toBe(true); - if (result.success) { - expect(result.data.returnAll).toBe(false); - } - }); - - test("accepts omitted returnAll (optional)", () => { - const input = { query: "find notes", salientTerms: ["notes"] }; - const result = localSearchSchema.safeParse(input); - expect(result.success).toBe(true); - if (result.success) { - expect(result.data.returnAll).toBeUndefined(); - } }); }); diff --git a/src/tools/SearchTools.ts b/src/tools/SearchTools.ts index db33dd4b..5684770d 100644 --- a/src/tools/SearchTools.ts +++ b/src/tools/SearchTools.ts @@ -9,7 +9,6 @@ import { getSettings } from "@/settings/model"; import { z } from "zod"; import { deduplicateSources } from "@/LLMProviders/chainRunner/utils/toolExecution"; import { createLangChainTool } from "./createLangChainTool"; -import { RETURN_ALL_LIMIT } from "@/search/v3/SearchCore"; import { getWebSearchCitationInstructions } from "@/LLMProviders/chainRunner/utils/citationUtils"; import { TieredLexicalRetriever } from "@/search/v3/TieredLexicalRetriever"; import { FilterRetriever } from "@/search/v3/FilterRetriever"; @@ -76,19 +75,6 @@ const localSearchSchema = z.object({ }) .optional() .describe("Optional time range filter. Use epoch milliseconds from getTimeRangeMs result."), - returnAll: z - .preprocess((val) => { - if (typeof val === "string") { - return val.toLowerCase() === "true"; - } - return val; - }, z.boolean().optional()) - .describe( - "Set to true when the user wants ALL matching notes, not just the best few. " + - "Use for requests like 'find all my X', 'list every Y', 'show me all my Z', " + - "'how many notes about W'. Returns up to 100 results instead of default 30. " + - "Leave false/undefined for normal questions." - ), _preExpandedQuery: z .object({ originalQuery: z.string(), @@ -107,26 +93,19 @@ async function performLexicalSearch({ salientTerms, forceLexical = false, preExpandedQuery, - returnAll = false, }: { timeRange?: { startTime: number; endTime: number }; query: string; salientTerms: string[]; forceLexical?: boolean; preExpandedQuery?: QueryExpansionInfo; - /** Caller-requested return-all (from LLM tool schema). Combined with implicit triggers. */ - returnAll?: boolean; }) { // Extract tag terms for self-host retriever (server-side tag filtering) const tagTerms = salientTerms.filter((term) => term.startsWith("#")); - // Time-range and tag-focused queries always use expanded result limits - const useExpandedLimits = returnAll || timeRange !== undefined || tagTerms.length > 0; - const effectiveMaxK = useExpandedLimits ? RETURN_ALL_LIMIT : DEFAULT_MAX_SOURCE_CHUNKS; + const effectiveMaxK = DEFAULT_MAX_SOURCE_CHUNKS; - logInfo( - `lexicalSearch useExpandedLimits: ${useExpandedLimits} (timeRange: ${!!timeRange}, tags: ${tagTerms.length > 0}, explicit: ${returnAll}), forceLexical: ${forceLexical}` - ); + logInfo(`lexicalSearch effectiveMaxK: ${effectiveMaxK}, forceLexical: ${forceLexical}`); // Convert QueryExpansionInfo to ExpandedQuery format (adding queries field) const convertedPreExpansion = preExpandedQuery @@ -145,7 +124,6 @@ async function performLexicalSearch({ salientTerms, timeRange, maxK: effectiveMaxK, - returnAll: useExpandedLimits, }); const filterDocs = await filterRetriever.getRelevantDocuments(query); @@ -158,11 +136,10 @@ async function performLexicalSearch({ if (!filterRetriever.hasTimeRange()) { const retrieverOptions = { - minSimilarityScore: useExpandedLimits ? 0.0 : 0.1, + minSimilarityScore: 0.1, maxK: effectiveMaxK, salientTerms, textWeight: TEXT_WEIGHT, - returnAll: useExpandedLimits, useRerankerThreshold: 0.5, tagTerms, // Used by SelfHostRetriever for server-side tag filtering preExpandedQuery: convertedPreExpansion, // Pass pre-expanded data to skip double expansion @@ -260,13 +237,12 @@ const lexicalSearchTool = createLangChainTool({ name: "lexicalSearch", description: "Search for notes using lexical/keyword-based search", schema: localSearchSchema, - func: async ({ timeRange: rawTimeRange, query, salientTerms, returnAll }) => { + func: async ({ timeRange: rawTimeRange, query, salientTerms }) => { const timeRange = validateTimeRange(rawTimeRange); return await performLexicalSearch({ timeRange, query, salientTerms, - returnAll: returnAll === true, }); }, }); @@ -276,28 +252,20 @@ const semanticSearchTool = createLangChainTool({ name: "semanticSearch", description: "Search for notes using semantic/meaning-based search with embeddings", schema: localSearchSchema, - func: async ({ timeRange: rawTimeRange, query, salientTerms, returnAll }) => { + func: async ({ timeRange: rawTimeRange, query, salientTerms }) => { const timeRange = validateTimeRange(rawTimeRange); - // Time-range and tag-focused queries always use expanded result limits - const tagTerms = salientTerms.filter((term) => term.startsWith("#")); - const useExpandedLimits = returnAll === true || timeRange !== undefined || tagTerms.length > 0; - const effectiveMaxK = useExpandedLimits - ? Math.max(DEFAULT_MAX_SOURCE_CHUNKS, 200) - : DEFAULT_MAX_SOURCE_CHUNKS; + const effectiveMaxK = DEFAULT_MAX_SOURCE_CHUNKS; - logInfo( - `semanticSearch useExpandedLimits: ${useExpandedLimits} (timeRange: ${!!timeRange}, tags: ${tagTerms.length > 0}, explicit: ${returnAll === true})` - ); + logInfo(`semanticSearch effectiveMaxK: ${effectiveMaxK}`); // Always use HybridRetriever for semantic search const retriever = new (await import("@/search/hybridRetriever")).HybridRetriever({ - minSimilarityScore: useExpandedLimits ? 0.0 : 0.1, + minSimilarityScore: 0.1, maxK: effectiveMaxK, salientTerms, timeRange, textWeight: TEXT_WEIGHT, - returnAll: useExpandedLimits, useRerankerThreshold: 0.5, }); @@ -387,24 +355,20 @@ function validateTimeRange(timeRange?: { async function performMiyoSearch({ query, salientTerms, - returnAll = false, timeRange, }: { query: string; salientTerms: string[]; - returnAll?: boolean; timeRange?: { startTime: number; endTime: number }; }) { const tagTerms = salientTerms.filter((term) => term.startsWith("#")); - const useExpandedLimits = returnAll || timeRange !== undefined || tagTerms.length > 0; - const effectiveMaxK = useExpandedLimits ? RETURN_ALL_LIMIT : DEFAULT_MAX_SOURCE_CHUNKS; + const effectiveMaxK = DEFAULT_MAX_SOURCE_CHUNKS; // FilterRetriever for local tag/title/time-range matches const filterRetriever = new FilterRetriever(app, { salientTerms, timeRange, maxK: effectiveMaxK, - returnAll: useExpandedLimits, }); const filterDocs = await filterRetriever.getRelevantDocuments(query); @@ -413,11 +377,10 @@ async function performMiyoSearch({ let miyoDocs: import("@langchain/core/documents").Document[] = []; if (!filterRetriever.hasTimeRange()) { const miyoRetriever = RetrieverFactory.createMiyoRetriever(app, { - minSimilarityScore: useExpandedLimits ? 0.0 : 0.1, + minSimilarityScore: 0.1, maxK: effectiveMaxK, salientTerms, textWeight: TEXT_WEIGHT, - returnAll: useExpandedLimits, useRerankerThreshold: 0.5, tagTerms, }); @@ -462,7 +425,7 @@ const localSearchTool = createLangChainTool({ description: "Search for notes in the vault based on query, salient terms, and optional time range", schema: localSearchSchema, - func: async ({ timeRange: rawTimeRange, query, salientTerms, returnAll, _preExpandedQuery }) => { + func: async ({ timeRange: rawTimeRange, query, salientTerms, _preExpandedQuery }) => { // Validate time range to prevent LLM hallucinations (e.g., {startTime: 0, endTime: 0}) const timeRange = validateTimeRange(rawTimeRange); @@ -472,7 +435,6 @@ const localSearchTool = createLangChainTool({ return await performMiyoSearch({ query, salientTerms, - returnAll: returnAll === true, timeRange, }); } @@ -490,7 +452,6 @@ const localSearchTool = createLangChainTool({ salientTerms, forceLexical: true, preExpandedQuery: _preExpandedQuery, - returnAll: returnAll === true, }); } @@ -504,7 +465,6 @@ const localSearchTool = createLangChainTool({ query, salientTerms, preExpandedQuery: _preExpandedQuery, - returnAll: returnAll === true, }); }, }); diff --git a/src/tools/builtinTools.ts b/src/tools/builtinTools.ts index 8acda5dc..91778f66 100644 --- a/src/tools/builtinTools.ts +++ b/src/tools/builtinTools.ts @@ -70,10 +70,8 @@ For time-based searches with meaningful terms (e.g., "python debugging notes fro 1. First call getTimeRangeMs with timeExpression: "yesterday" 2. Then use localSearch with the returned timeRange, query: "python debugging notes", salientTerms: ["python", "debugging", "notes"] -For exhaustive "find all" searches: -- Set returnAll: true when the user wants ALL matching notes (e.g., "find all my X", "list every Y", "show me all Z", "how many notes about W") -- Keep returnAll: false (or omit) for normal questions seeking specific information -- When setting returnAll: true, also call getFileTree to get all note titles as reference. This helps verify search completeness and identify notes the search may have missed.`, +For broad searches: +- If the user wants a comprehensive list, use getFileTree to get all note titles as reference. This helps verify search completeness and identify notes the search may have missed.`, }, }, {