Skip to content

feat: Move VertexAiExampleStore to the Gen AI example_stores module #6985

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
+185 −91
Open

feat: Move VertexAiExampleStore to the Gen AI example_stores module #6985

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
111 changes: 61 additions & 50 deletions src/google/adk/examples/vertex_ai_example_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,94 +14,105 @@

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(
input=types.Content(
role="user",
parts=[
types.Part.from_text(
text=result.example.stored_contents_example.search_key
text=stored_contents_example.search_key or ""
)
],
),
Expand Down
165 changes: 124 additions & 41 deletions tests/unittests/examples/test_vertex_ai_example_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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",
Expand All @@ -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):
Expand All @@ -93,14 +127,22 @@ 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=[
_result(
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")],
)
],
)
Expand Down Expand Up @@ -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"}
)
],
)
Expand All @@ -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},
)
],
)
Expand All @@ -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"]
Loading