fix: update Spanner query tools to async functions

PiperOrigin-RevId: 874318392

Author: google-genai-bot

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

@@ -14,9 +14,9 @@
 
 from __future__ import annotations
 
+import asyncio
 import functools
 import textwrap
-import types
 from typing import Callable
 
 from google.auth.credentials import Credentials
@@ -27,7 +27,7 @@
 from .settings import SpannerToolSettings
 
 
-def execute_sql(
+async def execute_sql(
     project_id: str,
     instance_id: str,
     database_id: str,
@@ -82,7 +82,8 @@ def execute_sql(
   Note:
     This is running with Read-Only Transaction for query that only read data.
   """
-  return utils.execute_sql(
+  return await asyncio.to_thread(
+      utils.execute_sql,
       project_id,
       instance_id,
       database_id,
@@ -179,15 +180,10 @@ def get_execute_sql(settings: SpannerToolSettings) -> Callable[..., dict]:
 
   if settings and settings.query_result_mode is QueryResultMode.DICT_LIST:
 
-    execute_sql_wrapper = types.FunctionType(
-        execute_sql.__code__,
-        execute_sql.__globals__,
-        execute_sql.__name__,
-        execute_sql.__defaults__,
-        execute_sql.__closure__,
-    )
-    functools.update_wrapper(execute_sql_wrapper, execute_sql)
-    # Update with the new docstring
+    @functools.wraps(execute_sql)
+    async def execute_sql_wrapper(*args, **kwargs) -> dict:
+      return await execute_sql(*args, **kwargs)
+
     execute_sql_wrapper.__doc__ = _EXECUTE_SQL_DICT_LIST_MODE_DOCSTRING
     return execute_sql_wrapper
 

src/google/adk/tools/spanner/search_tool.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import json
 from typing import Any
 from typing import Dict
@@ -230,7 +231,7 @@ def _generate_sql_for_ann(
   """
 
 
-def similarity_search(
+async def similarity_search(
     project_id: str,
     instance_id: str,
     database_id: str,
@@ -462,13 +463,16 @@ def similarity_search(
 
     # Generate embedding for the query according to the embedding options.
     if vertex_ai_embedding_model_name:
-      embedding = utils.embed_contents(
-          vertex_ai_embedding_model_name,
-          [query],
-          output_dimensionality,
+      embedding = (
+          await utils.embed_contents_async(
+              vertex_ai_embedding_model_name,
+              [query],
+              output_dimensionality,
+          )
       )[0]
     else:
-      embedding = _get_embedding_for_query(
+      embedding = await asyncio.to_thread(
+          _get_embedding_for_query,
           database,
           database.database_dialect,
           spanner_gsql_embedding_model_name,
@@ -507,30 +511,28 @@ def similarity_search(
     else:
       params = {_GOOGLESQL_PARAMETER_QUERY_EMBEDDING: embedding}
 
-    with database.snapshot() as snapshot:
-      result_set = snapshot.execute_sql(sql, params=params)
-      rows = []
-      result = {}
-      for row in result_set:
-        try:
-          # if the json serialization of the row succeeds, use it as is
-          json.dumps(row)
-        except (TypeError, ValueError, OverflowError):
-          row = str(row)
-
-        rows.append(row)
-
-      result["status"] = "SUCCESS"
-      result["rows"] = rows
-      return result
+    def _execute_sql():
+      with database.snapshot() as snapshot:
+        result_set = snapshot.execute_sql(sql, params=params)
+        rows = []
+        for row in result_set:
+          try:
+            # If the json serialization of the row succeeds, use it as is
+            json.dumps(row)
+          except (TypeError, ValueError, OverflowError):
+            row = str(row)
+          rows.append(row)
+        return {"status": "SUCCESS", "rows": rows}
+
+    return await asyncio.to_thread(_execute_sql)
   except Exception as ex:
     return {
         "status": "ERROR",
         "error_details": repr(ex),
     }
 
 
-def vector_store_similarity_search(
+async def vector_store_similarity_search(
     query: str,
     credentials: Credentials,
     settings: SpannerToolSettings,
@@ -605,7 +607,7 @@ def vector_store_similarity_search(
           settings.vector_store_settings.num_leaves_to_search
       )
 
-    return similarity_search(
+    return await similarity_search(
         project_id=settings.vector_store_settings.project_id,
         instance_id=settings.vector_store_settings.instance_id,
         database_id=settings.vector_store_settings.database_id,

tests/unittests/tools/spanner/test_search_tool.py

@@ -54,11 +54,12 @@ def mock_spanner_ids():
         ),
     ],
 )
-@mock.patch.object(utils, "embed_contents")
+@pytest.mark.asyncio
+@mock.patch.object(utils, "embed_contents_async", autospec=True)
 @mock.patch.object(client, "get_spanner_client")
-def test_similarity_search_knn_success(
+async def test_similarity_search_knn_success(
     mock_get_spanner_client,
-    mock_embed_contents,
+    mock_embed_contents_async,
     mock_spanner_ids,
     mock_credentials,
     embedding_option_key,
@@ -77,7 +78,7 @@ def test_similarity_search_knn_success(
   mock_get_spanner_client.return_value = mock_spanner_client
 
   if embedding_option_key == "vertex_ai_embedding_model_name":
-    mock_embed_contents.return_value = [expected_embedding]
+    mock_embed_contents_async.return_value = [expected_embedding]
     # execute_sql is called once for the kNN search
     mock_snapshot.execute_sql.return_value = iter([("result1",), ("result2",)])
   else:
@@ -90,7 +91,7 @@ def test_similarity_search_knn_success(
         iter([("result1",), ("result2",)]),
     ]
 
-  result = search_tool.similarity_search(
+  result = await search_tool.similarity_search(
       project_id=mock_spanner_ids["project_id"],
       instance_id=mock_spanner_ids["instance_id"],
       database_id=mock_spanner_ids["database_id"],
@@ -111,13 +112,14 @@ def test_similarity_search_knn_success(
   assert "@embedding" in sql
   assert call_args.kwargs == {"params": {"embedding": expected_embedding}}
   if embedding_option_key == "vertex_ai_embedding_model_name":
-    mock_embed_contents.assert_called_once_with(
+    mock_embed_contents_async.assert_called_once_with(
         embedding_option_value, ["test query"], None
     )
 
 
+@pytest.mark.asyncio
 @mock.patch.object(client, "get_spanner_client")
-def test_similarity_search_ann_success(
+async def test_similarity_search_ann_success(
     mock_get_spanner_client, mock_spanner_ids, mock_credentials
 ):
   """Test similarity_search function with ANN success."""
@@ -139,7 +141,7 @@ def test_similarity_search_ann_success(
   mock_spanner_client.instance.return_value = mock_instance
   mock_get_spanner_client.return_value = mock_spanner_client
 
-  result = search_tool.similarity_search(
+  result = await search_tool.similarity_search(
       project_id=mock_spanner_ids["project_id"],
       instance_id=mock_spanner_ids["instance_id"],
       database_id=mock_spanner_ids["database_id"],
@@ -164,13 +166,14 @@ def test_similarity_search_ann_success(
   assert call_args.kwargs == {"params": {"embedding": [0.1, 0.2, 0.3]}}
 
 
+@pytest.mark.asyncio
 @mock.patch.object(client, "get_spanner_client")
-def test_similarity_search_error(
+async def test_similarity_search_error(
     mock_get_spanner_client, mock_spanner_ids, mock_credentials
 ):
   """Test similarity_search function with a generic error."""
   mock_get_spanner_client.side_effect = Exception("Test Exception")
-  result = search_tool.similarity_search(
+  result = await search_tool.similarity_search(
       project_id=mock_spanner_ids["project_id"],
       instance_id=mock_spanner_ids["instance_id"],
       database_id=mock_spanner_ids["database_id"],
@@ -187,11 +190,12 @@ def test_similarity_search_error(
   assert "Test Exception" in result["error_details"]
 
 
-@mock.patch.object(utils, "embed_contents")
+@pytest.mark.asyncio
+@mock.patch.object(utils, "embed_contents_async")
 @mock.patch.object(client, "get_spanner_client")
-def test_similarity_search_circular_row_fallback_to_string(
+async def test_similarity_search_circular_row_fallback_to_string(
     mock_get_spanner_client,
-    mock_embed_contents,
+    mock_embed_contents_async,
     mock_spanner_ids,
     mock_credentials,
 ):
@@ -202,15 +206,15 @@ def test_similarity_search_circular_row_fallback_to_string(
   mock_snapshot = MagicMock()
   circular_row = []
   circular_row.append(circular_row)
-  mock_embed_contents.return_value = [[0.1, 0.2, 0.3]]
+  mock_embed_contents_async.return_value = [[0.1, 0.2, 0.3]]
   mock_snapshot.execute_sql.return_value = iter([circular_row])
   mock_database.snapshot.return_value.__enter__.return_value = mock_snapshot
   mock_database.database_dialect = DatabaseDialect.GOOGLE_STANDARD_SQL
   mock_instance.database.return_value = mock_database
   mock_spanner_client.instance.return_value = mock_instance
   mock_get_spanner_client.return_value = mock_spanner_client
 
-  result = search_tool.similarity_search(
+  result = await search_tool.similarity_search(
       project_id=mock_spanner_ids["project_id"],
       instance_id=mock_spanner_ids["instance_id"],
       database_id=mock_spanner_ids["database_id"],
@@ -228,8 +232,9 @@ def test_similarity_search_circular_row_fallback_to_string(
   assert result["rows"] == [str(circular_row)]
 
 
+@pytest.mark.asyncio
 @mock.patch.object(client, "get_spanner_client")
-def test_similarity_search_postgresql_knn_success(
+async def test_similarity_search_postgresql_knn_success(
     mock_get_spanner_client, mock_spanner_ids, mock_credentials
 ):
   """Test similarity_search with PostgreSQL dialect for kNN."""
@@ -249,7 +254,7 @@ def test_similarity_search_postgresql_knn_success(
   mock_spanner_client.instance.return_value = mock_instance
   mock_get_spanner_client.return_value = mock_spanner_client
 
-  result = search_tool.similarity_search(
+  result = await search_tool.similarity_search(
       project_id=mock_spanner_ids["project_id"],
       instance_id=mock_spanner_ids["instance_id"],
       database_id=mock_spanner_ids["database_id"],
@@ -273,8 +278,9 @@ def test_similarity_search_postgresql_knn_success(
   assert call_args.kwargs == {"params": {"p1": [0.1, 0.2, 0.3]}}
 
 
+@pytest.mark.asyncio
 @mock.patch.object(client, "get_spanner_client")
-def test_similarity_search_postgresql_ann_unsupported(
+async def test_similarity_search_postgresql_ann_unsupported(
     mock_get_spanner_client, mock_spanner_ids, mock_credentials
 ):
   """Test similarity_search with unsupported ANN for PostgreSQL dialect."""
@@ -286,7 +292,7 @@ def test_similarity_search_postgresql_ann_unsupported(
   mock_spanner_client.instance.return_value = mock_instance
   mock_get_spanner_client.return_value = mock_spanner_client
 
-  result = search_tool.similarity_search(
+  result = await search_tool.similarity_search(
       project_id=mock_spanner_ids["project_id"],
       instance_id=mock_spanner_ids["instance_id"],
       database_id=mock_spanner_ids["database_id"],
@@ -311,8 +317,9 @@ def test_similarity_search_postgresql_ann_unsupported(
   )
 
 
+@pytest.mark.asyncio
 @mock.patch.object(client, "get_spanner_client")
-def test_similarity_search_gsql_missing_embedding_model_error(
+async def test_similarity_search_gsql_missing_embedding_model_error(
     mock_get_spanner_client, mock_spanner_ids, mock_credentials
 ):
   """Test similarity_search with missing embedding_options for GoogleSQL dialect."""
@@ -324,7 +331,7 @@ def test_similarity_search_gsql_missing_embedding_model_error(
   mock_spanner_client.instance.return_value = mock_instance
   mock_get_spanner_client.return_value = mock_spanner_client
 
-  result = search_tool.similarity_search(
+  result = await search_tool.similarity_search(
       project_id=mock_spanner_ids["project_id"],
       instance_id=mock_spanner_ids["instance_id"],
       database_id=mock_spanner_ids["database_id"],
@@ -348,8 +355,9 @@ def test_similarity_search_gsql_missing_embedding_model_error(
   )
 
 
+@pytest.mark.asyncio
 @mock.patch.object(client, "get_spanner_client")
-def test_similarity_search_pg_missing_embedding_model_error(
+async def test_similarity_search_pg_missing_embedding_model_error(
     mock_get_spanner_client, mock_spanner_ids, mock_credentials
 ):
   """Test similarity_search with missing embedding_options for PostgreSQL dialect."""
@@ -361,7 +369,7 @@ def test_similarity_search_pg_missing_embedding_model_error(
   mock_spanner_client.instance.return_value = mock_instance
   mock_get_spanner_client.return_value = mock_spanner_client
 
-  result = search_tool.similarity_search(
+  result = await search_tool.similarity_search(
       project_id=mock_spanner_ids["project_id"],
       instance_id=mock_spanner_ids["instance_id"],
       database_id=mock_spanner_ids["database_id"],
@@ -427,8 +435,9 @@ def test_similarity_search_pg_missing_embedding_model_error(
         ),
     ],
 )
+@pytest.mark.asyncio
 @mock.patch.object(client, "get_spanner_client")
-def test_similarity_search_multiple_embedding_options_error(
+async def test_similarity_search_multiple_embedding_options_error(
     mock_get_spanner_client,
     mock_spanner_ids,
     mock_credentials,
@@ -443,7 +452,7 @@ def test_similarity_search_multiple_embedding_options_error(
   mock_spanner_client.instance.return_value = mock_instance
   mock_get_spanner_client.return_value = mock_spanner_client
 
-  result = search_tool.similarity_search(
+  result = await search_tool.similarity_search(
       project_id=mock_spanner_ids["project_id"],
       instance_id=mock_spanner_ids["instance_id"],
       database_id=mock_spanner_ids["database_id"],
@@ -461,8 +470,9 @@ def test_similarity_search_multiple_embedding_options_error(
   )
 
 
+@pytest.mark.asyncio
 @mock.patch.object(client, "get_spanner_client")
-def test_similarity_search_output_dimensionality_gsql_error(
+async def test_similarity_search_output_dimensionality_gsql_error(
     mock_get_spanner_client, mock_spanner_ids, mock_credentials
 ):
   """Test similarity_search with output_dimensionality and spanner_googlesql_embedding_model_name."""
@@ -474,7 +484,7 @@ def test_similarity_search_output_dimensionality_gsql_error(
   mock_spanner_client.instance.return_value = mock_instance
   mock_get_spanner_client.return_value = mock_spanner_client
 
-  result = search_tool.similarity_search(
+  result = await search_tool.similarity_search(
       project_id=mock_spanner_ids["project_id"],
       instance_id=mock_spanner_ids["instance_id"],
       database_id=mock_spanner_ids["database_id"],
@@ -492,8 +502,9 @@ def test_similarity_search_output_dimensionality_gsql_error(
   assert "is not supported when" in result["error_details"]
 
 
+@pytest.mark.asyncio
 @mock.patch.object(client, "get_spanner_client")
-def test_similarity_search_unsupported_algorithm_error(
+async def test_similarity_search_unsupported_algorithm_error(
     mock_get_spanner_client, mock_spanner_ids, mock_credentials
 ):
   """Test similarity_search with an unsupported nearest neighbors algorithm."""
@@ -505,7 +516,7 @@ def test_similarity_search_unsupported_algorithm_error(
   mock_spanner_client.instance.return_value = mock_instance
   mock_get_spanner_client.return_value = mock_spanner_client
 
-  result = search_tool.similarity_search(
+  result = await search_tool.similarity_search(
       project_id=mock_spanner_ids["project_id"],
       instance_id=mock_spanner_ids["instance_id"],
       database_id=mock_spanner_ids["database_id"],