ai-enhance: respect BigQuery cost guards in agent prompt#2057
Open
arsalann wants to merge 6 commits into
Open
Conversation
Add a Query Cost Safety section to the enhance system prompt instructing the agent to check `.bruin.yml` for max_query_cost / max_billable_bytes (and their _soft variants) before issuing queries against BigQuery assets, dry-run when limits exist, and require explicit user confirmation when they don't. Scoped to BigQuery only. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Mirror the new prompt-level Query Cost Safety section in the user documentation, listing the four configurable keys and explaining what the agent does when limits are configured versus absent. Scoped to BigQuery. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Contributor
|
Address review feedback: - Add an explicit asset-to-connection resolution step so the agent looks up cost guards on the exact connection the asset uses, not the first google_cloud_platform connection it finds. - Move the Query Cost Safety section above Context Discovery so it acts as a precondition the agent evaluates before planning any query. - Add prompt_test cases asserting the cost guard phrases are present and that the section precedes Context Discovery. - Mirror the resolution step in docs/commands/ai-enhance.md. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…proach The earlier approach relied on instructing the AI agent (via prompt text) to respect cost limits. That was unreliable — the agent could ignore or misinterpret the precondition. Move the check into the command itself so it runs deterministically before any AI work begins. When the input scope includes any BigQuery asset (bq.sql/seed/source/sensor), the command now resolves each asset's connection (asset.connection, falling back to pipeline default_connections), reads .bruin.yml, and checks for any of max_query_cost, max_query_cost_soft, max_billable_bytes, max_billable_bytes_soft. If a connection has none, the user is prompted interactively for a USD limit; pressing Enter applies a $5 default. The value is persisted to .bruin.yml so subsequent runs don't re-prompt. Non-TTY / --output json modes apply the default silently for CI safety. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
- gocritic: rewrite if-else chain in collectBigQueryConnections as switch. - unparam: drop unused afero.Fs parameter from collectBigQueryConnections — the function delegates to DefaultPipelineBuilder, which uses its own fs. - noctx: use exec.CommandContext in the gitInit test helper. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
golangci-lint's unparam flagged the error return as always nil — every code path either returns no connections or appends to the slice. Drop the error return and inline the check at the call site. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Summary
## Query Cost Safety (BigQuery only)section to thebruin ai-enhancesystem prompt (pkg/enhance/prompt.go)..bruin.ymlformax_query_cost/max_query_cost_soft/max_billable_bytes/max_billable_bytes_softon the asset'sgoogle_cloud_platformconnection before issuing any query.bruin query --dry-runfirst, rewrite if the estimate exceeds the limit, and never pass--dangerously-bypass-soft-limitswithout explicit user confirmation.No changes to BigQuery enforcement code or the enhance command itself — prompt-only update.
Test plan
make format— clean (golangci-lint not installed locally; unrelated)make test— full unit suite passes, includingpkg/enhance(TestBuildEnhancePrompt)bruin ai-enhance <bq-asset>against a pipeline with (a) no cost guard configured and (b) amax_query_cost_softconfigured; confirm the agent surfaces the no-limit warning in (a) and dry-runs before querying in (b)