google-adk 1.2.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

google/adk/tools/bigquery/query_tool.py

@@ -12,14 +12,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
+import types
+from typing import Callable
+
+from google.cloud import bigquery
 from google.oauth2.credentials import Credentials
 
-from ...tools.bigquery import client
+from . import client
+from .config import BigQueryToolConfig
+from .config import WriteMode
 
 MAX_DOWNLOADED_QUERY_RESULT_ROWS = 50
 
 
-def execute_sql(project_id: str, query: str, credentials: Credentials) -> dict:
+def execute_sql(
+    project_id: str,
+    query: str,
+    credentials: Credentials,
+    config: BigQueryToolConfig,
+) -> dict:
   """Run a BigQuery SQL query in the project and return the result.
 
   Args:
@@ -35,34 +47,51 @@ def execute_sql(project_id: str, query: str, credentials: Credentials) -> dict:
             query not returned in the result.
 
   Examples:
-      >>> execute_sql("bigframes-dev",
-      ... "SELECT island, COUNT(*) AS population "
-      ... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
-      {
-        "rows": [
-            {
-                "island": "Dream",
-                "population": 124
-            },
-            {
-                "island": "Biscoe",
-                "population": 168
-            },
-            {
-                "island": "Torgersen",
-                "population": 52
-            }
-        ]
-      }
+      Fetch data or insights from a table:
+
+          >>> execute_sql("bigframes-dev",
+          ... "SELECT island, COUNT(*) AS population "
+          ... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
+          {
+            "status": "ERROR",
+            "rows": [
+                {
+                    "island": "Dream",
+                    "population": 124
+                },
+                {
+                    "island": "Biscoe",
+                    "population": 168
+                },
+                {
+                    "island": "Torgersen",
+                    "population": 52
+                }
+            ]
+          }
   """
 
   try:
-    bq_client = client.get_bigquery_client(credentials=credentials)
+    bq_client = client.get_bigquery_client(
+        project=project_id, credentials=credentials
+    )
+    if not config or config.write_mode == WriteMode.BLOCKED:
+      query_job = bq_client.query(
+          query,
+          project=project_id,
+          job_config=bigquery.QueryJobConfig(dry_run=True),
+      )
+      if query_job.statement_type != "SELECT":
+        return {
+            "status": "ERROR",
+            "error_details": "Read-only mode only supports SELECT statements.",
+        }
+
     row_iterator = bq_client.query_and_wait(
         query, project=project_id, max_results=MAX_DOWNLOADED_QUERY_RESULT_ROWS
     )
     rows = [{key: val for key, val in row.items()} for row in row_iterator]
-    result = {"rows": rows}
+    result = {"status": "SUCCESS", "rows": rows}
     if (
         MAX_DOWNLOADED_QUERY_RESULT_ROWS is not None
         and len(rows) == MAX_DOWNLOADED_QUERY_RESULT_ROWS
@@ -74,3 +103,92 @@ def execute_sql(project_id: str, query: str, credentials: Credentials) -> dict:
         "status": "ERROR",
         "error_details": str(ex),
     }
