Added Artifact Handling for Chat Messages

- Enhanced the `Chat` model to include an `artifacts` field to hold associated artifacts.
- Updated the `AgentMessageTile` widget to display the number of artifacts and added functionality to trigger artifact downloads upon button press.
- Introduced a method in `ChatViewModel` to leverage the `ChatService` for artifact downloads.

5 files changed, 75 insertions, 42 deletions

diff --git a/frontend/lib/models/chat.dart b/frontend/lib/models/chat.dart
index 7120014fa..e9a281f41 100644
--- a/frontend/lib/models/chat.dart
+++ b/frontend/lib/models/chat.dart
@@ -8,6 +8,7 @@ class Chat {
   final DateTime timestamp;
   final MessageType messageType;
   final Map<String, dynamic>? jsonResponse;
+  final List<dynamic> artifacts;
 
   Chat({
     required this.id,
@@ -16,6 +17,7 @@ class Chat {
     required this.timestamp,
     required this.messageType,
     this.jsonResponse,
+    required this.artifacts,
   });
 
   // Convert a Map (usually from JSON) to a Chat object
@@ -27,6 +29,7 @@ class Chat {
       timestamp: DateTime.parse(map['timestamp']),
       messageType: MessageType.values.firstWhere(
           (e) => e.toString() == 'MessageType.${map['messageType']}'),
+      artifacts: List<dynamic>.from(map['artifacts'] ?? []),
     );
   }
 
@@ -39,7 +42,8 @@ class Chat {
           taskId == other.taskId &&
           message == other.message &&
           timestamp == other.timestamp &&
-          messageType == other.messageType;
+          messageType == other.messageType &&
+          artifacts == other.artifacts;
 
   @override
   int get hashCode =>
@@ -47,9 +51,10 @@ class Chat {
       taskId.hashCode ^
       message.hashCode ^
       timestamp.hashCode ^
-      messageType.hashCode;
+      messageType.hashCode ^
+      artifacts.hashCode;
 
   @override
   String toString() =>
-      'Chat(id: $id, taskId: $taskId, message: $message, timestamp: $timestamp, messageType: $messageType)';
+      'Chat(id: $id, taskId: $taskId, message: $message, timestamp: $timestamp, messageType: $messageType, artifacts: $artifacts)'; // Added artifacts in toString method
 }
diff --git a/frontend/lib/models/step.dart b/frontend/lib/models/step.dart
index 21b7e47cc..951c5eea6 100644
--- a/frontend/lib/models/step.dart
+++ b/frontend/lib/models/step.dart
@@ -8,6 +8,7 @@ class Step {
   final String status;
   final String output;
   final Map<String, dynamic> additionalOutput;
+  // TODO: Create an actual artifact object
   final List<dynamic> artifacts;
   final bool isLast;
 
diff --git a/frontend/lib/viewmodels/chat_viewmodel.dart b/frontend/lib/viewmodels/chat_viewmodel.dart
index 345c599d5..87c1e9d96 100644
--- a/frontend/lib/viewmodels/chat_viewmodel.dart
+++ b/frontend/lib/viewmodels/chat_viewmodel.dart
@@ -68,24 +68,23 @@ class ChatViewModel with ChangeNotifier {
         // Create a Chat object for 'input' if it exists and is not empty
         if (step.input.isNotEmpty) {
           chats.add(Chat(
-            id: step.stepId,
-            taskId: step.taskId,
-            message: step.input,
-            timestamp: currentTimestamp,
-            messageType: MessageType.user,
-          ));
+              id: step.stepId,
+              taskId: step.taskId,
+              message: step.input,
+              timestamp: currentTimestamp,
+              messageType: MessageType.user,
+              artifacts: step.artifacts));
         }
 
         // Create a Chat object for 'output'
         chats.add(Chat(
-          id: step.stepId,
-          taskId: step.taskId,
-          message: step.output,
-          timestamp: currentTimestamp,
-          messageType: MessageType.agent,
-          jsonResponse:
-              stepsJsonList[i], // Include the specific step's JSON here
-        ));
+            id: step.stepId,
+            taskId: step.taskId,
+            message: step.output,
+            timestamp: currentTimestamp,
+            messageType: MessageType.agent,
+            jsonResponse: stepsJsonList[i],
+            artifacts: step.artifacts));
       }
 
       // Assign the chats list
