fix(core): address PR feedback on tool preselection robustness

This commit is contained in:
mkorwel
2026-02-19 23:42:40 -06:00
parent 173cb56643
commit ab717bdb24
3 changed files with 59 additions and 5 deletions
+1 -1
View File
@@ -286,7 +286,7 @@ export class GeminiClient {
);
const selectedSet = new Set(selectedNames);
toolDeclarations = toolDeclarations.filter((t) =>
selectedSet.has(t.name!),
t.name ? selectedSet.has(t.name) : false,
);
}
@@ -124,4 +124,49 @@ describe('ToolPreselectionService', () => {
'tool6',
]);
});
it('falls back to all tools if LLM returns invalid format', async () => {
const tools: FunctionDeclaration[] = [
{ name: 'tool1', description: 'desc1' },
{ name: 'tool2', description: 'desc2' },
{ name: 'tool3', description: 'desc3' },
{ name: 'tool4', description: 'desc4' },
{ name: 'tool5', description: 'desc5' },
{ name: 'tool6', description: 'desc6' },
];
// Missing 'relevant_tools' key
mockLlmClient['generateJson'].mockResolvedValue({
something_else: ['tool1'],
});
const result = await service.selectTools(
'query',
tools,
new AbortController().signal,
);
expect(result).toEqual([
'tool1',
'tool2',
'tool3',
'tool4',
'tool5',
'tool6',
]);
});
it('safely handles tools with undefined names', async () => {
const tools: FunctionDeclaration[] = [
{ name: 'tool1', description: 'desc1' },
{ name: undefined, description: 'desc2' },
];
const result = await service.selectTools(
'query',
tools,
new AbortController().signal,
{ maxTools: 10 }, // Force threshold bypass
);
expect(result).toEqual(['tool1']);
});
});
@@ -43,7 +43,7 @@ export class ToolPreselectionService {
// Threshold below which we don't bother with pre-selection.
const threshold = options.maxTools ?? 5;
if (tools.length <= threshold) {
return tools.map((t) => t.name!);
return tools.filter((t) => !!t.name).map((t) => t.name!);
}
const schema = {
@@ -84,8 +84,17 @@ ${toolsList}`;
role: LlmRole.UTILITY_TOOL,
});
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
const selectedTools = result['relevant_tools'] as string[];
const selectedTools = result['relevant_tools'];
if (
!Array.isArray(selectedTools) ||
!selectedTools.every((item): item is string => typeof item === 'string')
) {
throw new Error(
'Tool preselection returned invalid data format. Expected an array of strings.',
);
}
debugLogger.debug(
`ToolPreselectionService: Selected ${selectedTools.length} tools out of ${tools.length} for query: "${query.substring(0, 50)}..."`,
);
@@ -93,7 +102,7 @@ ${toolsList}`;
} catch (error) {
debugLogger.error('ToolPreselectionService failed:', error);
// Fallback: return all tools if pre-selection fails.
return tools.map((t) => t.name!);
return tools.filter((t) => !!t.name).map((t) => t.name!);
}
}
}