Skip to content

Commit 167182b

Browse files
google-genai-botcopybara-github
authored andcommitted
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
1 parent 3f3aa7b commit 167182b

1 file changed

Lines changed: 86 additions & 16 deletions

File tree

src/google/adk/tools/bigquery/query_tool.py

Lines changed: 86 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,6 @@
1616

1717
import functools
1818
import json
19-
import sys
20-
import textwrap
2119
import types
2220
from typing import Callable
2321

@@ -79,7 +77,6 @@ def execute_sql(
7977
]
8078
}
8179
"""
82-
8380
try:
8481
# Get BigQuery client
8582
bq_client = client.get_bigquery_client(
@@ -189,7 +186,47 @@ def execute_sql(
189186
}
190187

191188

192-
_execute_sql_write_examples = """
189+
def _execute_sql_write_mode(*args, **kwargs) -> dict:
190+
"""Run a BigQuery or BigQuery ML SQL query in the project and return the result.
191+
192+
Args:
193+
project_id (str): The GCP project id in which the query should be
194+
executed.
195+
query (str): The BigQuery SQL query to be executed.
196+
credentials (Credentials): The credentials to use for the request.
197+
settings (BigQueryToolConfig): The settings for the tool.
198+
tool_context (ToolContext): The context for the tool.
199+
200+
Returns:
201+
dict: Dictionary representing the result of the query.
202+
If the result contains the key "result_is_likely_truncated" with
203+
value True, it means that there may be additional rows matching the
204+
query not returned in the result.
205+
206+
Examples:
207+
Fetch data or insights from a table:
208+
209+
>>> execute_sql("my_project",
210+
... "SELECT island, COUNT(*) AS population "
211+
... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
212+
{
213+
"status": "SUCCESS",
214+
"rows": [
215+
{
216+
"island": "Dream",
217+
"population": 124
218+
},
219+
{
220+
"island": "Biscoe",
221+
"population": 168
222+
},
223+
{
224+
"island": "Torgersen",
225+
"population": 52
226+
}
227+
]
228+
}
229+
193230
Create a table with schema prescribed:
194231
195232
>>> execute_sql("my_project",
@@ -328,9 +365,50 @@ def execute_sql(
328365
- Use "CREATE OR REPLACE MODEL" instead of "CREATE MODEL".
329366
- First run "DROP MODEL", followed by "CREATE MODEL".
330367
"""
368+
return execute_sql(*args, **kwargs)
331369

332370

333-
_execute_sql_protecetd_write_examples = """
371+
def _execute_sql_protected_write_mode(*args, **kwargs) -> dict:
372+
"""Run a BigQuery or BigQuery ML SQL query in the project and return the result.
373+
374+
Args:
375+
project_id (str): The GCP project id in which the query should be
376+
executed.
377+
query (str): The BigQuery SQL query to be executed.
378+
credentials (Credentials): The credentials to use for the request.
379+
settings (BigQueryToolConfig): The settings for the tool.
380+
tool_context (ToolContext): The context for the tool.
381+
382+
Returns:
383+
dict: Dictionary representing the result of the query.
384+
If the result contains the key "result_is_likely_truncated" with
385+
value True, it means that there may be additional rows matching the
386+
query not returned in the result.
387+
388+
Examples:
389+
Fetch data or insights from a table:
390+
391+
>>> execute_sql("my_project",
392+
... "SELECT island, COUNT(*) AS population "
393+
... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
394+
{
395+
"status": "SUCCESS",
396+
"rows": [
397+
{
398+
"island": "Dream",
399+
"population": 124
400+
},
401+
{
402+
"island": "Biscoe",
403+
"population": 168
404+
},
405+
{
406+
"island": "Torgersen",
407+
"population": 52
408+
}
409+
]
410+
}
411+
334412
Create a temporary table with schema prescribed:
335413
336414
>>> execute_sql("my_project",
@@ -460,6 +538,7 @@ def execute_sql(
460538
- Only temporary models can be created or deleted. Please do not try
461539
creating a permanent model (non-TEMP model) or deleting one.
462540
"""
541+
return execute_sql(*args, **kwargs)
463542

464543

465544
def get_execute_sql(settings: BigQueryToolConfig) -> Callable[..., dict]:
@@ -496,17 +575,8 @@ def get_execute_sql(settings: BigQueryToolConfig) -> Callable[..., dict]:
496575

497576
# Now, set the new docstring
498577
if settings.write_mode == WriteMode.PROTECTED:
499-
examples = _execute_sql_protecetd_write_examples
578+
execute_sql_wrapper.__doc__ = _execute_sql_protected_write_mode.__doc__
500579
else:
501-
examples = _execute_sql_write_examples
502-
503-
# Handle Python 3.13+ inspect.cleandoc behavior change
504-
# Python 3.13 changed inspect.cleandoc from lstrip() to lstrip(' '), making it
505-
# more conservative. The appended examples have inconsistent indentation that
506-
# Python 3.11/3.12's aggressive cleandoc would fix, but 3.13+ needs help.
507-
if sys.version_info >= (3, 13):
508-
examples = textwrap.dedent(examples)
509-
510-
execute_sql_wrapper.__doc__ += examples
580+
execute_sql_wrapper.__doc__ = _execute_sql_write_mode.__doc__
511581

512582
return execute_sql_wrapper

0 commit comments

Comments
 (0)