Ling0925
/
buster
mirror of https://github.com/buster-so/buster.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Refactor analyst agent tools integration and SQL execution context. Update the analyst agent to utilize new tool functions for metrics and dashboards, and modify SQL execution to accept context directly from the analyst agent options, enhancing clarity and maintainability.

This commit is contained in:
dal 2025-08-06 14:31:30 -06:00
parent 6b72213992
commit f694c534c7
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
2 changed files with 8 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
packages/ai/src/agents/analyst-agent/analyst-agent.ts
View File

@ -53,7 +53,7 @@ export function createAnalystAgent(analystAgentOptions: AnalystAgentOptions) {
toolChoice: 'required',
maxOutputTokens: 10000,
temperature: 0,
// No longer need experimental_context since it's baked into tools
experimental_context: analystAgentOptions,
}),
{
name: 'Analyst Agent',

12
packages/ai/src/tools/database-tools/execute-sql.ts
View File

@ -2,6 +2,7 @@ import { type DataSource, withRateLimit } from '@buster/data-source';
import { tool } from 'ai';
import { wrapTraced } from 'braintrust';
import { z } from 'zod';
import type { AnalystAgentOptions } from '../../agents/analyst-agent/analyst-agent';
import { getWorkflowDataSourceManager } from '../../utils/data-source-manager';
import { createPermissionErrorMessage, validateSqlPermissions } from '../../utils/sql-permissions';
import type { AnalystRuntimeContext } from '../../workflows/analyst-workflow';
@ -87,7 +88,8 @@ const executeSqlStatementOutputSchema = z.object({
const executeSqlStatement = wrapTraced(
async (
params: z.infer<typeof executeSqlStatementInputSchema>
params: z.infer<typeof executeSqlStatementInputSchema>,
context: AnalystAgentOptions
): Promise<z.infer<typeof executeSqlStatementOutputSchema>> => {
let { statements } = params;
@ -190,8 +192,8 @@ const executeSqlStatement = wrapTraced(
};
}
const dataSourceId = runtimeContext.get('dataSourceId');
const workflowStartTime = runtimeContext.get('workflowStartTime') as number | undefined;
const dataSourceId = context.dataSourceId;
const workflowStartTime = context.get('workflowStartTime') as number | undefined;
// Generate a unique workflow ID using start time and data source
const workflowId = workflowStartTime
@ -394,8 +396,8 @@ export const executeSql = tool({
Otherwise the queries wont run successfully.`,
inputSchema: executeSqlStatementInputSchema,
outputSchema: executeSqlStatementOutputSchema,
execute: async (input) => {
return await executeSqlStatement(input);
execute: async (input, { experimental_context: context }) => {
return await executeSqlStatement(input, context as AnalystAgentOptions);
},
});