From 52482409f21fadbfa062b7d5c6c95e92309e8ac5 Mon Sep 17 00:00:00 2001 From: Amy Wu Date: Thu, 20 Aug 2026 15:25:19 -0700 Subject: [PATCH] feat(agentplatform): Support ExampleStores module for few-shot example management PiperOrigin-RevId: 968089190 --- agentplatform/_genai/_transformers.py | 25 + agentplatform/_genai/client.py | 23 + agentplatform/_genai/example_stores.py | 1445 +++++++++++++++++ agentplatform/_genai/types/__init__.py | 204 +++ agentplatform/_genai/types/common.py | 1169 +++++++++++++ .../genai/replays/test_example_stores.py | 167 ++ 6 files changed, 3033 insertions(+) create mode 100644 agentplatform/_genai/example_stores.py create mode 100644 tests/unit/agentplatform/genai/replays/test_example_stores.py diff --git a/agentplatform/_genai/_transformers.py b/agentplatform/_genai/_transformers.py index 7f58b961c6..f4ec56bdb4 100644 --- a/agentplatform/_genai/_transformers.py +++ b/agentplatform/_genai/_transformers.py @@ -601,3 +601,28 @@ def t_strict_endpoint(endpoint: str) -> str: f"Invalid endpoint format: {endpoint}. Must be in the format of" " projects/.../locations/.../endpoints/... or endpoints/..." ) + + +_EXAMPLE_STORE_RES_NAME_RES = ( + re.compile(r"^projects/[^/]+/locations/[^/]+/exampleStores/[^/]+$"), + re.compile(r"^exampleStores/[^/]+$"), +) + + +def is_example_store_resource_name(name: str) -> bool: + """Returns whether the name addresses an ExampleStore resource.""" + return any(pattern.match(name) for pattern in _EXAMPLE_STORE_RES_NAME_RES) + + +def t_example_store(example_store: str) -> str: + """Validates a name that must address an ExampleStore resource.""" + if not example_store: + raise ValueError("example_store is required.") + + if is_example_store_resource_name(example_store): + return example_store + + raise ValueError( + f"Invalid example store format: {example_store}. Must be in the format" + " of projects/.../locations/.../exampleStores/... or exampleStores/..." + ) diff --git a/agentplatform/_genai/client.py b/agentplatform/_genai/client.py index 00a875655a..57e043c7f4 100644 --- a/agentplatform/_genai/client.py +++ b/agentplatform/_genai/client.py @@ -50,6 +50,9 @@ from agentplatform._genai import ( endpoints as endpoints_module, ) + from agentplatform._genai import ( + example_stores as example_stores_module, + ) _GENAI_MODULES_TELEMETRY_HEADER = "vertex-genai-modules" @@ -96,6 +99,7 @@ def __init__(self, api_client: genai_client.BaseApiClient): # type: ignore[name self._model_garden: Optional[ModuleType] = None self._feedback_entries: Optional[ModuleType] = None self._endpoints: Optional[ModuleType] = None + self._example_stores: Optional[ModuleType] = None @property @_common.experimental_warning( @@ -195,6 +199,15 @@ def endpoints(self) -> "endpoints_module.AsyncEndpoints": ) return self._endpoints.AsyncEndpoints(self._api_client) # type: ignore[no-any-return] + @property + def example_stores(self) -> "example_stores_module.AsyncExampleStores": + if self._example_stores is None: + self._example_stores = importlib.import_module( + ".example_stores", + __package__, + ) + return self._example_stores.AsyncExampleStores(self._api_client) # type: ignore[no-any-return] + @property @_common.experimental_warning( "The Vertex SDK GenAI async rag module is experimental, " @@ -328,6 +341,7 @@ def __init__( self._model_garden: Optional[ModuleType] = None self._feedback_entries: Optional[ModuleType] = None self._endpoints: Optional[ModuleType] = None + self._example_stores: Optional[ModuleType] = None @property def evals(self) -> "evals_module.Evals": @@ -452,6 +466,15 @@ def endpoints(self) -> "endpoints_module.Endpoints": ) return self._endpoints.Endpoints(self._api_client) # type: ignore[no-any-return] + @property + def example_stores(self) -> "example_stores_module.ExampleStores": + if self._example_stores is None: + self._example_stores = importlib.import_module( + ".example_stores", + __package__, + ) + return self._example_stores.ExampleStores(self._api_client) # type: ignore[no-any-return] + @property @_common.experimental_warning( "The Vertex SDK GenAI rag module is experimental, " diff --git a/agentplatform/_genai/example_stores.py b/agentplatform/_genai/example_stores.py new file mode 100644 index 0000000000..54c8754dba --- /dev/null +++ b/agentplatform/_genai/example_stores.py @@ -0,0 +1,1445 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +# Code generated by the Google Gen AI SDK generator DO NOT EDIT. + +import json +import logging +from typing import Any, Optional, Union +from urllib.parse import urlencode + +from google.genai import _api_module +from google.genai import _common +from google.genai._common import get_value_by_path as getv +from google.genai._common import set_value_by_path as setv + +from . import _transformers as t +from . import types + +logger = logging.getLogger("agentplatform_genai.examplestores") + + +def _CreateExampleStoreConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["description"]) is not None: + setv( + to_object, + ["_query", "exampleStore", "description"], + getv(from_object, ["description"]), + ) + + if getv(from_object, ["vertex_embedding_model"]) is not None: + setv( + to_object, + ["_query", "exampleStore", "exampleStoreConfig", "vertexEmbeddingModel"], + getv(from_object, ["vertex_embedding_model"]), + ) + + return to_object + + +def _CreateExampleStoreParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["display_name"]) is not None: + setv( + to_object, + ["_query", "exampleStore", "displayName"], + getv(from_object, ["display_name"]), + ) + + if getv(from_object, ["config"]) is not None: + _CreateExampleStoreConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + +def _DeleteExampleStoreConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + return to_object + + +def _DeleteExampleStoreRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv( + to_object, ["_url", "name"], t.t_example_store(getv(from_object, ["name"])) + ) + + if getv(from_object, ["config"]) is not None: + _DeleteExampleStoreConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + +def _FetchExamplesConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["example_ids"]) is not None: + setv(parent_object, ["exampleIds"], getv(from_object, ["example_ids"])) + + if getv(from_object, ["stored_contents_example_filter"]) is not None: + setv( + parent_object, + ["storedContentsExampleFilter"], + getv(from_object, ["stored_contents_example_filter"]), + ) + + if getv(from_object, ["page_size"]) is not None: + setv(parent_object, ["pageSize"], getv(from_object, ["page_size"])) + + if getv(from_object, ["page_token"]) is not None: + setv(parent_object, ["pageToken"], getv(from_object, ["page_token"])) + + return to_object + + +def _FetchExamplesParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv( + to_object, + ["_url", "exampleStore"], + t.t_example_store(getv(from_object, ["name"])), + ) + + if getv(from_object, ["config"]) is not None: + _FetchExamplesConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + +def _GetExampleStoreOperationParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["operation_name"]) is not None: + setv( + to_object, ["_url", "operationName"], getv(from_object, ["operation_name"]) + ) + + return to_object + + +def _GetExampleStoreParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv( + to_object, ["_url", "name"], t.t_example_store(getv(from_object, ["name"])) + ) + + return to_object + + +def _RemoveExamplesConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["example_ids"]) is not None: + setv(parent_object, ["exampleIds"], getv(from_object, ["example_ids"])) + + if getv(from_object, ["stored_contents_example_filter"]) is not None: + setv( + parent_object, + ["storedContentsExampleFilter"], + getv(from_object, ["stored_contents_example_filter"]), + ) + + return to_object + + +def _RemoveExamplesParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv( + to_object, + ["_url", "exampleStore"], + t.t_example_store(getv(from_object, ["name"])), + ) + + if getv(from_object, ["config"]) is not None: + _RemoveExamplesConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + +def _SearchExamplesConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["top_k"]) is not None: + setv(parent_object, ["topK"], getv(from_object, ["top_k"])) + + return to_object + + +def _SearchExamplesParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv( + to_object, + ["_url", "exampleStore"], + t.t_example_store(getv(from_object, ["name"])), + ) + + if getv(from_object, ["stored_contents_example_parameters"]) is not None: + setv( + to_object, + ["storedContentsExampleParameters"], + getv(from_object, ["stored_contents_example_parameters"]), + ) + + if getv(from_object, ["config"]) is not None: + _SearchExamplesConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + +def _UpsertExamplesConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["overwrite"]) is not None: + setv(parent_object, ["overwrite"], getv(from_object, ["overwrite"])) + + return to_object + + +def _UpsertExamplesParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv( + to_object, + ["_url", "exampleStore"], + t.t_example_store(getv(from_object, ["name"])), + ) + + if getv(from_object, ["examples"]) is not None: + setv( + to_object, ["examples"], [item for item in getv(from_object, ["examples"])] + ) + + if getv(from_object, ["config"]) is not None: + _UpsertExamplesConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + +class ExampleStores(_api_module.BaseModule): + """Class for managing Example Stores of few-shot examples.""" + + def _create( + self, + *, + display_name: str, + config: Optional[types.CreateExampleStoreConfigOrDict] = None, + ) -> types.ExampleStoreOperation: + """ + Creates an Example Store. + """ + + parameter_model = types._CreateExampleStoreParameters( + display_name=display_name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CreateExampleStoreParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{parent}/exampleStores:create".format_map(request_url_dict) + else: + path = "{parent}/exampleStores:create" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ExampleStoreOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def get( + self, *, name: str, config: Optional[types.GetExampleStoreConfigOrDict] = None + ) -> types.ExampleStore: + """ + Retrieves a specific Example Store resource by its name. + """ + + parameter_model = types._GetExampleStoreParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetExampleStoreParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ExampleStore._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _delete( + self, + *, + name: str, + config: Optional[types.DeleteExampleStoreConfigOrDict] = None, + ) -> types.DeleteExampleStoreOperation: + """ + Deletes an Example Store. + """ + + parameter_model = types._DeleteExampleStoreRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _DeleteExampleStoreRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("delete", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.DeleteExampleStoreOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def upsert_examples( + self, + *, + name: str, + examples: list[types.ExampleOrDict], + config: Optional[types.UpsertExamplesConfigOrDict] = None, + ) -> types.UpsertExamplesResponse: + """ + Creates or updates examples in an Example Store. + """ + + parameter_model = types._UpsertExamplesParameters( + name=name, + examples=examples, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _UpsertExamplesParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{exampleStore}:upsertExamples".format_map(request_url_dict) + else: + path = "{exampleStore}:upsertExamples" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.UpsertExamplesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def search_examples( + self, + *, + name: str, + stored_contents_example_parameters: Optional[ + types.StoredContentsExampleParametersOrDict + ] = None, + config: Optional[types.SearchExamplesConfigOrDict] = None, + ) -> types.SearchExamplesResponse: + """ + Searches an Example Store for similar examples. + """ + + parameter_model = types._SearchExamplesParameters( + name=name, + stored_contents_example_parameters=stored_contents_example_parameters, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _SearchExamplesParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{exampleStore}:searchExamples".format_map(request_url_dict) + else: + path = "{exampleStore}:searchExamples" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.SearchExamplesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def fetch_examples( + self, *, name: str, config: Optional[types.FetchExamplesConfigOrDict] = None + ) -> types.FetchExamplesResponse: + """ + Fetches examples from an Example Store. + """ + + parameter_model = types._FetchExamplesParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _FetchExamplesParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{exampleStore}:fetchExamples".format_map(request_url_dict) + else: + path = "{exampleStore}:fetchExamples" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.FetchExamplesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def remove_examples( + self, *, name: str, config: Optional[types.RemoveExamplesConfigOrDict] = None + ) -> types.RemoveExamplesResponse: + """ + Removes examples from an Example Store. + """ + + parameter_model = types._RemoveExamplesParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _RemoveExamplesParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{exampleStore}:removeExamples".format_map(request_url_dict) + else: + path = "{exampleStore}:removeExamples" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.RemoveExamplesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _get_example_store_operation( + self, + *, + operation_name: str, + config: Optional[types.GetExampleStoreOperationConfigOrDict] = None, + ) -> types.ExampleStoreOperation: + parameter_model = types._GetExampleStoreOperationParameters( + operation_name=operation_name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetExampleStoreOperationParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ExampleStoreOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + +class AsyncExampleStores(_api_module.BaseModule): + """Class for managing Example Stores of few-shot examples.""" + + async def _create( + self, + *, + display_name: str, + config: Optional[types.CreateExampleStoreConfigOrDict] = None, + ) -> types.ExampleStoreOperation: + """ + Creates an Example Store. + """ + + parameter_model = types._CreateExampleStoreParameters( + display_name=display_name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CreateExampleStoreParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{parent}/exampleStores:create".format_map(request_url_dict) + else: + path = "{parent}/exampleStores:create" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ExampleStoreOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def get( + self, *, name: str, config: Optional[types.GetExampleStoreConfigOrDict] = None + ) -> types.ExampleStore: + """ + Retrieves a specific Example Store resource by its name. + """ + + parameter_model = types._GetExampleStoreParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetExampleStoreParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ExampleStore._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _delete( + self, + *, + name: str, + config: Optional[types.DeleteExampleStoreConfigOrDict] = None, + ) -> types.DeleteExampleStoreOperation: + """ + Deletes an Example Store. + """ + + parameter_model = types._DeleteExampleStoreRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _DeleteExampleStoreRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "delete", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.DeleteExampleStoreOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def upsert_examples( + self, + *, + name: str, + examples: list[types.ExampleOrDict], + config: Optional[types.UpsertExamplesConfigOrDict] = None, + ) -> types.UpsertExamplesResponse: + """ + Creates or updates examples in an Example Store. + """ + + parameter_model = types._UpsertExamplesParameters( + name=name, + examples=examples, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _UpsertExamplesParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{exampleStore}:upsertExamples".format_map(request_url_dict) + else: + path = "{exampleStore}:upsertExamples" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.UpsertExamplesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def search_examples( + self, + *, + name: str, + stored_contents_example_parameters: Optional[ + types.StoredContentsExampleParametersOrDict + ] = None, + config: Optional[types.SearchExamplesConfigOrDict] = None, + ) -> types.SearchExamplesResponse: + """ + Searches an Example Store for similar examples. + """ + + parameter_model = types._SearchExamplesParameters( + name=name, + stored_contents_example_parameters=stored_contents_example_parameters, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _SearchExamplesParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{exampleStore}:searchExamples".format_map(request_url_dict) + else: + path = "{exampleStore}:searchExamples" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.SearchExamplesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def fetch_examples( + self, *, name: str, config: Optional[types.FetchExamplesConfigOrDict] = None + ) -> types.FetchExamplesResponse: + """ + Fetches examples from an Example Store. + """ + + parameter_model = types._FetchExamplesParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _FetchExamplesParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{exampleStore}:fetchExamples".format_map(request_url_dict) + else: + path = "{exampleStore}:fetchExamples" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.FetchExamplesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def remove_examples( + self, *, name: str, config: Optional[types.RemoveExamplesConfigOrDict] = None + ) -> types.RemoveExamplesResponse: + """ + Removes examples from an Example Store. + """ + + parameter_model = types._RemoveExamplesParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _RemoveExamplesParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{exampleStore}:removeExamples".format_map(request_url_dict) + else: + path = "{exampleStore}:removeExamples" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.RemoveExamplesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _get_example_store_operation( + self, + *, + operation_name: str, + config: Optional[types.GetExampleStoreOperationConfigOrDict] = None, + ) -> types.ExampleStoreOperation: + parameter_model = types._GetExampleStoreOperationParameters( + operation_name=operation_name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetExampleStoreOperationParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ExampleStoreOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index 96231a3a2b..24a8359701 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -41,6 +41,7 @@ from .common import _CreateEvaluationMetricParameters from .common import _CreateEvaluationRunParameters from .common import _CreateEvaluationSetParameters +from .common import _CreateExampleStoreParameters from .common import _CreateMultimodalDatasetParameters from .common import _CreateRagCorpusRequestParameters from .common import _CreateRuntimeFeedbackEntryRequestParameters @@ -60,6 +61,7 @@ from .common import _DeleteEvaluationExperimentParameters from .common import _DeleteEvaluationMetricParameters from .common import _DeleteEvaluationSetParameters +from .common import _DeleteExampleStoreRequestParameters from .common import _DeleteMultimodalDatasetRequestParameters from .common import _DeletePromptVersionRequestParameters from .common import _DeleteRagCorpusRequestParameters @@ -72,6 +74,7 @@ from .common import _EvaluateInstancesRequestParameters from .common import _ExecuteCodeAgentEngineSandboxRequestParameters from .common import _ExportPublisherModelRequestParameters +from .common import _FetchExamplesParameters from .common import _GenerateAgentEngineMemoriesRequestParameters from .common import _GenerateInstanceRubricsRequest from .common import _GenerateLossClustersParameters @@ -104,6 +107,8 @@ from .common import _GetEvaluationMetricParameters from .common import _GetEvaluationRunParameters from .common import _GetEvaluationSetParameters +from .common import _GetExampleStoreOperationParameters +from .common import _GetExampleStoreParameters from .common import _GetExportPublisherModelOperationParameters from .common import _GetImportFilesOperationParameters from .common import _GetMultimodalDatasetOperationParameters @@ -157,6 +162,7 @@ from .common import _QueryAgentEngineRequestParameters from .common import _QueryAgentEngineRuntimeRevisionRequestParameters from .common import _RecommendSpecRequestParameters +from .common import _RemoveExamplesParameters from .common import _RestoreVersionRequestParameters from .common import _RetrieveAgentEngineMemoriesRequestParameters from .common import _RetrieveMemoryProfilesRequestParameters @@ -167,6 +173,7 @@ from .common import _RunQueryJobAgentEngineConfigDict from .common import _RunQueryJobAgentEngineConfigOrDict from .common import _RunQueryJobAgentEngineRequestParameters +from .common import _SearchExamplesParameters from .common import _UndeployModelRequestParameters from .common import _UpdateAgentEngineMemoryRequestParameters from .common import _UpdateAgentEngineRequestParameters @@ -180,6 +187,7 @@ from .common import _UpdateRuntimeFeedbackEntryRequestParameters from .common import _UpdateSkillRequestParameters from .common import _UploadRagFileParameters +from .common import _UpsertExamplesParameters from .common import A2aPart from .common import A2aPartDict from .common import A2aPartOrDict @@ -255,6 +263,7 @@ from .common import AppendAgentEngineTaskEventResponse from .common import AppendAgentEngineTaskEventResponseDict from .common import AppendAgentEngineTaskEventResponseOrDict +from .common import ArrayOperator from .common import AskContextsConfig from .common import AskContextsConfigDict from .common import AskContextsConfigOrDict @@ -273,6 +282,12 @@ from .common import AttackCategoryResult from .common import AttackCategoryResultDict from .common import AttackCategoryResultOrDict +from .common import AudioTranscription +from .common import AudioTranscriptionDict +from .common import AudioTranscriptionOrDict +from .common import AudioTranscriptionWordInfo +from .common import AudioTranscriptionWordInfoDict +from .common import AudioTranscriptionWordInfoOrDict from .common import AutomaticResources from .common import AutomaticResourcesDict from .common import AutomaticResourcesOrDict @@ -349,6 +364,12 @@ from .common import ContentMapContentsOrDict from .common import ContentMapDict from .common import ContentMapOrDict +from .common import ContentsExample +from .common import ContentsExampleDict +from .common import ContentsExampleExpectedContent +from .common import ContentsExampleExpectedContentDict +from .common import ContentsExampleExpectedContentOrDict +from .common import ContentsExampleOrDict from .common import CorpusOperation from .common import CorpusOperationDict from .common import CorpusOperationOrDict @@ -391,6 +412,9 @@ from .common import CreateEvaluationSetConfig from .common import CreateEvaluationSetConfigDict from .common import CreateEvaluationSetConfigOrDict +from .common import CreateExampleStoreConfig +from .common import CreateExampleStoreConfigDict +from .common import CreateExampleStoreConfigOrDict from .common import CreateMultimodalDatasetConfig from .common import CreateMultimodalDatasetConfigDict from .common import CreateMultimodalDatasetConfigOrDict @@ -498,6 +522,12 @@ from .common import DeleteEvaluationSetOperation from .common import DeleteEvaluationSetOperationDict from .common import DeleteEvaluationSetOperationOrDict +from .common import DeleteExampleStoreConfig +from .common import DeleteExampleStoreConfigDict +from .common import DeleteExampleStoreConfigOrDict +from .common import DeleteExampleStoreOperation +from .common import DeleteExampleStoreOperationDict +from .common import DeleteExampleStoreOperationOrDict from .common import DeletePromptConfig from .common import DeletePromptConfigDict from .common import DeletePromptConfigOrDict @@ -715,12 +745,27 @@ from .common import ExactMatchSpec from .common import ExactMatchSpecDict from .common import ExactMatchSpecOrDict +from .common import Example +from .common import ExampleDict +from .common import ExampleOrDict from .common import Examples +from .common import ExamplesArrayFilter +from .common import ExamplesArrayFilterDict +from .common import ExamplesArrayFilterOrDict from .common import ExamplesDict from .common import ExamplesExampleGcsSource from .common import ExamplesExampleGcsSourceDict from .common import ExamplesExampleGcsSourceOrDict from .common import ExamplesOrDict +from .common import ExampleStore +from .common import ExampleStoreConfig +from .common import ExampleStoreConfigDict +from .common import ExampleStoreConfigOrDict +from .common import ExampleStoreDict +from .common import ExampleStoreOperation +from .common import ExampleStoreOperationDict +from .common import ExampleStoreOperationOrDict +from .common import ExampleStoreOrDict from .common import ExecuteCodeAgentEngineSandboxConfig from .common import ExecuteCodeAgentEngineSandboxConfigDict from .common import ExecuteCodeAgentEngineSandboxConfigOrDict @@ -779,6 +824,12 @@ from .common import FeedbackEntryDict from .common import FeedbackEntryOrDict from .common import FeedbackType +from .common import FetchExamplesConfig +from .common import FetchExamplesConfigDict +from .common import FetchExamplesConfigOrDict +from .common import FetchExamplesResponse +from .common import FetchExamplesResponseDict +from .common import FetchExamplesResponseOrDict from .common import FlexStart from .common import FlexStartDict from .common import FlexStartOrDict @@ -786,6 +837,7 @@ from .common import FullFineTunedResources from .common import FullFineTunedResourcesDict from .common import FullFineTunedResourcesOrDict +from .common import FunctionResponseScheduling from .common import GdcConfig from .common import GdcConfigDict from .common import GdcConfigOrDict @@ -907,6 +959,12 @@ from .common import GetEvaluationSetConfig from .common import GetEvaluationSetConfigDict from .common import GetEvaluationSetConfigOrDict +from .common import GetExampleStoreConfig +from .common import GetExampleStoreConfigDict +from .common import GetExampleStoreConfigOrDict +from .common import GetExampleStoreOperationConfig +from .common import GetExampleStoreOperationConfigDict +from .common import GetExampleStoreOperationConfigOrDict from .common import GetExportPublisherModelOperationConfig from .common import GetExportPublisherModelOperationConfigDict from .common import GetExportPublisherModelOperationConfigOrDict @@ -1207,6 +1265,7 @@ from .common import MapInstance from .common import MapInstanceDict from .common import MapInstanceOrDict +from .common import MediaResolution from .common import Memory from .common import MemoryBankCustomizationConfig from .common import MemoryBankCustomizationConfigConsolidationConfig @@ -1336,6 +1395,7 @@ from .common import OptimizeResponseEndpointOrDict from .common import OptimizeResponseOrDict from .common import OptimizeTarget +from .common import Outcome from .common import OverlayType from .common import PairwiseMetricInput from .common import PairwiseMetricInputDict @@ -1733,6 +1793,12 @@ from .common import RedTeamingAnalysisResult from .common import RedTeamingAnalysisResultDict from .common import RedTeamingAnalysisResultOrDict +from .common import RemoveExamplesConfig +from .common import RemoveExamplesConfigDict +from .common import RemoveExamplesConfigOrDict +from .common import RemoveExamplesResponse +from .common import RemoveExamplesResponseDict +from .common import RemoveExamplesResponseOrDict from .common import ReservationAffinity from .common import ReservationAffinityDict from .common import ReservationAffinityOrDict @@ -1973,6 +2039,15 @@ from .common import SchemaTextPromptDatasetMetadata from .common import SchemaTextPromptDatasetMetadataDict from .common import SchemaTextPromptDatasetMetadataOrDict +from .common import SearchExamplesConfig +from .common import SearchExamplesConfigDict +from .common import SearchExamplesConfigOrDict +from .common import SearchExamplesResponse +from .common import SearchExamplesResponseDict +from .common import SearchExamplesResponseOrDict +from .common import SearchExamplesResponseSimilarExample +from .common import SearchExamplesResponseSimilarExampleDict +from .common import SearchExamplesResponseSimilarExampleOrDict from .common import SecretEnvVar from .common import SecretEnvVarDict from .common import SecretEnvVarOrDict @@ -2025,6 +2100,24 @@ from .common import SpeculativeDecodingSpecNgramSpeculationOrDict from .common import SpeculativeDecodingSpecOrDict from .common import State +from .common import StoredContentsExample +from .common import StoredContentsExampleDict +from .common import StoredContentsExampleFilter +from .common import StoredContentsExampleFilterDict +from .common import StoredContentsExampleFilterOrDict +from .common import StoredContentsExampleOrDict +from .common import StoredContentsExampleParameters +from .common import StoredContentsExampleParametersContentSearchKey +from .common import StoredContentsExampleParametersContentSearchKeyDict +from .common import StoredContentsExampleParametersContentSearchKeyOrDict +from .common import StoredContentsExampleParametersDict +from .common import StoredContentsExampleParametersOrDict +from .common import StoredContentsExampleSearchKeyGenerationMethod +from .common import StoredContentsExampleSearchKeyGenerationMethodDict +from .common import StoredContentsExampleSearchKeyGenerationMethodLastEntry +from .common import StoredContentsExampleSearchKeyGenerationMethodLastEntryDict +from .common import StoredContentsExampleSearchKeyGenerationMethodLastEntryOrDict +from .common import StoredContentsExampleSearchKeyGenerationMethodOrDict from .common import Strategy from .common import StructuredMemoryConfig from .common import StructuredMemoryConfigDict @@ -2195,6 +2288,15 @@ from .common import UploadRagFileResponse from .common import UploadRagFileResponseDict from .common import UploadRagFileResponseOrDict +from .common import UpsertExamplesConfig +from .common import UpsertExamplesConfigDict +from .common import UpsertExamplesConfigOrDict +from .common import UpsertExamplesResponse +from .common import UpsertExamplesResponseDict +from .common import UpsertExamplesResponseOrDict +from .common import UpsertExamplesResponseUpsertResult +from .common import UpsertExamplesResponseUpsertResultDict +from .common import UpsertExamplesResponseUpsertResultOrDict from .common import VersionState from .common import VertexAiSearchConfig from .common import VertexAiSearchConfigDict @@ -4043,6 +4145,96 @@ "EndpointOperation", "EndpointOperationDict", "EndpointOperationOrDict", + "CreateExampleStoreConfig", + "CreateExampleStoreConfigDict", + "CreateExampleStoreConfigOrDict", + "ExampleStoreConfig", + "ExampleStoreConfigDict", + "ExampleStoreConfigOrDict", + "ExampleStore", + "ExampleStoreDict", + "ExampleStoreOrDict", + "ExampleStoreOperation", + "ExampleStoreOperationDict", + "ExampleStoreOperationOrDict", + "GetExampleStoreConfig", + "GetExampleStoreConfigDict", + "GetExampleStoreConfigOrDict", + "DeleteExampleStoreConfig", + "DeleteExampleStoreConfigDict", + "DeleteExampleStoreConfigOrDict", + "DeleteExampleStoreOperation", + "DeleteExampleStoreOperationDict", + "DeleteExampleStoreOperationOrDict", + "UpsertExamplesConfig", + "UpsertExamplesConfigDict", + "UpsertExamplesConfigOrDict", + "AudioTranscriptionWordInfo", + "AudioTranscriptionWordInfoDict", + "AudioTranscriptionWordInfoOrDict", + "AudioTranscription", + "AudioTranscriptionDict", + "AudioTranscriptionOrDict", + "ContentsExampleExpectedContent", + "ContentsExampleExpectedContentDict", + "ContentsExampleExpectedContentOrDict", + "ContentsExample", + "ContentsExampleDict", + "ContentsExampleOrDict", + "StoredContentsExampleSearchKeyGenerationMethodLastEntry", + "StoredContentsExampleSearchKeyGenerationMethodLastEntryDict", + "StoredContentsExampleSearchKeyGenerationMethodLastEntryOrDict", + "StoredContentsExampleSearchKeyGenerationMethod", + "StoredContentsExampleSearchKeyGenerationMethodDict", + "StoredContentsExampleSearchKeyGenerationMethodOrDict", + "StoredContentsExample", + "StoredContentsExampleDict", + "StoredContentsExampleOrDict", + "Example", + "ExampleDict", + "ExampleOrDict", + "UpsertExamplesResponseUpsertResult", + "UpsertExamplesResponseUpsertResultDict", + "UpsertExamplesResponseUpsertResultOrDict", + "UpsertExamplesResponse", + "UpsertExamplesResponseDict", + "UpsertExamplesResponseOrDict", + "SearchExamplesConfig", + "SearchExamplesConfigDict", + "SearchExamplesConfigOrDict", + "StoredContentsExampleParametersContentSearchKey", + "StoredContentsExampleParametersContentSearchKeyDict", + "StoredContentsExampleParametersContentSearchKeyOrDict", + "ExamplesArrayFilter", + "ExamplesArrayFilterDict", + "ExamplesArrayFilterOrDict", + "StoredContentsExampleParameters", + "StoredContentsExampleParametersDict", + "StoredContentsExampleParametersOrDict", + "SearchExamplesResponseSimilarExample", + "SearchExamplesResponseSimilarExampleDict", + "SearchExamplesResponseSimilarExampleOrDict", + "SearchExamplesResponse", + "SearchExamplesResponseDict", + "SearchExamplesResponseOrDict", + "StoredContentsExampleFilter", + "StoredContentsExampleFilterDict", + "StoredContentsExampleFilterOrDict", + "FetchExamplesConfig", + "FetchExamplesConfigDict", + "FetchExamplesConfigOrDict", + "FetchExamplesResponse", + "FetchExamplesResponseDict", + "FetchExamplesResponseOrDict", + "RemoveExamplesConfig", + "RemoveExamplesConfigDict", + "RemoveExamplesConfigOrDict", + "RemoveExamplesResponse", + "RemoveExamplesResponseDict", + "RemoveExamplesResponseOrDict", + "GetExampleStoreOperationConfig", + "GetExampleStoreOperationConfigDict", + "GetExampleStoreOperationConfigOrDict", "PromptOptimizerConfig", "PromptOptimizerConfigDict", "PromptOptimizerConfigOrDict", @@ -4185,6 +4377,10 @@ "Modality", "DeploymentType", "ModelProvider", + "Outcome", + "FunctionResponseScheduling", + "MediaResolution", + "ArrayOperator", "EvaluationExperimentMergeStrategy", "EvaluationItemType", "SamplingMethod", @@ -4374,6 +4570,14 @@ "_DeleteEndpointRequestParameters", "_GetEndpointParameters", "_GetEndpointOperationParameters", + "_CreateExampleStoreParameters", + "_GetExampleStoreParameters", + "_DeleteExampleStoreRequestParameters", + "_UpsertExamplesParameters", + "_SearchExamplesParameters", + "_FetchExamplesParameters", + "_RemoveExamplesParameters", + "_GetExampleStoreOperationParameters", "evals", "agent_engines", "prompts", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index 9b4c9434d5..d849d5c68b 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -630,6 +630,58 @@ class ModelProvider(_common.CaseInSensitiveEnum): """Anthropic.""" +class Outcome(_common.CaseInSensitiveEnum): + """Outcome of the code execution.""" + + OUTCOME_UNSPECIFIED = "OUTCOME_UNSPECIFIED" + """Unspecified status. This value should not be used.""" + OUTCOME_OK = "OUTCOME_OK" + """Code execution completed successfully. `output` contains the stdout, if any.""" + OUTCOME_FAILED = "OUTCOME_FAILED" + """Code execution failed. `output` contains the stderr and stdout, if any.""" + OUTCOME_DEADLINE_EXCEEDED = "OUTCOME_DEADLINE_EXCEEDED" + """Code execution ran for too long, and was cancelled. There may or may not be a partial `output` present.""" + + +class FunctionResponseScheduling(_common.CaseInSensitiveEnum): + """Specifies how the response should be scheduled in the conversation. Only applicable to NON_BLOCKING function calls, is ignored otherwise. Defaults to WHEN_IDLE.""" + + SCHEDULING_UNSPECIFIED = "SCHEDULING_UNSPECIFIED" + """This value is unused.""" + SILENT = "SILENT" + """Only add the result to the conversation context, do not interrupt or trigger generation.""" + WHEN_IDLE = "WHEN_IDLE" + """Add the result to the conversation context, and prompt to generate output without interrupting ongoing generation.""" + INTERRUPT = "INTERRUPT" + """Add the result to the conversation context, interrupt ongoing generation and prompt to generate output.""" + + +class MediaResolution(_common.CaseInSensitiveEnum): + """The tokenization quality used for given media.""" + + MEDIA_RESOLUTION_UNSPECIFIED = "MEDIA_RESOLUTION_UNSPECIFIED" + """Media resolution has not been set.""" + MEDIA_RESOLUTION_LOW = "MEDIA_RESOLUTION_LOW" + """Media resolution set to low.""" + MEDIA_RESOLUTION_MEDIUM = "MEDIA_RESOLUTION_MEDIUM" + """Media resolution set to medium.""" + MEDIA_RESOLUTION_HIGH = "MEDIA_RESOLUTION_HIGH" + """Media resolution set to high.""" + MEDIA_RESOLUTION_ULTRA_HIGH = "MEDIA_RESOLUTION_ULTRA_HIGH" + """Media resolution set to ultra high. This is for image only.""" + + +class ArrayOperator(_common.CaseInSensitiveEnum): + """The operator logic to use for filtering.""" + + ARRAY_OPERATOR_UNSPECIFIED = "ARRAY_OPERATOR_UNSPECIFIED" + """Not specified. This value should not be used.""" + CONTAINS_ANY = "CONTAINS_ANY" + """The metadata array field in the example must contain at least one of the values.""" + CONTAINS_ALL = "CONTAINS_ALL" + """The metadata array field in the example must contain all of the values.""" + + class EvaluationExperimentMergeStrategy(_common.CaseInSensitiveEnum): """Merge strategy for the evaluation experiment.""" @@ -28702,6 +28754,1123 @@ class EndpointOperationDict(TypedDict, total=False): EndpointOperationOrDict = Union[EndpointOperation, EndpointOperationDict] +class CreateExampleStoreConfig(_common.BaseModel): + """Config for creating an Example Store.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Whether to wait for the long running operation to complete.""", + ) + description: Optional[str] = Field( + default=None, description="""Optional. The description of the Example Store.""" + ) + vertex_embedding_model: Optional[str] = Field( + default=None, + description="""Optional. The embedding model used to generate the search key for + stored examples. + """, + ) + + +class CreateExampleStoreConfigDict(TypedDict, total=False): + """Config for creating an Example Store.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + wait_for_completion: Optional[bool] + """Whether to wait for the long running operation to complete.""" + + description: Optional[str] + """Optional. The description of the Example Store.""" + + vertex_embedding_model: Optional[str] + """Optional. The embedding model used to generate the search key for + stored examples. + """ + + +CreateExampleStoreConfigOrDict = Union[ + CreateExampleStoreConfig, CreateExampleStoreConfigDict +] + + +class _CreateExampleStoreParameters(_common.BaseModel): + """Parameters for creating an Example Store.""" + + display_name: Optional[str] = Field( + default=None, description="""Required. The display name of the Example Store.""" + ) + config: Optional[CreateExampleStoreConfig] = Field( + default=None, description="""Used to override the default configuration.""" + ) + + +class _CreateExampleStoreParametersDict(TypedDict, total=False): + """Parameters for creating an Example Store.""" + + display_name: Optional[str] + """Required. The display name of the Example Store.""" + + config: Optional[CreateExampleStoreConfigDict] + """Used to override the default configuration.""" + + +_CreateExampleStoreParametersOrDict = Union[ + _CreateExampleStoreParameters, _CreateExampleStoreParametersDict +] + + +class ExampleStoreConfig(_common.BaseModel): + """Configuration for the Example Store.""" + + vertex_embedding_model: Optional[str] = Field( + default=None, + description="""Required. The embedding model to be used for vector embedding. Immutable. Supported models: * "text-embedding-005" * "text-multilingual-embedding-002".""", + ) + + +class ExampleStoreConfigDict(TypedDict, total=False): + """Configuration for the Example Store.""" + + vertex_embedding_model: Optional[str] + """Required. The embedding model to be used for vector embedding. Immutable. Supported models: * "text-embedding-005" * "text-multilingual-embedding-002".""" + + +ExampleStoreConfigOrDict = Union[ExampleStoreConfig, ExampleStoreConfigDict] + + +class ExampleStore(_common.BaseModel): + """A storage bucket for few-shot examples used to steer a model.""" + + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when this ExampleStore was created.""", + ) + description: Optional[str] = Field( + default=None, description="""Optional. Description of the ExampleStore.""" + ) + display_name: Optional[str] = Field( + default=None, description="""Required. Display name of the ExampleStore.""" + ) + example_store_config: Optional[ExampleStoreConfig] = Field( + default=None, description="""Required. Example Store config.""" + ) + name: Optional[str] = Field( + default=None, + description="""Identifier. The resource name of the ExampleStore. This is a unique identifier. Format: projects/{project}/locations/{location}/exampleStores/{example_store}""", + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when this ExampleStore was most recently updated.""", + ) + + +class ExampleStoreDict(TypedDict, total=False): + """A storage bucket for few-shot examples used to steer a model.""" + + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this ExampleStore was created.""" + + description: Optional[str] + """Optional. Description of the ExampleStore.""" + + display_name: Optional[str] + """Required. Display name of the ExampleStore.""" + + example_store_config: Optional[ExampleStoreConfigDict] + """Required. Example Store config.""" + + name: Optional[str] + """Identifier. The resource name of the ExampleStore. This is a unique identifier. Format: projects/{project}/locations/{location}/exampleStores/{example_store}""" + + update_time: Optional[datetime.datetime] + """Output only. Timestamp when this ExampleStore was most recently updated.""" + + +ExampleStoreOrDict = Union[ExampleStore, ExampleStoreDict] + + +class ExampleStoreOperation(_common.BaseModel): + """Operation that has an Example Store as a response.""" + + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", + ) + response: Optional[ExampleStore] = Field( + default=None, description="""The created Example Store.""" + ) + + +class ExampleStoreOperationDict(TypedDict, total=False): + """Operation that has an Example Store as a response.""" + + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" + + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + response: Optional[ExampleStoreDict] + """The created Example Store.""" + + +ExampleStoreOperationOrDict = Union[ExampleStoreOperation, ExampleStoreOperationDict] + + +class GetExampleStoreConfig(_common.BaseModel): + """Optional parameters for example_stores.get method.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetExampleStoreConfigDict(TypedDict, total=False): + """Optional parameters for example_stores.get method.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GetExampleStoreConfigOrDict = Union[GetExampleStoreConfig, GetExampleStoreConfigDict] + + +class _GetExampleStoreParameters(_common.BaseModel): + """Parameters for retrieving an Example Store.""" + + name: Optional[str] = Field( + default=None, + description="""Required. The resource name of the Example Store to get.""", + ) + config: Optional[GetExampleStoreConfig] = Field( + default=None, description="""Optional parameters for the request.""" + ) + + +class _GetExampleStoreParametersDict(TypedDict, total=False): + """Parameters for retrieving an Example Store.""" + + name: Optional[str] + """Required. The resource name of the Example Store to get.""" + + config: Optional[GetExampleStoreConfigDict] + """Optional parameters for the request.""" + + +_GetExampleStoreParametersOrDict = Union[ + _GetExampleStoreParameters, _GetExampleStoreParametersDict +] + + +class DeleteExampleStoreConfig(_common.BaseModel): + """Config for deleting an Example Store.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Whether to wait for the long running operation to complete.""", + ) + + +class DeleteExampleStoreConfigDict(TypedDict, total=False): + """Config for deleting an Example Store.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + wait_for_completion: Optional[bool] + """Whether to wait for the long running operation to complete.""" + + +DeleteExampleStoreConfigOrDict = Union[ + DeleteExampleStoreConfig, DeleteExampleStoreConfigDict +] + + +class _DeleteExampleStoreRequestParameters(_common.BaseModel): + """Parameters for deleting an Example Store.""" + + name: Optional[str] = Field( + default=None, + description="""Required. The resource name of the Example Store to delete.""", + ) + config: Optional[DeleteExampleStoreConfig] = Field( + default=None, description="""Used to override the default configuration.""" + ) + + +class _DeleteExampleStoreRequestParametersDict(TypedDict, total=False): + """Parameters for deleting an Example Store.""" + + name: Optional[str] + """Required. The resource name of the Example Store to delete.""" + + config: Optional[DeleteExampleStoreConfigDict] + """Used to override the default configuration.""" + + +_DeleteExampleStoreRequestParametersOrDict = Union[ + _DeleteExampleStoreRequestParameters, _DeleteExampleStoreRequestParametersDict +] + + +class DeleteExampleStoreOperation(_common.BaseModel): + """Operation for deleting an Example Store.""" + + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", + ) + + +class DeleteExampleStoreOperationDict(TypedDict, total=False): + """Operation for deleting an Example Store.""" + + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" + + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + +DeleteExampleStoreOperationOrDict = Union[ + DeleteExampleStoreOperation, DeleteExampleStoreOperationDict +] + + +class UpsertExamplesConfig(_common.BaseModel): + """Config for upserting examples into an Example Store.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + overwrite: Optional[bool] = Field( + default=None, + description="""Optional. Whether to overwrite an example that already exists. If + false, the request fails when an example with the same id already exists. + """, + ) + + +class UpsertExamplesConfigDict(TypedDict, total=False): + """Config for upserting examples into an Example Store.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + overwrite: Optional[bool] + """Optional. Whether to overwrite an example that already exists. If + false, the request fails when an example with the same id already exists. + """ + + +UpsertExamplesConfigOrDict = Union[UpsertExamplesConfig, UpsertExamplesConfigDict] + + +class AudioTranscriptionWordInfo(_common.BaseModel): + """Information about a single recognized word.""" + + end_offset: Optional[str] = Field( + default=None, + description="""Optional. End offset in time of the word relative to the start of the audio.""", + ) + start_offset: Optional[str] = Field( + default=None, + description="""Optional. Start offset in time of the word relative to the start of the audio.""", + ) + word: Optional[str] = Field( + default=None, description="""Required. Transcript of the word.""" + ) + + +class AudioTranscriptionWordInfoDict(TypedDict, total=False): + """Information about a single recognized word.""" + + end_offset: Optional[str] + """Optional. End offset in time of the word relative to the start of the audio.""" + + start_offset: Optional[str] + """Optional. Start offset in time of the word relative to the start of the audio.""" + + word: Optional[str] + """Required. Transcript of the word.""" + + +AudioTranscriptionWordInfoOrDict = Union[ + AudioTranscriptionWordInfo, AudioTranscriptionWordInfoDict +] + + +class AudioTranscription(_common.BaseModel): + """The transcription of an audio part. For multi-speaker audio, each speaker segment is a separate Part with its own AudioTranscription carrying the speaker_label.""" + + speaker_label: Optional[str] = Field( + default=None, + description="""Optional. A label identifying the speaker of this audio segment (e.g. "spk_1", "spk_2"). Present when diarization is set.""", + ) + text: Optional[str] = Field( + default=None, + description="""Required. The transcription text of this audio segment.""", + ) + words: Optional[list[AudioTranscriptionWordInfo]] = Field( + default=None, + description="""Optional. Detailed word-level transcriptions and timing details. Present when word_timestamp is set.""", + ) + + +class AudioTranscriptionDict(TypedDict, total=False): + """The transcription of an audio part. For multi-speaker audio, each speaker segment is a separate Part with its own AudioTranscription carrying the speaker_label.""" + + speaker_label: Optional[str] + """Optional. A label identifying the speaker of this audio segment (e.g. "spk_1", "spk_2"). Present when diarization is set.""" + + text: Optional[str] + """Required. The transcription text of this audio segment.""" + + words: Optional[list[AudioTranscriptionWordInfoDict]] + """Optional. Detailed word-level transcriptions and timing details. Present when word_timestamp is set.""" + + +AudioTranscriptionOrDict = Union[AudioTranscription, AudioTranscriptionDict] + + +class ContentsExampleExpectedContent(_common.BaseModel): + """A single step of the expected output.""" + + content: Optional[genai_types.Content] = Field( + default=None, description="""Required. A single step's content.""" + ) + + +class ContentsExampleExpectedContentDict(TypedDict, total=False): + """A single step of the expected output.""" + + content: Optional[genai_types.ContentDict] + """Required. A single step's content.""" + + +ContentsExampleExpectedContentOrDict = Union[ + ContentsExampleExpectedContent, ContentsExampleExpectedContentDict +] + + +class ContentsExample(_common.BaseModel): + """A single example of a conversation with the model.""" + + contents: Optional[list[genai_types.Content]] = Field( + default=None, + description="""Required. The content of the conversation with the model that resulted in the expected output.""", + ) + expected_contents: Optional[list[ContentsExampleExpectedContent]] = Field( + default=None, + description="""Required. The expected output for the given `contents`. To represent multi-step reasoning, this is a repeated field that contains the iterative steps of the expected output.""", + ) + + +class ContentsExampleDict(TypedDict, total=False): + """A single example of a conversation with the model.""" + + contents: Optional[list[genai_types.ContentDict]] + """Required. The content of the conversation with the model that resulted in the expected output.""" + + expected_contents: Optional[list[ContentsExampleExpectedContentDict]] + """Required. The expected output for the given `contents`. To represent multi-step reasoning, this is a repeated field that contains the iterative steps of the expected output.""" + + +ContentsExampleOrDict = Union[ContentsExample, ContentsExampleDict] + + +class StoredContentsExampleSearchKeyGenerationMethodLastEntry(_common.BaseModel): + """Configuration for using only the last entry of the conversation history as the search key.""" + + pass + + +class StoredContentsExampleSearchKeyGenerationMethodLastEntryDict( + TypedDict, total=False +): + """Configuration for using only the last entry of the conversation history as the search key.""" + + pass + + +StoredContentsExampleSearchKeyGenerationMethodLastEntryOrDict = Union[ + StoredContentsExampleSearchKeyGenerationMethodLastEntry, + StoredContentsExampleSearchKeyGenerationMethodLastEntryDict, +] + + +class StoredContentsExampleSearchKeyGenerationMethod(_common.BaseModel): + """Options for generating the search key from the conversation history.""" + + last_entry: Optional[StoredContentsExampleSearchKeyGenerationMethodLastEntry] = ( + Field( + default=None, + description="""Use only the last entry of the conversation history (`contents_example.contents`) as the search key.""", + ) + ) + + +class StoredContentsExampleSearchKeyGenerationMethodDict(TypedDict, total=False): + """Options for generating the search key from the conversation history.""" + + last_entry: Optional[StoredContentsExampleSearchKeyGenerationMethodLastEntryDict] + """Use only the last entry of the conversation history (`contents_example.contents`) as the search key.""" + + +StoredContentsExampleSearchKeyGenerationMethodOrDict = Union[ + StoredContentsExampleSearchKeyGenerationMethod, + StoredContentsExampleSearchKeyGenerationMethodDict, +] + + +class StoredContentsExample(_common.BaseModel): + """A ContentsExample to be used with GenerateContent alongside information required for storage and retrieval with Example Store.""" + + contents_example: Optional[ContentsExample] = Field( + default=None, + description="""Required. The example to be used with GenerateContent.""", + ) + search_key: Optional[str] = Field( + default=None, + description="""Optional. (Optional) the search key used for retrieval. If not provided at upload-time, the search key will be generated from `contents_example.contents` using the method provided by `search_key_generation_method`. The generated search key will be included in retrieved examples.""", + ) + search_key_generation_method: Optional[ + StoredContentsExampleSearchKeyGenerationMethod + ] = Field( + default=None, + description="""Optional. The method used to generate the search key from `contents_example.contents`. This is ignored when uploading an example if `search_key` is provided.""", + ) + + +class StoredContentsExampleDict(TypedDict, total=False): + """A ContentsExample to be used with GenerateContent alongside information required for storage and retrieval with Example Store.""" + + contents_example: Optional[ContentsExampleDict] + """Required. The example to be used with GenerateContent.""" + + search_key: Optional[str] + """Optional. (Optional) the search key used for retrieval. If not provided at upload-time, the search key will be generated from `contents_example.contents` using the method provided by `search_key_generation_method`. The generated search key will be included in retrieved examples.""" + + search_key_generation_method: Optional[ + StoredContentsExampleSearchKeyGenerationMethodDict + ] + """Optional. The method used to generate the search key from `contents_example.contents`. This is ignored when uploading an example if `search_key` is provided.""" + + +StoredContentsExampleOrDict = Union[StoredContentsExample, StoredContentsExampleDict] + + +class Example(_common.BaseModel): + + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when this Example was created.""", + ) + display_name: Optional[str] = Field( + default=None, description="""Optional. The display name for Example.""" + ) + example_id: Optional[str] = Field( + default=None, + description="""Optional. Immutable. Unique identifier of an example. If not specified when upserting new examples, the example_id will be generated.""", + ) + stored_contents_example: Optional[StoredContentsExample] = Field( + default=None, + description="""An example of chat history and its expected outcome to be used with GenerateContent.""", + ) + + +class ExampleDict(TypedDict, total=False): + + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this Example was created.""" + + display_name: Optional[str] + """Optional. The display name for Example.""" + + example_id: Optional[str] + """Optional. Immutable. Unique identifier of an example. If not specified when upserting new examples, the example_id will be generated.""" + + stored_contents_example: Optional[StoredContentsExampleDict] + """An example of chat history and its expected outcome to be used with GenerateContent.""" + + +ExampleOrDict = Union[Example, ExampleDict] + + +class _UpsertExamplesParameters(_common.BaseModel): + """Parameters for upserting examples.""" + + name: Optional[str] = Field( + default=None, + description="""Required. The Example Store to upsert examples into.""", + ) + examples: Optional[list[Example]] = Field( + default=None, description="""Required. The examples to upsert.""" + ) + config: Optional[UpsertExamplesConfig] = Field(default=None, description="""""") + + +class _UpsertExamplesParametersDict(TypedDict, total=False): + """Parameters for upserting examples.""" + + name: Optional[str] + """Required. The Example Store to upsert examples into.""" + + examples: Optional[list[ExampleDict]] + """Required. The examples to upsert.""" + + config: Optional[UpsertExamplesConfigDict] + """""" + + +_UpsertExamplesParametersOrDict = Union[ + _UpsertExamplesParameters, _UpsertExamplesParametersDict +] + + +class UpsertExamplesResponseUpsertResult(_common.BaseModel): + """The result for creating/updating a single example.""" + + example: Optional[Example] = Field( + default=None, description="""The example created/updated successfully.""" + ) + status: Optional[genai_types.GoogleRpcStatus] = Field( + default=None, + description="""The error message of the example that was not created/updated successfully.""", + ) + + +class UpsertExamplesResponseUpsertResultDict(TypedDict, total=False): + """The result for creating/updating a single example.""" + + example: Optional[ExampleDict] + """The example created/updated successfully.""" + + status: Optional[genai_types.GoogleRpcStatusDict] + """The error message of the example that was not created/updated successfully.""" + + +UpsertExamplesResponseUpsertResultOrDict = Union[ + UpsertExamplesResponseUpsertResult, UpsertExamplesResponseUpsertResultDict +] + + +class UpsertExamplesResponse(_common.BaseModel): + """Response message for ExampleStoreService.UpsertExamples.""" + + results: Optional[list[UpsertExamplesResponseUpsertResult]] = Field( + default=None, + description="""A list of results for creating/updating. It's either a successfully created/updated example or a status with an error message.""", + ) + + +class UpsertExamplesResponseDict(TypedDict, total=False): + """Response message for ExampleStoreService.UpsertExamples.""" + + results: Optional[list[UpsertExamplesResponseUpsertResultDict]] + """A list of results for creating/updating. It's either a successfully created/updated example or a status with an error message.""" + + +UpsertExamplesResponseOrDict = Union[UpsertExamplesResponse, UpsertExamplesResponseDict] + + +class SearchExamplesConfig(_common.BaseModel): + """Config for searching an Example Store.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + top_k: Optional[int] = Field( + default=None, + description="""Optional. The number of similar examples to return.""", + ) + + +class SearchExamplesConfigDict(TypedDict, total=False): + """Config for searching an Example Store.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + top_k: Optional[int] + """Optional. The number of similar examples to return.""" + + +SearchExamplesConfigOrDict = Union[SearchExamplesConfig, SearchExamplesConfigDict] + + +class StoredContentsExampleParametersContentSearchKey(_common.BaseModel): + """The chat history to use to generate the search key for retrieval.""" + + contents: Optional[list[genai_types.Content]] = Field( + default=None, + description="""Required. The conversation for generating a search key.""", + ) + search_key_generation_method: Optional[ + StoredContentsExampleSearchKeyGenerationMethod + ] = Field( + default=None, description="""Required. The method of generating a search key.""" + ) + + +class StoredContentsExampleParametersContentSearchKeyDict(TypedDict, total=False): + """The chat history to use to generate the search key for retrieval.""" + + contents: Optional[list[genai_types.ContentDict]] + """Required. The conversation for generating a search key.""" + + search_key_generation_method: Optional[ + StoredContentsExampleSearchKeyGenerationMethodDict + ] + """Required. The method of generating a search key.""" + + +StoredContentsExampleParametersContentSearchKeyOrDict = Union[ + StoredContentsExampleParametersContentSearchKey, + StoredContentsExampleParametersContentSearchKeyDict, +] + + +class ExamplesArrayFilter(_common.BaseModel): + """Filters for examples' array metadata fields. An array field is example metadata where multiple values are attributed to a single example.""" + + array_operator: Optional[ArrayOperator] = Field( + default=None, + description="""Required. The operator logic to use for filtering.""", + ) + values: Optional[list[str]] = Field( + default=None, + description="""Required. The values by which to filter examples.""", + ) + + +class ExamplesArrayFilterDict(TypedDict, total=False): + """Filters for examples' array metadata fields. An array field is example metadata where multiple values are attributed to a single example.""" + + array_operator: Optional[ArrayOperator] + """Required. The operator logic to use for filtering.""" + + values: Optional[list[str]] + """Required. The values by which to filter examples.""" + + +ExamplesArrayFilterOrDict = Union[ExamplesArrayFilter, ExamplesArrayFilterDict] + + +class StoredContentsExampleParameters(_common.BaseModel): + """The metadata filters that will be used to search StoredContentsExamples. If a field is unspecified, then no filtering for that field will be applied""" + + content_search_key: Optional[StoredContentsExampleParametersContentSearchKey] = ( + Field( + default=None, + description="""The chat history to use to generate the search key for retrieval.""", + ) + ) + function_names: Optional[ExamplesArrayFilter] = Field( + default=None, description="""Optional. The function names for filtering.""" + ) + search_key: Optional[str] = Field( + default=None, description="""The exact search key to use for retrieval.""" + ) + + +class StoredContentsExampleParametersDict(TypedDict, total=False): + """The metadata filters that will be used to search StoredContentsExamples. If a field is unspecified, then no filtering for that field will be applied""" + + content_search_key: Optional[StoredContentsExampleParametersContentSearchKeyDict] + """The chat history to use to generate the search key for retrieval.""" + + function_names: Optional[ExamplesArrayFilterDict] + """Optional. The function names for filtering.""" + + search_key: Optional[str] + """The exact search key to use for retrieval.""" + + +StoredContentsExampleParametersOrDict = Union[ + StoredContentsExampleParameters, StoredContentsExampleParametersDict +] + + +class _SearchExamplesParameters(_common.BaseModel): + """Parameters for searching examples.""" + + name: Optional[str] = Field( + default=None, description="""Required. The Example Store to search.""" + ) + stored_contents_example_parameters: Optional[StoredContentsExampleParameters] = ( + Field( + default=None, + description="""Optional. The parameters that determine which examples to + retrieve. + """, + ) + ) + config: Optional[SearchExamplesConfig] = Field(default=None, description="""""") + + +class _SearchExamplesParametersDict(TypedDict, total=False): + """Parameters for searching examples.""" + + name: Optional[str] + """Required. The Example Store to search.""" + + stored_contents_example_parameters: Optional[StoredContentsExampleParametersDict] + """Optional. The parameters that determine which examples to + retrieve. + """ + + config: Optional[SearchExamplesConfigDict] + """""" + + +_SearchExamplesParametersOrDict = Union[ + _SearchExamplesParameters, _SearchExamplesParametersDict +] + + +class SearchExamplesResponseSimilarExample(_common.BaseModel): + """The result of the similar example.""" + + example: Optional[Example] = Field( + default=None, + description="""The example that is similar to the searched query.""", + ) + similarity_score: Optional[float] = Field( + default=None, description="""The similarity score of this example.""" + ) + + +class SearchExamplesResponseSimilarExampleDict(TypedDict, total=False): + """The result of the similar example.""" + + example: Optional[ExampleDict] + """The example that is similar to the searched query.""" + + similarity_score: Optional[float] + """The similarity score of this example.""" + + +SearchExamplesResponseSimilarExampleOrDict = Union[ + SearchExamplesResponseSimilarExample, SearchExamplesResponseSimilarExampleDict +] + + +class SearchExamplesResponse(_common.BaseModel): + """Response message for ExampleStoreService.SearchExamples.""" + + results: Optional[list[SearchExamplesResponseSimilarExample]] = Field( + default=None, description="""The results of searching for similar examples.""" + ) + + +class SearchExamplesResponseDict(TypedDict, total=False): + """Response message for ExampleStoreService.SearchExamples.""" + + results: Optional[list[SearchExamplesResponseSimilarExampleDict]] + """The results of searching for similar examples.""" + + +SearchExamplesResponseOrDict = Union[SearchExamplesResponse, SearchExamplesResponseDict] + + +class StoredContentsExampleFilter(_common.BaseModel): + """The metadata filters that will be used to remove or fetch StoredContentsExamples. If a field is unspecified, then no filtering for that field will be applied.""" + + function_names: Optional[ExamplesArrayFilter] = Field( + default=None, description="""Optional. The function names for filtering.""" + ) + search_keys: Optional[list[str]] = Field( + default=None, + description="""Optional. The search keys for filtering. Only examples with one of the specified search keys (StoredContentsExample.search_key) are eligible to be returned.""", + ) + + +class StoredContentsExampleFilterDict(TypedDict, total=False): + """The metadata filters that will be used to remove or fetch StoredContentsExamples. If a field is unspecified, then no filtering for that field will be applied.""" + + function_names: Optional[ExamplesArrayFilterDict] + """Optional. The function names for filtering.""" + + search_keys: Optional[list[str]] + """Optional. The search keys for filtering. Only examples with one of the specified search keys (StoredContentsExample.search_key) are eligible to be returned.""" + + +StoredContentsExampleFilterOrDict = Union[ + StoredContentsExampleFilter, StoredContentsExampleFilterDict +] + + +class FetchExamplesConfig(_common.BaseModel): + """Config for fetching examples from an Example Store.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + example_ids: Optional[list[str]] = Field( + default=None, + description="""Optional. Example IDs to fetch. If both this and the filter are + set, the fetched examples must match every example id AND the filter. + """, + ) + stored_contents_example_filter: Optional[StoredContentsExampleFilter] = Field( + default=None, + description="""Optional. The filter to apply to the fetched examples.""", + ) + page_size: Optional[int] = Field( + default=None, + description="""Optional. The maximum number of examples to return.""", + ) + page_token: Optional[str] = Field( + default=None, + description="""Optional. A page token from a previous FetchExamples call.""", + ) + + +class FetchExamplesConfigDict(TypedDict, total=False): + """Config for fetching examples from an Example Store.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + example_ids: Optional[list[str]] + """Optional. Example IDs to fetch. If both this and the filter are + set, the fetched examples must match every example id AND the filter. + """ + + stored_contents_example_filter: Optional[StoredContentsExampleFilterDict] + """Optional. The filter to apply to the fetched examples.""" + + page_size: Optional[int] + """Optional. The maximum number of examples to return.""" + + page_token: Optional[str] + """Optional. A page token from a previous FetchExamples call.""" + + +FetchExamplesConfigOrDict = Union[FetchExamplesConfig, FetchExamplesConfigDict] + + +class _FetchExamplesParameters(_common.BaseModel): + """Parameters for fetching examples.""" + + name: Optional[str] = Field( + default=None, + description="""Required. The Example Store to fetch examples from.""", + ) + config: Optional[FetchExamplesConfig] = Field(default=None, description="""""") + + +class _FetchExamplesParametersDict(TypedDict, total=False): + """Parameters for fetching examples.""" + + name: Optional[str] + """Required. The Example Store to fetch examples from.""" + + config: Optional[FetchExamplesConfigDict] + """""" + + +_FetchExamplesParametersOrDict = Union[ + _FetchExamplesParameters, _FetchExamplesParametersDict +] + + +class FetchExamplesResponse(_common.BaseModel): + """Response message for ExampleStoreService.FetchExamples.""" + + examples: Optional[list[Example]] = Field( + default=None, + description="""The examples in the Example Store that satisfy the metadata filters.""", + ) + next_page_token: Optional[str] = Field( + default=None, + description="""A token, which can be sent as FetchExamplesRequest.page_token to retrieve the next page. Absence of this field indicates there are no subsequent pages.""", + ) + + +class FetchExamplesResponseDict(TypedDict, total=False): + """Response message for ExampleStoreService.FetchExamples.""" + + examples: Optional[list[ExampleDict]] + """The examples in the Example Store that satisfy the metadata filters.""" + + next_page_token: Optional[str] + """A token, which can be sent as FetchExamplesRequest.page_token to retrieve the next page. Absence of this field indicates there are no subsequent pages.""" + + +FetchExamplesResponseOrDict = Union[FetchExamplesResponse, FetchExamplesResponseDict] + + +class RemoveExamplesConfig(_common.BaseModel): + """Config for removing examples from an Example Store.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + example_ids: Optional[list[str]] = Field( + default=None, + description="""Optional. Example IDs to remove. If both this and the filter are + set, the removed examples must match every example id AND the filter. + """, + ) + stored_contents_example_filter: Optional[StoredContentsExampleFilter] = Field( + default=None, + description="""Optional. The filter selecting which examples to remove.""", + ) + + +class RemoveExamplesConfigDict(TypedDict, total=False): + """Config for removing examples from an Example Store.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + example_ids: Optional[list[str]] + """Optional. Example IDs to remove. If both this and the filter are + set, the removed examples must match every example id AND the filter. + """ + + stored_contents_example_filter: Optional[StoredContentsExampleFilterDict] + """Optional. The filter selecting which examples to remove.""" + + +RemoveExamplesConfigOrDict = Union[RemoveExamplesConfig, RemoveExamplesConfigDict] + + +class _RemoveExamplesParameters(_common.BaseModel): + """Parameters for removing examples.""" + + name: Optional[str] = Field( + default=None, + description="""Required. The Example Store to remove examples from.""", + ) + config: Optional[RemoveExamplesConfig] = Field(default=None, description="""""") + + +class _RemoveExamplesParametersDict(TypedDict, total=False): + """Parameters for removing examples.""" + + name: Optional[str] + """Required. The Example Store to remove examples from.""" + + config: Optional[RemoveExamplesConfigDict] + """""" + + +_RemoveExamplesParametersOrDict = Union[ + _RemoveExamplesParameters, _RemoveExamplesParametersDict +] + + +class RemoveExamplesResponse(_common.BaseModel): + """Response message for ExampleStoreService.RemoveExamples.""" + + example_ids: Optional[list[str]] = Field( + default=None, description="""The IDs for the removed examples.""" + ) + + +class RemoveExamplesResponseDict(TypedDict, total=False): + """Response message for ExampleStoreService.RemoveExamples.""" + + example_ids: Optional[list[str]] + """The IDs for the removed examples.""" + + +RemoveExamplesResponseOrDict = Union[RemoveExamplesResponse, RemoveExamplesResponseDict] + + +class GetExampleStoreOperationConfig(_common.BaseModel): + """Optional parameters for example_stores.get_example_store_operation.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetExampleStoreOperationConfigDict(TypedDict, total=False): + """Optional parameters for example_stores.get_example_store_operation.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GetExampleStoreOperationConfigOrDict = Union[ + GetExampleStoreOperationConfig, GetExampleStoreOperationConfigDict +] + + +class _GetExampleStoreOperationParameters(_common.BaseModel): + """Parameters for getting an operation.""" + + operation_name: Optional[str] = Field( + default=None, description="""The server-assigned name for the operation.""" + ) + config: Optional[GetExampleStoreOperationConfig] = Field( + default=None, description="""Used to override the default configuration.""" + ) + + +class _GetExampleStoreOperationParametersDict(TypedDict, total=False): + """Parameters for getting an operation.""" + + operation_name: Optional[str] + """The server-assigned name for the operation.""" + + config: Optional[GetExampleStoreOperationConfigDict] + """Used to override the default configuration.""" + + +_GetExampleStoreOperationParametersOrDict = Union[ + _GetExampleStoreOperationParameters, _GetExampleStoreOperationParametersDict +] + + class PromptOptimizerConfig(_common.BaseModel): """VAPO Prompt Optimizer Config.""" diff --git a/tests/unit/agentplatform/genai/replays/test_example_stores.py b/tests/unit/agentplatform/genai/replays/test_example_stores.py new file mode 100644 index 0000000000..f9517e38fb --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_example_stores.py @@ -0,0 +1,167 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# pylint: disable=protected-access,bad-continuation,missing-function-docstring + +from agentplatform._genai import types +from tests.unit.agentplatform.genai.replays import pytest_helper +import pytest + +pytestmark = pytest_helper.setup( + file=__file__, +) + +pytest_plugins = ("pytest_asyncio",) + +# Real resources in the recording project (GOOGLE_CLOUD_PROJECT). Hard-coded so +# the recorded request URL is stable across runs. Re-recording against a +# different store changes every fixture, so keep these pinned. +EXAMPLE_STORE = "exampleStores/7241040532904869888" + +# Deliberately a non-existent example. removeExamples is a no-op for an unknown +# id, so this records a stable request/response pair without deleting the +# fixture example the search and fetch tests depend on. +EXAMPLE_ID = ( + "exampleTypes/stored_contents_example/examples/" "ffffffffffffffffffffffffffffffff" +) + +# Not an Example Store resource name, so the transformer must reject it before +# any HTTP request. No fixture is recorded for this test. +NOT_AN_EXAMPLE_STORE = "projects/p/locations/us-central1/endpoints/123" + + +def _stored_contents_example(search_key): + # search_key is required unless search_key_generation_method is set. With + # neither, upsertExamples reports INVALID_ARGUMENT per example inside an + # HTTP 200 and stores nothing. + return { + "stored_contents_example": { + "search_key": search_key, + "contents_example": { + "contents": [{"role": "user", "parts": [{"text": search_key}]}], + "expected_contents": [ + {"content": {"role": "model", "parts": [{"text": "hi"}]}} + ], + }, + } + } + + +def _assert_upserted(response): + assert isinstance(response, types.UpsertExamplesResponse) + # upsertExamples is partial-success: a result carries either an example or + # an error status, and the whole batch can fail inside a 200. Asserting the + # response type alone passes on an all-errors response. + assert response.results + assert response.results[0].status is None, response.results[0].status + assert response.results[0].example + + +def test_get(client): + response = client.example_stores.get(name=EXAMPLE_STORE) + assert isinstance(response, types.ExampleStore) + assert response.name + + +def test_upsert(client): + response = client.example_stores.upsert_examples( + name=EXAMPLE_STORE, + examples=[_stored_contents_example("hello")], + ) + _assert_upserted(response) + + +def test_search(client): + response = client.example_stores.search_examples( + name=EXAMPLE_STORE, + stored_contents_example_parameters={"search_key": "hello"}, + config={"top_k": 1}, + ) + assert isinstance(response, types.SearchExamplesResponse) + + +def test_fetch(client): + response = client.example_stores.fetch_examples( + name=EXAMPLE_STORE, + config={"page_size": 1}, + ) + assert isinstance(response, types.FetchExamplesResponse) + + +def test_remove(client): + response = client.example_stores.remove_examples( + name=EXAMPLE_STORE, + config={"example_ids": [EXAMPLE_ID]}, + ) + assert isinstance(response, types.RemoveExamplesResponse) + # EXAMPLE_ID does not exist, so this must remove nothing. Guards the config + # binding: while example_ids was dropped before the wire, removeExamples + # was sent an empty body and cleared the whole store instead. + assert not response.example_ids + + +@pytest.mark.parametrize( + "method", ["get", "upsert_examples", "search_examples", "fetch_examples"] +) +def test_rejects_non_example_store_name(client, method): + # The replay session is opened lazily on the first HTTP request, so a call + # that raises in the transformer never touches the wire and needs no fixture. + kwargs = {"examples": []} if method == "upsert_examples" else {} + with pytest.raises(ValueError, match="Invalid example store format"): + getattr(client.example_stores, method)(name=NOT_AN_EXAMPLE_STORE, **kwargs) + + +@pytest.mark.asyncio +async def test_get_async(client): + response = await client.aio.example_stores.get(name=EXAMPLE_STORE) + assert isinstance(response, types.ExampleStore) + assert response.name + + +@pytest.mark.asyncio +async def test_fetch_async(client): + response = await client.aio.example_stores.fetch_examples( + name=EXAMPLE_STORE, + config={"page_size": 1}, + ) + assert isinstance(response, types.FetchExamplesResponse) + + +@pytest.mark.asyncio +async def test_upsert_async(client): + response = await client.aio.example_stores.upsert_examples( + name=EXAMPLE_STORE, + examples=[_stored_contents_example("hello from the async client")], + ) + _assert_upserted(response) + + +@pytest.mark.asyncio +async def test_search_async(client): + response = await client.aio.example_stores.search_examples( + name=EXAMPLE_STORE, + stored_contents_example_parameters={"search_key": "hello"}, + config={"top_k": 1}, + ) + assert isinstance(response, types.SearchExamplesResponse) + + +@pytest.mark.asyncio +async def test_remove_async(client): + response = await client.aio.example_stores.remove_examples( + name=EXAMPLE_STORE, + config={"example_ids": [EXAMPLE_ID]}, + ) + assert isinstance(response, types.RemoveExamplesResponse) + assert not response.example_ids