You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-03-30 10:57:20 -07:00
fix: make tool description for bigquery execute_sql for various write modes self contained
So far we had a default docstring for the `execute-sql` tool and for non-default write modes we were concatenating more content to it. This was working fine in Python 3.9-3.12 but broke in Python 3.13 because of a nuanced difference in the string concatenation to the `__doc__` property of a function b/433914562#comment4. This change makes the docstring management more robust and readable. PiperOrigin-RevId: 797843736
This commit is contained in:
committed by
Copybara-Service
parent
3f3aa7b32d
commit
167182be01
@@ -16,8 +16,6 @@ from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import json
|
||||
import sys
|
||||
import textwrap
|
||||
import types
|
||||
from typing import Callable
|
||||
|
||||
@@ -79,7 +77,6 @@ def execute_sql(
|
||||
]
|
||||
}
|
||||
"""
|
||||
|
||||
try:
|
||||
# Get BigQuery client
|
||||
bq_client = client.get_bigquery_client(
|
||||
@@ -189,7 +186,47 @@ def execute_sql(
|
||||
}
|
||||
|
||||
|
||||
_execute_sql_write_examples = """
|
||||
def _execute_sql_write_mode(*args, **kwargs) -> dict:
|
||||
"""Run a BigQuery or BigQuery ML SQL query in the project and return the result.
|
||||
|
||||
Args:
|
||||
project_id (str): The GCP project id in which the query should be
|
||||
executed.
|
||||
query (str): The BigQuery SQL query to be executed.
|
||||
credentials (Credentials): The credentials to use for the request.
|
||||
settings (BigQueryToolConfig): The settings for the tool.
|
||||
tool_context (ToolContext): The context for the tool.
|
||||
|
||||
Returns:
|
||||
dict: Dictionary representing the result of the query.
|
||||
If the result contains the key "result_is_likely_truncated" with
|
||||
value True, it means that there may be additional rows matching the
|
||||
query not returned in the result.
|
||||
|
||||
Examples:
|
||||
Fetch data or insights from a table:
|
||||
|
||||
>>> execute_sql("my_project",
|
||||
... "SELECT island, COUNT(*) AS population "
|
||||
... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
|
||||
{
|
||||
"status": "SUCCESS",
|
||||
"rows": [
|
||||
{
|
||||
"island": "Dream",
|
||||
"population": 124
|
||||
},
|
||||
{
|
||||
"island": "Biscoe",
|
||||
"population": 168
|
||||
},
|
||||
{
|
||||
"island": "Torgersen",
|
||||
"population": 52
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
Create a table with schema prescribed:
|
||||
|
||||
>>> execute_sql("my_project",
|
||||
@@ -328,9 +365,50 @@ _execute_sql_write_examples = """
|
||||
- Use "CREATE OR REPLACE MODEL" instead of "CREATE MODEL".
|
||||
- First run "DROP MODEL", followed by "CREATE MODEL".
|
||||
"""
|
||||
return execute_sql(*args, **kwargs)
|
||||
|
||||
|
||||
_execute_sql_protecetd_write_examples = """
|
||||
def _execute_sql_protected_write_mode(*args, **kwargs) -> dict:
|
||||
"""Run a BigQuery or BigQuery ML SQL query in the project and return the result.
|
||||
|
||||
Args:
|
||||
project_id (str): The GCP project id in which the query should be
|
||||
executed.
|
||||
query (str): The BigQuery SQL query to be executed.
|
||||
credentials (Credentials): The credentials to use for the request.
|
||||
settings (BigQueryToolConfig): The settings for the tool.
|
||||
tool_context (ToolContext): The context for the tool.
|
||||
|
||||
Returns:
|
||||
dict: Dictionary representing the result of the query.
|
||||
If the result contains the key "result_is_likely_truncated" with
|
||||
value True, it means that there may be additional rows matching the
|
||||
query not returned in the result.
|
||||
|
||||
Examples:
|
||||
Fetch data or insights from a table:
|
||||
|
||||
>>> execute_sql("my_project",
|
||||
... "SELECT island, COUNT(*) AS population "
|
||||
... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
|
||||
{
|
||||
"status": "SUCCESS",
|
||||
"rows": [
|
||||
{
|
||||
"island": "Dream",
|
||||
"population": 124
|
||||
},
|
||||
{
|
||||
"island": "Biscoe",
|
||||
"population": 168
|
||||
},
|
||||
{
|
||||
"island": "Torgersen",
|
||||
"population": 52
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
Create a temporary table with schema prescribed:
|
||||
|
||||
>>> execute_sql("my_project",
|
||||
@@ -460,6 +538,7 @@ _execute_sql_protecetd_write_examples = """
|
||||
- Only temporary models can be created or deleted. Please do not try
|
||||
creating a permanent model (non-TEMP model) or deleting one.
|
||||
"""
|
||||
return execute_sql(*args, **kwargs)
|
||||
|
||||
|
||||
def get_execute_sql(settings: BigQueryToolConfig) -> Callable[..., dict]:
|
||||
@@ -496,17 +575,8 @@ def get_execute_sql(settings: BigQueryToolConfig) -> Callable[..., dict]:
|
||||
|
||||
# Now, set the new docstring
|
||||
if settings.write_mode == WriteMode.PROTECTED:
|
||||
examples = _execute_sql_protecetd_write_examples
|
||||
execute_sql_wrapper.__doc__ = _execute_sql_protected_write_mode.__doc__
|
||||
else:
|
||||
examples = _execute_sql_write_examples
|
||||
|
||||
# Handle Python 3.13+ inspect.cleandoc behavior change
|
||||
# Python 3.13 changed inspect.cleandoc from lstrip() to lstrip(' '), making it
|
||||
# more conservative. The appended examples have inconsistent indentation that
|
||||
# Python 3.11/3.12's aggressive cleandoc would fix, but 3.13+ needs help.
|
||||
if sys.version_info >= (3, 13):
|
||||
examples = textwrap.dedent(examples)
|
||||
|
||||
execute_sql_wrapper.__doc__ += examples
|
||||
execute_sql_wrapper.__doc__ = _execute_sql_write_mode.__doc__
|
||||
|
||||
return execute_sql_wrapper
|
||||
|
||||
Reference in New Issue
Block a user