From 01b6e767f0a8046f27fa87307da2742b9938e1d6 Mon Sep 17 00:00:00 2001 From: mikemikimike <13286568797@163.com> Date: Tue, 8 Sep 2026 14:05:17 +0800 Subject: [PATCH] fix: accept a2a agent framework in config --- tests/unit/vertexai/test_genai_types.py | 27 ++++++++++ vertexai/_genai/types/common.py | 72 ++++++++++++++++++++----- 2 files changed, 87 insertions(+), 12 deletions(-) create mode 100644 tests/unit/vertexai/test_genai_types.py diff --git a/tests/unit/vertexai/test_genai_types.py b/tests/unit/vertexai/test_genai_types.py new file mode 100644 index 0000000000..510d4d5112 --- /dev/null +++ b/tests/unit/vertexai/test_genai_types.py @@ -0,0 +1,27 @@ +# Copyright 2026 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. + +import pytest + +from vertexai._genai import types + + +@pytest.mark.parametrize( + "config_type", + [types.AgentEngineConfig, types.CreateAgentEngineConfig, types.UpdateAgentEngineConfig], +) +def test_agent_engine_configs_accept_a2a_framework(config_type): + config = config_type(agent_framework="a2a") + + assert config.agent_framework == "a2a" diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index 5bb3b8811b..41d29ddd64 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -8885,13 +8885,21 @@ class CreateAgentEngineConfig(_common.BaseModel): """, ) agent_framework: Optional[ - Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + Literal[ + "google-adk", + "langchain", + "langgraph", + "ag2", + "llama-index", + "custom", + "a2a", + ] ] = Field( default=None, description="""The agent framework to be used for the Agent Engine. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", - "ag2", "llama-index", "custom". + "ag2", "llama-index", "custom", "a2a". If not specified: - If `agent` is specified, the agent framework will be auto-detected. - If `source_packages` is specified, the agent framework will @@ -9031,12 +9039,20 @@ class CreateAgentEngineConfigDict(TypedDict, total=False): """ agent_framework: Optional[ - Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + Literal[ + "google-adk", + "langchain", + "langgraph", + "ag2", + "llama-index", + "custom", + "a2a", + ] ] """The agent framework to be used for the Agent Engine. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", - "ag2", "llama-index", "custom". + "ag2", "llama-index", "custom", "a2a". If not specified: - If `agent` is specified, the agent framework will be auto-detected. - If `source_packages` is specified, the agent framework will @@ -9552,13 +9568,21 @@ class UpdateAgentEngineConfig(_common.BaseModel): """, ) agent_framework: Optional[ - Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + Literal[ + "google-adk", + "langchain", + "langgraph", + "ag2", + "llama-index", + "custom", + "a2a", + ] ] = Field( default=None, description="""The agent framework to be used for the Agent Engine. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", - "ag2", "llama-index", "custom". + "ag2", "llama-index", "custom", "a2a". If not specified: - If `agent` is specified, the agent framework will be auto-detected. - If `source_packages` is specified, the agent framework will @@ -9707,12 +9731,20 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False): """ agent_framework: Optional[ - Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + Literal[ + "google-adk", + "langchain", + "langgraph", + "ag2", + "llama-index", + "custom", + "a2a", + ] ] """The agent framework to be used for the Agent Engine. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", - "ag2", "llama-index", "custom". + "ag2", "llama-index", "custom", "a2a". If not specified: - If `agent` is specified, the agent framework will be auto-detected. - If `source_packages` is specified, the agent framework will @@ -20352,13 +20384,21 @@ class AgentEngineConfig(_common.BaseModel): """, ) agent_framework: Optional[ - Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + Literal[ + "google-adk", + "langchain", + "langgraph", + "ag2", + "llama-index", + "custom", + "a2a", + ] ] = Field( default=None, description="""The agent framework to be used for the Agent Engine. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", - "ag2", "llama-index", "custom". + "ag2", "llama-index", "custom", "a2a". If not specified: - If `agent` is specified, the agent framework will be auto-detected. - If `source_packages` is specified, the agent framework will @@ -20548,12 +20588,20 @@ class AgentEngineConfigDict(TypedDict, total=False): """ agent_framework: Optional[ - Literal["google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"] + Literal[ + "google-adk", + "langchain", + "langgraph", + "ag2", + "llama-index", + "custom", + "a2a", + ] ] """The agent framework to be used for the Agent Engine. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", - "ag2", "llama-index", "custom". + "ag2", "llama-index", "custom", "a2a". If not specified: - If `agent` is specified, the agent framework will be auto-detected. - If `source_packages` is specified, the agent framework will