Skip to content

Commit

Permalink
[Security assistant] Update custom tool name regex (elastic#206487)
Browse files Browse the repository at this point in the history
  • Loading branch information
stephmilovic authored Jan 13, 2025
1 parent 95b76dc commit ef5977e
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -221,4 +221,31 @@ describe('getStructuredToolForIndexEntry', () => {
);
expect(result).toContain(`I'm sorry, but I was unable to find any information`);
});

it('should match the name regex correctly', () => {
const tool = getStructuredToolForIndexEntry({
indexEntry: getCreateKnowledgeBaseEntrySchemaMock({
type: 'index',
name: `1bad-name?`,
}) as IndexEntry,
esClient: mockEsClient,
logger: mockLogger,
});

const nameRegex = /^[a-zA-Z0-9_-]+$/;
expect(tool.lc_kwargs.name).toMatch(nameRegex);
});

it('dashes get removed before `a` is prepended', () => {
const tool = getStructuredToolForIndexEntry({
indexEntry: getCreateKnowledgeBaseEntrySchemaMock({
type: 'index',
name: `-testing`,
}) as IndexEntry,
esClient: mockEsClient,
logger: mockLogger,
});

expect(tool.lc_kwargs.name).toMatch('testing');
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,11 @@ export const getStructuredToolForIndexEntry = ({
}, {});

return new DynamicStructuredTool({
name: indexEntry.name.replace(/[^a-zA-Z0-9-]/g, ''), // // Tool names expects a string that matches the pattern '^[a-zA-Z0-9-]+$'
name: indexEntry.name
// Replace invalid characters with an empty string
.replace(/[^a-zA-Z0-9_]/g, '')
// Ensure it starts with a letter. If not, prepend 'a'
.replace(/^[^a-zA-Z]/, 'a'),
description: indexEntry.description,
schema: z.object({
query: z.string().describe(indexEntry.queryDescription),
Expand Down

0 comments on commit ef5977e

Please sign in to comment.