Merge pull request #6206 from continuedev/dallin/limit-grep-output

Limit and Truncate Grep Search output

Author: sestinj

core/config/types.ts

@@ -719,7 +719,7 @@ declare global {
   
     getPinnedFiles(): Promise<string[]>;
   
-    getSearchResults(query: string): Promise<string>;
+    getSearchResults(query: string, maxResults?: number): Promise<string>;
   
     subprocess(command: string, cwd?: string): Promise<[string, string]>;
   

core/context/providers/SearchContextProvider.ts

@@ -6,6 +6,7 @@ import {
 import { formatGrepSearchResults } from "../../util/grepSearch.js";
 import { BaseContextProvider } from "../index.js";
 
+const DEFAULT_MAX_SEARCH_CONTEXT_RESULTS = 200;
 class SearchContextProvider extends BaseContextProvider {
   static description: ContextProviderDescription = {
     title: "search",
@@ -19,8 +20,12 @@ class SearchContextProvider extends BaseContextProvider {
     query: string,
     extras: ContextProviderExtras,
   ): Promise<ContextItem[]> {
-    const results = await extras.ide.getSearchResults(query);
-    const formatted = formatGrepSearchResults(results);
+    const results = await extras.ide.getSearchResults(
+      query,
+      this.options?.maxResults ?? DEFAULT_MAX_SEARCH_CONTEXT_RESULTS,
+    );
+    // Note, search context provider will not truncate result chars, but will limit number of results
+    const { formatted } = formatGrepSearchResults(results);
     return [
       {
         description: "Search results",

core/index.d.ts

@@ -781,7 +781,7 @@ export interface IDE {
 
   getPinnedFiles(): Promise<string[]>;
 
-  getSearchResults(query: string): Promise<string>;
+  getSearchResults(query: string, maxResults?: number): Promise<string>;
 
   getFileResults(pattern: string): Promise<string[]>;
 

core/protocol/ide.ts

@@ -29,7 +29,7 @@ export type ToIdeFromWebviewOrCoreProtocol = {
   openFile: [{ path: string }, void];
   openUrl: [string, void];
   runCommand: [{ command: string; options?: TerminalOptions }, void];
-  getSearchResults: [{ query: string }, string];
+  getSearchResults: [{ query: string; maxResults?: number }, string];
   getFileResults: [{ pattern: string }, string[]];
   subprocess: [{ command: string; cwd?: string }, [string, string]];
   saveFile: [{ filepath: string }, void];

core/protocol/messenger/messageIde.ts

@@ -185,8 +185,8 @@ export class MessageIde implements IDE {
     return this.request("getPinnedFiles", undefined);
   }
 
-  getSearchResults(query: string): Promise<string> {
-    return this.request("getSearchResults", { query });
+  getSearchResults(query: string, maxResults?: number): Promise<string> {
+    return this.request("getSearchResults", { query, maxResults });
   }
 
   getFileResults(pattern: string): Promise<string[]> {

core/protocol/messenger/reverseMessageIde.ts

@@ -159,7 +159,7 @@ export class ReverseMessageIde {
     });
 
     this.on("getSearchResults", (data) => {
-      return this.ide.getSearchResults(data.query);
+      return this.ide.getSearchResults(data.query, data.maxResults);
     });
 
     this.on("getFileResults", (data) => {

core/tools/definitions/grepSearch.ts

@@ -12,7 +12,8 @@ export const grepSearchTool: Tool = {
   group: BUILT_IN_GROUP_NAME,
   function: {
     name: BuiltInToolNames.GrepSearch,
-    description: "Perform a search over the repository using ripgrep.",
+    description:
+      "Perform a search over the repository using ripgrep. Output may be truncated, so use targeted queries",
     parameters: {
       type: "object",
       required: ["query"],

core/tools/implementations/grepSearch.ts

@@ -1,13 +1,44 @@
 import { ToolImpl } from ".";
+import { ContextItem } from "../..";
 import { formatGrepSearchResults } from "../../util/grepSearch";
 
+const DEFAULT_GREP_SEARCH_RESULTS_LIMIT = 100;
+const DEFAULT_GREP_SEARCH_CHAR_LIMIT = 5000; // ~1000 tokens, will keep truncation simply for now
+
 export const grepSearchImpl: ToolImpl = async (args, extras) => {
-  const results = await extras.ide.getSearchResults(args.query);
-  return [
+  const results = await extras.ide.getSearchResults(
+    args.query,
+    DEFAULT_GREP_SEARCH_RESULTS_LIMIT,
+  );
+  const { formatted, numResults, truncated } = formatGrepSearchResults(
+    results,
+    DEFAULT_GREP_SEARCH_CHAR_LIMIT,
+  );
+  const truncationReasons: string[] = [];
+  if (numResults === DEFAULT_GREP_SEARCH_RESULTS_LIMIT) {
+    truncationReasons.push(
+      `the number of results exceeded ${DEFAULT_GREP_SEARCH_RESULTS_LIMIT}`,
+    );
+  }
+  if (truncated) {
+    truncationReasons.push(
+      `the number of characters exceeded ${DEFAULT_GREP_SEARCH_CHAR_LIMIT}`,
+    );
+  }
+
+  const contextItems: ContextItem[] = [
     {
       name: "Search results",
       description: "Results from grep search",
-      content: formatGrepSearchResults(results),
+      content: formatted,
     },
   ];
+  if (truncationReasons.length > 0) {
+    contextItems.push({
+      name: "Search truncation warning",
+      description: "Informs the model that search results were truncated",
+      content: `The above search results were truncated because ${truncationReasons.join(" and ")}. If the results are not satisfactory, try refining your search query.`,
+    });
+  }
+  return contextItems;
 };

core/util/filesystem.ts

@@ -234,7 +234,7 @@ class FileSystemIde implements IDE {
     return Promise.resolve([]);
   }
 
-  async getSearchResults(query: string): Promise<string> {
+  async getSearchResults(query: string, maxResults?: number): Promise<string> {
     return "";
   }
 

core/util/grepSearch.ts

@@ -1,4 +1,19 @@
-export function formatGrepSearchResults(results: string): string {
+/*
+  Formats the output of a grep search to reduce unnecessary indentation, lines, etc
+  Assumes a command with these params
+    ripgrep -i --ignore-file .continueignore --ignore-file .gitignore -C 2 --heading -m 100 -e <query> .
+  
+  Also can truncate the output to a specified number of characters
+*/
+export function formatGrepSearchResults(
+  results: string,
+  maxChars?: number,
+): {
+  formatted: string;
+  numResults: number;
+  truncated: boolean;
+} {
+  let numResults = 0;
   const keepLines: string[] = [];
 
   function countLeadingSpaces(line: string) {
@@ -45,6 +60,7 @@ export function formatGrepSearchResults(results: string): string {
     if (line.startsWith("./") || line === "--") {
       processResult(resultLines); // process previous result
       resultLines = [line];
+      numResults++;
       continue;
     }
 
@@ -57,5 +73,18 @@ export function formatGrepSearchResults(results: string): string {
   }
   processResult(resultLines);
 
-  return keepLines.join("\n");
+  const formatted = keepLines.join("\n");
+  if (maxChars && formatted.length > maxChars) {
+    return {
+      formatted: formatted.substring(0, maxChars),
+      numResults,
+      truncated: true,
+    };
+  } else {
+    return {
+      formatted,
+      numResults,
+      truncated: false,
+    };
+  }
 }