Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import type { StructuredTool } from '@langchain/core/tools';
import type { BrowserApiToolMetadata } from '@kbn/agent-builder-common';
import type { Logger } from '@kbn/logging';
import type { ToolReturnSummarizerFn } from '../tools/builtin';
import type { AgentEventEmitterFn, ExecutableTool } from '..';

export interface ToolManagerParams {
Expand Down Expand Up @@ -83,4 +84,10 @@ export interface ToolManager {
* @returns array of internal tool IDs
*/
getDynamicToolIds(): string[];

/**
* Gets the summarizer function for a tool by its internal tool ID.
* Returns undefined if the tool is not found or has no summarizer.
*/
getSummarizer(toolId: string): ToolReturnSummarizerFn | undefined;
}
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import type { MaybePromise } from '@kbn/utility-types';
import type { KibanaRequest } from '@kbn/core-http-server';
import type { ToolDefinition, ToolType } from '@kbn/agent-builder-common';
import type { RunToolReturn, ScopedRunnerRunToolsParams } from './runner';
import type { ToolReturnSummarizerFn } from '../tools/builtin';

/**
* Common interface shared across all tool providers.
Expand Down Expand Up @@ -50,6 +51,12 @@ export interface ExecutableTool<
* When provided, will replace the description when converting to llm tool.
*/
getLlmDescription?: LlmDescriptionHandler<TConfig>;
/**
* Optional function to summarize a tool return for conversation history.
* When provided, this function will be called when processing conversation history
* to replace large tool results with compact summaries.
*/
summarizeToolReturn?: ToolReturnSummarizerFn;
}

export interface LLmDescriptionHandlerParams<TConfig extends object = {}> {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,12 @@ export interface BuiltinToolDefinition<RunInput extends ZodObject<any> = ZodObje
* This helps prevent context bloat in long conversations.
*/
summarizeToolReturn?: ToolReturnSummarizerFn;
/**
* When true, results from this tool will not be stored in the filestore.
* Tools with this flag must also define {@link summarizeToolReturn} to provide
* a placeholder for the removed results in conversation history.
*/
excludeFromFilestore?: boolean;
}

type StaticToolRegistrationMixin<T extends ToolDefinition> = Omit<T, 'readonly'> &
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,10 @@ export interface InternalToolDefinition<
* Tool call policy to control tool call confirmation behavior
*/
confirmation?: ToolConfirmationPolicy;
/**
* When true, results from this tool will not be stored in the filestore.
*/
excludeFromFilestore?: boolean;
}

