From 23b93af3de225ea3ba2fcf43915a214f165704fb Mon Sep 17 00:00:00 2001 From: Amy Wu Date: Wed, 19 Aug 2026 11:13:52 -0700 Subject: [PATCH] feat: Move VertexAiExampleStore to the Gen AI example_stores module PiperOrigin-RevId: 967309371 --- .../adk/examples/vertex_ai_example_store.py | 111 ++++++------ .../examples/test_vertex_ai_example_store.py | 165 +++++++++++++----- 2 files changed, 185 insertions(+), 91 deletions(-) diff --git a/src/google/adk/examples/vertex_ai_example_store.py b/src/google/adk/examples/vertex_ai_example_store.py index e988454d5a8..029f60d3bf7 100644 --- a/src/google/adk/examples/vertex_ai_example_store.py +++ b/src/google/adk/examples/vertex_ai_example_store.py @@ -14,86 +14,97 @@ from __future__ import annotations +import os +from typing import Optional + from google.genai import types from typing_extensions import override from .base_example_provider import BaseExampleProvider from .example import Example +_TOP_K = 10 + +# Below this an example is more likely to mislead the model than help it. +_SIMILARITY_THRESHOLD = 0.5 + class VertexAiExampleStore(BaseExampleProvider): """Provides examples from Vertex example store.""" - def __init__(self, examples_store_name: str): + def __init__( + self, + examples_store_name: str, + *, + project: Optional[str] = None, + location: Optional[str] = None, + ): """Initializes the VertexAiExampleStore. Args: examples_store_name: The resource name of the vertex example store, in the format of ``projects/{project}/locations/{location}/exampleStores/{example_store}``. + project: The project to use for the Agent Platform client. If not set, + the GOOGLE_CLOUD_PROJECT environment variable is used, falling back to + the project in ``examples_store_name``. + location: The location to use for the Agent Platform client. If not set, + the GOOGLE_CLOUD_LOCATION environment variable is used, falling back + to the location in ``examples_store_name``. """ + try: + import agentplatform # noqa: F401 + except ImportError as e: + from ..utils._dependency import missing_extra + + raise missing_extra("google-cloud-aiplatform", "gcp") from e + self.examples_store_name = examples_store_name + self._project = project or os.environ.get("GOOGLE_CLOUD_PROJECT") + self._location = location or os.environ.get("GOOGLE_CLOUD_LOCATION") + + # Fallback: a fully-qualified store name already carries both, so a caller + # that passed one should not also have to set the environment. + if (not self._project or not self._location) and ( + examples_store_name.startswith("projects/") + ): + parts = examples_store_name.split("/") + if len(parts) >= 4 and parts[0] == "projects" and parts[2] == "locations": + self._project = self._project or parts[1] + self._location = self._location or parts[3] @override def get_examples(self, query: str) -> list[Example]: - from ..dependencies.vertexai import example_stores + import agentplatform - example_store = example_stores.ExampleStore(self.examples_store_name) - # Retrieve relevant examples. - request = { - "stored_contents_example_parameters": { + client = agentplatform.Client( + project=self._project, location=self._location + ) + response = client.example_stores.search_examples( + name=self.examples_store_name, + stored_contents_example_parameters={ "content_search_key": { "contents": [{"role": "user", "parts": [{"text": query}]}], "search_key_generation_method": {"last_entry": {}}, } }, - "top_k": 10, - "example_store": self.examples_store_name, - } - response = example_store.api_client.search_examples(request) + config={"top_k": _TOP_K}, + ) returned_examples = [] - # Convert results to genai formats - for result in response.results: - if result.similarity_score < 0.5: + for result in response.results or []: + if (result.similarity_score or 0.0) < _SIMILARITY_THRESHOLD: continue - expected_contents = [ - content.content - for content in ( - result.example.stored_contents_example.contents_example.expected_contents - ) + stored_contents_example = result.example.stored_contents_example + contents_example = stored_contents_example.contents_example + + # The module hands back google.genai Content already, so the expected + # output needs no part-by-part rebuilding. + expected_output = [ + expected.content + for expected in contents_example.expected_contents or [] + if expected.content ] - expected_output = [] - for content in expected_contents: - expected_parts = [] - for part in content.parts: - if part.text: - expected_parts.append(types.Part.from_text(text=part.text)) - elif part.function_call: - expected_parts.append( - types.Part.from_function_call( - name=part.function_call.name, - args={ - key: value - for key, value in part.function_call.args.items() - }, - ) - ) - elif part.function_response: - expected_parts.append( - types.Part.from_function_response( - name=part.function_response.name, - response={ - key: value - for key, value in ( - part.function_response.response.items() - ) - }, - ) - ) - expected_output.append( - types.Content(role=content.role, parts=expected_parts) - ) returned_examples.append( Example( @@ -101,7 +112,7 @@ def get_examples(self, query: str) -> list[Example]: role="user", parts=[ types.Part.from_text( - text=result.example.stored_contents_example.search_key + text=stored_contents_example.search_key or "" ) ], ), diff --git a/tests/unittests/examples/test_vertex_ai_example_store.py b/tests/unittests/examples/test_vertex_ai_example_store.py index 16086b36104..9f4d4c05142 100644 --- a/tests/unittests/examples/test_vertex_ai_example_store.py +++ b/tests/unittests/examples/test_vertex_ai_example_store.py @@ -15,24 +15,17 @@ """Tests for vertex_ai_example_store.""" from types import SimpleNamespace -from unittest import mock from google.adk.examples.vertex_ai_example_store import VertexAiExampleStore +from google.genai import types import pytest +from pytest_mock import MockerFixture _STORE_NAME = "projects/p/locations/l/exampleStores/s" -def _part(*, text=None, function_call=None, function_response=None): - return SimpleNamespace( - text=text, - function_call=function_call, - function_response=function_response, - ) - - def _expected_content(*, role, parts): - return SimpleNamespace(content=SimpleNamespace(role=role, parts=parts)) + return SimpleNamespace(content=types.Content(role=role, parts=parts)) def _result(*, search_key="search key", expected_contents=(), score=1.0): @@ -50,30 +43,21 @@ def _result(*, search_key="search key", expected_contents=(), score=1.0): @pytest.fixture -def mock_example_stores(): - with mock.patch( - "google.adk.dependencies.vertexai.example_stores" - ) as example_stores: - yield example_stores +def search_examples(mocker: MockerFixture): + """Patches agentplatform.Client and returns its search_examples mock.""" + client = mocker.Mock() + mocker.patch("agentplatform.Client", return_value=client) + return client.example_stores.search_examples -@pytest.fixture -def search_examples(mock_example_stores): - return ( - mock_example_stores.ExampleStore.return_value.api_client.search_examples - ) - - -def test_get_examples_searches_the_configured_store( - mock_example_stores, search_examples -): +def test_get_examples_searches_the_configured_store(search_examples): search_examples.return_value = SimpleNamespace(results=[]) VertexAiExampleStore(_STORE_NAME).get_examples("what is the weather?") - mock_example_stores.ExampleStore.assert_called_once_with(_STORE_NAME) - search_examples.assert_called_once_with({ - "stored_contents_example_parameters": { + search_examples.assert_called_once_with( + name=_STORE_NAME, + stored_contents_example_parameters={ "content_search_key": { "contents": [{ "role": "user", @@ -82,9 +66,59 @@ def test_get_examples_searches_the_configured_store( "search_key_generation_method": {"last_entry": {}}, } }, - "top_k": 10, - "example_store": _STORE_NAME, - }) + config={"top_k": 10}, + ) + + +def test_get_examples_derives_project_and_location_from_the_store_name( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch +): + # The suite's conftest exports both, so the fallback is only reachable once + # they are unset. + monkeypatch.delenv("GOOGLE_CLOUD_PROJECT", raising=False) + monkeypatch.delenv("GOOGLE_CLOUD_LOCATION", raising=False) + client_factory = mocker.patch("agentplatform.Client") + client_factory.return_value.example_stores.search_examples.return_value = ( + SimpleNamespace(results=[]) + ) + + VertexAiExampleStore(_STORE_NAME).get_examples("query") + + client_factory.assert_called_once_with(project="p", location="l") + + +def test_get_examples_prefers_the_environment_over_the_store_name( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "env-project") + monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "env-location") + client_factory = mocker.patch("agentplatform.Client") + client_factory.return_value.example_stores.search_examples.return_value = ( + SimpleNamespace(results=[]) + ) + + VertexAiExampleStore(_STORE_NAME).get_examples("query") + + client_factory.assert_called_once_with( + project="env-project", location="env-location" + ) + + +def test_get_examples_prefers_explicit_project_and_location( + mocker: MockerFixture, +): + client_factory = mocker.patch("agentplatform.Client") + client_factory.return_value.example_stores.search_examples.return_value = ( + SimpleNamespace(results=[]) + ) + + VertexAiExampleStore( + _STORE_NAME, project="other-project", location="other-location" + ).get_examples("query") + + client_factory.assert_called_once_with( + project="other-project", location="other-location" + ) def test_get_examples_returns_empty_list_without_results(search_examples): @@ -93,6 +127,13 @@ def test_get_examples_returns_empty_list_without_results(search_examples): assert VertexAiExampleStore(_STORE_NAME).get_examples("query") == [] +def test_get_examples_tolerates_unset_results(search_examples): + # The response field is optional, so an empty search omits it entirely. + search_examples.return_value = SimpleNamespace(results=None) + + assert VertexAiExampleStore(_STORE_NAME).get_examples("query") == [] + + def test_get_examples_converts_text_part(search_examples): search_examples.return_value = SimpleNamespace( results=[ @@ -100,7 +141,8 @@ def test_get_examples_converts_text_part(search_examples): search_key="what is the weather?", expected_contents=[ _expected_content( - role="model", parts=[_part(text="it is sunny")] + role="model", + parts=[types.Part.from_text(text="it is sunny")], ) ], ) @@ -144,10 +186,8 @@ def test_get_examples_converts_function_call_part(search_examples): _expected_content( role="model", parts=[ - _part( - function_call=SimpleNamespace( - name="get_weather", args={"city": "London"} - ) + types.Part.from_function_call( + name="get_weather", args={"city": "London"} ) ], ) @@ -171,11 +211,9 @@ def test_get_examples_converts_function_response_part(search_examples): _expected_content( role="user", parts=[ - _part( - function_response=SimpleNamespace( - name="get_weather", - response={"temperature": 12}, - ) + types.Part.from_function_response( + name="get_weather", + response={"temperature": 12}, ) ], ) @@ -189,3 +227,48 @@ def test_get_examples_converts_function_response_part(search_examples): function_response = examples[0].output[0].parts[0].function_response assert function_response.name == "get_weather" assert function_response.response == {"temperature": 12} + + +def test_get_examples_preserves_multi_step_expected_output(search_examples): + # expected_contents is repeated to represent iterative reasoning steps; all + # of them belong in the example's output, in order. + search_examples.return_value = SimpleNamespace( + results=[ + _result( + expected_contents=[ + _expected_content( + role="model", parts=[types.Part.from_text(text="step 1")] + ), + _expected_content( + role="model", parts=[types.Part.from_text(text="step 2")] + ), + ], + ) + ] + ) + + examples = VertexAiExampleStore(_STORE_NAME).get_examples("query") + + assert [content.parts[0].text for content in examples[0].output] == [ + "step 1", + "step 2", + ] + + +def test_get_examples_skips_expected_contents_without_content(search_examples): + search_examples.return_value = SimpleNamespace( + results=[ + _result( + expected_contents=[ + SimpleNamespace(content=None), + _expected_content( + role="model", parts=[types.Part.from_text(text="kept")] + ), + ], + ) + ] + ) + + examples = VertexAiExampleStore(_STORE_NAME).get_examples("query") + + assert [content.parts[0].text for content in examples[0].output] == ["kept"]