From e017184cc6998d6b829fbb850fff9168fc0a66ee Mon Sep 17 00:00:00 2001 From: Artem Niehrieiev Date: Thu, 26 Feb 2026 14:14:12 +0000 Subject: [PATCH] Refactor AI panel generation to remove table_name parameter and enhance prompt clarity --- .../generate-panel-position-with-ai.ds.ts | 1 - .../panel-position.controller.ts | 9 +- ...enerate-panel-position-with-ai.use.case.ts | 234 ++++++++++++------ .../dashboard-ai-widget-e2e.test.ts | 158 +++++------- 4 files changed, 233 insertions(+), 169 deletions(-) diff --git a/backend/src/entities/visualizations/panel-position/data-structures/generate-panel-position-with-ai.ds.ts b/backend/src/entities/visualizations/panel-position/data-structures/generate-panel-position-with-ai.ds.ts index 64b49d504..46c5cea33 100644 --- a/backend/src/entities/visualizations/panel-position/data-structures/generate-panel-position-with-ai.ds.ts +++ b/backend/src/entities/visualizations/panel-position/data-structures/generate-panel-position-with-ai.ds.ts @@ -4,7 +4,6 @@ export class GeneratePanelPositionWithAiDs { masterPassword: string; userId: string; chart_description: string; - table_name: string; name?: string; position_x?: number; position_y?: number; diff --git a/backend/src/entities/visualizations/panel-position/panel-position.controller.ts b/backend/src/entities/visualizations/panel-position/panel-position.controller.ts index 03c95fd90..0ebb72333 100644 --- a/backend/src/entities/visualizations/panel-position/panel-position.controller.ts +++ b/backend/src/entities/visualizations/panel-position/panel-position.controller.ts @@ -7,11 +7,10 @@ import { Param, Post, Put, - Query, UseGuards, UseInterceptors, } from '@nestjs/common'; -import { ApiBearerAuth, ApiBody, ApiOperation, ApiParam, ApiQuery, ApiResponse, ApiTags } from '@nestjs/swagger'; +import { ApiBearerAuth, ApiBody, ApiOperation, ApiParam, ApiResponse, ApiTags } from '@nestjs/swagger'; import { UseCaseType } from '../../../common/data-injection.tokens.js'; import { MasterPassword } from '../../../decorators/master-password.decorator.js'; import { SlugUuid } from '../../../decorators/slug-uuid.decorator.js'; @@ -157,7 +156,7 @@ export class DashboardWidgetController { @ApiOperation({ summary: 'Generate a panel using AI', description: - 'Creates a new panel with SQL query and chart configuration generated by AI based on natural language description', + 'Creates a new panel with SQL query and chart configuration generated by AI based on natural language description. AI autonomously discovers relevant tables.', }) @ApiResponse({ status: 201, @@ -167,13 +166,12 @@ export class DashboardWidgetController { @ApiBody({ type: GeneratePanelPositionWithAiDto }) @ApiParam({ name: 'dashboardId', required: true }) @ApiParam({ name: 'connectionId', required: true }) - @ApiQuery({ name: 'tableName', required: true, description: 'The table name to generate the panel for' }) @UseGuards(ConnectionEditGuard) + @Timeout(TimeoutDefaults.AI) @Post('/dashboard/:dashboardId/widget/generate/:connectionId') async generateWidgetWithAi( @SlugUuid('connectionId') connectionId: string, @Param('dashboardId') dashboardId: string, - @Query('tableName') tableName: string, @MasterPassword() masterPwd: string, @UserId() userId: string, @Body() generateDto: GeneratePanelPositionWithAiDto, @@ -184,7 +182,6 @@ export class DashboardWidgetController { masterPassword: masterPwd, userId, chart_description: generateDto.chart_description, - table_name: tableName, name: generateDto.name, }; return await this.generateWidgetWithAiUseCase.execute(inputData, InTransactionEnum.OFF); diff --git a/backend/src/entities/visualizations/panel-position/use-cases/generate-panel-position-with-ai.use.case.ts b/backend/src/entities/visualizations/panel-position/use-cases/generate-panel-position-with-ai.use.case.ts index 583b58166..1ba9e1aec 100644 --- a/backend/src/entities/visualizations/panel-position/use-cases/generate-panel-position-with-ai.use.case.ts +++ b/backend/src/entities/visualizations/panel-position/use-cases/generate-panel-position-with-ai.use.case.ts @@ -3,12 +3,23 @@ import { getDataAccessObject } from '@rocketadmin/shared-code/dist/src/data-acce import { ConnectionTypesEnum } from '@rocketadmin/shared-code/dist/src/shared/enums/connection-types-enum.js'; import { IDataAccessObject } from '@rocketadmin/shared-code/dist/src/shared/interfaces/data-access-object.interface.js'; import { IDataAccessObjectAgent } from '@rocketadmin/shared-code/dist/src/shared/interfaces/data-access-object-agent.interface.js'; -import { AICoreService, AIProviderType, cleanAIJsonResponse } from '../../../../ai-core/index.js'; +import { + AICoreService, + AIProviderType, + AIToolCall, + AIToolDefinition, + cleanAIJsonResponse, + createDashboardGenerationTools, + encodeError, + encodeToToon, + MessageBuilder, +} from '../../../../ai-core/index.js'; import AbstractUseCase from '../../../../common/abstract-use.case.js'; import { IGlobalDatabaseContext } from '../../../../common/application/global-database-context.interface.js'; import { BaseType } from '../../../../common/data-injection.tokens.js'; import { DashboardWidgetTypeEnum } from '../../../../enums/dashboard-widget-type.enum.js'; import { Messages } from '../../../../exceptions/text/messages.js'; +import { isConnectionTypeAgent } from '../../../../helpers/is-connection-entity-agent.js'; import { validateQuerySafety } from '../../panel/utils/check-query-is-safe.util.js'; import { GeneratePanelPositionWithAiDs } from '../data-structures/generate-panel-position-with-ai.ds.js'; import { GeneratedPanelWithPositionDto } from '../dto/generated-panel-with-position.dto.js'; @@ -38,6 +49,7 @@ interface AIGeneratedPanelResponse { }; } +const MAX_TOOL_ITERATIONS = 10; const MAX_FEEDBACK_ITERATIONS = 3; const EXPLAIN_SUPPORTED_TYPES: ReadonlySet = new Set([ @@ -49,11 +61,6 @@ const EXPLAIN_SUPPORTED_TYPES: ReadonlySet = new Set([ ConnectionTypesEnum.agent_clickhouse, ]); -interface TableInfo { - table_name: string; - columns: Array<{ name: string; type: string; nullable: boolean }>; -} - @Injectable({ scope: Scope.REQUEST }) export class GeneratePanelPositionWithAiUseCase extends AbstractUseCase @@ -74,8 +81,8 @@ export class GeneratePanelPositionWithAiUseCase dashboardId, connectionId, masterPassword, + userId, chart_description, - table_name, name, position_x, position_y, @@ -103,31 +110,25 @@ export class GeneratePanelPositionWithAiUseCase const dao = getDataAccessObject(foundConnection); - let tableInfo: TableInfo; - - try { - const structure = await dao.getTableStructure(table_name, null); - tableInfo = { - table_name: table_name, - columns: structure.map((col) => ({ - name: col.column_name, - type: col.data_type, - nullable: col.allow_null, - })), - }; - } catch (error) { - throw new BadRequestException(`Failed to get table structure for "${table_name}": ${error.message}`); + let userEmail: string; + if (isConnectionTypeAgent(foundConnection.type)) { + userEmail = await this._dbContext.userRepository.getUserEmailOrReturnNull(userId); } - if (tableInfo.columns.length === 0) { - throw new BadRequestException(`The specified table "${table_name}" does not have any columns or does not exist.`); - } + const tools = createDashboardGenerationTools(); - const prompt = this.buildPrompt(chart_description, tableInfo, foundConnection.type); + const systemPrompt = this.buildSystemPrompt(foundConnection.type as ConnectionTypesEnum); - const aiResponse = await this.aiCoreService.completeWithProvider(AIProviderType.BEDROCK, prompt, { - temperature: 0.3, - }); + const messages = new MessageBuilder() + .system(systemPrompt) + .human( + `Generate a single panel based on this description: "${chart_description}". ` + + `Start by listing the available tables, then inspect the ones that are relevant to this request. ` + + `Generate the panel configuration.`, + ) + .build(); + + const aiResponse = await this.runToolLoop(messages, tools, dao, userEmail); const generatedPanel = this.parseAIResponse(aiResponse); @@ -136,9 +137,7 @@ export class GeneratePanelPositionWithAiUseCase const refinedPanel = await this.validateAndRefineQueryWithExplain( dao, generatedPanel, - tableInfo, foundConnection.type as ConnectionTypesEnum, - chart_description, ); return { @@ -161,33 +160,26 @@ export class GeneratePanelPositionWithAiUseCase }; } - private buildPrompt( - chartDescription: string, - tableInfo: { table_name: string; columns: Array<{ name: string; type: string; nullable: boolean }> }, - databaseType: string, - ): string { - const schemaDescription = `Table: ${tableInfo.table_name}\n Columns:\n${tableInfo.columns.map((col) => ` - ${col.name}: ${col.type}${col.nullable ? ' (nullable)' : ''}`).join('\n')}`; - - return `You are a database analytics assistant. Based on the user's chart description and the database schema, generate the SQL query and chart configuration. + private buildSystemPrompt(databaseType: ConnectionTypesEnum): string { + return `You are a database analytics assistant that generates a single panel. You have access to two tools: getTablesList and getTableStructure. You do NOT have the ability to execute queries. DATABASE TYPE: ${databaseType} -DATABASE SCHEMA: -${schemaDescription} - -USER'S CHART DESCRIPTION: -"${chartDescription}" +YOUR WORKFLOW: +1. Call getTablesList to see all available tables +2. Call getTableStructure for tables that are relevant to the user's request +3. Based on the table schemas, generate a single panel configuration -Generate a JSON response with the following structure: +CRITICAL: Your final response MUST be a raw JSON object only — no explanations, no markdown, no code fences, no text before or after. Just the JSON object in this format: { - "name": "Short descriptive name for the chart (max 50 chars)", - "description": "Brief description of what the chart shows", - "query_text": "SELECT query that returns data for the chart. The query should return columns that can be used for labels and values. Use appropriate aggregations (COUNT, SUM, AVG, etc.) and GROUP BY clauses as needed. Always use the table name '${tableInfo.table_name}'.", + "name": "Short descriptive name for the panel (max 50 chars)", + "description": "Brief description of what the panel shows", + "query_text": "SELECT query that returns data for the visualization", "panel_type": "chart" | "table" | "counter" | "text", "chart_type": "bar" | "line" | "pie" | "doughnut" | "polarArea", "panel_options": { "label_column": "column name for labels/categories (x-axis)", - "value_column": "column name for values (y-axis) - use this for single series", + "value_column": "column name for values (y-axis) - use for single series", "series": [ { "value_column": "column name", @@ -221,13 +213,115 @@ IMPORTANT GUIDELINES: 6. Set panel_options.label_column to the column that should be used for labels 7. For single value series, use value_column directly 8. For multiple series, use the series array +9. You may use multiple tables in queries (JOINs) if they are related +10. Your final response must start with { and end with } — no text, no markdown, no explanation`; + } + + private async runToolLoop( + messages: ReturnType, + tools: AIToolDefinition[], + dao: IDataAccessObject | IDataAccessObjectAgent, + userEmail: string, + ): Promise { + let currentMessages = [...messages]; + let depth = 0; + + while (depth < MAX_TOOL_ITERATIONS) { + const stream = await this.aiCoreService.streamChatWithToolsAndProvider( + AIProviderType.BEDROCK, + currentMessages, + tools, + { temperature: 0.3 }, + ); + + let accumulatedContent = ''; + const pendingToolCalls: AIToolCall[] = []; + + for await (const chunk of stream) { + if (chunk.type === 'text' && chunk.content) { + accumulatedContent += chunk.content; + } + if (chunk.type === 'tool_call' && chunk.toolCall) { + pendingToolCalls.push(chunk.toolCall); + } + } + + this.logger.log( + `Tool loop iteration ${depth + 1}: toolCalls=${pendingToolCalls.map((tc) => tc.name).join(', ') || 'none'}, ` + + `contentLength=${accumulatedContent.length}`, + ); + + if (pendingToolCalls.length === 0) { + return accumulatedContent; + } -Respond ONLY with the JSON object, no additional text or explanation.`; + const toolResults = await this.executeToolCalls(pendingToolCalls, dao, userEmail); + + const continuationBuilder = MessageBuilder.fromMessages(currentMessages); + continuationBuilder.ai(accumulatedContent, pendingToolCalls); + for (const tr of toolResults) { + continuationBuilder.toolResult(tr.toolCallId, tr.result); + } + currentMessages = continuationBuilder.build(); + + depth++; + } + + throw new BadRequestException('AI tool loop exceeded maximum iterations. Please try again.'); + } + + private async executeToolCalls( + toolCalls: AIToolCall[], + dao: IDataAccessObject | IDataAccessObjectAgent, + userEmail: string, + ): Promise> { + const results: Array<{ toolCallId: string; result: string }> = []; + + for (const toolCall of toolCalls) { + let result: string; + + try { + switch (toolCall.name) { + case 'getTablesList': { + const tables = await dao.getTablesFromDB(userEmail); + result = encodeToToon(tables); + break; + } + + case 'getTableStructure': { + const tableName = toolCall.arguments.tableName as string; + if (!tableName) { + throw new Error('Missing required argument "tableName"'); + } + const structure = await dao.getTableStructure(tableName, userEmail); + result = encodeToToon({ + tableName, + columns: structure.map((col) => ({ + name: col.column_name, + type: col.data_type, + nullable: col.allow_null, + })), + }); + break; + } + + default: + result = encodeError({ error: `Unknown tool: ${toolCall.name}` }); + } + } catch (error) { + result = encodeError({ error: error.message }); + } + + results.push({ toolCallId: toolCall.id, result }); + } + + return results; } private parseAIResponse(aiResponse: string): AIGeneratedPanelResponse { try { - const cleanedResponse = cleanAIJsonResponse(aiResponse); + const extracted = this.extractJsonFromResponse(aiResponse); + const cleanedResponse = cleanAIJsonResponse(extracted); const parsed = JSON.parse(cleanedResponse) as AIGeneratedPanelResponse; if (!parsed.name || !parsed.query_text || !parsed.panel_type) { @@ -236,20 +330,31 @@ Respond ONLY with the JSON object, no additional text or explanation.`; return parsed; } catch (error) { - console.error('Error parsing AI response:', error.message); - console.error('AI Response:', aiResponse); + this.logger.error('Error parsing AI response:', error.message); throw new BadRequestException( 'Failed to generate panel configuration from AI. Please try again with a different description.', ); } } + private extractJsonFromResponse(response: string): string { + const jsonBlockMatch = response.match(/```(?:json)?\s*\n?([\s\S]*?)```/); + if (jsonBlockMatch) { + return jsonBlockMatch[1].trim(); + } + + const firstBrace = response.indexOf('{'); + if (firstBrace !== -1) { + return response.slice(firstBrace); + } + + return response; + } + private async validateAndRefineQueryWithExplain( dao: IDataAccessObject | IDataAccessObjectAgent, generatedPanel: AIGeneratedPanelResponse, - tableInfo: TableInfo, connectionType: ConnectionTypesEnum, - chartDescription: string, ): Promise { if (!EXPLAIN_SUPPORTED_TYPES.has(connectionType)) { return generatedPanel; @@ -258,15 +363,14 @@ Respond ONLY with the JSON object, no additional text or explanation.`; let currentQuery = generatedPanel.query_text; for (let iteration = 0; iteration < MAX_FEEDBACK_ITERATIONS; iteration++) { - const explainResult = await this.runExplainQuery(dao, currentQuery, tableInfo.table_name); + const explainResult = await this.runExplainQuery(dao, currentQuery); const correctionPrompt = this.buildQueryCorrectionPrompt( currentQuery, explainResult.success ? explainResult.result : explainResult.error, !explainResult.success, - tableInfo, connectionType, - chartDescription, + generatedPanel.name, ); const aiResponse = await this.aiCoreService.completeWithProvider(AIProviderType.BEDROCK, correctionPrompt, { @@ -296,11 +400,10 @@ Respond ONLY with the JSON object, no additional text or explanation.`; private async runExplainQuery( dao: IDataAccessObject | IDataAccessObjectAgent, query: string, - tableName: string, ): Promise<{ success: boolean; result?: string; error?: string }> { try { const explainQuery = `EXPLAIN ${query.replace(/;\s*$/, '')}`; - const result = await (dao as IDataAccessObject).executeRawQuery(explainQuery, tableName); + const result = await (dao as IDataAccessObject).executeRawQuery(explainQuery, ''); return { success: true, result: JSON.stringify(result, null, 2) }; } catch (error) { return { success: false, error: error.message }; @@ -311,14 +414,9 @@ Respond ONLY with the JSON object, no additional text or explanation.`; currentQuery: string, explainResultOrError: string, isError: boolean, - tableInfo: TableInfo, connectionType: ConnectionTypesEnum, - chartDescription: string, + panelName: string, ): string { - const schemaDescription = `Table: ${tableInfo.table_name}\n Columns:\n${tableInfo.columns - .map((col) => ` - ${col.name}: ${col.type}${col.nullable ? ' (nullable)' : ''}`) - .join('\n')}`; - const feedbackSection = isError ? `The query FAILED with the following error:\n${explainResultOrError}\n\nPlease fix the query to resolve this error.` : `The EXPLAIN output for the query is:\n${explainResultOrError}\n\nReview the execution plan. If the query has performance issues (full table scans on large datasets, inefficient joins, etc.), optimize it if possible. If the query is already acceptable, return it unchanged.`; @@ -327,11 +425,7 @@ Respond ONLY with the JSON object, no additional text or explanation.`; DATABASE TYPE: ${connectionType} -DATABASE SCHEMA: -${schemaDescription} - -ORIGINAL USER REQUEST: -"${chartDescription}" +PANEL NAME: "${panelName}" CURRENT QUERY: ${currentQuery} diff --git a/backend/test/ava-tests/saas-tests/dashboard-ai-widget-e2e.test.ts b/backend/test/ava-tests/saas-tests/dashboard-ai-widget-e2e.test.ts index c65b91a5c..78a2bfd5c 100644 --- a/backend/test/ava-tests/saas-tests/dashboard-ai-widget-e2e.test.ts +++ b/backend/test/ava-tests/saas-tests/dashboard-ai-widget-e2e.test.ts @@ -60,26 +60,69 @@ const MOCK_AI_RESPONSE_UNSAFE_QUERY = JSON.stringify({ panel_type: 'table', }); +const MOCK_TABLES_LIST = [ + { tableName: 'test_table', isView: false }, +]; + +const MOCK_TABLE_STRUCTURE = [ + { column_name: 'id', data_type: 'integer', allow_null: false }, + { column_name: 'name', data_type: 'varchar', allow_null: true }, +]; + let mockResponse = MOCK_AI_RESPONSE_CHART; +let toolCallCounter = 0; const mockAICoreService = { - completeWithProvider: async () => mockResponse, - complete: async () => mockResponse, - chat: async () => ({ content: mockResponse, responseId: faker.string.uuid() }), + streamChatWithToolsAndProvider: async () => { + toolCallCounter++; + // First call: AI requests getTablesList + if (toolCallCounter === 1) { + return { + *[Symbol.asyncIterator]() { + yield { type: 'tool_call', toolCall: { id: 'tc_1', name: 'getTablesList', arguments: {} } }; + yield { type: 'done' }; + }, + }; + } + // Second call: AI requests getTableStructure + if (toolCallCounter === 2) { + return { + *[Symbol.asyncIterator]() { + yield { + type: 'tool_call', + toolCall: { id: 'tc_2', name: 'getTableStructure', arguments: { tableName: 'test_table' } }, + }; + yield { type: 'done' }; + }, + }; + } + // Third call: AI returns final panel JSON + return { + *[Symbol.asyncIterator]() { + yield { type: 'text', content: mockResponse }; + yield { type: 'done' }; + }, + }; + }, + completeWithProvider: async (_provider: string, prompt: string) => { + if (prompt.includes('query optimization assistant')) { + const match = prompt.match(/CURRENT QUERY:\n([\s\S]*?)\n\n/); + return match ? match[1].trim() : 'SELECT 1'; + } + return 'SELECT 1'; + }, + complete: async () => 'SELECT 1', + chat: async () => ({ content: '{}', responseId: faker.string.uuid() }), + chatWithToolsAndProvider: async () => ({ content: '{}', toolCalls: [] }), streamChat: async () => ({ *[Symbol.asyncIterator]() { - yield { type: 'text', content: mockResponse, responseId: faker.string.uuid() }; + yield { type: 'text', content: '{}', responseId: faker.string.uuid() }; }, }), - chatWithTools: async () => ({ content: mockResponse, responseId: faker.string.uuid() }), + chatWithTools: async () => ({ content: '{}', responseId: faker.string.uuid() }), streamChatWithTools: async () => ({ *[Symbol.asyncIterator]() { - yield { type: 'text', content: mockResponse, responseId: faker.string.uuid() }; - }, - }), - streamChatWithToolsAndProvider: async () => ({ - *[Symbol.asyncIterator]() { - yield { type: 'text', content: mockResponse, responseId: faker.string.uuid() }; + yield { type: 'text', content: '{}', responseId: faker.string.uuid() }; }, }), getDefaultProvider: () => 'bedrock', @@ -125,6 +168,7 @@ currentTest = 'POST /dashboard/:dashboardId/widget/generate/:connectionId'; test.serial(`${currentTest} should generate a widget with AI for chart type`, async (t) => { mockResponse = MOCK_AI_RESPONSE_CHART; + toolCallCounter = 0; const connectionToTestDB = getTestData(mockFactory).connectionToPostgres; const { token } = await registerUserAndReturnUserInfo(app); @@ -151,7 +195,7 @@ test.serial(`${currentTest} should generate a widget with AI for chart type`, as t.is(createDashboard.status, 201); const generateWidget = await request(app.getHttpServer()) - .post(`/dashboard/${dashboardId}/widget/generate/${connectionId}?tableName=${testTableName}`) + .post(`/dashboard/${dashboardId}/widget/generate/${connectionId}`) .send({ chart_description: 'Show total sales by category as a bar chart', }) @@ -194,6 +238,7 @@ test.serial(`${currentTest} should generate a widget with AI for chart type`, as test.serial(`${currentTest} should generate a counter widget with AI`, async (t) => { mockResponse = MOCK_AI_RESPONSE_COUNTER; + toolCallCounter = 0; const connectionToTestDB = getTestData(mockFactory).connectionToPostgres; const { token } = await registerUserAndReturnUserInfo(app); @@ -219,7 +264,7 @@ test.serial(`${currentTest} should generate a counter widget with AI`, async (t) const dashboardId = JSON.parse(createDashboard.text).id; const generateWidget = await request(app.getHttpServer()) - .post(`/dashboard/${dashboardId}/widget/generate/${connectionId}?tableName=${testTableName}`) + .post(`/dashboard/${dashboardId}/widget/generate/${connectionId}`) .send({ chart_description: 'Show total count of orders', }) @@ -243,6 +288,7 @@ test.serial(`${currentTest} should generate a counter widget with AI`, async (t) test.serial(`${currentTest} should reject AI-generated unsafe query`, async (t) => { mockResponse = MOCK_AI_RESPONSE_UNSAFE_QUERY; + toolCallCounter = 0; const connectionToTestDB = getTestData(mockFactory).connectionToPostgres; const { token } = await registerUserAndReturnUserInfo(app); @@ -267,7 +313,7 @@ test.serial(`${currentTest} should reject AI-generated unsafe query`, async (t) const dashboardId = JSON.parse(createDashboard.text).id; const generateWidget = await request(app.getHttpServer()) - .post(`/dashboard/${dashboardId}/widget/generate/${connectionId}?tableName=${testTableName}`) + .post(`/dashboard/${dashboardId}/widget/generate/${connectionId}`) .send({ chart_description: 'Delete all data', }) @@ -284,6 +330,7 @@ test.serial(`${currentTest} should reject AI-generated unsafe query`, async (t) test.serial(`${currentTest} should generate widget with custom name`, async (t) => { mockResponse = MOCK_AI_RESPONSE_CHART; + toolCallCounter = 0; const connectionToTestDB = getTestData(mockFactory).connectionToPostgres; const { token } = await registerUserAndReturnUserInfo(app); @@ -309,7 +356,7 @@ test.serial(`${currentTest} should generate widget with custom name`, async (t) const customName = 'My Custom Widget Name'; const generateWidget = await request(app.getHttpServer()) - .post(`/dashboard/${dashboardId}/widget/generate/${connectionId}?tableName=${testTableName}`) + .post(`/dashboard/${dashboardId}/widget/generate/${connectionId}`) .send({ chart_description: 'Show sales data', name: customName, @@ -324,46 +371,9 @@ test.serial(`${currentTest} should generate widget with custom name`, async (t) t.is(generateWidgetRO.name, customName); }); -test.serial(`${currentTest} should fail without tableName query parameter`, async (t) => { - mockResponse = MOCK_AI_RESPONSE_CHART; - - const { token } = await registerUserAndReturnUserInfo(app); - const newConnection = getTestData(mockFactory).newEncryptedConnection; - - const createConnectionResponse = await request(app.getHttpServer()) - .post('/connection') - .send(newConnection) - .set('Cookie', token) - .set('masterpwd', 'ahalaimahalai') - .set('Content-Type', 'application/json') - .set('Accept', 'application/json'); - const connectionId = JSON.parse(createConnectionResponse.text).id; - - const createDashboard = await request(app.getHttpServer()) - .post(`/dashboards/${connectionId}`) - .send({ name: 'No Table Dashboard' }) - .set('Cookie', token) - .set('masterpwd', 'ahalaimahalai') - .set('Content-Type', 'application/json') - .set('Accept', 'application/json'); - - const dashboardId = JSON.parse(createDashboard.text).id; - - const generateWidget = await request(app.getHttpServer()) - .post(`/dashboard/${dashboardId}/widget/generate/${connectionId}`) - .send({ - chart_description: 'Show some data', - }) - .set('Cookie', token) - .set('masterpwd', 'ahalaimahalai') - .set('Content-Type', 'application/json') - .set('Accept', 'application/json'); - - t.is(generateWidget.status, 400); -}); - test.serial(`${currentTest} should fail without chart_description`, async (t) => { mockResponse = MOCK_AI_RESPONSE_CHART; + toolCallCounter = 0; const connectionToTestDB = getTestData(mockFactory).connectionToPostgres; const { token } = await registerUserAndReturnUserInfo(app); @@ -388,7 +398,7 @@ test.serial(`${currentTest} should fail without chart_description`, async (t) => const dashboardId = JSON.parse(createDashboard.text).id; const generateWidget = await request(app.getHttpServer()) - .post(`/dashboard/${dashboardId}/widget/generate/${connectionId}?tableName=${testTableName}`) + .post(`/dashboard/${dashboardId}/widget/generate/${connectionId}`) .send({}) .set('Cookie', token) .set('masterpwd', 'ahalaimahalai') @@ -400,6 +410,7 @@ test.serial(`${currentTest} should fail without chart_description`, async (t) => test.serial(`${currentTest} should fail for non-existent dashboard`, async (t) => { mockResponse = MOCK_AI_RESPONSE_CHART; + toolCallCounter = 0; const connectionToTestDB = getTestData(mockFactory).connectionToPostgres; const { token } = await registerUserAndReturnUserInfo(app); @@ -416,7 +427,7 @@ test.serial(`${currentTest} should fail for non-existent dashboard`, async (t) = const fakeDashboardId = faker.string.uuid(); const generateWidget = await request(app.getHttpServer()) - .post(`/dashboard/${fakeDashboardId}/widget/generate/${connectionId}?tableName=${testTableName}`) + .post(`/dashboard/${fakeDashboardId}/widget/generate/${connectionId}`) .send({ chart_description: 'Show some data', }) @@ -427,40 +438,3 @@ test.serial(`${currentTest} should fail for non-existent dashboard`, async (t) = t.is(generateWidget.status, 404); }); - -test.serial(`${currentTest} should fail for non-existent table`, async (t) => { - mockResponse = MOCK_AI_RESPONSE_CHART; - - const connectionToTestDB = getTestData(mockFactory).connectionToPostgres; - const { token } = await registerUserAndReturnUserInfo(app); - - const createConnectionResponse = await request(app.getHttpServer()) - .post('/connection') - .send(connectionToTestDB) - .set('Cookie', token) - .set('Content-Type', 'application/json') - .set('Accept', 'application/json'); - const connectionId = JSON.parse(createConnectionResponse.text).id; - - const createDashboard = await request(app.getHttpServer()) - .post(`/dashboards/${connectionId}`) - .send({ name: 'Non-existent Table Dashboard' }) - .set('Cookie', token) - .set('masterpwd', 'ahalaimahalai') - .set('Content-Type', 'application/json') - .set('Accept', 'application/json'); - - const dashboardId = JSON.parse(createDashboard.text).id; - - const generateWidget = await request(app.getHttpServer()) - .post(`/dashboard/${dashboardId}/widget/generate/${connectionId}?tableName=non_existent_table_xyz`) - .send({ - chart_description: 'Show some data', - }) - .set('Cookie', token) - .set('masterpwd', 'ahalaimahalai') - .set('Content-Type', 'application/json') - .set('Accept', 'application/json'); - - t.is(generateWidget.status, 400); -});