diff --git a/.release-please-manifest.json b/.release-please-manifest.json index 3d04639af4..e275cc3c25 100644 --- a/.release-please-manifest.json +++ b/.release-please-manifest.json @@ -1,3 +1,3 @@ { - ".": "1.115.0" + ".": "1.116.0" } diff --git a/CHANGELOG.md b/CHANGELOG.md index e358f6a503..a5e64e4e5c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,25 @@ # Changelog +## [1.116.0](https://github.com/googleapis/python-aiplatform/compare/v1.115.0...v1.116.0) (2025-09-22) + + +### Features + +* Add experimental async list_prompts and list_version methods to prompt management ([13a626b](https://github.com/googleapis/python-aiplatform/commit/13a626bb38f5e8f9ba5c0b7db41b41c46ecf4f58)) +* GenAI SDK client - add experimental async delete_prompt and delete_version method to Prompt Management ([45e616a](https://github.com/googleapis/python-aiplatform/commit/45e616a94f1b14646930cb9d598ac20776326038)) +* Update gapic utils to allows creating RAG clients with api_override ([9851905](https://github.com/googleapis/python-aiplatform/commit/98519056af7f3f95ac15c0b585b788a82ab548a7)) + + +### Bug Fixes + +* GenAI Client(evals) - Add support for `context` in Grounding metric ([e75d91f](https://github.com/googleapis/python-aiplatform/commit/e75d91fc928358299ed4d6cd1a210a2448cbea4e)) +* Pass plugins to the Runner in ADK template. ([b43b7b4](https://github.com/googleapis/python-aiplatform/commit/b43b7b4f259cea295fe447fa0c751ff5c5a34c6a)) + + +### Documentation + +* Update AgentEngine with ADK template to the version in GA ([0d600fd](https://github.com/googleapis/python-aiplatform/commit/0d600fdf93a3e70e9a6b0d43913147b100f116ce)) + ## [1.115.0](https://github.com/googleapis/python-aiplatform/compare/v1.114.0...v1.115.0) (2025-09-18) diff --git a/gemini_docs/README.md b/gemini_docs/README.md index 51992d9c17..3ffe4ffd1a 100644 --- a/gemini_docs/README.md +++ b/gemini_docs/README.md @@ -318,7 +318,7 @@ result = EvalTask( Before you begin, install the packages with ```shell -pip3 install --upgrade --user "google-cloud-aiplatform[agent_engines,adk]>=1.95.1" +pip3 install --upgrade --user "google-cloud-aiplatform[agent_engines,adk]>=1.111" ``` First, define a function that looks up the exchange rate: @@ -351,7 +351,7 @@ Next, define an ADK Agent: ```python from google.adk.agents import Agent -from vertexai.preview.reasoning_engines import AdkApp +from vertexai.agent_engines import AdkApp app = AdkApp(agent=Agent( model="gemini-2.0-flash", # Required. @@ -363,7 +363,7 @@ app = AdkApp(agent=Agent( Test the agent locally using US dollars and Swedish Krona: ```python -for event in app.stream_query( +async for event in app.async_stream_query( user_id="user-id", message="What is the exchange rate from US dollars to SEK today?", ): @@ -388,7 +388,7 @@ remote_app = vertexai.agent_engines.create( You can also run queries against the deployed agent: ```python -for event in remote_app.stream_query( +async for event in remote_app.async_stream_query( user_id="user-id", message="What is the exchange rate from US dollars to SEK today?", ): diff --git a/google/cloud/aiplatform/gapic_version.py b/google/cloud/aiplatform/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/gapic_version.py +++ b/google/cloud/aiplatform/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py b/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/version.py b/google/cloud/aiplatform/version.py index 7d272003ae..46dd2a2baa 100644 --- a/google/cloud/aiplatform/version.py +++ b/google/cloud/aiplatform/version.py @@ -15,4 +15,4 @@ # limitations under the License. # -__version__ = "1.115.0" +__version__ = "1.116.0" diff --git a/google/cloud/aiplatform_v1/gapic_version.py b/google/cloud/aiplatform_v1/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform_v1/gapic_version.py +++ b/google/cloud/aiplatform_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform_v1beta1/gapic_version.py b/google/cloud/aiplatform_v1beta1/gapic_version.py index 2a17984cf3..a299e2bf67 100644 --- a/google/cloud/aiplatform_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.115.0" # {x-release-please-version} +__version__ = "1.116.0" # {x-release-please-version} diff --git a/pypi/_vertex_ai_placeholder/version.py b/pypi/_vertex_ai_placeholder/version.py index 5e53e7f921..22b8338ce1 100644 --- a/pypi/_vertex_ai_placeholder/version.py +++ b/pypi/_vertex_ai_placeholder/version.py @@ -15,4 +15,4 @@ # limitations under the License. # -__version__ = "1.115.0" +__version__ = "1.116.0" diff --git a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json index f1717f0f64..04beb40ab9 100644 --- a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json +++ b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json @@ -8,7 +8,7 @@ ], "language": "PYTHON", "name": "google-cloud-aiplatform", - "version": "1.115.0" + "version": "1.116.0" }, "snippets": [ { diff --git a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json index c45c62cbf7..bef7b2a179 100644 --- a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json +++ b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json @@ -8,7 +8,7 @@ ], "language": "PYTHON", "name": "google-cloud-aiplatform", - "version": "1.115.0" + "version": "1.116.0" }, "snippets": [ { diff --git a/setup.py b/setup.py index a7ad9f14b0..c645ae442b 100644 --- a/setup.py +++ b/setup.py @@ -170,7 +170,7 @@ "jsonschema", "ruamel.yaml", "pyyaml", - "litellm >= 1.72.4", + "litellm >= 1.72.4, <= 1.76.3", ] langchain_extra_require = [ diff --git a/tests/unit/vertexai/genai/replays/test_delete_prompt.py b/tests/unit/vertexai/genai/replays/test_delete_prompt.py index 151dd010d7..acf5204202 100644 --- a/tests/unit/vertexai/genai/replays/test_delete_prompt.py +++ b/tests/unit/vertexai/genai/replays/test_delete_prompt.py @@ -19,6 +19,8 @@ from tests.unit.vertexai.genai.replays import pytest_helper from vertexai._genai import types from google.genai import types as genai_types +import pytest + TEST_PROMPT_DATASET_ID = "6550997480673116160" TEST_PROMPT_VERSION_ID = "2" @@ -79,3 +81,57 @@ def test_delete_dataset_version(client, caplog): globals_for_file=globals(), test_method="prompt_management.delete", ) + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_delete_dataset_async(client, caplog): + caplog.set_level(logging.INFO) + prompt = client.prompt_management.create_version( + prompt=types.Prompt( + prompt_data=types.PromptData( + model="gemini-2.5-flash", + contents=[ + genai_types.Content( + parts=[genai_types.Part(text="What is the capital of France?")], + ) + ], + ) + ), + config=types.CreatePromptConfig( + prompt_display_name="test_delete_prompt_dataset", + version_display_name="test_delete_prompt_dataset_version", + ), + ) + await client.aio.prompt_management.delete_prompt( + prompt_id=prompt.prompt_id, + ) + assert "Deleted prompt with id: " in caplog.text + + +@pytest.mark.asyncio +async def test_delete_version_async(client, caplog): + caplog.set_level(logging.INFO) + prompt = client.prompt_management.create_version( + prompt=types.Prompt( + prompt_data=types.PromptData( + model="gemini-2.5-flash", + contents=[ + genai_types.Content( + parts=[genai_types.Part(text="What is the capital of France?")], + ) + ], + ) + ), + config=types.CreatePromptConfig( + prompt_display_name="test_delete_prompt_dataset", + version_display_name="test_delete_prompt_dataset_version", + ), + ) + version_id = prompt.dataset_version.name.split("/")[-1] + await client.aio.prompt_management.delete_version( + prompt_id=prompt.prompt_id, + version_id=version_id, + ) + assert "Deleted prompt version" in caplog.text diff --git a/tests/unit/vertexai/genai/replays/test_evaluate_predefined_metrics.py b/tests/unit/vertexai/genai/replays/test_evaluate_predefined_metrics.py index 54ee4cba07..a982681656 100644 --- a/tests/unit/vertexai/genai/replays/test_evaluate_predefined_metrics.py +++ b/tests/unit/vertexai/genai/replays/test_evaluate_predefined_metrics.py @@ -37,7 +37,7 @@ def test_evaluation_result(client): ) predefined_metrics = [ - types.PrebuiltMetric.GENERAL_QUALITY, + types.RubricMetric.GENERAL_QUALITY, ] evaluation_result = client.evals.evaluate( @@ -201,6 +201,53 @@ def test_multi_turn_predefined_metric(client): assert case_result.response_candidate_results is not None +def test_evaluation_grounding_metric(client): + """Tests that grounding metric produces a correctly structured EvaluationResult.""" + prompts_df = pd.DataFrame( + { + "prompt": ["Explain the concept of machine learning in simple terms."], + "response": [ + "Machine learning is a type of artificial intelligence that allows" + " computers to learn from data without being explicitly programmed." + ], + "context": [ + "Article: 'Intro to AI', Section 2.1\n" + "Machine learning (ML) is a subfield of artificial intelligence (AI). " + "The core idea of machine learning is that it allows computer systems to " + "learn from and adapt to new data without being explicitly programmed. " + "Instead of a developer writing code for every possible scenario, the " + "system builds a model based on patterns in training data." + ], + } + ) + + eval_dataset = types.EvaluationDataset( + eval_dataset_df=prompts_df, + candidate_name="gemini-2.5-flash", + ) + + evaluation_result = client.evals.evaluate( + dataset=eval_dataset, + metrics=[ + types.RubricMetric.GROUNDING, + ], + ) + + assert isinstance(evaluation_result, types.EvaluationResult) + + assert evaluation_result.summary_metrics is not None + for summary in evaluation_result.summary_metrics: + assert isinstance(summary, types.AggregatedMetricResult) + assert summary.metric_name is not None + assert summary.mean_score is not None + + assert evaluation_result.eval_case_results is not None + for case_result in evaluation_result.eval_case_results: + assert isinstance(case_result, types.EvalCaseResult) + assert case_result.eval_case_index is not None + assert case_result.response_candidate_results is not None + + pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), diff --git a/tests/unit/vertexai/genai/replays/test_list_prompts.py b/tests/unit/vertexai/genai/replays/test_list_prompts.py index 0e64e601c5..d72569563f 100644 --- a/tests/unit/vertexai/genai/replays/test_list_prompts.py +++ b/tests/unit/vertexai/genai/replays/test_list_prompts.py @@ -16,6 +16,8 @@ from tests.unit.vertexai.genai.replays import pytest_helper from vertexai._genai import types +import pytest + def test_list_returns_prompts(client): prompt_refs = client.prompt_management.list_prompts() @@ -71,3 +73,26 @@ def test_list_versions(client): globals_for_file=globals(), test_method="prompt_management.list_prompts", ) + +pytest_plugins = ("pytest_asyncio",) + + +@pytest.mark.asyncio +async def test_list_returns_prompts_async(client): + prompt_refs = client.aio.prompt_management.list_prompts() + async for prompt in prompt_refs: + assert isinstance(prompt, types.PromptRef) + assert prompt.prompt_id is not None + assert prompt.model is not None + + +@pytest.mark.asyncio +async def test_list_versions_async(client): + prompt_version_refs = client.aio.prompt_management.list_versions( + prompt_id="3331020504126455808" + ) + async for prompt_version in prompt_version_refs: + assert isinstance(prompt_version, types.PromptVersionRef) + assert prompt_version.prompt_id is not None + assert prompt_version.version_id is not None + assert prompt_version.model is not None diff --git a/vertexai/_genai/_evals_metric_handlers.py b/vertexai/_genai/_evals_metric_handlers.py index cb50d1f755..8f97f37f92 100644 --- a/vertexai/_genai/_evals_metric_handlers.py +++ b/vertexai/_genai/_evals_metric_handlers.py @@ -871,6 +871,19 @@ def _build_request_payload( eval_case.prompt ) + other_data_map = {} + if hasattr(eval_case, "context") and eval_case.context: + if isinstance(eval_case.context, str): + other_data_map["context"] = types.InstanceData(text=eval_case.context) + elif isinstance(eval_case.context, genai_types.Content): + other_data_map["context"] = ( + PredefinedMetricHandler._content_to_instance_data(eval_case.context) + ) + else: + logger.warning( + f"Unsupported type for context: {type(eval_case.context)}" + ) + instance_payload = types.EvaluationInstance( prompt=prompt_instance_data, response=PredefinedMetricHandler._content_to_instance_data( @@ -878,6 +891,11 @@ def _build_request_payload( ), reference=reference_instance_data, rubric_groups=eval_case.rubric_groups, + other_data=( + types.MapInstance(map_instance=other_data_map) + if other_data_map + else None + ), ) return { diff --git a/vertexai/_genai/prompt_management.py b/vertexai/_genai/prompt_management.py index 4dec340fb8..ef65de259e 100644 --- a/vertexai/_genai/prompt_management.py +++ b/vertexai/_genai/prompt_management.py @@ -19,7 +19,7 @@ import json import logging import time -from typing import Any, Iterator, Optional, Union +from typing import Any, AsyncIterator, Iterator, Optional, Union from urllib.parse import urlencode from google.genai import _api_module @@ -28,7 +28,7 @@ from google.genai import types as genai_types from google.genai._common import get_value_by_path as getv from google.genai._common import set_value_by_path as setv -from google.genai.pagers import Pager +from google.genai.pagers import AsyncPager, Pager from . import _prompt_management_utils from . import types @@ -2397,3 +2397,210 @@ async def get( prompt._dataset_version = prompt_version_resource return prompt + + async def _wait_for_project_operation( + self, + operation: types.DatasetOperation, + timeout: int, + ) -> None: + """Waits for a dataset deletion operation to complete. + + Delete operations are project level operations and are separate from dataset resource operations, for example: projects/123/locations/us-central1/operations/789. + + Args: + operation: The project operation to wait for. + timeout: The maximum time to wait for the operation to complete. + Raises: + TimeoutError: If the operation does not complete within the timeout. + ValueError: If the operation fails. + """ + done = False + + start_time = time.time() + sleep_duration = 5 + wait_multiplier = 2 + max_wait_time = 60 + previous_time = time.time() + while not done: + if (time.time() - start_time) > timeout: + raise TimeoutError( + f"Delete operation did not complete within the" + f" specified timeout of {timeout} seconds." + ) + current_time = time.time() + if current_time - previous_time >= sleep_duration: + sleep_duration = min(sleep_duration * wait_multiplier, max_wait_time) + previous_time = current_time + await asyncio.sleep(sleep_duration) + operations_module = operations.AsyncOperations(api_client_=self._api_client) + + operation = await operations_module._get( + operation_id=operation.name.split("/")[-1], + ) + done = operation.done if hasattr(operation, "done") else False + if hasattr(operation, "error") and operation.error is not None: + raise ValueError(f"Error in delete operation: {operation.error}") + + async def delete_prompt( + self, + *, + prompt_id: str, + config: Optional[types.DeletePromptConfig] = None, + ) -> None: + """Deletes a prompt resource. + + Args: + prompt_id: The id of the prompt resource to delete. + + Raises: + TimeoutError: If the delete operation does not complete within the timeout. + ValueError: If the delete operation fails. + """ + + delete_prompt_operation = await self._delete_dataset( + prompt_id=prompt_id, + config=config, + ) + await self._wait_for_project_operation( + operation=delete_prompt_operation, timeout=config.timeout if config else 90 + ) + logger.info(f"Deleted prompt with id: {prompt_id}") + + async def delete_version( + self, + *, + prompt_id: str, + version_id: str, + config: Optional[types.DeletePromptConfig] = None, + ) -> None: + """Deletes a prompt version resource. + + Args: + prompt_id: The id of the prompt resource to delete. + version_id: The id of the prompt version resource to delete. + + Raises: + TimeoutError: If the delete operation does not complete within the timeout. + ValueError: If the delete operation fails. + """ + delete_version_operation = await self._delete_dataset_version( + prompt_id=prompt_id, + version_id=version_id, + config=config, + ) + + await self._wait_for_project_operation( + operation=delete_version_operation, timeout=config.timeout if config else 90 + ) + logger.info( + f"Deleted prompt version {version_id} from prompt with id: {prompt_id}" + ) + + async def _list_prompts_pager( + self, + *, + config: Optional[types.ListPromptsConfigOrDict] = None, + ) -> AsyncPager[types.Dataset]: + return AsyncPager( + "datasets", + self._list_prompts, + await self._list_prompts(config=config), + config, + ) + + async def _list_versions_pager( + self, + *, + prompt_id: str, + config: Optional[types.ListPromptsConfigOrDict] = None, + ) -> AsyncPager[types.DatasetVersion]: + return AsyncPager( + "dataset_versions", + self._list_versions, + await self._list_versions(config=config, dataset_id=prompt_id), + config, + ) + + async def list_prompts( + self, + *, + config: Optional[types.ListPromptsConfigOrDict] = None, + ) -> AsyncIterator[types.PromptRef]: + """Lists prompt resources in a project. + + This method retrieves all the prompts from the project provided in the + vertexai.Client constructor and returns a list of prompt references containing the prompt_id and model for the prompt. + + To get the full types.Prompt resource for a PromptRef after calling this method, use the get() method with the prompt_id as the prompt_id argument. + Example usage: + + ``` + prompt_refs = client.aio.prompt_management.list_prompts() + async for prompt_ref in prompt_refs: + await client.prompt_management.get(prompt_id=prompt_ref.prompt_id) + ``` + + Args: + config: Optional configuration for listing prompts. + + Returns: + An async iterator of types.PromptRef objects. + """ + if isinstance(config, dict): + config = types.ListPromptsConfig(**config) + elif not config: + config = types.ListPromptsConfig() + async for dataset in await self._list_prompts_pager(config=config): + if not dataset or not dataset.model_reference or not dataset.name: + continue + prompt_ref = types.PromptRef( + model=dataset.model_reference, prompt_id=dataset.name.split("/")[-1] + ) + yield prompt_ref + + async def list_versions( + self, + *, + prompt_id: str, + config: Optional[types.ListPromptsConfigOrDict] = None, + ) -> AsyncIterator[types.PromptVersionRef]: + """Lists prompt version resources for a provided prompt_id. + + This method retrieves all the prompt versions for a provided prompt_id. + + To get the full types.Prompt resource for a PromptVersionRef after calling this method, use the get() method with the returned prompt_id and version_id. + Example usage: + + ``` + prompt_version_refs = await client.prompt_management.list_versions(prompt_id="123") + async for version_ref in prompt_version_refs: + await client.aio.prompt_management.get(prompt_id=version_ref.prompt_id, version_id=version_ref.version_id) + ``` + + Args: + prompt_id: The id of the Vertex Dataset resource containing the prompt. For example, if the prompt resource name is "projects/123/locations/us-central1/datasets/456", then the prompt_id is "456". + config: Optional configuration for listing prompts. + + Returns: + An async iterator of types.PromptVersionRef objects representing the prompt version resources for the provided prompt_id. + + """ + if isinstance(config, dict): + config = types.ListPromptsConfig(**config) + elif not config: + config = types.ListPromptsConfig() + async for dataset_version in await self._list_versions_pager( + config=config, prompt_id=prompt_id + ): + if ( + not dataset_version + or not dataset_version.model_reference + or not dataset_version.name + ): + continue + prompt_version_ref = types.PromptVersionRef( + model=dataset_version.model_reference, + version_id=dataset_version.name.split("/")[-1], + prompt_id=prompt_id, + ) + yield prompt_version_ref diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index 57afaf440b..aaee9095d0 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -569,6 +569,7 @@ def set_up(self): self._tmpl_attrs["runner"] = Runner( agent=self._tmpl_attrs.get("agent"), + plugins=self._tmpl_attrs.get("plugins"), session_service=self._tmpl_attrs.get("session_service"), artifact_service=self._tmpl_attrs.get("artifact_service"), memory_service=self._tmpl_attrs.get("memory_service"), diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index 22e9ee0bb7..3f5c12fa3f 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -571,6 +571,7 @@ def set_up(self): self._tmpl_attrs["in_memory_memory_service"] = InMemoryMemoryService() self._tmpl_attrs["in_memory_runner"] = Runner( agent=self._tmpl_attrs.get("agent"), + plugins=self._tmpl_attrs.get("plugins"), session_service=self._tmpl_attrs.get("in_memory_session_service"), artifact_service=self._tmpl_attrs.get("in_memory_artifact_service"), memory_service=self._tmpl_attrs.get("in_memory_memory_service"), diff --git a/vertexai/rag/rag_retrieval.py b/vertexai/rag/rag_retrieval.py index fd5638b5a7..8448644e18 100644 --- a/vertexai/rag/rag_retrieval.py +++ b/vertexai/rag/rag_retrieval.py @@ -27,6 +27,8 @@ def retrieval_query( text: str, + parent_override: Optional[str] = None, + api_path_override: Optional[str] = None, rag_resources: Optional[List[resources.RagResource]] = None, rag_retrieval_config: Optional[resources.RagRetrievalConfig] = None, ) -> aiplatform_v1.RetrieveContextsResponse: @@ -62,6 +64,8 @@ def retrieval_query( Args: text: The query in text format to get relevant contexts. + parent_override: Optional. The resource path of the parent. + api_path_override: Optional. The base API endpoint to use for the request. rag_resources: A list of RagResource. It can be used to specify corpus only or ragfiles. Currently only support one corpus or multiple files from one corpus. In the future we may open up multiple corpora support. @@ -72,8 +76,10 @@ def retrieval_query( RetrieveContextsResonse. """ parent = initializer.global_config.common_location_path() + if parent_override: + parent = parent_override - client = _gapic_utils.create_rag_service_client() + client = _gapic_utils.create_rag_service_client(api_path_override) if rag_resources: if len(rag_resources) > 1: @@ -82,7 +88,7 @@ def retrieval_query( else: raise ValueError("rag_resources must be specified.") - data_client = _gapic_utils.create_rag_data_service_client() + data_client = _gapic_utils.create_rag_data_service_client(api_path_override) if data_client.parse_rag_corpus_path(name): rag_corpus_name = name elif re.match("^{}$".format(_gapic_utils._VALID_RESOURCE_NAME_REGEX), name): diff --git a/vertexai/rag/utils/_gapic_utils.py b/vertexai/rag/utils/_gapic_utils.py index 565bd845ff..3ee39a7a0f 100644 --- a/vertexai/rag/utils/_gapic_utils.py +++ b/vertexai/rag/utils/_gapic_utils.py @@ -73,21 +73,30 @@ ) -def create_rag_data_service_client(): +def create_rag_data_service_client( + api_path_override: Optional[str] = None, +): return initializer.global_config.create_client( client_class=VertexRagDataClientWithOverride, + api_path_override=api_path_override, ).select_version("v1") -def create_rag_data_service_async_client(): +def create_rag_data_service_async_client( + api_path_override: Optional[str] = None, +): return initializer.global_config.create_client( client_class=VertexRagDataAsyncClientWithOverride, + api_path_override=api_path_override, ).select_version("v1") -def create_rag_service_client(): +def create_rag_service_client( + api_path_override: Optional[str] = None, +): return initializer.global_config.create_client( client_class=VertexRagClientWithOverride, + api_path_override=api_path_override, ).select_version("v1")