-
Notifications
You must be signed in to change notification settings - Fork 43
fix: copilot by using bucketId and adjust agent models #1691
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from 9 commits
Commits
Show all changes
13 commits
Select commit
Hold shift + click to select a range
18db504
chore: bump bedrock model for data copilot
joanagmaia db49be8
chore: add region prefix to model id
joanagmaia d5dac30
chore: test different models
joanagmaia 5f5c7a4
Merge remote-tracking branch 'origin/main' into chore/bump-bedrock-model
joanagmaia 25595b2
chore: revert timeout
joanagmaia 01c2483
fix: models
joanagmaia b83314d
chore: further improvements
joanagmaia 76da4d8
fix: readme
joanagmaia 22e0e71
fix: imports
joanagmaia 88593d0
fix: address PR comments
joanagmaia cb5ca71
fix: tests model and documentation
joanagmaia cef9536
chore: revert model
joanagmaia 5932c7b
Merge remote-tracking branch 'origin/main' into chore/bump-bedrock-model
joanagmaia File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -5,6 +5,7 @@ import { createAmazonBedrock } from '@ai-sdk/amazon-bedrock'; | |
| import { experimental_createMCPClient as createMCPClient, type LanguageModelV1 } from 'ai'; | ||
| import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'; | ||
| import type { Pool } from 'pg'; | ||
| import { ofetch } from 'ofetch'; | ||
| import type { ChatResponse, IChatResponseDb } from '../../server/repo/chat.repo'; | ||
| import { ChatRepository } from '../../server/repo/chat.repo'; | ||
|
|
||
|
|
@@ -56,14 +57,23 @@ export class DataCopilot { | |
| /** Human-readable overview of tools for router agent decision making */ | ||
| private toolsOverview: string = ''; | ||
|
|
||
| /** Tinybird bucketId for the current project — fetched once and reused across all pipe calls */ | ||
| private bucketId: number | null = null; | ||
|
|
||
| /** Tinybird MCP server URL */ | ||
| private tbMcpUrl: string = ''; | ||
|
|
||
| /** Amazon Bedrock language model instance */ | ||
| private model: LanguageModelV1; | ||
| /** Amazon Bedrock language model instance for routing and auditing (Sonnet) */ | ||
| private sonnetModel: LanguageModelV1; | ||
|
|
||
| /** Amazon Bedrock language model instance for text-to-SQL, pipe, and chart agents (Opus) */ | ||
joanagmaia marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| private opusModel: LanguageModelV1; | ||
|
|
||
| /** Bedrock model identifier for general agents */ | ||
| private readonly BEDROCK_SONNET_MODEL_ID = 'us.anthropic.claude-sonnet-4-20250514-v1:0'; | ||
|
|
||
| /** Bedrock model identifier */ | ||
| private readonly BEDROCK_MODEL_ID = 'us.anthropic.claude-sonnet-4-20250514-v1:0'; | ||
| /** Bedrock model identifier for text-to-SQL and pipe agent */ | ||
| private readonly BEDROCK_OPUS_MODEL_ID = 'us.anthropic.claude-opus-4-6-v1'; | ||
|
|
||
| /** Maximum number of auditor retry attempts */ | ||
| private readonly MAX_AUDITOR_RETRIES = 1; | ||
|
|
@@ -72,7 +82,8 @@ export class DataCopilot { | |
| private readonly MAX_SQL_RETRIES = 2; | ||
|
|
||
| constructor() { | ||
| this.model = bedrock(this.BEDROCK_MODEL_ID); | ||
| this.sonnetModel = bedrock(this.BEDROCK_SONNET_MODEL_ID); | ||
| this.opusModel = bedrock(this.BEDROCK_OPUS_MODEL_ID); | ||
| this.tbMcpUrl = `https://mcp.tinybird.co?token=${process.env.NUXT_INSIGHTS_DATA_COPILOT_TINYBIRD_TOKEN}&host=${process.env.NUXT_TINYBIRD_BASE_URL}`; | ||
| } | ||
|
|
||
|
|
@@ -88,10 +99,43 @@ export class DataCopilot { | |
| }), | ||
| }); | ||
|
|
||
| this.tbTools = await this.mcpClient.tools({}); | ||
| const allTools = await this.mcpClient.tools({}); | ||
|
|
||
| // Filter out tools with empty descriptions — Bedrock rejects them with a validation error | ||
| this.tbTools = Object.fromEntries( | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Anything important is missing descriptions? Should we report these to Tinybird support? |
||
| Object.entries(allTools).filter(([_, tool]: [string, any]) => { | ||
| const description = | ||
| (tool?.description as string) || (tool?.meta?.description as string) || ''; | ||
| return description.trim().length > 0; | ||
| }), | ||
| ); | ||
|
|
||
| this.buildToolsOverview(); | ||
| } | ||
|
|
||
| /** | ||
| * Fetch and store the Tinybird bucketId for a project. | ||
| * Uses ofetch directly to stay outside the Nuxt server context (no useStorage/createError). | ||
| */ | ||
| private async fetchBucketId(project: string): Promise<void> { | ||
| if (!project) return; | ||
| const tinybirdBaseUrl = | ||
| process.env.NUXT_TINYBIRD_BASE_URL || 'https://api.us-west-2.aws.tinybird.co'; | ||
| const tinybirdToken = process.env.NUXT_INSIGHTS_DATA_COPILOT_TINYBIRD_TOKEN; | ||
| if (!tinybirdToken) return; | ||
| try { | ||
| const response = await ofetch( | ||
| `${tinybirdBaseUrl}/v0/pipes/project_buckets.json?project=${encodeURIComponent(project)}`, | ||
| { headers: { Authorization: `Bearer ${tinybirdToken}` }, timeout: 10_000 }, | ||
| ); | ||
| this.bucketId = response.data?.[0]?.bucketId ?? null; | ||
| console.warn(`🪣 [DataCopilot] bucketId for "${project}": ${this.bucketId}`); | ||
joanagmaia marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| } catch (error) { | ||
| console.error(`[DataCopilot] Failed to fetch bucketId for "${project}":`, error); | ||
| this.bucketId = null; | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Build human-readable overview of available tools for the router agent | ||
| */ | ||
|
|
@@ -170,10 +214,20 @@ export class DataCopilot { | |
| instructions?: string, | ||
| ): Promise<void> { | ||
| const chatRepo = new ChatRepository(insightsDbPool); | ||
|
|
||
| let model: string | undefined = this.BEDROCK_SONNET_MODEL_ID; | ||
|
|
||
| if (agent === 'EXECUTE_INSTRUCTIONS') { | ||
| model = undefined; | ||
| // For now, we use the Opus model for text-to-SQL and auditor. | ||
| // The model is currently too slow for both the pipe and router agents. | ||
joanagmaia marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| } else if (agent === 'TEXT_TO_SQL' || agent === 'AUDITOR') { | ||
| model = this.BEDROCK_OPUS_MODEL_ID; | ||
| } | ||
| await chatRepo.saveAgentStep({ | ||
| chatResponseId, | ||
| agent, | ||
| model: agent === 'EXECUTE_INSTRUCTIONS' ? undefined : this.BEDROCK_MODEL_ID, | ||
| model, | ||
| response, | ||
| inputTokens: response?.usage?.promptTokens || 0, | ||
| outputTokens: response?.usage?.completionTokens || 0, | ||
|
|
@@ -207,7 +261,7 @@ export class DataCopilot { | |
| }: Omit<RouterAgentInput, 'toolsOverview' | 'model' | 'tools'>) { | ||
| const agent = new RouterAgent(); | ||
| return agent.execute({ | ||
| model: this.model, | ||
| model: this.sonnetModel, | ||
| messages, | ||
| tools: this.tbTools, | ||
| toolsOverview: this.toolsOverview, | ||
|
|
@@ -248,7 +302,7 @@ export class DataCopilot { | |
|
|
||
| const agent = new TextToSqlAgent(); | ||
| return agent.execute({ | ||
| model: this.model, | ||
| model: this.opusModel, | ||
| messages, | ||
| tools: followUpTools, | ||
| date, | ||
|
|
@@ -288,12 +342,21 @@ export class DataCopilot { | |
| const followUpTools: Record<string, unknown> = {}; | ||
| for (const toolName of toolNames) { | ||
| if (this.tbTools[toolName]) { | ||
| followUpTools[toolName] = this.tbTools[toolName]; | ||
| const tool = this.tbTools[toolName] as any; | ||
| // Wrap execute to inject bucketId into every MCP tool call during planning | ||
| followUpTools[toolName] = | ||
| this.bucketId !== null && this.tbTools[toolName]?.execute | ||
| ? { | ||
| ...this.tbTools[toolName], | ||
| execute: async (params: any) => | ||
| tool.execute({ bucketId: this.bucketId, ...params }), | ||
| } | ||
| : tool; | ||
| } | ||
| } | ||
| const agent = new PipeAgent(); | ||
| return agent.execute({ | ||
| model: this.model, | ||
| model: this.sonnetModel, | ||
| messages, | ||
| tools: followUpTools, | ||
| date, | ||
|
|
@@ -329,7 +392,7 @@ export class DataCopilot { | |
| const dataSummary = generateDataSummary(data); | ||
| const agent = new AuditorAgent(); | ||
| return agent.execute({ | ||
| model: this.model, | ||
| model: this.opusModel, | ||
| messages, | ||
| originalQuestion, | ||
| reformulatedQuestion, | ||
|
|
@@ -606,7 +669,7 @@ export class DataCopilot { | |
| } | ||
|
|
||
| // Prepare for retry - add feedback to messages and loop | ||
| previousFeedback = auditorResult.feedback_to_router; | ||
| previousFeedback = auditorResult.feedback_to_router || undefined; | ||
| attemptNumber++; | ||
|
|
||
| dataStream.writeData({ | ||
|
|
@@ -742,6 +805,12 @@ export class DataCopilot { | |
| const parametersString = JSON.stringify(parameters || {}); | ||
| const date = new Date().toISOString().slice(0, 10); | ||
|
|
||
| // Fetch bucketId once upfront — required by all Tinybird pipes for data partitioning | ||
| const project = (parameters as any)?.project; | ||
| if (project) { | ||
| await this.fetchBucketId(project); | ||
| } | ||
|
|
||
| // Build messages from conversation history | ||
| const { messages, previousWasClarification } = await this.buildMessagesFromConversation( | ||
| currentQuestion, | ||
|
|
@@ -768,7 +837,7 @@ export class DataCopilot { | |
| userPrompt: currentQuestion, | ||
| inputTokens: 0, | ||
| outputTokens: 0, | ||
| model: this.BEDROCK_MODEL_ID, | ||
| model: this.BEDROCK_SONNET_MODEL_ID, | ||
| conversationId: conversationId || '', | ||
| routerResponse: RouterDecisionAction.STOP, | ||
| routerReason: '', | ||
|
|
@@ -866,7 +935,7 @@ export class DataCopilot { | |
| routerReason: routerOutput.reasoning, | ||
| pipeInstructions: undefined, | ||
| sqlQuery: undefined, | ||
| model: this.BEDROCK_MODEL_ID, | ||
| model: this.BEDROCK_SONNET_MODEL_ID, | ||
| conversationId: conversationId, | ||
| }, | ||
| insightsDbPool, | ||
|
|
@@ -901,7 +970,7 @@ export class DataCopilot { | |
| clarificationQuestion: routerOutput.clarification_question || undefined, | ||
| pipeInstructions: undefined, | ||
| sqlQuery: undefined, | ||
| model: this.BEDROCK_MODEL_ID, | ||
| model: this.BEDROCK_SONNET_MODEL_ID, | ||
| conversationId: conversationId, | ||
| }, | ||
| insightsDbPool, | ||
|
|
@@ -1186,7 +1255,7 @@ export class DataCopilot { | |
| // Execute the pipes according to the instructions and combine results (don't stream data yet, auditor will do it) | ||
| const pipeExecutionStart = Date.now(); | ||
| try { | ||
| const combinedData = await executePipeInstructions(pipeOutput.instructions); | ||
| const combinedData = await executePipeInstructions(pipeOutput.instructions, this.bucketId); | ||
| const pipeExecutionTime = (Date.now() - pipeExecutionStart) / 1000; | ||
|
|
||
| // Track successful pipe execution step | ||
|
|
@@ -1246,7 +1315,7 @@ export class DataCopilot { | |
| routerReason: routerOutput.reasoning, | ||
| pipeInstructions, | ||
| sqlQuery, | ||
| model: this.BEDROCK_MODEL_ID, | ||
| model: this.BEDROCK_SONNET_MODEL_ID, | ||
| conversationId: conversationId, | ||
| }, | ||
| insightsDbPool, | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.