diff --git a/src/google/adk/tools/spanner/query_tool.py b/src/google/adk/tools/spanner/query_tool.py index 3cdede43..24c1be60 100644 --- a/src/google/adk/tools/spanner/query_tool.py +++ b/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 QueryResultMode 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 diff --git a/src/google/adk/tools/spanner/search_tool.py b/src/google/adk/tools/spanner/search_tool.py index 03f695b8..6fb4a93f 100644 --- a/src/google/adk/tools/spanner/search_tool.py +++ b/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,22 +511,20 @@ 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) + 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} - rows.append(row) - - result["status"] = "SUCCESS" - result["rows"] = rows - return result + return await asyncio.to_thread(_execute_sql) except Exception as ex: return { "status": "ERROR", @@ -530,7 +532,7 @@ def similarity_search( } -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, diff --git a/tests/unittests/tools/spanner/test_search_tool.py b/tests/unittests/tools/spanner/test_search_tool.py index 4532dd56..c6a6c742 100644 --- a/tests/unittests/tools/spanner/test_search_tool.py +++ b/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,7 +206,7 @@ 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 @@ -210,7 +214,7 @@ def test_similarity_search_circular_row_fallback_to_string( 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"], diff --git a/tests/unittests/tools/spanner/test_spanner_query_tool.py b/tests/unittests/tools/spanner/test_spanner_query_tool.py index 6c75a3ea..928c207d 100644 --- a/tests/unittests/tools/spanner/test_spanner_query_tool.py +++ b/tests/unittests/tools/spanner/test_spanner_query_tool.py @@ -191,8 +191,9 @@ async def test_execute_sql_query_result( assert tool.description == expected_description +@pytest.mark.asyncio @mock.patch.object(query_tool.utils, "execute_sql", spec_set=True) -def test_execute_sql(mock_utils_execute_sql): +async def test_execute_sql(mock_utils_execute_sql): """Test execute_sql function in query result default mode.""" mock_credentials = mock.create_autospec( Credentials, instance=True, spec_set=True @@ -202,7 +203,7 @@ def test_execute_sql(mock_utils_execute_sql): ) mock_utils_execute_sql.return_value = {"status": "SUCCESS", "rows": [[1]]} - result = query_tool.execute_sql( + result = await query_tool.execute_sql( project_id="test-project", instance_id="test-instance", database_id="test-database",