@@ -123,12 +122,12 @@ class ChatViewModel with ChangeNotifier {
       // Create a Chat object for the user message
       if (executedStep.input.isNotEmpty) {
         final userChat = Chat(
-          id: executedStep.stepId,
-          taskId: executedStep.taskId,
-          message: executedStep.input,
-          timestamp: DateTime.now(),
-          messageType: MessageType.user,
-        );
+            id: executedStep.stepId,
+            taskId: executedStep.taskId,
+            message: executedStep.input,
+            timestamp: DateTime.now(),
+            messageType: MessageType.user,
+            artifacts: executedStep.artifacts);
 
         _chats.add(userChat);
       }
@@ -140,7 +139,8 @@ class ChatViewModel with ChangeNotifier {
           message: executedStep.output,
           timestamp: DateTime.now(),
           messageType: MessageType.agent,
-          jsonResponse: executedStepResponse);
+          jsonResponse: executedStepResponse,
+          artifacts: executedStep.artifacts);
 
       _chats.add(agentChat);
 
@@ -165,4 +165,20 @@ class ChatViewModel with ChangeNotifier {
       // TODO: Handle additional error scenarios or log them as required
     }
   }
+
+  /// Downloads an artifact associated with a specific chat.
+  ///
+  /// [taskId] is the ID of the task.
+  /// [artifactId] is the ID of the artifact to be downloaded.
+  Future<void> downloadArtifact(String taskId, String artifactId) async {
+    try {
+      // Call the downloadArtifact method from the ChatService class
+      await _chatService.downloadArtifact(taskId, artifactId);
+
+      print("Artifact $artifactId downloaded successfully for task $taskId!");
+    } catch (error) {
+      print("Error downloading artifact: $error");
+      // TODO: Handle the error appropriately, perhaps notify the user
+    }
+  }
 }
