diff --git a/library/src/androidTest/java/com/owncloud/android/lib/resources/assistant/v2/AssistantV2Tests.kt b/library/src/androidTest/java/com/owncloud/android/lib/resources/assistant/v2/AssistantV2Tests.kt
index 92afbf401..50222020f 100644
--- a/library/src/androidTest/java/com/owncloud/android/lib/resources/assistant/v2/AssistantV2Tests.kt
+++ b/library/src/androidTest/java/com/owncloud/android/lib/resources/assistant/v2/AssistantV2Tests.kt
@@ -9,6 +9,7 @@
 package com.owncloud.android.lib.resources.assistant.v2
 
 import com.owncloud.android.AbstractIT
+import com.owncloud.android.lib.resources.assistant.v2.model.Shape
 import com.owncloud.android.lib.resources.assistant.v2.model.TaskInputShape
 import com.owncloud.android.lib.resources.assistant.v2.model.TaskOutputShape
 import com.owncloud.android.lib.resources.assistant.v2.model.TaskTypeData
@@ -29,20 +30,24 @@ class AssistantV2Tests : AbstractIT() {
             "core:text2text",
             "Free text to text prompt",
             "Runs an arbitrary prompt through a language model that returns a reply",
-            listOf(
+            inputShape =
                 TaskInputShape(
-                    "Prompt",
-                    "Describe a task that you want the assistant to do or ask a question",
-                    "Text"
-                )
-            ),
-            listOf(
+                    input =
+                        Shape(
+                            "Prompt",
+                            "Describe a task that you want the assistant to do or ask a question",
+                            "Text"
+                        )
+                ),
+            outputShape =
                 TaskOutputShape(
-                    "Generated reply",
-                    "The generated text from the assistant",
-                    "Text"
+                    output =
+                        Shape(
+                            "Generated reply",
+                            "The generated text from the assistant",
+                            "Text"
+                        )
                 )
-            )
         )
 
     private fun getSelectedTaskType(): String = "core:text2text"
diff --git a/library/src/main/java/com/owncloud/android/lib/resources/assistant/v2/GetTaskTypesRemoteOperationV2.kt b/library/src/main/java/com/owncloud/android/lib/resources/assistant/v2/GetTaskTypesRemoteOperationV2.kt
index e4b4bd862..b3f10305e 100644
--- a/library/src/main/java/com/owncloud/android/lib/resources/assistant/v2/GetTaskTypesRemoteOperationV2.kt
+++ b/library/src/main/java/com/owncloud/android/lib/resources/assistant/v2/GetTaskTypesRemoteOperationV2.kt
@@ -45,12 +45,8 @@ class GetTaskTypesRemoteOperationV2 : OCSRemoteOperation<List<TaskTypeData>>() {
 
                 val supportedTaskTypeList =
                     taskTypeList?.filter { taskType ->
-                        taskType.inputShape?.any { inputShape ->
-                            inputShape.type == supportedTaskType
-                        } == true &&
-                            taskType.outputShape?.any { outputShape ->
-                                outputShape.type == supportedTaskType
-                            } == true
+                        taskType.inputShape?.input?.type == supportedTaskType &&
+                            taskType.outputShape?.output?.type == supportedTaskType
                     }
 
                 result = RemoteOperationResult(true, getMethod)
diff --git a/library/src/main/java/com/owncloud/android/lib/resources/assistant/v2/model/TaskTypes.kt b/library/src/main/java/com/owncloud/android/lib/resources/assistant/v2/model/TaskTypes.kt
index bd33d3972..e0a3861aa 100644
--- a/library/src/main/java/com/owncloud/android/lib/resources/assistant/v2/model/TaskTypes.kt
+++ b/library/src/main/java/com/owncloud/android/lib/resources/assistant/v2/model/TaskTypes.kt
@@ -16,18 +16,20 @@ data class TaskTypeData(
     val id: String?,
     val name: String?,
     val description: String?,
-    val inputShape: List<TaskInputShape>?,
-    val outputShape: List<TaskOutputShape>?
+    val inputShape: TaskInputShape?,
+    val outputShape: TaskOutputShape?
 )
 
 data class TaskInputShape(
-    val name: String?,
-    val description: String?,
-    val type: String?
+    val input: Shape?
 )
 
 data class TaskOutputShape(
-    val name: String?,
-    val description: String?,
-    val type: String?
+    val output: Shape?
+)
+
+data class Shape(
+    val name: String,
+    val description: String,
+    val type: String
 )