export type InternalToolAvailabilityHandler = (
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,7 @@ export const runDefaultAgentMode: RunChatAgentFn = async (
// Create unified result transformer for tool result optimization
const resultTransformer = createResultTransformer({
toolRegistry,
toolManager,
filestore,
filestoreEnabled: experimentalFeatures.filestore,
});
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import type { ToolCallWithResult, ToolResult } from '@kbn/agent-builder-common';
import { ToolResultType } from '@kbn/agent-builder-common';
import type { IFileStore, FileEntry } from '@kbn/agent-builder-server/runner/filestore';
import type { ToolManager } from '@kbn/agent-builder-server/runner/tool_manager';
import type { ToolRegistry } from '@kbn/agent-builder-server';
import {
createResultTransformer,
Expand Down Expand Up @@ -61,6 +62,16 @@ describe('createResultTransformer', () => {
list: jest.fn(async () => []),
} as unknown as ToolRegistry);

const createMockToolManager = (
summarizers: Map<
string,
(step: ToolCallWithResult) => ToolResult[] | null | undefined
> = new Map()
): ToolManager =>
({
getSummarizer: jest.fn((toolId: string) => summarizers.get(toolId)),
} as unknown as ToolManager);

describe('tool-specific summarization', () => {
it('applies summarizeToolReturn when tool has it defined', async () => {
const toolWithSummarizer = {
Expand All @@ -78,6 +89,7 @@ describe('createResultTransformer', () => {

const transformer = createResultTransformer({
toolRegistry,
toolManager: createMockToolManager(),
filestore,
filestoreEnabled: false,
});
Expand Down Expand Up @@ -106,6 +118,7 @@ describe('createResultTransformer', () => {

const transformer = createResultTransformer({
toolRegistry,
toolManager: createMockToolManager(),
filestore,
filestoreEnabled: false,
});
Expand All @@ -130,6 +143,7 @@ describe('createResultTransformer', () => {

const transformer = createResultTransformer({
toolRegistry,
toolManager: createMockToolManager(),
filestore,
filestoreEnabled: false,
});
Expand All @@ -150,6 +164,7 @@ describe('createResultTransformer', () => {

const transformer = createResultTransformer({
toolRegistry,
toolManager: createMockToolManager(),
filestore,
filestoreEnabled: false,
});
Expand All @@ -163,6 +178,82 @@ describe('createResultTransformer', () => {
expect(result).toHaveLength(1);
expect(result[0].data).toEqual({ original: 'data' });
});

it('prefers toolManager summarizer over toolRegistry', async () => {
const managerSummarizer = jest.fn((step: ToolCallWithResult) => [
{
tool_result_id: 'manager-summarized',
type: ToolResultType.other,
data: { summary: 'From manager' },
},
]);

const registrySummarizer = {
summarizeToolReturn: jest.fn((step: ToolCallWithResult) => [
{
tool_result_id: 'registry-summarized',
type: ToolResultType.other,
data: { summary: 'From registry' },
},
]),
};

const toolManager = createMockToolManager(new Map([['search', managerSummarizer]]));
const toolRegistry = createMockToolRegistry(new Map([['search', registrySummarizer as any]]));
const filestore = createMockFileStore(new Map());

const transformer = createResultTransformer({
toolRegistry,
toolManager,
filestore,
filestoreEnabled: false,
});

const toolCall = makeToolCallWithResult('call-1', 'search', [
{ tool_result_id: 'result-1', type: ToolResultType.other, data: { original: 'data' } },
]);

const result = await transformer(toolCall);

expect(managerSummarizer).toHaveBeenCalledWith(toolCall);
expect(registrySummarizer.summarizeToolReturn).not.toHaveBeenCalled();
expect(result).toHaveLength(1);
expect(result[0].data).toEqual({ summary: 'From manager', _summary: true });
});

it('falls back to toolRegistry when toolManager has no summarizer', async () => {
const registrySummarizer = {
summarizeToolReturn: jest.fn((step: ToolCallWithResult) => [
{
tool_result_id: 'registry-summarized',
type: ToolResultType.other,
data: { summary: 'From registry' },
},
]),
};

const toolManager = createMockToolManager(new Map());
const toolRegistry = createMockToolRegistry(new Map([['search', registrySummarizer as any]]));
const filestore = createMockFileStore(new Map());

const transformer = createResultTransformer({
toolRegistry,
toolManager,
filestore,
filestoreEnabled: false,
});

const toolCall = makeToolCallWithResult('call-1', 'search', [
{ tool_result_id: 'result-1', type: ToolResultType.other, data: { original: 'data' } },
]);

const result = await transformer(toolCall);

expect(toolManager.getSummarizer).toHaveBeenCalledWith('search');
expect(registrySummarizer.summarizeToolReturn).toHaveBeenCalledWith(toolCall);
expect(result).toHaveLength(1);
expect(result[0].data).toEqual({ summary: 'From registry', _summary: true });
});
});

describe('filestore substitution', () => {
Expand All @@ -182,6 +273,7 @@ describe('createResultTransformer', () => {

const transformer = createResultTransformer({
toolRegistry,
toolManager: createMockToolManager(),
filestore,
filestoreEnabled: true,
});
Expand Down Expand Up @@ -212,6 +304,7 @@ describe('createResultTransformer', () => {

const transformer = createResultTransformer({
toolRegistry,
toolManager: createMockToolManager(),
filestore,
filestoreEnabled: true,
});
Expand All @@ -233,6 +326,7 @@ describe('createResultTransformer', () => {

const transformer = createResultTransformer({
toolRegistry,
toolManager: createMockToolManager(),
filestore,
filestoreEnabled: true,
});
Expand Down Expand Up @@ -263,6 +357,7 @@ describe('createResultTransformer', () => {

const transformer = createResultTransformer({
toolRegistry,
toolManager: createMockToolManager(),
filestore,
filestoreEnabled: false,
});
Expand Down Expand Up @@ -307,6 +402,7 @@ describe('createResultTransformer', () => {

const transformer = createResultTransformer({
toolRegistry,
toolManager: createMockToolManager(),
filestore,
filestoreEnabled: true,
});
Expand Down Expand Up @@ -335,6 +431,7 @@ describe('createResultTransformer', () => {

const transformer = createResultTransformer({
toolRegistry,
toolManager: createMockToolManager(),
filestore,
filestoreEnabled: true,
});
Expand Down Expand Up @@ -371,6 +468,7 @@ describe('createResultTransformer', () => {

const transformer = createResultTransformer({
toolRegistry,
toolManager: createMockToolManager(),
filestore,
filestoreEnabled: false,
});
Expand All @@ -395,6 +493,7 @@ describe('createResultTransformer', () => {

const transformer = createResultTransformer({
toolRegistry,
toolManager: createMockToolManager(),
filestore,
filestoreEnabled: true,
});
Expand All @@ -419,6 +518,7 @@ describe('createResultTransformer', () => {

const transformer = createResultTransformer({
toolRegistry,
toolManager: createMockToolManager(),
filestore,
filestoreEnabled: true,
tokenThreshold: 100,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import type { ToolCallWithResult, ToolResult } from '@kbn/agent-builder-common';
import { ToolResultType } from '@kbn/agent-builder-common';
import type { IFileStore } from '@kbn/agent-builder-server/runner/filestore';
import type { ToolRegistry } from '@kbn/agent-builder-server';
import type { ToolManager } from '@kbn/agent-builder-server/runner/tool_manager';
import { getToolCallEntryPath } from '../../../runner/store/volumes/tool_results/utils';

/**
Expand Down Expand Up @@ -75,8 +76,16 @@ export type ToolCallResultTransformer = (toolCall: ToolCallWithResult) => Promis
export interface CreateResultTransformerOptions {
/**
* Tool registry to look up tool-specific summarization functions.
* Used as a fallback when the tool is not found in the tool manager
* (e.g. for evicted dynamic tools from previous rounds).
*/
toolRegistry: ToolRegistry;
/**
* Tool manager to look up tool-specific summarization functions.
* Checked first, as it contains all active tools including internal ones
* (filestore tools, attachment tools) that may not be in the registry.
*/
toolManager: ToolManager;
/**
* Filestore to check token counts for file reference substitution.
*/
Expand All @@ -102,6 +111,7 @@ export interface CreateResultTransformerOptions {
*/
export const createResultTransformer = ({
toolRegistry,
toolManager,
filestore,
filestoreEnabled,
tokenThreshold = FILE_REFERENCE_TOKEN_THRESHOLD,
Expand All @@ -113,7 +123,7 @@ export const createResultTransformer = ({
}

// Step 1: Try tool-specific summarization
const summarized = await tryToolSummarization(toolCall, toolRegistry);
const summarized = await tryToolSummarization(toolCall, toolManager, toolRegistry);
if (summarized) {
return summarized.map(markResultAsCleaned);
}
Expand Down Expand Up @@ -141,12 +151,21 @@ export const createResultTransformer = ({

/**
* Attempts to apply tool-specific summarization using the tool's `summarizeToolReturn` function.
* Checks the tool manager first (has all active tools including internal ones like filestore tools),
* then falls back to the tool registry (for evicted dynamic tools from previous rounds).
* Returns the summarized results, or undefined if no summarization is available/applicable.
*/
const tryToolSummarization = async (
toolCall: ToolCallWithResult,
toolManager: ToolManager,
toolRegistry: ToolRegistry
): Promise<ToolResult[] | undefined> => {
const managerSummarizer = toolManager.getSummarizer(toolCall.tool_id);
if (managerSummarizer) {
const summarizedResults = managerSummarizer(toolCall);
return summarizedResults ?? undefined;
}
Comment on lines +163 to +167
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Internal tools such as the filestore and attachment tools aren't registered in the tool registry, so that whole logic of result summarization wasnt working properly for them. Had to adapt a bunch of things, and use the tool manager as a second source of truth for the tool definitions.

Note that this also fixes a bug with attachment tools, because they currently define summarizeToolReturn which aren't used without that PR.


try {
const tool = await toolRegistry.get(toolCall.tool_id);
if (!tool?.summarizeToolReturn) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,7 @@ const builtinToolToExecutable = ({
configuration: {},
readonly: true,
getSchema: () => tool.schema,
summarizeToolReturn: tool.summarizeToolReturn,
execute: async (params) => {
return runner.runInternalTool({
...params,
Expand All @@ -175,6 +176,7 @@ const builtinToolToExecutable = ({
configuration: {},
readonly: true,
confirmation: { askUser: 'never' },
excludeFromFilestore: tool.excludeFromFilestore,
isAvailable: async () => ({ status: 'available' as const }),
getSchema: () => tool.schema,
getHandler: () => tool.handler,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -183,7 +183,7 @@ export const runInternalTool = async <TParams = Record<string, unknown>>({
const afterToolHooksResult = await hooks.run(HookLifecycle.afterToolCall, postContext);
runToolReturn = afterToolHooksResult.toolReturn;

if (runToolReturn.results) {
if (runToolReturn.results && !tool.excludeFromFilestore) {
runToolReturn.results.forEach((result) => {
resultStore.add({
tool_id: tool.id,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import { filestoreTools } from '@kbn/agent-builder-common/tools';
import { createOtherResult } from '@kbn/agent-builder-server';
import type { IFileStore, FileEntry } from '@kbn/agent-builder-server/runner/filestore';
import type { BuiltinToolDefinition } from '@kbn/agent-builder-server/tools';
import { summarizeFilestoreToolReturn } from './summarize';

const schema = z.object({
pattern: z.string().describe('Glob pattern to match files (e.g., "/**/*.json", "/logs/*.log")'),
Expand All @@ -27,6 +28,8 @@ export const globTool = ({
type: ToolType.builtin,
schema,
tags: ['filestore'],
excludeFromFilestore: true,
summarizeToolReturn: summarizeFilestoreToolReturn,
handler: async ({ pattern }, context) => {
const entries = await filestore.glob(pattern);
const summaries = entries.map(toSummary);
Expand Down
Loading