From e97ec9c00269f9f7ab5e7bb2d8ab972fc579d1a5 Mon Sep 17 00:00:00 2001 From: dal Date: Mon, 25 Aug 2025 12:48:19 -0600 Subject: [PATCH] Enhance S3 integration functionality - Added bucketName to S3 integration responses and handlers to support bucket-specific operations. - Updated tests to reflect the inclusion of bucketName in integration creation and retrieval processes. - Improved error handling for fetching bucket names from the vault during S3 integration retrieval. These changes improve the S3 integration management by allowing for more granular control and visibility of storage buckets. --- .../create-s3-integration.test.ts | 1 + .../s3-integrations/create-s3-integration.ts | 1 + .../v2/s3-integrations/get-s3-integration.ts | 24 ++++- .../analyst-agent-task/analyst-agent-task.ts | 2 +- .../integrations/StorageIntegrations.tsx | 2 +- packages/ai/src/llm/sonnet-4.ts | 18 ---- .../done-tool/done-tool-start.ts | 49 ++++++++- .../done-tool/done-tool-streaming.test.ts | 5 + .../done-tool/done-tool.int.test.ts | 3 + .../done-tool/done-tool.ts | 1 + .../helpers/done-tool-file-selection.ts | 100 ++++++++++++++++++ .../src/s3-integrations/responses.ts | 1 + 12 files changed, 184 insertions(+), 23 deletions(-) diff --git a/apps/server/src/api/v2/s3-integrations/create-s3-integration.test.ts b/apps/server/src/api/v2/s3-integrations/create-s3-integration.test.ts index a9bfc2d35..87454fd7f 100644 --- a/apps/server/src/api/v2/s3-integrations/create-s3-integration.test.ts +++ b/apps/server/src/api/v2/s3-integrations/create-s3-integration.test.ts @@ -60,6 +60,7 @@ describe('createS3IntegrationHandler', () => { id: 'integration-123', provider: 's3', organizationId: 'org-123', + bucketName: 'test-bucket', createdAt: new Date('2024-01-15'), updatedAt: new Date('2024-01-15'), deletedAt: null, diff --git a/apps/server/src/api/v2/s3-integrations/create-s3-integration.ts b/apps/server/src/api/v2/s3-integrations/create-s3-integration.ts index 0802014dc..233ad5e06 100644 --- a/apps/server/src/api/v2/s3-integrations/create-s3-integration.ts +++ b/apps/server/src/api/v2/s3-integrations/create-s3-integration.ts @@ -72,6 +72,7 @@ export async function createS3IntegrationHandler( id: integration.id, provider: integration.provider, organizationId: integration.organizationId, + bucketName: request.bucket, createdAt: integration.createdAt, updatedAt: integration.updatedAt, deletedAt: integration.deletedAt, diff --git a/apps/server/src/api/v2/s3-integrations/get-s3-integration.ts b/apps/server/src/api/v2/s3-integrations/get-s3-integration.ts index 586250840..041eba5b7 100644 --- a/apps/server/src/api/v2/s3-integrations/get-s3-integration.ts +++ b/apps/server/src/api/v2/s3-integrations/get-s3-integration.ts @@ -1,6 +1,10 @@ import type { User } from '@buster/database'; -import { getS3IntegrationByOrganizationId, getUserOrganizationId } from '@buster/database'; -import type { GetS3IntegrationResponse } from '@buster/server-shared'; +import { + getS3IntegrationByOrganizationId, + getSecretByName, + getUserOrganizationId, +} from '@buster/database'; +import type { CreateS3IntegrationRequest, GetS3IntegrationResponse } from '@buster/server-shared'; import { HTTPException } from 'hono/http-exception'; /** @@ -32,10 +36,26 @@ export async function getS3IntegrationHandler(user: User): Promise => { const taskStartTime = Date.now(); const resourceTracker = new ResourceTracker(); diff --git a/apps/web/src/components/features/integrations/StorageIntegrations.tsx b/apps/web/src/components/features/integrations/StorageIntegrations.tsx index bc06398fe..6ec2a6bcb 100644 --- a/apps/web/src/components/features/integrations/StorageIntegrations.tsx +++ b/apps/web/src/components/features/integrations/StorageIntegrations.tsx @@ -140,7 +140,7 @@ const StorageConfiguration = React.memo(() => {
- {providerLabels[s3Integration.provider]} + {s3Integration.bucketName || providerLabels[s3Integration.provider]}
diff --git a/packages/ai/src/llm/sonnet-4.ts b/packages/ai/src/llm/sonnet-4.ts index 89a67f879..526e5c75b 100644 --- a/packages/ai/src/llm/sonnet-4.ts +++ b/packages/ai/src/llm/sonnet-4.ts @@ -37,15 +37,6 @@ function initializeSonnet4(): ReturnType { } } - if (process.env.ANTHROPIC_API_KEY) { - try { - models.push(anthropicModel('claude-opus-4-1-20250805')); - console.info('Opus41: Anthropic model added to fallback chain'); - } catch (error) { - console.warn('Opus41: Failed to initialize Anthropic model:', error); - } - } - // Ensure we have at least one model if (models.length === 0) { throw new Error( @@ -53,15 +44,6 @@ function initializeSonnet4(): ReturnType { ); } - if (process.env.OPENAI_API_KEY) { - try { - models.push(openaiModel('gpt-5')); - console.info('Sonnet4: OpenAI model added to fallback chain'); - } catch (error) { - console.warn('Sonnet4: Failed to initialize OpenAI model:', error); - } - } - console.info(`Sonnet4: Initialized with ${models.length} model(s) in fallback chain`); _sonnet4Instance = createFallback({ diff --git a/packages/ai/src/tools/communication-tools/done-tool/done-tool-start.ts b/packages/ai/src/tools/communication-tools/done-tool/done-tool-start.ts index 71924c8b6..1584389b8 100644 --- a/packages/ai/src/tools/communication-tools/done-tool/done-tool-start.ts +++ b/packages/ai/src/tools/communication-tools/done-tool/done-tool-start.ts @@ -1,9 +1,10 @@ -import { updateMessage, updateMessageEntries } from '@buster/database'; +import { updateChat, updateMessage, updateMessageEntries } from '@buster/database'; import type { ToolCallOptions } from 'ai'; import type { UpdateMessageEntriesParams } from '../../../../../database/src/queries/messages/update-message-entries'; import type { DoneToolContext, DoneToolState } from './done-tool'; import { createFileResponseMessages, + extractAllFilesForChatUpdate, extractFilesFromToolCalls, } from './helpers/done-tool-file-selection'; import { @@ -26,13 +27,24 @@ export function createDoneToolStart(context: DoneToolContext, doneToolState: Don toolCallId: options.toolCallId, }); + // Extract files for response messages (filtered to avoid duplicates) const extractedFiles = extractFilesFromToolCalls(options.messages); + // Extract ALL files for updating the chat's most recent file (includes reports) + const allFilesForChatUpdate = extractAllFilesForChatUpdate(options.messages); + console.info('[done-tool-start] Files extracted', { extractedCount: extractedFiles.length, files: extractedFiles.map((f) => ({ id: f.id, type: f.fileType, name: f.fileName })), + allFilesCount: allFilesForChatUpdate.length, + allFiles: allFilesForChatUpdate.map((f) => ({ + id: f.id, + type: f.fileType, + name: f.fileName, + })), }); + // Add extracted files as response messages (these are filtered to avoid duplicates) if (extractedFiles.length > 0 && context.messageId) { const fileResponses = createFileResponseMessages(extractedFiles); @@ -50,6 +62,41 @@ export function createDoneToolStart(context: DoneToolContext, doneToolState: Don console.error('[done-tool] Failed to add file response entries:', error); } } + + // Update the chat with the most recent file (using ALL files, including reports) + if (context.chatId && allFilesForChatUpdate.length > 0) { + // Sort files by version number (descending) to get the most recent + const sortedFiles = allFilesForChatUpdate.sort((a, b) => { + const versionA = a.versionNumber || 1; + const versionB = b.versionNumber || 1; + return versionB - versionA; + }); + + // Prefer reports over other file types for the chat's most recent file + const reportFile = sortedFiles.find((f) => f.fileType === 'report'); + const mostRecentFile = reportFile || sortedFiles[0]; + + if (mostRecentFile) { + console.info('[done-tool-start] Updating chat with most recent file', { + chatId: context.chatId, + fileId: mostRecentFile.id, + fileType: mostRecentFile.fileType, + fileName: mostRecentFile.fileName, + versionNumber: mostRecentFile.versionNumber, + isReport: mostRecentFile.fileType === 'report', + }); + + try { + await updateChat(context.chatId, { + mostRecentFileId: mostRecentFile.id, + mostRecentFileType: mostRecentFile.fileType as 'metric' | 'dashboard' | 'report', + mostRecentVersionNumber: mostRecentFile.versionNumber || 1, + }); + } catch (error) { + console.error('[done-tool] Failed to update chat with most recent file:', error); + } + } + } } const doneToolResponseEntry = createDoneToolResponseMessage(doneToolState, options.toolCallId); diff --git a/packages/ai/src/tools/communication-tools/done-tool/done-tool-streaming.test.ts b/packages/ai/src/tools/communication-tools/done-tool/done-tool-streaming.test.ts index 15a436e4a..a16212577 100644 --- a/packages/ai/src/tools/communication-tools/done-tool/done-tool-streaming.test.ts +++ b/packages/ai/src/tools/communication-tools/done-tool/done-tool-streaming.test.ts @@ -8,11 +8,13 @@ import { createDoneToolStart } from './done-tool-start'; vi.mock('@buster/database', () => ({ updateMessageEntries: vi.fn().mockResolvedValue({ success: true }), updateMessage: vi.fn().mockResolvedValue({ success: true }), + updateChat: vi.fn().mockResolvedValue({ success: true }), })); describe('Done Tool Streaming Tests', () => { const mockContext: DoneToolContext = { messageId: 'test-message-id-123', + chatId: 'test-chat-id-456', workflowStartTime: Date.now(), }; @@ -112,6 +114,7 @@ describe('Done Tool Streaming Tests', () => { test('should handle context without messageId', async () => { const contextWithoutMessageId: DoneToolContext = { messageId: '', + chatId: 'test-chat-id-456', workflowStartTime: Date.now(), }; const state: DoneToolState = { @@ -364,11 +367,13 @@ The following items were processed: test('should enforce DoneToolContext type requirements', () => { const validContext: DoneToolContext = { messageId: 'message-123', + chatId: 'test-chat-id-456', workflowStartTime: Date.now(), }; const extendedContext = { messageId: 'message-456', + chatId: 'test-chat-id-456', workflowStartTime: Date.now(), additionalField: 'extra-data', }; diff --git a/packages/ai/src/tools/communication-tools/done-tool/done-tool.int.test.ts b/packages/ai/src/tools/communication-tools/done-tool/done-tool.int.test.ts index 005af0147..f74f92765 100644 --- a/packages/ai/src/tools/communication-tools/done-tool/done-tool.int.test.ts +++ b/packages/ai/src/tools/communication-tools/done-tool/done-tool.int.test.ts @@ -28,6 +28,7 @@ describe('Done Tool Integration Tests', () => { mockContext = { messageId: testMessageId, + chatId: testChatId, workflowStartTime: Date.now(), }; }); @@ -233,6 +234,7 @@ All operations completed successfully.`; test('should handle database errors gracefully', async () => { const invalidContext: DoneToolContext = { messageId: 'non-existent-message-id', + chatId: testChatId, workflowStartTime: Date.now(), }; @@ -258,6 +260,7 @@ All operations completed successfully.`; const invalidContext: DoneToolContext = { messageId: 'invalid-id', + chatId: testChatId, workflowStartTime: Date.now(), }; diff --git a/packages/ai/src/tools/communication-tools/done-tool/done-tool.ts b/packages/ai/src/tools/communication-tools/done-tool/done-tool.ts index 98e4d4979..2b0699126 100644 --- a/packages/ai/src/tools/communication-tools/done-tool/done-tool.ts +++ b/packages/ai/src/tools/communication-tools/done-tool/done-tool.ts @@ -22,6 +22,7 @@ const DoneToolOutputSchema = z.object({ const DoneToolContextSchema = z.object({ messageId: z.string().describe('The message ID of the message that triggered the done tool'), + chatId: z.string().describe('The chat ID that this message belongs to'), workflowStartTime: z.number().describe('The start time of the workflow'), }); diff --git a/packages/ai/src/tools/communication-tools/done-tool/helpers/done-tool-file-selection.ts b/packages/ai/src/tools/communication-tools/done-tool/helpers/done-tool-file-selection.ts index e26522883..66771d826 100644 --- a/packages/ai/src/tools/communication-tools/done-tool/helpers/done-tool-file-selection.ts +++ b/packages/ai/src/tools/communication-tools/done-tool/helpers/done-tool-file-selection.ts @@ -66,6 +66,106 @@ interface ReportInfo { operation: 'created' | 'modified'; } +/** + * Extract ALL files from tool calls for updating the chat's most recent file + * This includes reports that would normally be filtered out + */ +export function extractAllFilesForChatUpdate(messages: ModelMessage[]): ExtractedFile[] { + const files: ExtractedFile[] = []; + + console.info('[done-tool-file-selection] Extracting ALL files for chat update', { + messageCount: messages.length, + }); + + // First pass: extract create report content from assistant messages + const createReportContents: Map = new Map(); + + for (const message of messages) { + if (message.role === 'assistant' && Array.isArray(message.content)) { + for (const content of message.content) { + if ( + content && + typeof content === 'object' && + 'type' in content && + content.type === 'tool-call' && + 'toolName' in content && + content.toolName === CREATE_REPORTS_TOOL_NAME + ) { + const contentObj = content as { toolCallId?: string; input?: unknown }; + const toolCallId = contentObj.toolCallId; + const input = contentObj.input as { + files?: Array<{ yml_content?: string; content?: string }>; + }; + if (toolCallId && input && input.files && Array.isArray(input.files)) { + for (const file of input.files) { + const reportContent = file.yml_content || file.content; + if (reportContent) { + createReportContents.set(toolCallId, reportContent); + } + } + } + } + } + } + } + + // Second pass: process tool results + for (const message of messages) { + if (message.role === 'tool') { + const toolContent = message.content; + + if (Array.isArray(toolContent)) { + for (const content of toolContent) { + if (content && typeof content === 'object') { + if ('type' in content && content.type === 'tool-result') { + const toolName = (content as unknown as Record).toolName; + const output = (content as unknown as Record).output; + const contentWithCallId = content as { toolCallId?: string }; + const toolCallId = contentWithCallId.toolCallId; + + const outputObj = output as Record; + if (outputObj && outputObj.type === 'json' && outputObj.value) { + try { + const parsedOutput = + typeof outputObj.value === 'string' + ? JSON.parse(outputObj.value) + : outputObj.value; + + processToolOutput( + toolName as string, + parsedOutput, + files, + toolCallId, + createReportContents + ); + } catch (error) { + console.warn('[done-tool-file-selection] Failed to parse JSON output', { + error, + }); + } + } + } else if ('files' in content || 'file' in content) { + processDirectFileContent(content, files); + } + } + } + } + } + } + + // Deduplicate files by ID, keeping highest version + const deduplicatedFiles = deduplicateFilesByVersion(files); + + console.info('[done-tool-file-selection] All extracted files for chat update', { + totalFiles: deduplicatedFiles.length, + metrics: deduplicatedFiles.filter((f) => f.fileType === 'metric').length, + dashboards: deduplicatedFiles.filter((f) => f.fileType === 'dashboard').length, + reports: deduplicatedFiles.filter((f) => f.fileType === 'report').length, + }); + + return deduplicatedFiles; +} + /** * Extract files from tool call responses in the conversation messages * Focuses on tool result messages that contain file information diff --git a/packages/server-shared/src/s3-integrations/responses.ts b/packages/server-shared/src/s3-integrations/responses.ts index 18d37a985..33e1c8264 100644 --- a/packages/server-shared/src/s3-integrations/responses.ts +++ b/packages/server-shared/src/s3-integrations/responses.ts @@ -8,6 +8,7 @@ export const S3IntegrationResponseSchema = z.object({ id: z.string().uuid(), provider: StorageProviderSchema, organizationId: z.string().uuid(), + bucketName: z.string().optional(), createdAt: z.string(), updatedAt: z.string(), deletedAt: z.string().nullable(),