diff --git a/frontend/lib/views/chat/agent_message_tile.dart b/frontend/lib/views/chat/agent_message_tile.dart
index 7ea81f6ff..0587175b8 100644
--- a/frontend/lib/views/chat/agent_message_tile.dart
+++ b/frontend/lib/views/chat/agent_message_tile.dart
@@ -6,10 +6,12 @@ import 'package:flutter/material.dart';
 
 class AgentMessageTile extends StatefulWidget {
   final Chat chat;
+  final VoidCallback onArtifactsButtonPressed;
 
   const AgentMessageTile({
     Key? key,
-    required this.chat, // The agent message to be displayed
+    required this.chat,
+    required this.onArtifactsButtonPressed,
   }) : super(key: key);
 
   @override
@@ -22,12 +24,12 @@ class _AgentMessageTileState extends State<AgentMessageTile> {
   @override
   Widget build(BuildContext context) {
     String jsonString = jsonEncode(widget.chat.jsonResponse);
+    int artifactsCount = widget.chat.artifacts.length;
+
     return LayoutBuilder(
       builder: (context, constraints) {
-        double chatViewWidth = constraints.maxWidth; // Get the chat view width
-        double tileWidth = (chatViewWidth >= 1000)
-            ? 900
-            : chatViewWidth - 40; // Determine tile width
+        double chatViewWidth = constraints.maxWidth;
+        double tileWidth = (chatViewWidth >= 1000) ? 900 : chatViewWidth - 40;
 
         return Align(
           alignment: Alignment.center,
@@ -43,13 +45,11 @@ class _AgentMessageTileState extends State<AgentMessageTile> {
             child: Column(
               mainAxisAlignment: MainAxisAlignment.start,
               children: [
-                // Container for Agent title, message, and controls
                 Container(
                   constraints: const BoxConstraints(minHeight: 50),
                   child: Row(
                     crossAxisAlignment: CrossAxisAlignment.center,
                     children: [
-                      // Agent title
                       const Text(
                         "Agent",
                         style: TextStyle(
@@ -59,7 +59,6 @@ class _AgentMessageTileState extends State<AgentMessageTile> {
                         ),
                       ),
                       const SizedBox(width: 20),
-                      // Message content
                       Expanded(
                         child: Container(
                           padding: const EdgeInsets.fromLTRB(0, 10, 20, 10),
@@ -69,9 +68,8 @@ class _AgentMessageTileState extends State<AgentMessageTile> {
                           ),
                         ),
                       ),
-                      // Artifacts button (static for now)
                       ElevatedButton(
-                        onPressed: () {},
+                        onPressed: widget.onArtifactsButtonPressed,
                         style: ElevatedButton.styleFrom(
                           backgroundColor: Colors.white,
                           foregroundColor: Colors.black,
@@ -80,7 +78,7 @@ class _AgentMessageTileState extends State<AgentMessageTile> {
                             borderRadius: BorderRadius.circular(8),
                           ),
                         ),
-                        child: const Text("2 Artifacts"),
+                        child: Text('$artifactsCount Artifacts'),
                       ),
                       const SizedBox(width: 20),
                       // Expand/Collapse button
@@ -91,7 +89,7 @@ class _AgentMessageTileState extends State<AgentMessageTile> {
                             : Icons.keyboard_arrow_down),
                         onPressed: () {
                           setState(() {
-                            isExpanded = !isExpanded; // Toggle expanded view
+                            isExpanded = !isExpanded;
                           });
                         },
                       ),
@@ -105,10 +103,8 @@ class _AgentMessageTileState extends State<AgentMessageTile> {
                     child: SizedBox(
                       height: 200,
                       child: Padding(
-                        padding: const EdgeInsets.only(
-                            right: 20), // Padding for the right side
+                        padding: const EdgeInsets.only(right: 20),
                         child: JsonCodeSnippetView(
-                          // JSON code snippet view
                           jsonString: jsonString,
                         ),
                       ),
diff --git a/frontend/lib/views/chat/chat_view.dart b/frontend/lib/views/chat/chat_view.dart
index a03cb6827..44b93a66d 100644
--- a/frontend/lib/views/chat/chat_view.dart
+++ b/frontend/lib/views/chat/chat_view.dart
@@ -8,8 +8,6 @@ import 'package:flutter/material.dart';
 import 'package:auto_gpt_flutter_client/viewmodels/chat_viewmodel.dart';
 import 'package:provider/provider.dart';
 
-// TODO: Implement artifacts
-
 class ChatView extends StatefulWidget {
   final ChatViewModel viewModel;
 
@@ -84,7 +82,24 @@ class _ChatViewState extends State<ChatView> {
                 if (chat.messageType == MessageType.user) {
                   return UserMessageTile(message: chat.message);
                 } else {
-                  return AgentMessageTile(chat: chat);
+                  return AgentMessageTile(
+                    chat: chat,
+                    onArtifactsButtonPressed: () {
+                      // TODO: Create an actual artifact object
+                      // Loop through each artifact and download it using the artifact_id
+                      for (var artifact in chat.artifacts) {
+                        if (artifact is Map) {
+                          final artifactMap = artifact.cast<String,
+                              dynamic>(); // Cast each item to Map<String, dynamic>
+
+                          final artifactId = artifactMap['artifact_id']
+                              .toString(); // Get the artifact_id
+                          widget.viewModel.downloadArtifact(
+                              chat.taskId, artifactId); // Download the artifact
+                        }
+                      }
+                    },
+                  );
                 }
               },
             ),