feat: add more fields to the GenAI client for the Agent Engine Task Store Service

tests/unit/vertexai/genai/replays/test_create_agent_engine_a2a_task.py

@@ -34,13 +34,58 @@ 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",
+                            )
+                        ],
+                    )
+                ],
+            ),
         ),
     )
 
     assert isinstance(task, types.A2aTask)
     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)

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
 
 

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