Website: update query generator (#26926)
Related to: https://github.com/fleetdm/confidential/issues/9884 Changes: - Updated the prompt in the `get-llm-generated-sql` action to include a note about using wildcard characters when generating queries that use the LIKE operator. - Improved error handling in the `get-llm-generated-sql` action
This commit is contained in:
@@ -18,8 +18,9 @@ module.exports = {
|
||||
description: 'A SQL query was generated'
|
||||
},
|
||||
|
||||
errorFromOpenAi: {
|
||||
description: 'The Open AI API reutrned an error.'
|
||||
couldNotGenerateQueries: {
|
||||
description: 'A set of queries could not be generated for a user using the provided question.',
|
||||
responseType: 'badRequest'
|
||||
}
|
||||
},
|
||||
|
||||
@@ -57,17 +58,34 @@ module.exports = {
|
||||
return lighterTable;}))}
|
||||
\`\`\`
|
||||
|
||||
Please respond in JSON, with the same data shape as the provided context, but with the array filtered to include only relevant tables.`;
|
||||
Please respond in JSON, with the same data shape as the provided context, but with the array filtered to include only relevant tables.
|
||||
|
||||
|
||||
If no queries can be generated from the provided instructions do not return the datashape above and instead return this JSON in this exact data shape:
|
||||
|
||||
{
|
||||
"couldNotGenerateQueries": true
|
||||
}`;
|
||||
|
||||
let filteredTables = await sails.helpers.ai.prompt(schemaFiltrationPrompt, 'gpt-4o-mini-2024-07-18', true, 'Please only respond in valid JSON with no codefences or backticks.')
|
||||
.intercept((err)=>{
|
||||
sails.log.warn(`When trying to get a subset of tables to use to generate a query for a user, an error occurred. Full error: ${require('util').inspect(err, {depth: 2})}`);
|
||||
if(this.req.isSocket){
|
||||
// If this request was from a socket and an error occurs, broadcast an 'error' event and unsubscribe the socket from this room.
|
||||
sails.sockets.broadcast(roomId, 'error', {error: err});
|
||||
sails.sockets.leave(this.req, roomId);
|
||||
}
|
||||
return new Error(`When trying to get a subset of tables to use to generate a query for an Admin user, an error occurred. Full error: ${require('util').inspect(err, {depth: 2})}`);
|
||||
return 'couldNotGenerateQueries';
|
||||
});
|
||||
|
||||
if(filteredTables.couldNotGenerateQueries){
|
||||
if(this.req.isSocket){
|
||||
sails.sockets.broadcast(roomId, 'error', {error: 'couldNotGenerateQueries'});
|
||||
sails.sockets.leave(this.req, roomId);
|
||||
} else {
|
||||
throw 'couldNotGenerateQueries';
|
||||
}
|
||||
}
|
||||
|
||||
// 2024-02-26: Testing using a system prompt with a single API request.
|
||||
// let systemPrompt = `You are an AI that generates osquery SQL queries for IT admin questions. Use the following osquery schema as context:
|
||||
@@ -120,11 +138,12 @@ module.exports = {
|
||||
|
||||
When generating the SQL:
|
||||
1. Please do not use the SQL "AS" operator, nor alias tables. Always reference tables by their full name.
|
||||
2. If this question is related to an application or program, consider using LIKE instead of something verbatim.
|
||||
3. If this question is not possible to ask given the tables and columns available in the provided context (the osquery schema) for a particular operating system, then use empty string.
|
||||
4. If this question is a "yes" or "no" question, or a "how many people" question, or a "how many hosts" question, then build the query such that a "yes" returns exactly one row and a "no" returns zero rows. In other words, if this question is about finding out which hosts match a "yes" or "no" question, then if a host does not match, do not include any rows for it.
|
||||
5. Use only tables that are supported for each target platform, as documented in the provided context, considering the examples if they exist, and the available columns.
|
||||
6. For each table that you use, only use columns that are documented for that table, as documented in the provided context.
|
||||
2. When generating a query that uses the "LIKE" operator, you should include wildcard characters.
|
||||
3. If this question is related to an application or program, consider using LIKE instead of something verbatim.
|
||||
4. If this question is not possible to ask given the tables and columns available in the provided context (the osquery schema) for a particular operating system, then use empty string.
|
||||
5. If this question is a "yes" or "no" question, or a "how many people" question, or a "how many hosts" question, then build the query such that a "yes" returns exactly one row and a "no" returns zero rows. In other words, if this question is about finding out which hosts match a "yes" or "no" question, then if a host does not match, do not include any rows for it.
|
||||
6. Use only tables that are supported for each target platform, as documented in the provided context, considering the examples if they exist, and the available columns.
|
||||
7. For each table that you use, only use columns that are documented for that table, as documented in the provided context.
|
||||
|
||||
Provided context:
|
||||
\`\`\`
|
||||
@@ -142,6 +161,13 @@ module.exports = {
|
||||
"windowsCaveats": "TODO",
|
||||
"linuxCaveats": "TODO",
|
||||
"chromeOSCaveats": "TODO",
|
||||
}
|
||||
|
||||
|
||||
If no queries can be generated from the provided instructions do not return the datashape above and instead return this JSON in this exact data shape:
|
||||
|
||||
{
|
||||
"couldNotGenerateQueries": true
|
||||
}`;
|
||||
|
||||
let sqlReport = await sails.helpers.ai.prompt.with({prompt:sqlPrompt, baseModel:'o3-mini-2025-01-31', expectJson: true})
|
||||
@@ -151,9 +177,19 @@ module.exports = {
|
||||
sails.sockets.broadcast(roomId, 'error', {error: err});
|
||||
sails.sockets.leave(this.req, roomId);
|
||||
}
|
||||
return new Error(`When trying to generate a query for an Admin user, an error occurred. Full error: ${require('util').inspect(err, {depth: 2})}`);
|
||||
sails.log.warn(`When trying to generate a query for a user, an error occurred. Full error: ${require('util').inspect(err, {depth: 2})}`);
|
||||
return 'couldNotGenerateQueries';
|
||||
});
|
||||
|
||||
if(sqlReport.couldNotGenerateQueries){
|
||||
if(this.req.isSocket){
|
||||
sails.sockets.broadcast(roomId, 'error', {error: 'couldNotGenerateQueries'});
|
||||
sails.sockets.leave(this.req, roomId);
|
||||
} else {
|
||||
throw 'couldNotGenerateQueries';
|
||||
}
|
||||
}
|
||||
|
||||
// If this request was from a socket, we'll broadcast a 'queryGenerated' event with the sqlReport and unsubscribe the socket
|
||||
if(this.req.isSocket){
|
||||
sails.sockets.broadcast(roomId, 'queryGenerated', {result: sqlReport});
|
||||
|
||||
Reference in New Issue
Block a user