Skip to content

Conversation

@Faraaz1994
Copy link

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

2. Or, if no issue exists, describe the change:

If applicable, please follow the issue templates to provide as much detail as
possible.

Problem:
Currently, the BigQuery tool in ADK does not provide a way for developers to add custom labels to BigQuery jobs created by their agents. This makes it difficult to:

Track and monitor BigQuery costs associated with specific agents or use cases
Organize and filter BigQuery jobs in the Google Cloud Console
Implement billing attribution and resource organization strategies
Differentiate between jobs from different environments (dev, staging, production)
While the tool automatically adds an internal adk-bigquery-tool label with the caller_id, there's no mechanism for users to add their own custom labels for tracking and monitoring purposes.

Solution:
Add a labels configuration field to BigQueryToolConfig that allows users to specify custom key-value pairs to be applied to all BigQuery jobs executed by the agent. The solution should:

Configuration Option: Add an optional labels parameter to BigQueryToolConfig accepting a dictionary of string key-value pairs
Validation: Ensure labels follow BigQuery's requirements (non-empty string keys, string values)
Job Application: Automatically apply configured labels to all BigQuery jobs alongside the existing internal labels Documentation: Provide clear documentation on how to use labels for tracking and monitoring

Testing Plan

Please describe the tests that you ran to verify your changes. This is required
for all PRs that are not small documentation or typo fixes.

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

Please include a summary of passed pytest results.

pytest tests/unittests/tools/bigquery/test_bigquery_tool_config.py -v --tb=line -W ignore::UserWarning
========================================= test session starts ==========================================
platform darwin -- Python 3.11.14, pytest-9.0.1, pluggy-1.6.0 -- *****redacted******
cachedir: .pytest_cache
rootdir: *****redacted******
configfile: pyproject.toml
plugins: mock-3.15.1, anyio-4.11.0, xdist-3.8.0, langsmith-0.4.43, asyncio-1.3.0
asyncio: mode=Mode.AUTO, debug=False, asyncio_default_fixture_loop_scope=function, asyncio_default_test_loop_scope=function
collected 14 items                                                                                                                                                                                                         

tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_experimental_warning PASSED                                                                                                   [  7%]
tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_invalid_property PASSED                                                                                                       [ 14%]
tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_invalid_application_name PASSED                                                                                               [ 21%]
tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_max_query_result_rows_default PASSED                                                                                          [ 28%]
tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_max_query_result_rows_custom PASSED                                                                                           [ 35%]
tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_valid_maximum_bytes_billed PASSED                                                                                             [ 42%]
tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_invalid_maximum_bytes_billed PASSED                                                                                           [ 50%]
tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_valid_labels PASSED                                                                                                           [ 57%]
tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_empty_labels PASSED                                                                                                           [ 64%]
tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_none_labels PASSED                                                                                                            [ 71%]
tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_invalid_labels_type PASSED                                                                                                    [ 78%]
tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_invalid_label_key_type PASSED                                                                                                 [ 85%]
tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_invalid_label_value_type PASSED                                                                                               [ 92%]
tests/unittests/tools/bigquery/test_bigquery_tool_config.py::test_bigquery_tool_config_empty_label_key PASSED                                                                                                        [100%]

==================================================================================================== 14 passed in 2.02s ====================================================================================================

Manual End-to-End (E2E) Tests:

Please provide instructions on how to manually test your changes, including any
necessary setup or configuration. Please provide logs or screenshots to help
reviewers better understand the fix.

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

Additional context

Add any other context or screenshots about the feature request here.

@google-cla
Copy link

google-cla bot commented Nov 17, 2025

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello @Faraaz1994, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the BigQuery tool within the ADK by enabling developers to attach custom labels to their BigQuery jobs. Previously, users lacked the ability to add their own labels, hindering effective cost tracking, job organization, and environment differentiation. The new 'labels' configuration option in 'BigQueryToolConfig' addresses this by allowing the specification of key-value pairs, which are then automatically applied to all BigQuery jobs. This change provides greater flexibility for resource management and improves operational visibility within Google Cloud.

Highlights

  • Custom BigQuery Labels: Introduced an optional 'labels' field in 'BigQueryToolConfig' to allow users to specify custom key-value pairs for BigQuery jobs.
  • Label Validation: Implemented robust validation for the new 'labels' field, ensuring keys and values are strings and keys are not empty, aligning with BigQuery's requirements.
  • Automatic Label Application: Configured the BigQuery tool to automatically apply these custom labels to all executed BigQuery jobs, alongside existing internal labels.
  • Enhanced Test Coverage: Added comprehensive unit tests to verify the correct behavior and validation of the new 'labels' configuration.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

@Faraaz1994 Faraaz1994 changed the title Add custom bq labels feat(bigquery): Add labels support to BigQueryToolConfig for job tracking and monitoring Nov 17, 2025
@adk-bot adk-bot added the tools [Component] This issue is related to tools label Nov 17, 2025
@adk-bot
Copy link
Collaborator

adk-bot commented Nov 17, 2025

Response from ADK Triaging Agent

Hello @Faraaz1994, thank you for creating this PR!

Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). You can visit https://cla.developers.google.com/ to see your current agreements or to sign a new one.

This information will help reviewers to review your PR more efficiently. Thanks!

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a valuable feature by allowing users to add custom labels to BigQuery jobs, which will significantly improve tracking and cost management. The implementation is well-structured, adding a labels field to BigQueryToolConfig, along with corresponding validation and integration into the query execution logic. The unit tests are comprehensive, covering a good range of scenarios.

I have a few suggestions to enhance the code's robustness and maintainability. In src/google/adk/tools/bigquery/config.py, the label validation can be simplified by removing redundant checks already handled by Pydantic. In src/google/adk/tools/bigquery/query_tool.py, it's crucial to create a copy of the labels dictionary to prevent unintended side effects from mutating the configuration object. Lastly, for consistency and test robustness, a pytest.warns context manager should be added to the new tests in tests/unittests/tools/bigquery/test_bigquery_tool_config.py.

Overall, this is a solid contribution. Addressing these points will further strengthen the implementation.

@ryanaiagent ryanaiagent self-assigned this Nov 18, 2025
@Faraaz1994
Copy link
Author

@shobsi thanks for the review. I have addressed the review comments please take a look.

@Faraaz1994 Faraaz1994 requested a review from shobsi November 19, 2025 06:21
@ryanaiagent
Copy link
Collaborator

Hi @Faraaz1994 , Thank you for your contribution! It appears you haven't yet signed the Contributor License Agreement (CLA). Please visit https://cla.developers.google.com/ to complete the signing process. Once the CLA is signed, we'll be able to proceed with the review of your PR. Thank you!

@ryanaiagent ryanaiagent added the request clarification [Status] The maintainer need clarification or more information from the author label Nov 20, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

request clarification [Status] The maintainer need clarification or more information from the author tools [Component] This issue is related to tools

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat(bigquery): Add labels support to BigQueryToolConfig for job tracking and monitoring

4 participants