From ba93c25ff4d5044d901a0bc2405649f1aee71edd Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Thu, 29 Feb 2024 23:16:45 -0300
Subject: [PATCH] Add Decision class for decision making components

---
 .../langflow/components/utilities/Branch.py     |  4 +++-
 .../components/utilities/ShouldRunNext.py       |  4 +++-
 src/backend/langflow/schema.py                  | 17 +++++++++++++++++
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/src/backend/langflow/components/utilities/Branch.py b/src/backend/langflow/components/utilities/Branch.py
index daa103e09..3e786b7b9 100644
--- a/src/backend/langflow/components/utilities/Branch.py
+++ b/src/backend/langflow/components/utilities/Branch.py
@@ -1,5 +1,6 @@
 from langflow import CustomComponent
 from langflow.field_typing import Text
+from langflow.schema import Decision
 
 
 class BranchComponent(CustomComponent):
@@ -11,4 +12,5 @@ class BranchComponent(CustomComponent):
         return {"param": {"display_name": "Parameter"}}
 
     def build(self, param: Text) -> Text:
-        return {"path": True, "result": param}
+
+        return Decision(path=True, result=param)
diff --git a/src/backend/langflow/components/utilities/ShouldRunNext.py b/src/backend/langflow/components/utilities/ShouldRunNext.py
index 995f1cb43..bdc328884 100644
--- a/src/backend/langflow/components/utilities/ShouldRunNext.py
+++ b/src/backend/langflow/components/utilities/ShouldRunNext.py
@@ -5,6 +5,7 @@ from langchain_core.prompts import PromptTemplate
 
 from langflow import CustomComponent
 from langflow.field_typing import BaseLanguageModel, Prompt
+from langflow.schema import Decision
 
 
 class ShouldRunNext(CustomComponent):
@@ -49,4 +50,5 @@ class ShouldRunNext(CustomComponent):
         # The string should be the words true or false
         # if not raise an error
         bool_result = result.lower() == "true"
-        return {"path": bool_result, "result": kwargs}
+
+        return Decision(path=bool_result, result=kwargs)
diff --git a/src/backend/langflow/schema.py b/src/backend/langflow/schema.py
index e9f437038..bfc8f8efa 100644
--- a/src/backend/langflow/schema.py
+++ b/src/backend/langflow/schema.py
@@ -68,3 +68,20 @@ def docs_to_records(documents: list[Document]) -> list[Record]:
         list[Record]: The converted list of Records.
     """
     return [Record.from_document(document) for document in documents]
+
+
+#  {"path": bool_result, "result": kwargs}
+# Create a class for the above dictionary
+# with a good name that fits the context of
+# a decision making component
+class Decision(BaseModel):
+    """
+    Represents a decision made in the Graph.
+
+    Attributes:
+        path (bool): The path taken in the Graph.
+        result (dict): The result of the decision.
+    """
+
+    path: bool
+    result: Any