fix: accept "a2a" as an agent_framework in the agent engine config types

PiperOrigin-RevId: 974373620

Author: vertex-sdk-bot

tests/unit/agentplatform/genai/test_genai_client.py

@@ -16,6 +16,7 @@
 # pylint: disable=protected-access,bad-continuation
 
 import importlib
+import pydantic
 import pytest
 import warnings
 from unittest import mock
@@ -26,6 +27,7 @@
 from google.cloud.aiplatform import initializer as aiplatform_initializer
 import vertexai
 from vertexai._genai import client as vertexai_client
+from vertexai._genai import types as vertexai_types
 
 
 _TEST_PROJECT = "test-project"
@@ -122,3 +124,23 @@ def test_vertexai_client_deprecation_warning(self):
                 warnings.simplefilter("error", FutureWarning)
                 _ = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
                 _ = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
+
+
+_AGENT_ENGINE_CONFIG_TYPES = [
+    vertexai_types.AgentEngineConfig,
+    vertexai_types.CreateAgentEngineConfig,
+    vertexai_types.UpdateAgentEngineConfig,
+]
+
+
+@pytest.mark.parametrize("config_type", _AGENT_ENGINE_CONFIG_TYPES)
+def test_vertexai_agent_framework_accepts_a2a(config_type):
+    # googleapis/python-aiplatform#7097: the backend and the SDK's own
+    # _SUPPORTED_AGENT_FRAMEWORKS both accept a2a, so these types must too.
+    assert config_type(agent_framework="a2a").agent_framework == "a2a"
+
+
+@pytest.mark.parametrize("config_type", _AGENT_ENGINE_CONFIG_TYPES)
+def test_vertexai_agent_framework_rejects_unknown(config_type):
+    with pytest.raises(pydantic.ValidationError):
+        config_type(agent_framework="not-a-framework")

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