feat(tools): Centralize shell tool summarization (#4009)
This commit is contained in:
parent
09a3b7d5e1
commit
44ef0408f3
|
@ -656,39 +656,15 @@ export class CoreToolScheduler {
|
|||
return;
|
||||
}
|
||||
|
||||
let resultForDisplay: ToolResult = toolResult;
|
||||
let summary: string | undefined;
|
||||
if (scheduledCall.tool.summarizer) {
|
||||
try {
|
||||
const toolSignal = new AbortController();
|
||||
summary = await scheduledCall.tool.summarizer(
|
||||
toolResult,
|
||||
this.config.getGeminiClient(),
|
||||
toolSignal.signal,
|
||||
);
|
||||
if (toolSignal.signal.aborted) {
|
||||
console.debug('aborted summarizing tool result');
|
||||
return;
|
||||
}
|
||||
if (scheduledCall.tool?.shouldSummarizeDisplay) {
|
||||
resultForDisplay = {
|
||||
...toolResult,
|
||||
returnDisplay: summary,
|
||||
};
|
||||
}
|
||||
} catch (e) {
|
||||
console.error('Error summarizing tool result:', e);
|
||||
}
|
||||
}
|
||||
const response = convertToFunctionResponse(
|
||||
toolName,
|
||||
callId,
|
||||
summary ? [summary] : toolResult.llmContent,
|
||||
toolResult.llmContent,
|
||||
);
|
||||
const successResponse: ToolCallResponseInfo = {
|
||||
callId,
|
||||
responseParts: response,
|
||||
resultDisplay: resultForDisplay.returnDisplay,
|
||||
resultDisplay: toolResult.returnDisplay,
|
||||
error: undefined,
|
||||
};
|
||||
|
||||
|
|
|
@ -68,17 +68,9 @@ export async function executeToolCall(
|
|||
// No live output callback for non-interactive mode
|
||||
);
|
||||
|
||||
const tool_output = tool.summarizer
|
||||
? await tool.summarizer(
|
||||
toolResult,
|
||||
config.getGeminiClient(),
|
||||
effectiveAbortSignal,
|
||||
)
|
||||
: toolResult.llmContent;
|
||||
const tool_output = toolResult.llmContent;
|
||||
|
||||
const tool_display = tool.shouldSummarizeDisplay
|
||||
? (tool_output as string)
|
||||
: toolResult.returnDisplay;
|
||||
const tool_display = toolResult.returnDisplay;
|
||||
|
||||
const durationMs = Date.now() - startTime;
|
||||
logToolCall(config, {
|
||||
|
|
|
@ -4,9 +4,11 @@
|
|||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { expect, describe, it } from 'vitest';
|
||||
import { expect, describe, it, vi, beforeEach } from 'vitest';
|
||||
import { ShellTool } from './shell.js';
|
||||
import { Config } from '../config/config.js';
|
||||
import * as summarizer from '../utils/summarizer.js';
|
||||
import { GeminiClient } from '../core/client.js';
|
||||
|
||||
describe('ShellTool', () => {
|
||||
it('should allow a command if no restrictions are provided', async () => {
|
||||
|
@ -396,3 +398,35 @@ describe('ShellTool', () => {
|
|||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('ShellTool Bug Reproduction', () => {
|
||||
let shellTool: ShellTool;
|
||||
let config: Config;
|
||||
|
||||
beforeEach(() => {
|
||||
config = {
|
||||
getCoreTools: () => undefined,
|
||||
getExcludeTools: () => undefined,
|
||||
getDebugMode: () => false,
|
||||
getGeminiClient: () => ({}) as GeminiClient,
|
||||
getTargetDir: () => '.',
|
||||
} as unknown as Config;
|
||||
shellTool = new ShellTool(config);
|
||||
});
|
||||
|
||||
it('should not let the summarizer override the return display', async () => {
|
||||
const summarizeSpy = vi
|
||||
.spyOn(summarizer, 'summarizeToolOutput')
|
||||
.mockResolvedValue('summarized output');
|
||||
|
||||
const abortSignal = new AbortController().signal;
|
||||
const result = await shellTool.execute(
|
||||
{ command: 'echo "hello"' },
|
||||
abortSignal,
|
||||
);
|
||||
|
||||
expect(result.returnDisplay).toBe('hello\n');
|
||||
expect(result.llmContent).toBe('summarized output');
|
||||
expect(summarizeSpy).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
|
|
@ -27,7 +27,7 @@ export interface ShellToolParams {
|
|||
directory?: string;
|
||||
}
|
||||
import { spawn } from 'child_process';
|
||||
import { llmSummarizer } from '../utils/summarizer.js';
|
||||
import { summarizeToolOutput } from '../utils/summarizer.js';
|
||||
|
||||
const OUTPUT_UPDATE_INTERVAL_MS = 1000;
|
||||
|
||||
|
@ -74,8 +74,6 @@ Process Group PGID: Process group started or \`(none)\``,
|
|||
},
|
||||
false, // output is not markdown
|
||||
true, // output can be updated
|
||||
llmSummarizer,
|
||||
true, // should summarize display output
|
||||
);
|
||||
}
|
||||
|
||||
|
@ -490,6 +488,16 @@ Process Group PGID: Process group started or \`(none)\``,
|
|||
// returnDisplayMessage will remain empty, which is fine.
|
||||
}
|
||||
}
|
||||
return { llmContent, returnDisplay: returnDisplayMessage };
|
||||
|
||||
const summary = await summarizeToolOutput(
|
||||
llmContent,
|
||||
this.config.getGeminiClient(),
|
||||
abortSignal,
|
||||
);
|
||||
|
||||
return {
|
||||
llmContent: summary,
|
||||
returnDisplay: returnDisplayMessage,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
|
|
@ -11,7 +11,6 @@ import { spawn } from 'node:child_process';
|
|||
import { StringDecoder } from 'node:string_decoder';
|
||||
import { discoverMcpTools } from './mcp-client.js';
|
||||
import { DiscoveredMCPTool } from './mcp-tool.js';
|
||||
import { defaultSummarizer } from '../utils/summarizer.js';
|
||||
import { parse } from 'shell-quote';
|
||||
|
||||
type ToolParams = Record<string, unknown>;
|
||||
|
@ -48,7 +47,6 @@ Signal: Signal number or \`(none)\` if no signal was received.
|
|||
parameterSchema,
|
||||
false, // isOutputMarkdown
|
||||
false, // canUpdateOutput
|
||||
defaultSummarizer,
|
||||
);
|
||||
}
|
||||
|
||||
|
|
|
@ -5,7 +5,6 @@
|
|||
*/
|
||||
|
||||
import { FunctionDeclaration, PartListUnion, Schema } from '@google/genai';
|
||||
import { Summarizer, defaultSummarizer } from '../utils/summarizer.js';
|
||||
|
||||
/**
|
||||
* Interface representing the base Tool functionality
|
||||
|
@ -44,16 +43,6 @@ export interface Tool<
|
|||
*/
|
||||
canUpdateOutput: boolean;
|
||||
|
||||
/**
|
||||
* A function that summarizes the result of the tool execution.
|
||||
*/
|
||||
summarizer?: Summarizer;
|
||||
|
||||
/**
|
||||
* Whether the tool's display output should be summarized
|
||||
*/
|
||||
shouldSummarizeDisplay?: boolean;
|
||||
|
||||
/**
|
||||
* Validates the parameters for the tool
|
||||
* Should be called from both `shouldConfirmExecute` and `execute`
|
||||
|
@ -109,8 +98,6 @@ export abstract class BaseTool<
|
|||
* @param isOutputMarkdown Whether the tool's output should be rendered as markdown
|
||||
* @param canUpdateOutput Whether the tool supports live (streaming) output
|
||||
* @param parameterSchema JSON Schema defining the parameters
|
||||
* @param summarizer Function to summarize the tool's output
|
||||
* @param shouldSummarizeDisplay Whether the tool's display output should be summarized
|
||||
*/
|
||||
constructor(
|
||||
readonly name: string,
|
||||
|
@ -119,8 +106,6 @@ export abstract class BaseTool<
|
|||
readonly parameterSchema: Schema,
|
||||
readonly isOutputMarkdown: boolean = true,
|
||||
readonly canUpdateOutput: boolean = false,
|
||||
readonly summarizer: Summarizer = defaultSummarizer,
|
||||
readonly shouldSummarizeDisplay: boolean = false,
|
||||
) {}
|
||||
|
||||
/**
|
||||
|
|
Loading…
Reference in New Issue