diff --git a/lib/viewmodels/chat_viewmodel.dart b/lib/viewmodels/chat_viewmodel.dart
index e2a0127c..6e12423a 100644
--- a/lib/viewmodels/chat_viewmodel.dart
+++ b/lib/viewmodels/chat_viewmodel.dart
@@ -38,9 +38,12 @@ class ChatViewModel with ChangeNotifier {
     }
     try {
       // Fetch task steps from the data source
-      final List<dynamic> stepsJsonList =
+      final Map<String, dynamic> stepsResponse =
           await _chatService.listTaskSteps(_currentTaskId!);
 
+      // Extract steps from the response
+      final List<dynamic> stepsJsonList = stepsResponse['steps'] ?? [];
+
       // Convert each map into a Step object
       List<Step> steps =
           stepsJsonList.map((stepMap) => Step.fromMap(stepMap)).toList();
diff --git a/lib/viewmodels/task_viewmodel.dart b/lib/viewmodels/task_viewmodel.dart
index e4041c63..d4085d87 100644
--- a/lib/viewmodels/task_viewmodel.dart
+++ b/lib/viewmodels/task_viewmodel.dart
@@ -1,4 +1,5 @@
 import 'package:auto_gpt_flutter_client/models/task.dart';
+import 'package:auto_gpt_flutter_client/models/task_response.dart';
 import 'package:flutter/foundation.dart';
 import 'package:collection/collection.dart';
 import 'package:auto_gpt_flutter_client/services/task_service.dart';
@@ -45,13 +46,9 @@ class TaskViewModel with ChangeNotifier {
   /// Fetches tasks from the data source.
   void fetchTasks() async {
     try {
-      final fetchedTasks = await _taskService.listAllTasks();
-      _tasks = fetchedTasks.map((taskMap) {
-        return Task(
-          id: taskMap['task_id'].toString(),
-          title: taskMap['output'].toString(),
-        );
-      }).toList();
+      final TaskResponse tasksResponse = await _taskService.listAllTasks();
+      _tasks = tasksResponse.tasks;
+
       notifyListeners();
       print("Tasks fetched successfully!");
     } catch (error) {