+
+
+_execute_sql_write_examples = """
+      Create a table from the result of a query:
+
+          >>> execute_sql("bigframes-dev",
+          ... "CREATE TABLE my_project.my_dataset.my_table AS "
+          ... "SELECT island, COUNT(*) AS population "
+          ... "FROM bigquery-public-data.ml_datasets.penguins GROUP BY island")
+          {
+            "status": "SUCCESS",
+            "rows": []
+          }
+
+      Delete a table:
+
+          >>> execute_sql("bigframes-dev",
+          ... "DROP TABLE my_project.my_dataset.my_table")
+          {
+            "status": "SUCCESS",
+            "rows": []
+          }
+
+      Copy a table to another table:
+
+          >>> execute_sql("bigframes-dev",
+          ... "CREATE TABLE my_project.my_dataset.my_table_clone "
+          ... "CLONE my_project.my_dataset.my_table")
+          {
+            "status": "SUCCESS",
+            "rows": []
+          }
+
+      Create a snapshot (a lightweight, read-optimized copy) of en existing
+      table:
+
+          >>> execute_sql("bigframes-dev",
+          ... "CREATE SNAPSHOT TABLE my_project.my_dataset.my_table_snapshot "
+          ... "CLONE my_project.my_dataset.my_table")
+          {
+            "status": "SUCCESS",
+            "rows": []
+          }
+
+  Notes:
+      - If a destination table already exists, there are a few ways to overwrite
+      it:
+          - Use "CREATE OR REPLACE TABLE" instead of "CREATE TABLE".
+          - First run "DROP TABLE", followed by "CREATE TABLE".
+      - To insert data into a table, use "INSERT INTO" statement.
+  """
+
+
+def get_execute_sql(config: BigQueryToolConfig) -> Callable[..., dict]:
+  """Get the execute_sql tool customized as per the given tool config.
+
+  Args:
+      config: BigQuery tool configuration indicating the behavior of the
+        execute_sql tool.
+
+  Returns:
+      callable[..., dict]: A version of the execute_sql tool respecting the tool
+      config.
+  """
+
+  if not config or config.write_mode == WriteMode.BLOCKED:
+    return execute_sql
+
+  # Create a new function object using the original function's code and globals.
+  # We pass the original code, globals, name, defaults, and closure.
+  # This creates a raw function object without copying other metadata yet.
+  execute_sql_wrapper = types.FunctionType(
+      execute_sql.__code__,
+      execute_sql.__globals__,
+      execute_sql.__name__,
+      execute_sql.__defaults__,
+      execute_sql.__closure__,
+  )
+
+  # Use functools.update_wrapper to copy over other essential attributes
+  # from the original function to the new one.
+  # This includes __name__, __qualname__, __module__, __annotations__, etc.
+  # It specifically allows us to then set __doc__ separately.
+  functools.update_wrapper(execute_sql_wrapper, execute_sql)
+
+  # Now, set the new docstring
+  execute_sql_wrapper.__doc__ += _execute_sql_write_examples
+
+  return execute_sql_wrapper

google/adk/tools/google_api_tool/googleapi_to_openapi_converter.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import argparse
 import json
 import logging
@@ -505,11 +507,12 @@ def main():
     converter = GoogleApiToOpenApiConverter(args.api_name, args.api_version)
     converter.convert()
     converter.save_openapi_spec(args.output)
-    print(
-        f"Successfully converted {args.api_name} {args.api_version} to"
-        " OpenAPI v3"
+    logger.info(
+        "Successfully converted %s %s to OpenAPI v3",
+        args.api_name,
+        args.api_version,
     )
-    print(f"Output saved to {args.output}")
+    logger.info("Output saved to %s", args.output)
   except Exception as e:
     logger.error("Conversion failed: %s", e)
     return 1

google/adk/tools/google_search_tool.py

@@ -48,7 +48,6 @@ class GoogleSearchTool(BaseTool):
     llm_request.config.tools = llm_request.config.tools or []
     if llm_request.model and 'gemini-1' in llm_request.model:
       if llm_request.config.tools:
-        print(llm_request.config.tools)
         raise ValueError(
             'Google search tool can not be used with other tools in Gemini 1.x.'
         )

google/adk/tools/mcp_tool/__init__.py

@@ -17,6 +17,9 @@ __all__ = []
 try:
   from .conversion_utils import adk_to_mcp_tool_type
   from .conversion_utils import gemini_to_json_schema
+  from .mcp_session_manager import SseConnectionParams
+  from .mcp_session_manager import StdioConnectionParams
+  from .mcp_session_manager import StreamableHTTPConnectionParams
   from .mcp_tool import MCPTool
   from .mcp_toolset import MCPToolset
 
@@ -25,6 +28,9 @@ try:
       'gemini_to_json_schema',
       'MCPTool',
       'MCPToolset',
+      'StdioConnectionParams',
+      'SseConnectionParams',
+      'StreamableHTTPConnectionParams',
   ])
 
 except ImportError as e: