diff --git a/tests/unit/vertexai/genai/replays/test_create_agent_engine_a2a_task.py b/tests/unit/vertexai/genai/replays/test_create_agent_engine_a2a_task.py index 6988dd6c2c..ea4640816d 100644 --- a/tests/unit/vertexai/genai/replays/test_create_agent_engine_a2a_task.py +++ b/tests/unit/vertexai/genai/replays/test_create_agent_engine_a2a_task.py @@ -34,6 +34,39 @@ def test_create_simple_a2a_task(client): a2a_task_id="task123", config=types.CreateAgentEngineTaskConfig( context_id="context123", + metadata={ + "key": "value", + "key2": [{"key3": "value3", "key4": "value4"}], + }, + state=types.A2aTaskState.STATE_UNSPECIFIED, + status_details=types.TaskStatusDetails( + task_message=types.TaskMessage( + role="user", + message_id="message123", + parts=[ + types.Part( + text="hello123", + ) + ], + metadata={ + "key42": "value42", + }, + ), + ), + output=types.TaskOutput( + artifacts=[ + types.TaskArtifact( + artifact_id="artifact123", + display_name="display_name123", + description="description123", + parts=[ + types.Part( + text="hello456", + ) + ], + ) + ], + ), ), ) @@ -41,6 +74,18 @@ def test_create_simple_a2a_task(client): assert task.name == f"{agent_engine.api_resource.name}/a2aTasks/task123" assert task.context_id == "context123" assert task.state == types.State.SUBMITTED + assert task.status_details.task_message.role == "user" + assert task.status_details.task_message.message_id == "message123" + assert task.status_details.task_message.parts[0].text == "hello123" + assert task.status_details.task_message.metadata["key42"] == "value42" + assert task.output.artifacts[0].artifact_id == "artifact123" + assert task.output.artifacts[0].display_name == "display_name123" + assert task.output.artifacts[0].description == "description123" + assert task.output.artifacts[0].parts[0].text == "hello456" + assert task.metadata == { + "key": "value", + "key2": [{"key3": "value3", "key4": "value4"}], + } # Clean up resources. client.agent_engines.delete(name=agent_engine.api_resource.name, force=True) diff --git a/vertexai/_genai/a2a_tasks.py b/vertexai/_genai/a2a_tasks.py index f3e4c1ca14..2b8a2ee26a 100644 --- a/vertexai/_genai/a2a_tasks.py +++ b/vertexai/_genai/a2a_tasks.py @@ -51,6 +51,18 @@ def _CreateAgentEngineTaskConfig_to_vertex( if getv(from_object, ["context_id"]) is not None: setv(parent_object, ["contextId"], getv(from_object, ["context_id"])) + if getv(from_object, ["metadata"]) is not None: + setv(parent_object, ["metadata"], getv(from_object, ["metadata"])) + + if getv(from_object, ["state"]) is not None: + setv(parent_object, ["state"], getv(from_object, ["state"])) + + if getv(from_object, ["status_details"]) is not None: + setv(parent_object, ["statusDetails"], getv(from_object, ["status_details"])) + + if getv(from_object, ["output"]) is not None: + setv(parent_object, ["output"], getv(from_object, ["output"])) + return to_object diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index 0903178153..c7f91953fe 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -1177,6 +1177,18 @@ class CreateAgentEngineTaskConfig(_common.BaseModel): context_id: Optional[str] = Field( default=None, description="""The context id of the task to create.""" ) + metadata: Optional[dict[str, Any]] = Field( + default=None, description="""The metadata of the task to create.""" + ) + state: Optional[A2aTaskState] = Field( + default=None, description="""The state of the task to create.""" + ) + status_details: Optional[TaskStatusDetails] = Field( + default=None, description="""The status details of the task to create.""" + ) + output: Optional[TaskOutput] = Field( + default=None, description="""The output of the task to create.""" + ) class CreateAgentEngineTaskConfigDict(TypedDict, total=False): @@ -1188,6 +1200,18 @@ class CreateAgentEngineTaskConfigDict(TypedDict, total=False): context_id: Optional[str] """The context id of the task to create.""" + metadata: Optional[dict[str, Any]] + """The metadata of the task to create.""" + + state: Optional[A2aTaskState] + """The state of the task to create.""" + + status_details: Optional[TaskStatusDetailsDict] + """The status details of the task to create.""" + + output: Optional[TaskOutputDict] + """The output of the task to create.""" + CreateAgentEngineTaskConfigOrDict = Union[ CreateAgentEngineTaskConfig, CreateAgentEngineTaskConfigDict