From 55af7e0f1569e92d34bdbe367e26d3d51e7ef507 Mon Sep 17 00:00:00 2001
From: Brandon Hancock <brandon@brandonhancock.io>
Date: Wed, 3 Jul 2024 11:09:19 -0400
Subject: [PATCH] WIP. Needing team to review change

---
 src/crewai/crews/crew_output.py | 20 ++++++++++++++------
 src/crewai/tasks/task_output.py |  2 +-
 tests/crew_test.py              |  7 +++++++
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/crewai/crews/crew_output.py b/src/crewai/crews/crew_output.py
index 4bbf91354..99a8e37bc 100644
--- a/src/crewai/crews/crew_output.py
+++ b/src/crewai/crews/crew_output.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Union
 
 from pydantic import BaseModel, Field
 
@@ -16,20 +16,28 @@ class CrewOutput(BaseModel):
     )
 
     # TODO: Ask @joao what is the desired behavior here
-    def result(self) -> Union[str, BaseModel, Dict[str, Any]]:
+    def result(
+        self,
+    ) -> List[str | BaseModel | Dict[str, Any]]]:
         """Return the result of the task based on the available output."""
         results = [output.result() for output in self.output]
-        return results if len(results) > 1 else results[0]
+        return results
 
     def raw_output(self) -> str:
         """Return the raw output of the task."""
         return aggregate_raw_outputs_from_task_outputs(self.output)
 
-    def to_output_dict(self) -> Dict[str, Any]:
-        self.output.to_output_dict()
+    def to_output_dict(self) -> List[Dict[str, Any]]:
+        output_dict = [output.to_output_dict() for output in self.output]
+        return output_dict
 
     def __getitem__(self, key: str) -> Any:
-        self.output[key]
+        if len(self.output) == 0:
+            return None
+        elif len(self.output) == 1:
+            return self.output[0][key]
+        else:
+            return [output[key] for output in self.output]
 
     # TODO: Confirm with Joao that we want to print the raw output and not the object
     def __str__(self):
diff --git a/src/crewai/tasks/task_output.py b/src/crewai/tasks/task_output.py
index 893b6c335..8dc50b0ae 100644
--- a/src/crewai/tasks/task_output.py
+++ b/src/crewai/tasks/task_output.py
@@ -49,7 +49,7 @@ class TaskOutput(BaseModel):
         if self.json_output:
             output_dict.update(self.json_output)
         if self.pydantic_output:
-            output_dict.update(self.pydantic_output.dict())
+            output_dict.update(self.pydantic_output.model_dump())
         return output_dict
 
     def __str__(self) -> str:
diff --git a/tests/crew_test.py b/tests/crew_test.py
index bc5c69d14..f1140630d 100644
--- a/tests/crew_test.py
+++ b/tests/crew_test.py
@@ -7,6 +7,7 @@ from unittest.mock import patch
 
 import pydantic_core
 import pytest
+
 from crewai.agent import Agent
 from crewai.agents.cache import CacheHandler
 from crewai.crew import Crew
@@ -146,6 +147,7 @@ def test_crew_creation():
     assert result.raw_output() == expected_string_output
     assert isinstance(result, CrewOutput)
     assert len(result.tasks_output) == len(tasks)
+    assert result.result() == [expected_string_output]
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
@@ -1736,3 +1738,8 @@ def test__setup_for_training():
 
     for agent in agents:
         assert agent.allow_delegation is False
+
+
+# TODO: TEST EXPORT OUTPUT TASK WITH PYDANTIC
+# TODO: TEST EXPORT OUTPUT TASK WITH JSON
+# TODO: TEST EXPORT OUTPUT TASK CALLBACK